技術的概要：Selective Access Transformer (SATFormer)

問題定義

深層トランスフォーマーアーキテクチャにおいて、初期層で計算される低レベルの語彙的および構造的特徴は、残差ストリームがアテンション、正規化、およびフィードフォワード層を介して繰り返し変換されるにつれて、回復が困難になる傾向があります。この「情報の希薄化」を緩和するため、最近の研究では、後続の層が初期表現（具体的には第 1 層の値投影 $V_1$）にアクセスできるようにする、層間パスウェイが導入されました。

既存の解決策は、明確なトレードオフを持つ 2 つのカテゴリに分類されます。

1. 静的値残差（例：ResFormer）： これらは学習されたスカラー係数を追加し、層内のすべてのトークンとアテンションヘッドに対して $V_1$ を均一に露出させます。計算コストは低いものの、粒度が粗く、文脈に関係なく同じ量の初期情報を適用してしまいます。
2. 高密度または動的な代替案（例：DenseFormer, MUDDFormer, HyperConnections）： これらの手法は、より広範な層の履歴を集約するか、より細粒度のアクセスを回復するための動的パスウェイを生成します。しかし、これらは著しく高いメモリコストと低いスループットを伴います。

本研究が扱う中心的な問いは、モデルがパスウェイの「いつ」および「どこで」使用するかを「制御」することを学習すれば、単一の初期値パスウェイが競争力を持つ可能性があるかどうか、すなわち、層間再利用を純粋な接続性の問題としてではなく、制御の問題として扱うことができるかどうかです。

手法：SATFormer

著者は、Selective Access Transformer (SATFormer) を提案します。これは、初期表現の再利用を、高密度なパスウェイを必要とする接続性の問題ではなく、文脈依存のゲートによって制御される検索問題として再定義します。

アーキテクチャ

SATFormer は、ResFormer の単一の初期値パスウェイを維持しつつ、静的な層ごとのスカラー係数をトークンごと、ヘッドごとのゲートに置き換えます。

• ゲート計算： 層 $n$ のトークン $t$ および KV ヘッド $j$ に対して、ゲート $\alpha^{(n)}_{t,j}$ は、現在の正規化された隠れ状態 $x^{(n)}_t$ の単一の線形投影を通じて計算されます。
  $$\alpha^{(n)}_{t,j} = \left[ \text{ReLU}\left( x^{(n)}_t W^{(n)}_\alpha \right) \right]_j$$
  ここで、$W^{(n)}_\alpha$ は学習可能な投影行列です。ReLU 活性化関数は非負性を保証し、ゲートを正確にゼロにすることを可能にします。
• 値の修正： アテンションで使用される値ベクトルは、現在の層の値 ($V_{t,j,r}$) と第 1 層の値 ($V^{(1)}_{t,j,r}$) の混合となります。
  $$V'_{t,j,r} = V_{t,j,r} + \alpha^{(n)}_{t,j} \cdot V^{(1)}_{t,j,r}$$

設計の根拠

この設計は、効率性を確保するために 3 つの最小限の選択を行います。

1. 粒度： ゲートはトークンごとかつヘッドごとであり、$V_1$ の有用性が両軸（例：特定のトークンは語彙情報が必要だが、他のトークンは不要など）で変化する可能性を認識しています。
2. パラメータ効率： ゲートは現在の隠れ状態からの単一の線形投影を用いて計算され、層あたり $d_{model} \times N_{kv}$ の無視できるパラメータと実行時コストしか追加しません。
3. 疎性： ReLU ゲートにより、モデルは初期情報が役立たない場合に $V_1$ との接続を完全に遮断（ゼロに設定）でき、疎なアクセスパターンを促進します。

主要な貢献

