Forget What Matters, Keep the Rest: Selective Unlearning of Informative Tokens

1. 문제 정의 (Problem Definition)

대규모 언어 모델 (LLM) 의 **기계적 망각 (Machine Unlearning)**은 특정 학습 데이터 (예: 저작권 침해 정보, 민감한 개인정보, 유해 지식) 를 모델에서 선택적으로 제거하면서도 모델의 전반적인 성능 (Utility) 을 유지하는 것을 목표로 합니다.

기존의 망각 기법들은 주로 전체 토큰에 균일하게 손실 함수 (Loss) 를 적용하거나, **정답 토큰의 신뢰도 (Confidence)**만을 기준으로 가중치를 조정하는 방식을 사용했습니다. 그러나 이러한 접근 방식에는 다음과 같은 한계가 있습니다:

• 불필요한 성능 저하: 문법적 역할을 하는 구조적 토큰 (예: 'the', 'a' 등) 과 핵심 정보를 담고 있는 정보적 토큰 (예: 고유명사, 핵심 동사) 을 구별하지 못해, 구조적 토큰까지 과도하게 망각시키려다 모델의 일반화 능력이 떨어집니다.
• 신뢰도의 한계: 정답 토큰의 신뢰도 (Ground-truth Confidence) 만으로는 토큰이 얼마나 '정보적 (Informative)'인지 판단하기 어렵습니다. 높은 신뢰도를 가진 토큰이라도 문맥상 핵심 정보일 수 있고, 낮은 신뢰도를 가진 토큰이 단순한 반복일 수도 있기 때문입니다.
• 외부 의존성: 일부 기존 연구는 spaCy 와 같은 외부 언어 파서를 사용하여 명사나 개체명 (NER) 만을 정보적 토큰으로 간주하는데, 이는 문맥적 정보를 반영하지 못합니다.

2. 제안 방법: 엔트로피 기반 토큰 가중치 (Entropy-guided Token Weighting, ETW)

저자들은 **예측 분포의 엔트로피 (Entropy)**를 토큰의 정보성 (Informativeness) 을 판단하는 지표로 활용하는 ETW를 제안합니다.

핵심 아이디어

• 구조적 토큰 (Structural Tokens): 문법적 역할을 하며 다음 토큰 예측이 매우 명확합니다. 따라서 모델의 예측 분포가 한쪽으로 치우쳐 엔트로피가 낮습니다.
• 정보적 토큰 (Informative Tokens): 핵심 의미를 담고 있으며, 여러 가지 가능한 대안이 존재할 수 있어 예측 불확실성이 높습니다. 따라서 예측 분포가 분산되어 엔트로피가 높습니다.

수식 및 알고리즘

1. 토큰별 엔트로피 계산: $i$번째 토큰 $y_i$에 대한 엔트로피 $H(y_i)$는 모델의 예측 분포 $p(v|y_{<i}, x)$를 사용하여 계산합니다.
   $$H(y_i|y_{<i}, x; \theta) = -\sum_{v \in V} p(v|y_{<i}, x; \theta) \log p(v|y_{<i}, x; \theta)$$
2. 가중치 부여 (Reweighting): 망각 손실 (Gradient Ascent Loss) 에 적용할 가중치 $\omega_i$를 엔트로피에 비례하도록 설정합니다.
   $$\omega_i^{ETW} = \frac{n \cdot H(y_i)}{\sum_{j=1}^n H(y_j)}$$
   • 여기서 $n$은 시퀀스 길이로, 전체 손실의 크기를 일정하게 유지하기 위해 정규화합니다.
   • Stop-Gradient: 가중치 계산 시 모델 파라미터 $\theta$의 그래디언트 흐름을 차단 ($\hat{\theta}$ 사용) 하여, 가중치 계산 자체가 학습을 방해하지 않도록 합니다.
