技術的概要：Cubit – カーネルリッジ回帰を備えたトークンミキサー

問題定義

2017 年にトランスフォーマーアーキテクチャが導入されて以来、それは自然言語処理、コンピュータビジョン、マルチモーダルタスクにおける深層学習の支配的なパラダイムとなっています。しかし、トランスフォーマーにおける核心的なトークンミキシング機構は自己アテンションに依存しており、これは理論的には指数関数カーネル（softmax 経由）と L1 正規化を用いたナダラヤ・ワトソン（NW）回帰を実行していると解釈できます。

効果的である一方で、この NW ベースのアプローチには本質的な限界があります：

1. 計算複雑性: 標準的なアテンションは、系列長に対して二次的な計算コスト（$O(N^2)$）を伴い、長い系列の拡張性を制限します。
2. 理論的制約: NW 回帰はバイアス - バリアンスのトレードオフに対する明示的な正則化を欠いており、ノイズや境界効果に対して敏感になり得ます。
3. 改善の拡張性: スパースアテンション、線形アテンション、エキスパート混合（Mixture-of-Experts）など多数の修正が提案されてきましたが、根本的な計算パラダイムは NW 回帰に根ざしたままです。

著者らは、NW 回帰の基盤をより堅牢な統計的枠組みに置き換えることで、より強力な数学的基盤を持ち、長い系列のモデリング能力が向上したアーキテクチャが得られると仮説を立てています。

手法：Cubit アーキテクチャ

本論文は、アテンションモジュールの NW 回帰をカーネルリッジ回帰（KRR）に置き換える新しいアーキテクチャCubitを提案します。

1. 理論的基盤：NW から KRR へ

• ナダラヤ・ワトソン（トランスフォーマー）: 標準的なアテンションにおいて、出力トークン $z_i$ は値ベクトル $v_j$ の重み付き和であり、重みはクエリとキー間のソフトマックス正規化された類似度によって決定されます。これは、グローバルなシステムを解くことなく局所的な重み付き平均に基づいて関数値を推定する NW 回帰に対応します。
• カーネルリッジ回帰（Cubit）: KRR は、正則化された経験的リスク最小化問題を解きます。新しい入力に対する予測には、カーネル行列の逆行列に正則化項を加えた $(K + \lambda I)^{-1}$ が関与します。
  • 主な相違点: NW と異なり、KRR はバイアス - バリアンスのトレードオフを制御するための明示的な正則化パラメータ $\lambda$ を導入し、グラム行列の逆行列を利用してすべてのデータサンプル間で情報を結合します。これにより、ノイズや境界効果に対する堅牢性が向上します。

2. アーキテクチャ設計

Cubit は、KRR の閉形式解を統合することで古典的なアテンション計算を変更します：

• カーネル構築: カーネル行列は、キー埋め込みから導出された特徴マップを使用して構築されます。著者らは、学習可能な射影行列 $W_K$ と規定された活性化関数 $\phi$ を用いて特徴マップ $\psi(x)$ をパラメータ化します。
• 逆行列による正規化: 行ごとのソフトマックス正規化の代わりに、Cubit は正規化行列 $\Sigma^{(h)} = (K^{(h)}K^{(h)\top} + \lambda I)^{-1}$ を計算します。この行列はグローバルな結合演算子として機能します。
• 出力定式化: 出力は $Z^{(h)} = A^{(h)} \Sigma^{(h)} V^{(h)}$ として計算されます。ここで、$A^{(h)}$ はカーネル類似度行列（クエリ - キー相互作用）を表し、$\Sigma^{(h)}$ は KRR に着想を得た正規化行列です。
• 因果的制約: 自己回帰生成を処理するために、カーネル構築が適応され、$i$ 番目のトークンが自分自身と過去のトークンのみにアテンションを向けるように調整されます。これにより、利用可能なコンテキスト上で回帰が実行されることを保証します。

3. 限定範囲リスケール（LRR）

大規模な行列の逆演算に内在する訓練の安定性の問題に対処するため、著者らは**限定範囲リスケール（LRR）**を導入します：

• 学習可能なスケーリングベクトル $s^{(h)}$ が適用され、カーネル精度を適応的に調整します。
• スケーリング因子はシグモイド関数を用いて有界範囲 $(\beta, \alpha + \beta)$ 内に制約されます。
• これにより、スケーリング行列の可逆性が保証され、逆行列の要素が有界化され、線形システムの解算中の数値的不安定性が防止されます。

