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. トークン帰属メカニズム

どのトークンが忘却すべき知識を符号化しているかを特定するため、このフレームワークは 2 つのシグナルを用いて各トークンの複合重要度スコア（$\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. トークン選択的最適化戦略

重要度スコアが生成された後、このフレームワークは 2 つの戦略を用いて、高重要度のトークンのみに対して忘却目的関数を適用します。

1. ハード選択:
   • $\phi_i$ に基づいて上位$r$分（例：上位 20%）のトークンを選択します。
   • 忘却損失をこれらのトークンのみ排他的に適用し、他のトークンには勾配更新を適用しません。
2. ソフト重み付け:
   • 正規化された重要度スコア（温度$\tau$を用いた softmax）に基づき、すべてのトークンの勾配寄与を調整します。
   • 一様更新と完全な選択的更新との間の滑らかな補間を可能にします。

C. 既存手法との統合

このフレームワークはモデル非依存であり、標準的な勾配ベースの忘却目的関数と統合されます。これには以下が含まれます：

• 勾配上昇（Gradient Ascent: GA）
• 重み付き勾配上昇（Weighted Gradient Ascent: WGA）
• 負の選好最適化（Negative Preference Optimization: NPO）
• 表現誤方向忘却（Representation Misdirection Unlearning: RMU）
• 保持: すべての手法は、一般的能力を維持するために「保持セット（retain set）」に対する KL 発散正則化と組み合わされます。

3. 理論的解析

この論文は、トークンレベルの選択が忘却を改善する理由を理論的に正当化しています。

• 信号対雑音比（SNR）: 著者は、忘却をパラメータの低次元部分空間（$U$）に沿った移動としてモデル化します。
• ノイズ低減: 一様更新はすべての$T$個のトークンからのノイズを蓄積します。トークンレベルの選択は、勾配重みを知識に重要なトークンのサブセット$K$に集中させます。
• 結果: 勾配推定量の SNR は、選択比率（$T/|K|$）の逆数に比例して改善します。重要でないトークンへの更新を抑制することで、この手法は勾配ノイズを低減し、無関係な知識をより効果的に保持します。

4. 実験結果

著者は、3 つのモデルアーキテクチャ（Phi-1.5、Llama-2-7B、Qwen-3-8B）にわたる 2 つのベンチマークで TokenUnlearn を評価しました。

ベンチマーク

1. TOFU（Task of Fictitious Unlearning）: 微細な知識忘却（架空の著者に関する事実の削除）。
2. WMDP（Weapons of Mass Destruction Proxy）: 安全性アライメント（危険なバイオ/サイバーセキュリティ知識の削除）。

主要な知見

• 優れた忘却: トークンレベルのバリアント（T-* および S-*）は、忘却の有效性（忘却セットでの精度が低いこと）において、シーケンスレベルのベースラインを一貫して上回りました。
  • 例: Qwen-3-8B における TOFU では、標準的な WGA に比べて T-WGA が忘却の有效性を大幅に改善しました。
• より良い有用性保持: トークンレベルの手法は、非ターゲットデータにおいてより高いモデル有用性（保持スコア）を維持しました。
  • 例: Qwen-3-8B において、忘却の改善を達成しながら、保持スコアはシーケンスレベルのベースラインと比較して最大**19.0%**向上しました。
• 汎用性: このアプローチは、異なるモデルサイズやアーキテクチャ間で一貫して機能しました。
• アブレーション研究:
  • 選択比率（$r$）: 20%（上位 20% のトークン）の比率が、忘却と保持の間の最適なトレードオフを提供しました。
  • エントロピー重み: エントロピーシグナルの組み合わせ（約 30% の重み）は、マスキング単独では見逃す可能性のある不確実な意思決定点を捉えることで、補完的な利益を提供しました。

5. 主要な貢献

1. トークンレベル帰属フレームワーク: 補助モデルを必要とせず、対照マスキングとエントロピーシグナルを用いて知識に重要なトークンを特定する原理的な手法を導入しました。
2. 選択的忘却戦略: 既存の忘却アルゴリズムをトークンレベルで動作するように拡張する「ハード選択」と「ソフト重み付け」の戦略を提案しました。
3. 理論的洞察: トークン選択的更新が勾配 SNR を改善することを証明し、忘却と保持という二重の目的を最適化するために一様更新が失敗する理由を説明しました。
4. 実証的検証: 3 つのモデルと 2 つの異なるベンチマーク（知識と安全性）にわたる一貫した改善を実証し、忘却において最大32.6%、有用性保持において**19.0%**の改善を達成しました。

6. 意義と影響

• プライバシーと規制: モデルの一般有用性を劣化させることなく特定のデータ痕跡を削除することで、より精密な「忘れられる権利（Right to be Forgotten）」への準拠（GDPR/CCPA）を可能にします。
• 安全性: 事前学習済みモデルから危険な能力（例：バイオセキュリティ脅威）を削除するためのより効果的なメカニズムを提供します。
• 効率性: 忘却には一様なシーケンス更新が必要であるという仮定に挑戦し、LLM 内の知識はスパースであり、正確にターゲットを絞って削除可能であることを示唆します。
• 限界: この手法は、マスキングのための追加のフォワードパスにより計算オーバーヘッドを導入します。これは近似手法であり、知識痕跡の 100% 削除を保証するものではなく（残留痕跡が存在する可能性があります）。

結論として、TokenUnlearnは、粗粒度のシーケンス更新から微細粒度のトークン帰属へと移行する機械的忘却におけるパラダイムシフトを表しており、効果的な忘却と有用性保持の間の根本的な緊張関係を解決します。