技術的サマリー：トランスフォーマー加速のためのゲート付き部分空間推論

1. 問題定義

大規模言語モデル（LLM）の推論における主要なボトルネックは、特にバッチサイズが 1 の場合、計算スループットではなくメモリ帯域幅である。屋根線モデル（roofline model）によって形式化される通り、トランスフォーマーのフォワードパスは、線形層の演算（$y = Wx$）の算術強度が約 1 FLOP/バイトであるため、完全にメモリ帯域幅に制限される。高帯域幅メモリ（HBM）からの重み行列の読み出しがレイテンシを支配する。

既存の加速戦略は重大な限界に直面している：

• 量子化は重み要素のサイズを縮小するが、読み取る要素の数を削減しない。
• プルーニングは要素の数を削減するが、通常は再学習または微調整を必要とする。
• 静的低ランク分解（LoRA、SVD-LLM など）はオフラインで重みを分解するが、推論時の特定の入力分布に適応できない。
• 活性化スパース性手法（CATS、PowerInfer など）は離散的なニューロン部分集合または要素ごとのスパース性を活用するが、活性化多様体の連続的な低次元構造を活用しない。

本論文は、任意の層における活性化ベクトルが、ほぼ低次元部分空間（$k \ll d$）に存在すると仮定する。したがって、重みと活性化の間の相互作用は、実効ランクが低い。入力分布が活性化空間のすべての方向を訪問するわけではないため、すべてのトークンに対して完全な重み行列を読み取ることは無駄である。

2. 手法：ゲート付き部分空間推論（GSI）

GSI は、再学習、アーキテクチャの変更、またはアテンションの近似なしに、重み - 活性化相互作用の実効低ランク性を活用する損失なし加速手法である。これは、Skyline Subspace Inference（SSI）フレームワークを MLP 層からトランスフォーマー全体（QKV および出力射影を含む）に拡張するものである。

2.1. 核心となる分解

重み行列 $W \in \mathbb{R}^{d_{out} \times d}$ と活性化 $x \in \mathbb{R}^d$ を持つ線形層において、活性化は部分空間成分と残差に正確に分解される：
$$x = V_k V_k^\top x + r$$
ここで、$V_k \in \mathbb{R}^{d \times k}$ は活性化部分空間を表すランク $k$ の正規直交基底であり、$r$ は残差である。出力 $y = Wx$ は以下のように計算される：
$$y = (W V_k)(V_k^\top x) + Wr = Mg + Wr$$
ここで、$M = W V_k$ は事前計算されキャッシュされた「低ランク重み画像」であり、$g = V_k^\top x$ は活性化の部分空間への射影である。

2.2. ゲート付き残差パススルー

決定的な革新は、残差補正 $Wr$ の計算を決定するトークンごとのゲートである。

• 高速パス：残差比率 $\rho_t = \|r\| / \|x\|$ が閾値 $\epsilon$ 未満の場合、出力は $y \approx Mg$ と近似される。これにより、完全な $d$ 列の代わりに、重み行列の $k$ 列（またはキャッシュされた画像 $M$）のみを読み出す必要がある。
• 低速パス：$\rho_t \ge \epsilon$ の場合、正確性を確保するために完全な計算 $y = Wx$ が実行される。

このゲートにより、出力分布は制御可能な許容誤差内で保持される。理論的分析（定理 4）は、高速パスにおけるトークンごとの誤差を $\|W\|_2 \cdot \epsilon \cdot \|x\|$ によって上限付けする。重要なのは、静的射影のように残差を完全に破棄すると壊滅的な出力劣化を招くが、ゲート付きパススルーは損失のない品質を維持することを論文が実証している点である。

2.3. 深度カスケードと較正

この手法は「深度の整合性」を活用し、連続する層の活性化部分空間が高度に相関している（深い層では平均コサイン類似度が 0.90 以上）ことを観察する。

• カスケード初期化：層 $l+1$ の基底 $V_k$ は、層 $l$ の基底から初期化され、すべての層に対して完全な特異値分解（SVD）を行うのではなく、わずかなランク 1 の補正（DGKS 更新）のみを必要とする。
• 較正：基底は、代表的な較正シーケンスでの単一のフォワードパスを通じて一度計算される。6B モデルの場合、これは約 30 秒を要する。
• ストレージオーバーヘッド：キャッシュされた画像 $M$ と基底 $V_k$ は、モデルのメモリフットプリントに 4% 未満（13 GB モデルの場合約 469 MB）を追加するのみである。