4. 関連研究との区別

本論文は、Cubit をDeltaFormer（これも逆行列を利用する）から明確に区別します：

• 動機: DeltaFormer は DeltaNet とトランスフォーマーを組み合わせますが、Cubit は NW から KRR への理論的転換に動機付けられています。
• 実装: DeltaFormer は逆行列を構築するために 2 つの異なる埋め込みを使用し、対角値を 1 に強制します。Cubit は、対角値を 1 に強制することなく自然に可逆性を保証するために、単一の特定の埋め込み（例えば $K$ のみ）を使用します。
• 安定性: Cubit は安定性を向上させるために LRR を提案しており、これは DeltaFormer には存在しない構成要素です。

主要な貢献

1. 統合された理論的枠組み: 著者らは、トークンミキシング機構と古典的な回帰手法との間の関連性を確立し、NW 回帰と KRR を体系的に分析することで、モデルアーキテクチャの開発に対する新たな視点を提供します。
2. Cubit アーキテクチャ: トランスフォーマーにおける NW 推定量を置き換える KRR ベースのトークンミキサーの提案。訓練の安定性を高めるための**限定範囲リスケール（LRR）**機構を備えています。
3. 実証的検証: Cubit が標準的なトランスフォーマーや DeltaFormer よりも強力な長い系列のモデリング能力を示し、系列長が増加するにつれて性能向上が増大することを示す包括的な実験。

実験結果

著者らは、Cubit をトランスフォーマーおよび DeltaFormer と比較し、多様なデータセット（Arxiv, Books3, FinWeb-Edu）、系列長（512 から 8192）、モデルサイズ（1.25 億から 13 億パラメータ）で評価しました。

• 言語モデリングのパフォーマンス: Cubit は、ベースラインよりも一貫して低いパープレキシティ（検証損失）を達成しました。例えば、Arxiv データセット（1.25 億パラメータモデル、512 トークン）において、Cubit は 1.8802 の損失を達成し、トランスフォーマーの 1.9003 と比較されました。
• 長い系列のモデリング: Cubit とトランスフォーマーの間の性能差は、訓練系列長が増加するにつれて拡大しました。8192 トークンにおいて、損失の差は 0.05〜0.06 で安定し、Cubit のカーネルベースの機構が長距離依存関係をより効果的に捉えていることを示唆しています。
• スケーリング特性: Cubit は好ましいスケーリング挙動を示しました。モデルサイズが増加する（3.5 億から 13 億へ）につれて DeltaFormer のトランスフォーマーに対する優位性が減少するのに対し、Cubit の性能向上は持続し、改善されました。
• 下流タスク: ゼロショット評価（ARC, HellaSwag, PIQA など）において、Cubit は異なるデータサイズ（100 億から 500 億トークン）およびモデルスケールにおいて、トランスフォーマーと DeltaFormer の両方を一貫して上回りました。
• アブレーション研究:
  • アーキテクチャ設計: 性能向上はパラメータの拡張ではなく、構造的な設計（KRR 定式化）に起因しました。
  • LRR の有効性: LRR の導入は測定可能な改善をもたらしました。「Cubit-NoLRR」は標準的なトランスフォーマーを上回っていましたが、「Cubit」（LRR 付き）が最良の結果を達成しました。

意義と主張

本論文は、Cubit が明示的正則化、再生核ヒルベルト空間（RKHS）におけるより速い収束率、およびノイズに対する堅牢性を含むカーネルリッジ回帰の理論的利点を活用することで、バニラトランスフォーマーよりも強力な数学的基盤を提供すると主張しています。

著者らは、Cubit が特に長いコンテキスト理解を必要とするタスクに対して、次世代の基盤モデルアーキテクチャを表す可能性があると示唆しています。性能向上は単に訓練トークンの増加によるものではなく、長い系列に対するアーキテクチャの有効性に由来すると指摘しています。また、より強力なモデルの悪用可能性は認めつつも、モデリング能力の向上がもたらす社会的利益の可能性を強調しています。

この研究は、境界バイアスを解決するが計算的に遅い局所線形回帰や、トークンミキシングのための他の回帰ベースの枠組みの探求など、将来の方向性にも触れており、アテンションに対する回帰ベースの視点がさらなるアーキテクチャ革新の肥沃な土壌であることを示唆しています。