Matryoshka attribution: Learning to attribute language model outputs to representations and weights
Abstractの概要
本論文では、言語モデルにおける帰属（アトリビューション）を、下流損失を最小化する内部コンポーネントの入れ子状サブセットの順序付けを学習する問題として定式化しています。著者らは、微分可能なsigmoid top-k演算子とランダムにサンプリングされたスパース性予算に基づく学習を中心に構築されたマスク学習手法であるMatryoshka Attribution（MAttr）を導入し、単一モデルで複数のスパース性レベルにわたる帰属を同時に学習できるようにしました。Mechanistic Interpretability Benchmark（MIB）および拡張スイートMIB+において内部表現への帰属を評価したほか、モデルチェックポイント間の補間によるパラメータ帰属にも本手法を適用しています。実践的なケーススタディとして、強化学習を用いてMAttrを訓練し、一般的な能力を維持しながら拒絶動作の原因となるファインチューニング起因の重み変化を特定しています。
新規性
主な新規性は、帰属を入れ子状サブセットの最適化問題として定式化し、ランダムな予算サンプリングとsigmoid top-kマスクを介したスパース性横断の共同学習によって実装した点にあります。また、このフレームワークを表現帰属からチェックポイント間のパラメータ帰属へと拡張し、微分不可能な行動ターゲットに対する強化学習ベースの訓練を含めた点も独自性を持っています。
成果
表現帰属において、MAttrは公式MIBリーダーボードのトップスコア獲得やノード・エッジレベルの局所化におけるベースラインに対する統計的に有意な向上など、MIBおよびMIB+で最高水準の性能を達成したと報告されています。また、本手法はタスクや基底間で転移可能なスパース回路を生成します。パラメータ帰属に関しては、Llama-3.2-1B-Instructの2%およびLlama-3.1-8B-Instructの1%のパラメータをベースモデルの値に戻すだけで、能力指標を概ね維持しつつ拒絶動作を排除できることが示されました。
論文の注目点
- MAttrは、スパース性ごとに個別のマスクを学習するのではなく、複数のスパース性予算にわたって入れ子状のマスクを最適化することで、コンポーネントに対する単一の帰属順序を学習する。
- 本手法は、MIB形式の回路局所化において勾配ベース、介入ベース、および従来のマスク学習ベースラインを上回り、スパース性に敏感なコンパクト性分析でも良好な結果を示す。
- MAttrはパラメータレベルの帰属のために強化学習で訓練可能であり、拒絶などのファインチューニングに起因する振る舞いをわずかな割合の重み変化へ局所化することを可能にする。