[論文レビュー] Circuit Breaking: Removing Model Behaviors with Targeted Ablation

Maximilian Xiling Li, Xander Davies+1 more|arXiv (Cornell University)|Sep 12, 2023

Adversarial Robustness in Machine Learning被引用数 5

ひとこと要約

この論文では、言語モデルの計算グラフ内の重要な因果経路を特定し、その一部を削除することで望ましくない行動を除去する、標的的なエッジアブレーションという手法を提案している。GPT-2では、わずか12のエッジをアブレーションするだけで、毒性生成が平均で28%低下し、最も毒性の強いプロンプトでは40%低下したが、他の入力ではパフォーマンスの低下は最小限に抑えられた。

ABSTRACT

Language models often exhibit behaviors that improve performance on a pre-training objective but harm performance on downstream tasks. We propose a novel approach to removing undesirable behaviors by ablating a small number of causal pathways between model components, with the intention of disabling the computational circuit responsible for the bad behavior. Given a small dataset of inputs where the model behaves poorly, we learn to ablate a small number of important causal pathways. In the setting of reducing GPT-2 toxic language generation, we find ablating just 12 of the 11.6K causal edges mitigates toxic generation with minimal degradation of performance on other inputs.

研究の動機と目的

• ファインチューニングによっても消えない言語モデルの有害行動を除去する課題に対処すること。
• 関連しないタスクのパフォーマンスを損なうことなく、特定の望ましくない行動を無効化する手法を開発すること。
• 計算グラフのアブレーションを用いて、悪質な行動を引き起こす最小限で因果的な経路を同定すること。
• パフォーマンスを保ちながら、不連続で関連のない行動に対して高い特異性を達成すること。
• 全ファインチューニングや重みの変更を回避する、ポストトレーニング編集アプローチを検討すること。

提案手法

• 所定の粒度でモデルを有向無閉路グラフ（DAG）として再定式化し、計算をノードとして、依存関係をエッジとして表現する。
• 有害行動入力での損失増加を促進し、アブレーションをペナルティ化する微分可能目的関数を用いて、エッジ上のバイナリーマスクを学習する。
• 推論時にゼロアブレーションまたは平均アブレーションを適用し、選択されたエッジに沿った因果的影響を無効化する。
• 有害行動入力での損失増加とトレーニングデータでのパフォーマンス維持の両立を図る損失関数を最適化する。
• 小さな有害行動データセットへの過学習を防ぐために、早期停止と正則化を適用する。
• 推論時にアブレーションを適用することで、モデル重みの変更なしに特定の経路を無効化する。

Figure 1: In targeted ablation, we (1) rewrite our model as a computation graph of a desired granularity, (2) learn a binary mask over edges while regularizing to penalize ablations, and (3) ablate edges at inference time to avoid the target bad behavior.

実験結果

リサーチクエスチョン

• RQ1 少数の標的的エッジアブレーションが、言語モデルにおける特定の望ましくない行動を効果的に緩和できるか？
• RQ2 有効性と特異性の観点から、タスク算術やファインチューニングに比べて、標的的エッジアブレーションはどのように異なるか？
• RQ3 因果的経路のアブレーションによって、関連のない不連続なタスクのパフォーマンスを損なわず、悪質な行動を解消できるか？
• RQ4 アブレーションのスパarsity（スパarsity）と、非ターゲット行動におけるパフォーマンス低下のトレードオフは何か？
• RQ5 重みベースの編集や活性化プロジェクション手法に比べ、エッジアブレーションは行動除去において優れているか？

主な発見

• GPT-2で12のエッジをアブレーションした結果、毒性プロンプトの平均毒性スコアが0.458から0.328に低下し、相対的に28%の低下を達成した。
• 最も毒性を誘発するプロンプトでは、毒性スコアが0.944から0.567に低下し、相対的に39.9%の低下を達成した。
• すべての有効性と特異性の指標において、タスク算術を上回った。特に、毒性プロンプト後の不整合性が低減された。
• 84エッジを用いた平均アブレーションでも、毒性生成の緩和効果は同等であり、他の入力でのパフォーマンスは維持された。
• 共同ファインチューニングと比較して、関連のない行動への劣化が顕著に少なかったため、特異性が非常に高かった。
• この手法は、毒性生成を引き起こす回路を最小限の因果カットとして効果的に同定できており、メカニズムに基づくアプローチの有効性を裏付けた。

Figure 2: Ablating GPT-2 Small to remove toxicity. Left: Grey nodes are attention heads, and purple nodes are MLPs. Computation proceeds upwards, with horizontal alignment corresponding to layers. The computational graph has 11,611 edges; red edges are the 12 ablations learned to remove toxicity. Ri