3. 최종 목적 함수:
   $$L_f = \sum_{i=1}^n \omega_i^{ETW} \log p(y_i | y_{<i}, x; \theta)$$
   • 높은 엔트로피 (정보적) 토큰에는 강한 페널티를, 낮은 엔트로피 (구조적) 토큰에는 약한 페널티를 부과하여 선택적 망각을 수행합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 정규화 기법 (ETW) 도입: 정답 토큰의 신뢰도나 외부 언어 파서에 의존하지 않고, 모델이 예측한 전체 어휘 분포의 엔트로피를 활용하여 토큰별 정보성을 정량화하는 새로운 정규화 기법을 제안했습니다.
2. 정보적/구조적 토큰 구분 능력 입증: 기존 방법들 (WGA, Imp, SatImp, TNPO, 언어 기반 방법 등) 보다 정보적 토큰과 구조적 토큰을 구분하는 능력 (ROC-AUC) 이 월등히 뛰어남을 실험적으로 증명했습니다.
3. 효율적인 망각과 성능 유지의 균형: ETW 를 적용한 모델은 망각 품질 (Forget Quality) 을 높이는 동시에, 재학습 (Retraining) 된 모델과 유사한 예측 확률 분포를 유지하여 모델의 유용성 (Utility) 저하를 최소화했습니다.

4. 실험 결과 (Results)

실험은 TOFU(가상 저자 망각 벤치마크) 와 WMDP(유해 지식 제거 벤치마크) 에서 수행되었습니다.

• TOFU 벤치마크 (LLaMA-3.2 1B, 3B, 7B):
  • 망각 품질 (Forget Quality): ETW 는 10%, 5%, 1% 망각 비율 모두에서 가장 우수한 망각 품질을 달성했습니다. 특히 10% 망각 설정에서 다른 방법들은 성능이 급격히 떨어지는 반면, ETW 는 안정적으로 성능을 유지했습니다.
  • 모델 유용성 (Model Utility): 기존 방법들은 망각 과정에서 모델의 일반화 능력이 크게 저하되었으나, ETW 는 구조적 토큰을 보호하여 유용성 손실을 최소화했습니다.
  • 정보적 토큰 분석: 망각 후 ETW 모델이 정보적 토큰에 부여하는 확률 분포가 '재학습 (Retrain) 모델'의 분포와 가장 유사하게 일치했습니다. 이는 ETW 가 핵심 정보를 정확히 제거하면서도 모델의 언어 능력을 보존했음을 의미합니다.
• WMDP 벤치마크:
  • 사이버 보안 및 바이오 보안 관련 유해 지식 제거 실험에서 ETW 는 다른 방법들보다 낮은 정확도 (망각 성공) 를 보이며, 유지해야 할 지식 (MMLU) 에 대한 성능은 잘 보존했습니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 망각 기술의 중요한 패러다임 전환을 제시합니다.

• 맥락 기반의 선택적 망각: 단순한 토큰 종류 (명사, 동사 등) 나 정답 확률에 의존하지 않고, **문맥적 예측 불확실성 (엔트로피)**을 통해 "무엇이 중요한지"를 모델 스스로 판단하게 함으로써 더 정교한 망각이 가능해졌습니다.
• 실용성: 기존 방법들이 겪던 "망각을 하면 성능이 떨어지고, 성능을 지키면 망각이 안 되는" 딜레마를 해결하여, 실제 응용 환경에서 신뢰할 수 있는 LLM 망각 솔루션을 제공합니다.
• 확장성: 다양한 모델 크기 (1B~7B) 와 다양한 망각 비율에서 일관된 성능을 보여주어, ETW 가 LLM 망각을 위한 강력한 표준 기법으로 자리 잡을 수 있음을 시사합니다.

요약하자면, "무엇이 중요한지 (정보적 토큰) 잊게 하고, 나머지는 (구조적 토큰) 기억하게 하는" 엔트로피 기반의 선택적 망각 기법인 ETW 는 기존 기법들의 한계를 극복하고 효율적인 LLM 망각을 실현하는 획기적인 방법론입니다.