技術的概要：LoRIF – スケーラブルなトレーニングデータ帰属のための低ランク影響関数

問題定義

トレーニングデータ帰属（TDA）は、特定のモデルが与えられたクエリに対して行った予測に、どのトレーニング例が最も影響を与えたかを特定することを目的としています。影響関数（IF）は、勾配と逆ヘッシアン行列を用いて「1 例除外再学習」を近似することで、TDA に対して理論的に裏付けられたアプローチを提供しますが、数百万の例で学習された数十億パラメータ規模の現代のモデルに対しては、実質的に利用不可能なままでした。

既存のスケーラブルな変種である TRAK、LoGRA、TrackStar は、投影された例ごとの勾配を保存し、減衰ガウス・ニュートンヘッシアン近似を使用することで計算上のボトルネックに対処しています。しかし、これらの手法は、帰属の質と計算の実現性の間でトレードオフを生み出す、2 つの決定的なスケーラビリティボトルネックに直面しています。

1. 勾配の保存と I/O: $N$ 個のトレーニング例に対する投影された例ごとの勾配を保存するには、$O(ND)$の空間が必要であり、ここで$D$ は有効な投影次元です。クエリ時にこれらの勾配を読み込むことがレイテンシを支配します（クエリ時間の約 96%）。
2. 逆ヘッシアンのスケーリング: $D \times D$ の逆ヘッシアン近似を形成・保存するには、$O(D^2)$ のメモリが必要です。

帰属の質を向上させるためには投影次元 $D$ を増やす必要がありますが、そうすると保存/I/O とメモリの両方のボトルネックが悪化し、最先端規模のモデルにおいて高品質な帰属を実用的にするのが不可能になります。

手法：LoRIF

LoRIF（Low-Rank Influence Functions）は、ニューラルネットワークの勾配に内在する低ランク構造を利用することで、これらのボトルネックに対処します。この手法は 3 段階のパイプラインに従います。

1. 低ランク勾配保存（I/O ボトルネックへの対応）

LoRIF は、完全な投影された例ごとの勾配行列 $\tilde{G}_i \in \mathbb{R}^{d_1 \times d_2}$（ここで $D = d_1 d_2$）を保存する代わりに、ランク-$c$ の因数分解を保存します。
$$\tilde{G}_i \approx u_i v_i^\top$$
ここで、$u_i \in \mathbb{R}^{d_1 \times c}$、$v_i \in \mathbb{R}^{d_2 \times c}$ です。

• メカニズム: この因数分解は、ブロックパワー反復法によって計算されます。
• 影響: 層あたり・サンプルあたりの保存量とクエリ時の I/O は、$O(D)$ から $O(c(\sqrt{D}))$ に削減されます。実証結果によると、$c=1$ であっても近似は意味のある帰属の質を保持し、高い $c$ よりも大きな $D$ を優先することが、より良いコスト効率をもたらします。

2. 切り捨て SVD による曲率近似（メモリボトルネックへの対応）

密な $D \times D$ 逆ヘッシアン $(G^\top G + \lambda I)^{-1}$ を形成することを避けるため、LoRIF は集約勾配行列 $G$ のランク-$r$ 切り捨て特異値分解（SVD）を計算します。
$$G \approx U_r \Sigma_r V_r^\top$$
ウッドベリー行列恒等式を用いることで、逆ヘッシアン項を $r$ 次元部分空間で近似します。
$$(V_r \Sigma_r^2 V_r^\top + \lambda I)^{-1} = \frac{1}{\lambda}I - \frac{1}{\lambda^2} V_r \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} V_r^\top$$

• メカニズム: 上位-$r$ 個の特異ベクトル（$V_r$）と特異値（$\Sigma_r$）のみが保存されます。
• 影響: 曲率項のメモリ複雑性は $O(D^2)$ から $O(Dr)$に削減されます。これにより、固定されたメモリ予算内で、はるかに大きな有効投影次元$D$ を可能にします。

3. 影響スコアの計算

