Lifelong Knowledge Editing requires Better Regularization 논문 리뷰

Lifelong Knowledge Editing requires Better Regularization

Summary

Knowledge editing에서는 다음과 같은 테스크를 목표로 한다:

• 새로운 지식 업데이트(혹은 수정) 과정에서 기존의 지식이 무너지지 않게 방지
• 새로운 지식 업데이트(혹은 수정) 과정에서 업데이트 된 지식이 overfit 되는 것을 방지
• 단순히 하나의 지식을 수정하는 것이 아닌, 여러 지식을 수정(sequential editing)하면서 학습 안정성 확보

기존의 논문 AlphaEdit, MEMIT 등은 새로운 지식을 업데이트하는 과정에서 기존의 지식이 무너지지 않게 방지하지만,

MEMIT의 경우 Loss Objective에 기존 지식에 관련된 term이 포함되어 기존 지식이 무너지는 것을 막기 위한 hyperparameter 튜닝이 어려우며, 업데이트 된 지식이 overfit 되는 경향이 있다. AlphaEdit 또한 해당 문제를 일부 해결했으나, sequential editing이 길어질수록 학습 안정성이 떨어지는 문제점을 보인다.

 

따라서 본 논문은 MPES(Most-Propable Early Stopping)와 NC(Norm Constraint)를 활용하며 MEMIT와 AlphaEdit이 가진 문제점을 해결하고자 한다. 또한, 기존 metric이 sequential editing 안정성을 충분히 반영하지 못한다는 점을 보완하기 위해 새로운 평가 지표를 제안, 각 기법의 편집 성능과 안정성을 종합적으로 검증하였다.

Evaluation

본 논문은 기존 방법론들에 대한 실험 및 평가, 이후 새로운 방안을 제안하는 방식으로 서술한다. 따라서 새로운 metric을 소개한 후 기존 방법론이 가지는 문제점에 대해서 분석한 것을 먼저 제시하고자 한다.

Metrics

• Efficacy Score (ES): Editing process가 효과적으로 이루어졌는지에 대한 측도. $P(\text{new fact}) > P(\text{old fact})$인 경우의 비율을 확인함
• Paraphrase Score (PS): 다른 시나리오에 대해서 모델이 수정된 사실을 얼마나 회상(recall)하는지에 대한 측도. ES에서 사실을 직접 묻는 질문으로 측정한다면, PS에서는 다른 시나리오를 프롬프트로 제공하여 $P(\text{new fact}) > P(\text{old fact})$를 확인함
• Neighborhood Score (NS): 수정하고자 하는 지식이 관련되지 않은 지식을 얼마나 방해하는지(disturb)에 관한 측도. Edited fact의 neighborhood에서 영향받지 않은 지식(fact)의 비율을 확인함
• Overall Score (S): ES/PS/NS의 조화평균. 낮을수록 editing이 실패했다고 볼 수 있음
• Generation Entropy (GE): 모델이 생성한 텍스트를 토큰 단위로 나누어, bi-gram과 tri-gram 엔트로피를 가중평균한다. 즉 낮은 GE를 가지고 있으면 엔트로피가 낮으니 부자연스럽고 의미 없는 반복을 하며 GE가 높을수록 단어 조합이 다양(fluent)해진다.

이전의 metric은 Downstream Performance Metrics이며, 좋은 방법이기는 하지만 여전히 Model editing을 직접적으로 평가하지는 못한다는 문제점이 있다. 다만 본 논문에서는 이 metric도 포함시켰다.

 

주의: AlphaEdit의 다운스트림 성능이 MEMIT보다 우수함에도 불구하고, 새로 제시된 Metric으로 측정한 경우에는 오히려 점수가 더 낮게 나타났다. 이는 해당 Metric이 아직 완전하지 않음을 시사하며, 따라서 기존 Metric과 새로운 Metric을 보완적으로 활용하는 것이 바람직하다.

 

본 논문이 기존 방법론에 대해서 지적하는 것은 다음과 같다:

Sequential Editing은 과적합을 불러일으킨다.

두 번째 그림에서 보이는 빨간 색 부분은, 편집된 레이어이다. 편집된 레이어가 전체 레이어의 대부분을 차지하여 실제 단어 생성에 관여하는 최종 representation이 편집된 레이어에 지나치게 관여되어 표현의 자유도가 떨어진다. 즉 편집된 레이어가 모델의 다른 부분에서 생성된 정보를 전부 덮어버리는(override) 문제가 있다.

 

논문에서는 이 현상을 단순히 LayerNorm이나 RMSNorm으로는 완화할 수 없다고 하며, 이들은 입력 데이터를 정규화할 뿐, 편집된 레이어가 최종 representation의 내용 자체를 바꾸어버리는 것을 막을 수 없다고 한다.

 

이 부분을 뒷받침하기 위해 AlphaEdit을 이용하여 추가적인 실험을 하였다. AlphaEdit은 이미 loss function에 Norm에 대한 constraint term이 존재한다.

본 논문은 AlphaEdit의 loss에서 해당 제약 조건을 없애보며 실험했는데,

Score가 거의 무너지는 것을 확인할 수 있다. (무엇보다 이 논문의 평가 지표에 대한 문제점이 드러나기도 하는데, 그렇게 무너진 모델이 또 GE에서는 성능이 준수하다. 따라서 평가 지표에 대한 재논의가 필요하다)

