Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning

다음은 논문 "Unlearning What Matters: Token-Level Attribution for Precise Language Model Unlearning"에 대한 상세한 기술 요약입니다.

1. 문제 제기

대규모 언어 모델 (LLM) 은 종종 민감하거나 저작권이 있거나 유해한 데이터를 암기하므로, 처음부터 재학습하지 않고 특정 지식을 제거하기 위해 **기계적 망각 (machine unlearning)**이 필요합니다.

• 현재의 한계: 기존 망각 방법들은 시퀀스 수준에서 작동하여 시퀀스 내 모든 토큰에 균일한 그래디언트 업데이트를 적용합니다.
• 핵심 문제: 지식은 균일하게 인코딩되지 않습니다. 오직 희소 부분집합의 토큰들 (예: 특정 개체, 날짜, 사실) 만이 표적 정보를 담고 있으며, 나머지는 구조적 또는 문맥적 역할을 수행합니다.
• 균일 업데이트의 결과:
  1. 그래디언트 노이즈: 비중요한 토큰을 업데이트하면 노이즈가 발생하여 망각 과정의 신호 대 잡음비 (SNR) 가 저하됩니다.
  2. 유용성 저하: 무분별한 업데이트는 관련 없는 지식에 대한 모델의 성능을 손상시킵니다 (부수적 피해).
  3. 비최적 망각: 그래디언트 신호가 관련 없는 토큰에 의해 희석되므로 모델이 표적 지식을 효과적으로 제거하지 못합니다.

2. 방법론: TokenUnlearn

저자들은 "지식-중요" 토큰을 식별하고 선택적으로 표적화함으로써 망각을 시퀀스 수준에서 토큰 수준으로 전환하는 TokenUnlearn 프레임워크를 제안합니다.

A. 토큰 귀속 메커니즘

망각해야 할 지식을 인코딩하는 토큰을 식별하기 위해, 프레임워크는 두 가지 신호를 사용하여 각 토큰에 대한 복합 중요도 점수 ($\phi_i$) 를 계산합니다.

1. 지식 인식 귀속 (마스킹):
   • 메커니즘: 모델이 원래 입력 ($z_{orig}$) 과 지식 관련 부분 (특히 프롬프트의 명사) 이 마스킹된 입력 ($z_{mask}$) 에 대한 예측을 생성합니다.
   • 지표: 절대 로그 확률 변화량 ($\Delta^{unlearn}_i$) 을 계산합니다. 큰 변화량은 해당 토큰의 예측이 표적 지식에 크게 의존함을 나타냅니다.
2. 엔트로피 인식 보완 신호:
   • 메커니즘: 토큰의 예측 엔트로피 ($H_i$) 를 계산합니다.
   • 근거: 높은 엔트로피는 모델이 지식 의존적 연속체 (예: 여러 가능한 사실 사이에서 선택) 에 대해 불확실한 결정 지점을 나타낼 수 있음을 의미합니다.
   • 결합: 최종 점수는 가중 합입니다: $\phi_i = \alpha \cdot \bar{\Delta}^{unlearn}_i + (1-\alpha) \cdot \bar{H}_i$.

B. 토큰 선택적 최적화 전략

중요도 점수가 생성되면, 프레임워크는 두 가지 전략을 사용하여 고중요도 토큰에만 망각 목적을 적용합니다.

1. 하드 선택 (Hard Selection):
   • $\phi_i$를 기반으로 상위-$r$ 비율의 토큰 (예: 상위 20%) 을 선택합니다.
   • 망각 손실을 오직 이러한 토큰에만 적용하며, 다른 토큰은 0 의 그래디언트 업데이트를 받습니다.
2. 소프트 가중치 (Soft Weighting):
   • 정규화된 중요도 점수 (온도 $\tau$를 사용하는 소프트맥스) 를 기반으로 모든 토큰의 그래디언트 기여도를 조절합니다.
   • 균일한 업데이트와 완전히 선택적인 업데이트 사이의 부드러운 보간을 가능하게 합니다.

C. 기존 방법과의 통합

이 프레임워크는 모델에 구애받지 않으며 표준 그래디언트 기반 망각 목적과 통합됩니다.

• 그래디언트 상승 (GA)
• 가중치 그래디언트 상승 (WGA)
• 부정적 선호도 최적화 (NPO)
• 표현 방향 전환 망각 (RMU)
• 유지: 모든 방법은 일반 능력을 보존하기 위해 "유지 집합 (retain set)"에 KL-발산 정규화와 결합됩니다.

3. 이론적 분석

이 논문은 토큰 수준 선택이 망각을 개선하는 이유에 대한 이론적 근거를 제공합니다.