3. 主要な貢献

1. ゲート付き部分空間推論（GSI）：活性化を学習された低ランク部分空間に射影し、条件付きで残差補正をスキップすることにより、トランスフォーマー内のすべての線形マップ（QKV、出力、および MLP 射影）を加速する手法。
2. ゲートによる損失なし品質：ゲート付き残差パススルーが不可欠であることを論文は確立している。これなしでは、静的射影は生成の崩壊を引き起こす。これにより、特定の動作点においてベースラインと文字単位で完全に同一の出力を達成する。
3. 深度カスケード：層間での部分空間の整合性を活用して較正コストを 96% 削減するメカニズム。これにより、ネットワークが深度を通じて伝播する一貫した低次元多様体上で動作していることが明らかになった。
4. 直交的加速：GSI は隠れ次元（$d$ 軸）で動作し、トークン次元（$T$ 軸）で動作する ADA などの既存のアテンション加速手法を補完する。これらはフォワードパス全体を網羅する。

4. 実験結果

実験は、AMD MI300X GPU を使用して、GPT-2（124M）、GPT-J（6B）、OPT（6.7B）モデルに対して実施された。

• 高速化：この手法は、線形層の重み読み出しにおいて 3.0 倍から 10.5 倍 の実効的な高速化を達成する。
  • GPT-J 6B（動作点 $k=256, \epsilon=0.10$）：重み読み出しにおいて 15.6 倍 の高速化。99.8% のトークンが高速パスをとる。
  • OPT 6.7B（$k=256, \epsilon=0.10$）：10.5 倍 の高速化。
  • GPT-2 124M（$k=256, \epsilon=0.10$）：3.0 倍 の高速化。
• 品質保持：
  • パープレキシティ：比率は一貫して 1.00 未満（GPT-J の場合 0.991 など）であり、加速モデルはベースラインと統計的に区別がつかないか、わずかに優れていることを示している。
  • トークン一致率：すべての設定においてトップ 1 トークンの一致率が 98% を超える。
  • 生成忠実度：目標動作点（$k=256, \epsilon=0.05$ または $0.10$）において、GPT-J と GPT-2 は 50 個の貪欲トークンにわたってベースラインと 100% 文字単位で完全に同一 の出力を生成する。
• 実効パラメータ数：推論コストは「実効パラメータ数」（活性化空間で訪問された方向に対応するパラメータ）に比例する。GPT-J 6B の場合、最適な動作点では実効数が 6B から約 375M に低下し（16 倍の削減）、達成される。
• ネガティブコントロール：$k=128$ における静的射影（残差なし）は、パープレキシティを 52% 増加させ、生成の崩壊を引き起こした。これはゲートの必要性を検証するものである。

5. 意義と主張

本論文は、GSI がトランスフォーマー推論に対する 損失なし加速 を提供し、コストモデルを総パラメータ数から 実効パラメータ数 に根本的にシフトさせることを主張する。

• 再学習不要：この手法は、微調整やアーキテクチャの変更なしに事前学習済みモデルで機能する。
• アテンションの近似なし：アテンション機構を変更しないため、FlashAttention や ADA などの他の加速手法と互換性がある。
• 構造的洞察：結果は、重み - 活性化相互作用の実効低ランク性が、異なるアーキテクチャやモデルサイズにわたって成立する、学習済みトランスフォーマーの構造的性質であることを確認している。
• 実用的影響：線形層のメモリ帯域幅要件を $d/k$ 倍（例：16 倍）削減することで、GSI はバッチサイズ 1 の推論における主要なボトルネックに対処し、大幅に高速なデコード時間を可能にする（例：アテンション加速と組み合わせることで、1.3 ms のフォワードパスを 0.22 ms に短縮）。

著者らは、ゲート付き残差パススルーが、数学的な出力の正確性を犠牲にすることなく高い圧縮率を可能にする決定的な要因であると結論付けており、これが損失なし推論加速の実用的な道筋となるとしている。