技術要約：プロトタイプ言語モデル (PRISM)

問題提起

現在の大規模言語モデル（LLM）は、**学習データ属性付け（Training Data Attribution: TDA）**における根本的な課題に直面しています。それは、どの特定の学習例がモデルの出力に影響を与えているかを特定することです。標準的な高密度（dense）アーキテクチャでは、学習データは、数十億のパラメータの中に蓄積された効果を通じて、高密度で絡み合った経路内の間接的な影響として予測に寄与します。このため、単一の例の影響は分散しており、高価で近似的、かつしばしば統計的に誤設定された事後解析（post-hoc analysis）なしには追跡不可能です。

既存の勾配ベースのTDA手法（例：影響関数/Influence Functions）は、スケールアップにおいて2つの複合的な問題に直面します：

1. 計算の非効率性： 巨大なパラメータを持つモデルに対して、逆ヘッセ行列・ベクトル積（iHVP）の正確な計算を行うことは不可能であり、忠実度を犠牲にしてスケーラビリティを確保する近似が必要となります。
2. 統計的誤設定： これらの手法は、決定論的で凸な目的関数が収束まで最適化されることを前提としていますが、確率的なLLMの学習においてはこれらの条件は成立しません。その結果、属性付けの推定値は、Leave-one-outによる再学習の結果と一致しないことがあり、また学習パイプラインのハイパーパラメータ（深さ、バッチサイズ、データの順序）に対して敏感になります。

本論文は、TDAを単なる事後的な推定問題としてではなく、アーキテクチャ設計の問題として扱うべきであると主張しています。

手法：PRISM アーキテクチャ

著者らは、**解釈可能なシーケンスモデリングのためのプロトタイプ（Prototypes for Interpretable Sequence Modeling: PRISM）**という、高密度な出力ヘッドではなく、学習されたプロトタイプの疎な非負混合によって予測を形成する言語モデルアーキテクチャを提案しています。

コア・アーキテクチャ

PRISMは標準的なTransformerデコーダーのバックボーンを保持しますが、最終的な高密度出力経路を**疎なプロトタイプ再構成（sparse prototype reconstruction）**に置き換えます。

1. プロトタイプ・バンク： 学習されたプロトタイプベクトル $P = [p_1, \dots, p_K] \in \mathbb{R}^{d \times K}$ の集合が保持されます。
2. 疎な活性化： 与えられた隠れ状態 $z_t$ に対して、すべてのプロトタイプとのコサイン類似度を計算し、ReLUを適用して非負性を強制し、上位 $k$ 個の活性化プロトタイプを選択します（$k \ll K$）。
3. 再構成： 隠れ状態は、プロトタイプ再構成 $\hat{z}_t$ と残差 $r_t$ に分解されます：
   $$z_t = \hat{z}_t + r_t = \sum_{i \in \mathcal{K}_t} a_{t,i} p_i + r_t$$
   ここで、$\mathcal{K}_t$ は活性化しているプロトタイプの集合であり、$a_{t,i}$ はそれらの活性化値です。
4. ロジット形成： 次のトークンのロジットは加算的に形成されます：
   $$\ell_t = W z_t = W r_t + \sum_{i \in \mathcal{K}_t} a_{t,i} (W p_i)$$
   ここで、$W p_i$ は各プロトタイプの固定された「トークン・ロジット・シグネチャ」を表します。

学習目的関数

PRISMは、言語モデリングの忠実度とプロトタイプの解釈可能性のバランスをとるために、4つの損失項を用いてエンドツーエンドで学習されます：

• 交差エントロピー ($L_{CE}$): 標準的な次トークン予測の品質を保証します。
• 再構成損失 ($L_{REC}$): 残差にペナルティを課し、プロトタイプが隠れ状態の大部分を説明するように促します。
• クラスタリング損失 ($L_{R1}, L_{R2}$): これら対称的な目的関数は、プロトタイプを訓練データのコヒーレントな近傍に固定します。$L_{R1}$ はプロトタイプをバッチ内の最も近いトークン表現へと引き寄せ、$L_{R2}$ はトークン状態を最も近いプロトタイプへと引き寄せます。これにより、プロトタイプが任意の疎な辞書方向へ漂流することなく、データ多様体上に位置することを保証します。

理論的洞察：曲率の局所化

主要な理論的貢献は、プロトタイプ部分空間における損失ヘッセ行列の分析です。著者らは、クラスタリング目的関数が**曲率を局所化（localize curvature）**することを示しています：

• ヘッセ行列は、グローバルな共使用項（プロトタイプがどのように共に使用されるかを捉える）と、ローカルなブロック対角項（各プロトタイプの周囲の幾何学を捉える）に分離されます。
• クラスタリングを強化すると、プロトタイプ間の結合が減少し、ヘッセ行列の条件数が改善されます。
• この構造により、スケーラブルなTDAが可能になります。曲率が局所化されているため、キャッシュされた疎な統計量を使用して影響スコアを高速に計算でき、フルパラメータ空間の勾配積を回避できます。