クエリ時に、影響スコアは、クエリ勾配を $r$ 次元部分空間に投影し、保存されたトレーニング勾配の低ランク因数を利用することで計算されます。影響スコア $I(x_{tr}, x_{te})$ は以下のように近似されます。
$$I(x_{tr}, x_{te}) \approx \frac{1}{\lambda} g_{te}^\top g_{tr} - \frac{1}{\lambda^2} (g'_{te})^\top \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} g'_{tr}$$
ここで、$g'$ は特異ベクトル $V_r$ 上に投影された勾配を表します。最初の項はランク-$c$ の因数を用いて効率的に計算され、2 番目の項は中央の行列が対角であるため、$O(r)$ の演算のみで済みます。

主要な貢献

1. ボトルネックの特定: 本論文は、投影次元 $D$ がスケーラブルな IF 型 TDA における主要な制約であることを特定し、正確な帰属（大きな $D$）が現在非現実的であるという、質とスケーラビリティのトレードオフを明らかにしました。
2. ランク-$c$ 因数分解: 投影された勾配をランク-$c$ の因数として保存する提案により、例あたり・層あたりの保存量を $O(D)$ から $O(c\sqrt{D})$ に削減しました。
3. ランク-$r$ 曲率近似: 切り捨て SVD とウッドベリー恒等式を用いて曲率項を近似する提案により、層あたりのメモリを $O(D^2)$ から $O(Dr)$ に削減しました。
4. スケーラビリティの実証: 数百万のトレーニング例を有する最大 700 億パラメータ（Apertus-70B）のモデルにおいて、実用的なコストとレイテンシを維持しながら、効果的なデータ帰属を初めて実証しました。

実験結果

LoRIF は、1 億 2400 万から 700 億パラメータのモデル（GPT2-small、OLMo-3-7B、Apertus-70B）と、最大 380 万例のデータセットを用いて評価されました。

• 保存量とレイテンシ: LoGRA と比較して、LoRIF は帰属の質を同等かそれ以上としながら、勾配保存量を 2.3〜20 倍削減し、クエリレイテンシを 1.3〜20 倍削減しました。
  • GPT2-small において、LoRIF は中程度の保存 regimes において、保存量を最大 9 倍、レイテンシを 5.6 倍削減しました。
  • OLMo-3-7B において、LoRIF は LoGRA とほぼ同等のテールパッチスコアを達成しつつ、保存量を 20.3 倍、レイテンシを 22.0 倍削減しました。
  • Apertus-70B において、LoRIF は LoGRA よりも 1.5〜2.2 倍高いテールパッチスコアを改善し、より低い保存量とレイテンシを実現しました。
• 帰属の質: Linear Datamodeling Score（LDS）とテールパッチスコアで測定したところ、LoRIF は一貫してパレートフロンティアを改善しました。メモリ制約により以前は非現実的だったより大きな有効投影次元 $D$ を可能にすることで、より小さな $D$ に制約されたベースラインよりも高い質を達成しました。
• 定性的評価: トップ 1 の検索例に対する LLM-as-a-judge 評価では、LoRIF は LoGRA に比べて、より高い意味的関連性を持つ例（OLMo-3-7B における平均スコア 3.04 対 2.49）を検索し、完全に無関係な検索結果が少なくなることが示されました。

意義と主張

本論文は、LoRIF によって勾配ベースのトレーニングデータ帰属が最先端規模で実用的になったと主張しています。勾配の低ランク構造を利用することで、LoRIF は大規模言語モデルへの影響関数の応用を妨げてきた質とスケーラビリティのトレードオフを解決しました。

著者らは、LoRIF が新しい低ランクプリミティブに依存するのではなく、例ごとおよび集約勾配の両方で観察される既存の低ランク構造を、投影勾配影響法において体系的にどのように利用できるかを実証している点を強調しています。結果は、帰属インデックスが多数のクエリで再利用される場合、従来の手法が要求する禁止的な保存量や I/O コストを必要とすることなく、大規模 LLM の監査、デバッグ、安全性チェックに勾配ベースの TDA が現在使用可能であることを示唆しています。

本論文は限界を認めつつ、帰属インデックスの構築にはコーパス全体での勾配パス（前処理コスト）が必要であること、および現在の大規模実験は完全な事前学習コーパスではなく、教師あり微調整データに焦点を当てていることを指摘しています。しかし、この研究はスケーラブルな帰属のための新たな基準を確立し、以前はこのような手法の適用範囲外と考えられていたモデルの分析を可能にしました。