위의 도표는 모델이 편집 방식에 따라서 업데이트된 fact를 얼마나 예측하는지에 대한 확률을 보여준다. MEMIT의 경우 그 확률이 지나치게 많이 올라 기존 지식이 붕괴되거나 특정 패턴을 반복하는 것에 빠질 수 있음을 보여주며, LTI(KL divergence constraint를 두어 편집 전과 편집 후의 예측 분포 차이 최소화)를 적용하는 경우 그 제약이 지나쳐 새로운 fact의 예측력 자체를 너무 낮추었다.

(이 숫자만으로는 예측 자체가 좋아졌다고 볼 수는 없고, 예측이 지나치게 심해졌거나 너무 약해졌거나 둘 중 하나를 확인하면 된다. 숫자가 미미하게 줄어들거나 늘어난 경우에 대해서는 해석할 수 없는 수치이다)

 

결론적으로, 기존 방법론들은 model editing 자체가 다소 위험한 경향이 있으며, sequential model editing에서 norm constraint는 반드시 필수적임을 알 수 있다. 또한 이렇게만 제약을 둘 것이 아니라 동적인 early stopping 알고리즘을 두면 빠른 학습과 과적합 방지를 도달할 수 있기에 본 논문은 MPES를 제안한다. 아래의 표로 MPES의 정당성을 부여하고자 한다:

Methods

사실 실험 결과가 더욱 중요하기는 해서... Norm constraint와 MPES 알고리즘만 보여주고 넘어가고자 한다.

Norm Constraint

보다보면 어? MEMIT에도 안정화 term($\Delta$)가 존재하는데 왜 Norm Constraint를 달아준다고 하지? 할 수 있다. 차이점은 행렬 전체의 크기 성장을 억제하는, frobenius norm을 이용한 제약조건이라는 것이다.

Frobenius Norm을 적용하는 이유는, 단순히 행렬의 특정 방향 벡터만 크기 제한을 하는 것(L2 regularization)이 아니라 행렬 크기 자체를 제한하여 앞서 설명한 특정 layer가 지나치게 커지는 것을 막을 수 있기 때문이다.

 

MPES(Most-Probable Early Stopping)

Code는 다음과 같다:

# inputs:
#   model, prompts (N개), target_token_ids (길이 T), optimizer
#   prob_cutoff (float)   # >0: 확률 임계치, <0: 성공 후 추가 스텝 수
#   max_steps (int)

correct_counter = 0
for step in range(max_steps):
    optimizer.zero_grad()

    # 1) 모델 통과하여 각 위치의 log_probs 계산 (shape: N x seq_len x V)
    log_probs = forward_and_get_log_probs(model, prompts)  # softmax(logits)

    # 2) 타깃 토큰 위치 마스크와 인덱스 만들기
    #    gather로 타깃 토큰의 log prob만 뽑아냄 → (N x T)
    target_log_probs = gather_target_log_probs(log_probs, target_token_positions, target_token_ids)

    # 3) 프롬프트별 평균 NLL 계산
    #    nll_i = -(1/T) * sum_t log p(y_t | prompt_i)
    nll_each = - target_log_probs.mean(dim=1)          # shape: (N,)

    # 4) 평균 예측 확률(프롬프트 평균) 계산
    #    avg_prob = mean_i exp(-nll_i)
    avg_prob = torch.exp(-nll_each).mean().item()

    # 5) top-1 달성 여부 체크 (모든 타깃 위치에서 argmax == 정답인지)
    #    current_top = argmax over vocab at 각 타깃 위치
    num_correct = count_top1_hits(log_probs, target_token_positions, target_token_ids)
    all_top1 = (num_correct == total_target_positions)

    # ---- MPES stopping rules ----
    if prob_cutoff > 0:
        # 확률 임계치 모드: avg_prob가 기준 넘으면 즉시 종료
        if avg_prob >= prob_cutoff:
            break

    elif prob_cutoff < 0:
        # 성공 후 추가 스텝 모드: 모든 위치 top-1 달성 후 |prob_cutoff| 스텝 더
        if all_top1:
            correct_counter += 1
            if correct_counter >= abs(prob_cutoff):
                break

    # 6) 아직 종료 아님 → 손실 계산/역전파/스텝
    loss = compute_locate_loss(target_log_probs)  # 예: nll_each.me_

 

이렇게 적용하는 경우 하드코딩된 기존의 Early Stop 방법보다 더욱 유연하게 조기 종료를 할 수 있다.

 

Result

다른 모델들이 Sequential editing 과정에서 모델이 무너지는 것에 반해, NC와 MPES를 적용하는 경우 downstream task에서 무너지지 않는 것을 확인할 수 있다.

MPES와 NC를 적용하는 것이 기존 방법론들의 성능을 향상시키는 것을 확인할 수 있다.

 

Conclusion

본 논문은 Model editing이 가져오는 결과에 대해서 면밀히 분석하여 Norm Constraint Term을 제공하는 것이 바람직하며, 학습 시간 및 성능에 있어서 MPES를 적용하는 것이 바람직함을 보여주고 있다.

다만 살짝 아쉬운 점은 새로 만들어진 평가 metric이 다소 일관적이지 않다는 점, 그리고 Downstream task에 대해서 다양하게 벤치마크를 보여주지 않았다는 점이다.

 

그럼에도 NC와 MPES의 중요성을 알린 점에서 매우 좋게 생각한다! 끝.