1. 制御に基づく枠組み： 本論文は、初期表現の再利用のパラダイムを「接続性」（いくつのパスウェイがあるか）から「制御」（単一のパスウェイをどのようにゲートするか）へと転換します。
2. SATFormer アーキテクチャ： 静的なスカラーを入力依存のゲーティングに置き換える計算コストの低いメカニズムであり、無視できるオーバーヘッドしか追加しません。
3. パレート有利なスケーリング： SATFormer は、1 億 3 千万から 13 億パラメータまでのスケールにおいて、標準的なトランスフォーマーおよび ResFormer 両方に対して、検証損失とゼロショット精度を向上させました。
4. 機械的証拠： 分析により、学習されたゲーティングは疎であり、深さに依存し、ヘッド固有であり、トークンカテゴリ（意味的特徴を好む）に敏感であることが明らかになりました。これにより、モデルが均一なコピーではなく、選択的な再利用を学習していることが確認されました。

実験結果

検索集約型ベンチマーク

文脈内検索を必要とするタスク（TriviaQA, SWDE, SQuAD, NQ, FDA, DROP）において、SATFormer は評価されたアーキテクチャの中で最高平均スコア（29.86）を達成しました。

• ResFormer を約 1.5 平均ポイント 上回りました。
• MUDDFormer (29.855) をわずかに上回りましたが、計算プロファイルは著しく効率的でした。

効率性とスループット

SATFormer は、ベースラインのトランスフォーマーおよび ResFormer に近い計算プロファイルを維持し、高密度な手法の遅延を回避しています。

• スループット： 7 億 6 千万パラメータのスケールにおいて、SATFormer は HyperConnections よりも約 1.75 倍、MUDDFormer よりも約 1.82 倍 高いスループットを達成しました。
• メモリ： HyperConnections よりも 31% 少ない メモリ、MUDDFormer よりも 5% 少ない メモリを使用しました。
• ウォールクロックトレーニング： 壁掛け時間に対して測定した場合、SATFormer は高密度な代替案よりも速く低いトレーニング損失に到達しました。

一般言語モデリングとスケーリング

• 検証損失： SATFormer は、すべてのスケール（1 億 3 千万、3 億 4 千万、7 億 6 千万、13 億）において、ResFormer に対して最終検証損失を一貫して改善しました。
• ゼロショット精度： SATFormer は、1 億 3 千万、3 億 4 千万、13 億のスケールにおいて、ResFormer に対する平均ゼロショット精度を向上させました。最大の改善は 13 億（XL）スケールで +0.67 の平均精度 向上でした。
• 高密度な手法との比較： SATFormer は高密度な代替案（例：HyperConnections や MUDDFormer は、いくつかの指標において Medium および Large スケールでより高いゼロショットスコアを達成）を均一に支配したわけではありませんが、適応性、モデル品質、トレーニング効率の間の優れたトレードオフを提供しました。

機械的分析

• 疎なアクセス： ゲート活性化は疎であることが判明し、初期層および中間層のほとんどのヘッドは $V_1$ にほぼゼロの重みを割り当てました。アクセスは後続の層のヘッドのサブセットに集中しました。
• 深さ依存性： 最も大きな改善はより深いモデル（13 億）で現れ、制御メカニズムを適用する層がより多いほど選択的アクセスが恩恵を受けることを示唆しています。
• カテゴリ感受性： 後続の層のヘッドは、「構造的」または「機能的」特徴と比較して「意味的」特徴に対してより強いアクセスを示しました。これは、モデルが内容を含む初期表現を選択的に再利用しているという仮説を支持します。
• 介入： ゲートをアブレーション（ゼロにする、または平均に置き換える）すると、パープレキシティが増加し、$V_1$ パスウェイと学習されたゲーティング構造の両方が機能的に重要であることを確認しました。

意義と主張

本論文は、初期表現が有用であるために、高密度な接続性を介して利用可能である必要はない と主張しています。SATFormer はアーキテクチャ設計空間における特定の点を占めています。それは高密度な層間ルーティングよりも著しく安価であり、静的な値残差よりも適応的であり、検索指向の環境で特に効果的です。

層間再利用を「最大接続性」ではなく「選択的アクセス」の問題として再定義することにより、SATFormer は将来のトランスフォーマーアーキテクチャのためのシンプルで計算効率の高い設計点を提供します。著者は、モデルが制御された実証的分析を目的としており、この研究は、実質的に高密度化または計算コストの高い接続性を必要とすることなく、トランスフォーマーにおける情報フローの効率性と解釈可能性を向上させることを目指していると強調しています。