Composing Linear Layers from Irreducibles

技術要約：既約要素からの線形層の構成

問題提起
現代の大規模言語モデル（LLM）は、低レベルのプリミティブが組み合わさることで、より豊かな機能を持つモジュールへと構成されることを示唆する挙動を示している。しかし、これらの線形変換の基礎となる構成要素については、依然として理解が進んでいない。先行研究では、パラメータ効率の高い近似（例：Low-Rank Adaptation、Block-Hadamard）やモジュール型の構成（例：Mixture-of-Experts）が探索されてきたが、複雑な線形挙動を生成する低レベルのプリミティブがいかにしてメカニズム的に構成されるかについての理解は不足している。著者らは具体的な問題を提示している。すなわち、通常数百万のパラメータ（$O(d^2)$）を含む線形層の機能を、最小限の既約な幾何学的オブジェクトの集合から合成できるか、という問いである。

手法
著者らは、線形変換を分解および再構成するために、**クリフォード代数（Clifford algebra）**を用いた代数的なフレームワークを提案している。

1. 代数的定式化:

   • 標準的な線形写像は伝統的に密な行列として表現される。著者らは代わりに、線形変換をクリフォード代数 $Cl(n)$ 内のバイベクタ（bivector）（向き付けられた平面を符号化する幾何学的オブジェクト）の合成として表現する。
   • 彼らは、特殊直交群 $SO(n)$ の二重被覆であるスピン群（Spin group） $Spin(n)$ を利用する。この群の要素であるローター（rotor） ($r$) は、「サンドイッチ積」 $x \mapsto rxr^\dagger$ を介してマルチベクタに作用する。
   • 任意のマルチベクタには $2^n$ 個のパラメータが必要となるが、著者らはローターによって生成されるバイベクタに構成を限定する。ローター $r$ は、バイベクタ $b$ の指数関数として定義される：$r = \exp(b)$。

2. パラメータ効率:

   • 密な $d \times d$ 行列には $O(d^2)$ のパラメータが必要である。
   • 変換を、部分空間に作用するローターの合成として表現することで、必要なスカラーパラメータの数は $O(\log_2 d)$ にスケールする。具体的には、幅 $w$ に対して、パラメータ数は約 $w \binom{\log_2 d}{2}$ となる。$d=2048$ の場合、これは数百万のパラメータを数百へと削減する。

3. 微分可能な分解アルゴリズム:

   • 一般的なバイベクタに対する指数写像は無限級数を伴うため、計算上の問題がある。しかし、単純なバイベクタ（simple bivectors）（単一の回転平面を表すもの）については、指数写像は閉形式の解を持つ：$\exp(b) = \cos(\|b\|) + \sin(\|b\|)\frac{b}{\|b\|}$。
   • 一般的なバイベクタを扱うために、著者らは不変分解（invariant decomposition）（Roelfs and Keninck [2021]に基づく）を採用し、あらゆるバイベクタを、互いに可換で直交する最大 $k = \lfloor n/2 \rfloor$ 個の単純なバイベクタの和として表現する。
   • 彼らは、クリフォード代数に適応させたクライロフ部分空間に触発されたパワーイテレーション法を用いて、これらの単純なバイベクタを反復的に抽出する微分可能な不変分解アルゴリズム（アルゴリズム1および2）を導入した。これにより、固有値分解を明示的に行うことなく、勾配ベースの最適化のための数値的安定性と微分可能性を確保している。

4. 汎用ローター・ガジェット:

   • $2$の累乗ではない任意の入出力次元（$d_{in}, d_{out}$）を扱うために、著者らは異なる部分空間上で動作する複数のローター・サンドイッチ・モジュールを集約し、プーリング演算子を適用することで、汎用的な変換$\psi$ を構築する。

主な貢献

1. 代数的表現: 著者らは、ローターのサンドイッチ積を用いることで、線形変換を幾何学的プリミティブ（バイベクタ）の合成として表現し、パラメータの複雑さを $O(d^2)$ から $O(\log_2 d)$ へと削減した。
2. 微分可能な分解: 任意のバイベクタを不変分解を介して対応するローターへと写す、閉形式の微分可能なアルゴリズムを提案し、autogradフレームワークへの統合を可能にした。
3. 実証的検証: 本手法をLLMのアテンション層におけるKey、Query、Value投影に適用した。ローターベースの層は、密な層を模倣するように（層単位で）訓練され、またFMNIST上でエンドツーエンドで訓練された。

結果

• パラメータ効率: LLaMa-3.2 1B および Qwen-2.5 1.5B における実験において、ローター層は劇的なパラメータ削減を達成した。例えば、LLaMa-3.2 1B のクエリ投影では、約4.19M個のパラメータ（密な層）から $\le 896$ 個（ローター）へと減少しており、約4700倍の削減を実現している。
• 性能: この極端な圧縮にもかかわらず、ローターベースの層は、強力なベースラインである Low-Rank (LR1, LR4) および Block-Hadamard (BH1) 近似と同等、あるいはそれ以上の性能を示した。
  • Wikitext2 (LLaMa-3.2 1B) において、単一のアテンション層を置き換えた場合、ローターの log-PPL は 2.629 であり、最良のベースライン (BH1) の 2.636 を上回った。
  • Arc Challenge において、ローターは 55.31% の精度を達成し、BH1 (54.83%) や LR4 (53.84%) を上回った。
• 堅牢性: ローターの性能は、様々な層の組み合わせ（1層、2層、または3層の置換）および様々なデータセット（Wikitext2, C4, PTB, Arc, HellaSwag）にわたって一貫していた。
• エンドツーエンド学習: FMNISTでゼロから訓練した場合、ローターベースのMLPは、最終的な精度（88.36%）において密なベースライン（89.67%）に匹敵する結果を得たが、収束はわずかに遅かった。

意義と主張
本論文は、コンパクトな幾何学的構成要素の集合から線形層を合成することの実現可能性を示すものであると主張している。著者らは、その目的が単に密な層のドロップイン・リプレースメント（現状ではハードウェア最適化されたカーネルの欠如により、壁時計時間の推論において不利になる）を提供することではなく、低レベルのプリミティブがいかにして高レベルの関数へと構成されるかという代数的視点を提供することにあると強調している。

これらの知見は、線形層の複雑な機能が、数百万の独立した重みではなく、数百の幾何学的オブジェクト（ローター）の構造化された合成を通じて実現できることを示唆している。これは以下の新たな道を切り開くものである：

• 解釈可能性: 線形層を、向き付けられた平面の合成として理解すること。
• パラメータ効率: パフォーマンスを維持しながら、パラメータ数を劇的に削減すること。
• 将来のハードウェア: 幾何学的パラメータからオンチップで重みを合成することで、計算量とメモリ帯域幅をトレードオフにし、大規模モデルの推論におけるメモリ帯域幅のボトルネックに対処すること。

著者らは、実用的な利益を得るには、本研究の範囲を超えた追加的なシステムレベルの統合とハードウェア認識の実装が必要であることを述べ、謙虚な姿勢を保っている。