• 신호 대 잡음비 (SNR): 저자들은 망각을 매개변수의 저차원 부분 공간 ($U$) 을 따른 이동으로 모델링합니다.
• 노이즈 감소: 균일한 업데이트는 모든 $T$개의 토큰에서 노이즈를 누적합니다. 토큰 수준 선택은 그래디언트 가중치를 지식-중요 토큰의 부분집합 $K$에 집중시킵니다.
• 결과: 그래디언트 추정기의 SNR 은 선택 비율 ($T/|K|$) 의 역수에 비례하여 개선됩니다. 비중요 토큰에 대한 업데이트를 억제함으로써, 이 방법은 그래디언트 노이즈를 줄이고 관련 없는 지식을 더 효과적으로 보존합니다.

4. 실험 결과

저자들은 세 가지 모델 아키텍처 (Phi-1.5, Llama-2-7B, Qwen-3-8B) 에 걸쳐 두 가지 벤치마크에서 TokenUnlearn 을 평가했습니다.

벤치마크

1. TOFU (가상 망각 작업): 세밀한 지식 망각 (가상 저자에 대한 사실 제거).
2. WMDP (대량 살상 무기 프록시): 안전 정렬 (유해한 바이오/사이버 보안 지식 제거).

주요 발견

• 우수한 망각: 토큰 수준 변형 (T-* 및 S-*) 은 망각 효과성 (망각 집합에서의 낮은 정확도) 에서 시퀀스 수준 베이스라인보다 일관되게 우수했습니다.
  • 예시: Qwen-3-8B 를 사용한 TOFU 에서 T-WGA 는 표준 WGA 에 비해 망각 효과성을 크게 개선했습니다.
• 더 나은 유용성 보존: 토큰 수준 방법은 비표적 데이터에서 더 높은 모델 유용성 (유지 점수) 을 유지했습니다.
  • 예시: Qwen-3-8B 에서 유지 점수는 더 나은 망각을 달성하는 동시에 시퀀스 수준 베이스라인 대비 최대 **19.0%**까지 개선되었습니다.
• 일반화: 이 접근법은 다양한 모델 크기와 아키텍처에서 일관되게 작동했습니다.
• 어블레이션 연구:
  • 선택 비율 ($r$): 20% 비율 (상위 20% 토큰) 이 망각과 유지 사이의 최적 균형을 제공했습니다.
  • 엔트로피 가중치: 엔트로피 신호를 결합 (약 30% 가중치) 하면 마스킹만으로는 놓칠 수 있는 불확실한 결정 지점을 포착하여 보완적인 이점을 제공했습니다.

5. 주요 기여

1. 토큰 수준 귀속 프레임워크: 보조 모델 없이 반사실적 마스킹과 엔트로피 신호를 사용하여 지식-중요 토큰을 식별하는 원칙적인 방법을 도입했습니다.
2. 선택적 망각 전략: 기존 망각 알고리즘을 토큰 수준에서 작동하도록 확장하는 "하드 선택" 및 "소프트 가중치" 전략을 제안했습니다.
3. 이론적 통찰: 토큰 선택적 업데이트가 그래디언트 SNR 을 개선하여 망각과 유지라는 이중 목적을 최적화하는 데 균일한 업데이트가 실패하는 이유를 설명했습니다.
4. 실증적 검증: 세 가지 모델과 두 가지 서로 다른 벤치마크 (지식 및 안전) 에서 일관된 개선을 입증하여, 최대 32.6% 의 망각 개선과 19.0% 의 유용성 보존 개선을 달성했습니다.

6. 중요성 및 영향

• 개인정보 보호 및 규제: 모델의 일반 유용성을 저하시키지 않고 특정 데이터 흔적을 제거함으로써 더 정밀한 "잊힐 권리" (GDPR/CCPA) 준수를 가능하게 합니다.
• 안전: 사전 학습된 모델에서 유해한 능력 (예: 바이오 보안 위협) 을 제거하는 더 효과적인 메커니즘을 제공합니다.
• 효율성: 망각이 균일한 시퀀스 업데이트를 필요로 한다는 가정에 도전하며, LLM 의 지식은 희소하며 정밀하게 표적화될 수 있음을 시사합니다.
• 한계: 이 방법은 추가적인 순방향 통과 (forward pass) 로 인해 계산 오버헤드를 도입합니다. 이는 근사 방법이며 지식 흔적의 100% 제거를 보장하지 않습니다 (잔여 흔적이 존재할 수 있음).

결론적으로, TokenUnlearn은 거친 시퀀스 업데이트에서 세밀한 토큰 귀속으로의 기계적 망각 패러다임 전환을 나타내며, 효과적인 망각과 유용성 보존 사이의 근본적인 긴장 관계를 해결합니다.