主な貢献

1. アーキテクチャ (PRISM): 次のトークン予測のための、初のLLMスケールのプロトタイプ言語モデル・ファミリーです。これは、予測のたびに、訓練に基づいた疎なプロトタイプの混合を露出させることで、事例ベース推論（case-based reasoning）を自己回帰モデリングへと拡張します。
2. 理論と属性付け: 疎なプロトタイプ構造とクラスタリング目的関数が、損失ランドスケープにおける曲率を局所化することを証明しました。これにより、構造化されたヘッセ行列が得られ、メモリ消費量を同等に保ちつつ、高密度な勾配ベースのベースライン（EK-FACなど）よりも約500倍高速に学習データ属性付けを実行できます。
3. 実証的性能と制御:
   • 性能: FineWebおよびNemotron-CCコーパス（最大50Bトークン）を用いて、130Mから1.6Bパラメータまでの範囲でPRISMをスケールさせました。
   • 精度: 1.6Bパラメータにおいて、10Bトークンで学習したPRISMは、高密度なGPTベースラインと同等の性能を示しました。20B〜50Bトークンを用いると、高密度ベースラインとの差を縮め、ダウンストリームの平均精度において2.5ポイント以内の差に収まりました。
   • 効率: プロトタイプ層によるオーバーヘッドは極めて小さく（XLスケールでパラメータ増加はわずか1.7%）、高密度GPTモデルの7.3%のトレーニングスループットを達成しています。
   • グラウンディング: プロトタイプの被覆指標（$\bar{R}_1, R_2$）は高い水準（0.94–1.00）を維持しており、プロトタイプが任意の疎な辞書方向へ漂流することなく、実際の訓練コンテキストに留まっていることを示しています。

実験結果

スケーリングと性能

• スケール: 130Mから1.6Bパラメータまで、FineWebおよびNemotron-CCコーパス（最大50Bトークン）を用いてスケールさせました。
• 精度: 1.6Bパラメータにおいて、10Bトークンで学習したPRISMは、高密度なGPTベースラインと同等の性能を実現しました。20B–50Bトークンを用いることで、高密度ベースラインとの差を詰め、ダウンストリーム精度の平均誤差を2.5ポイント以内に抑えました。
• 効率: プロトタイプ層による追加コストは最小限であり（XLスケールでパラメータ増加は1.7%）、高密度GPTモデルのスループットの97.3%を達成しています。
• グラウンディング: プロトタイプの被覆指標（$\bar{R}_1, R_2$）は高い値を維持しており（0.94–1.00）、プロトタイプが訓練データの現実的な文脈から逸脱していないことを示しています。

属性付けの効率

• 速度: キャッシュされたプロトタイプ属性付けは、同等のGPUメモリ制約下での124Mモデルにおいて、EK-FACに対して470倍の高速化を達成しました。
• 信号の保存: キャッシュされたスコアは、フルプロトタイプ空間の勾配と比較して、上位ランクの訓練シーケンスの約80%を保持しており、疎で局所的な幾何学が、必要な属性付け信号を保持していることを実証しました。

制御とステアリング

• コントローラー: 活性化プロトタイプに対する疎な線形コントローラーは、4つの非バイナリ多肢選択タスクにおいて、平均+2.72ポイントの精度向上を実現し、SAEベースのコントローラー（+0.32ポイント）を大幅に上回りました。
• 安全性への整列（Alignment）: 推論時に特定された86個のNSFW関連プロトタイプ（バンクの0.52%）を抑制することで、50,000回の生成における平均NSFWスコアを0.668から0.154に低下させました。この際、テキストの品質は変化しませんでした。最も困難なケース（ベースのNSFWスコア = 2）においても、スコアは0.144まで低下しました。

意義と主張

本論文は、PRISMが、解釈可能な言語モデルをどのように設計すべきかというパラダイムシフトを示していると主張しています。高密度なパラメータから訓練データの痕跡を回収しようとする困難な事後解析に頼るのではなく、PRISMは予測経路の中に、訓練に基づいた疎な座標を直接的に露出させています。

著者らは、予測性能が同等であるならば、その計算が属性付け、監査、および制御しやすいアーキテクチャの方が好ましいと主張しています。PRISMは以下の方法でこれを実現しています：

1. 学習データ属性付けを、モデルアーキテクチャのネイティブかつ計算コストの低い特性にすること。
2. ファインチューニングを行うことなく、モデルの挙動に対する直接的な介入（ステアリング、抑制）を可能にすること。
3. 特定のテスト時の挙動を、特定の訓練近傍へと遡ることができる「モデルの汎化の科学」を提供すること。

この研究は、「設計による解釈可能性（interpretable by design）」を備えたモデリングが大規模なスケールでも実現可能であることを示唆しており、標準的な高密度LLMに見られる効果の絡み合いに対する、実用的な代替案を提示しています。