技術的サマリー：Toeplitz MLP ミキサーを低複雑度かつ情報豊富なシーケンスモデルとして

1. 問題提起

トランスフォーマーベースの大規模言語モデル（LLM）は、自己注意メカニズムの二次的な時間計算量（$O(n^2)$）と空間計算量によって根本的に制約されています。この計算量を二次未満または線形レベルに削減する試み（例えば、注意の線形化や Mamba などの状態空間モデル）は数多く行われていますが、これらのアプローチはしばしば重大なトレードオフをもたらします。具体的には、線形計算量のモデルは、訓練効率の低下、数値的不安定性、および入力情報の保持能力の顕著な欠如に悩まされることが多いです。この情報保持の欠如は、正確なコピー、コンテキスト内学習、長文脈検索を必要とするタスクにおける性能低下として現れます。これらは現代の LLM 能力の印です。核心的な課題は、これらの創発的機能を劣化させるトークン情報へのアクセスに対する厳格な制限を課すことなく、計算複雑度を削減することです。

2. 手法：Toeplitz MLP ミキサー（TMM）

著者らは、シーケンス次元における三角マスク付き Toeplitz 行列乗算によって自己注意を置き換える、トランスフォーマーに似たアーキテクチャである**Toeplitz MLP ミキサー（TMM）**を提案します。

中核アーキテクチャ

• Toeplitz パラメータ化: 標準的な MLP ミキサーではトークン混合の重みが任意の行列であるのに対し、TMM はこれらの重みを Toeplitz 行列としてパラメータ化します。Toeplitz 行列では、与えられた対角線上のすべての値が同一であるため、行列全体を単一のベクトル（最初の行と列）で表現できます。
• 因果的マスク: 次のトークンの予測を因果的に強制するため、実体化された重み行列は三角マスクを施されます。実装では、活性化値がこのマスク付き Toeplitz 行列（$T'$、上三角行列）によって左から乗算されます。
• FFT による計算: Toeplitz 行列を巡回行列に埋め込むことで、著者らは高速フーリエ変換（FFT）を活用して行列 - ベクトル積を計算します。これにより、訓練および推論のプレフィル段階での計算複雑度が $O(n^2)$ から$O(dn \log n)$に削減され、空間計算量は**$O(dn)$**となります。
• マルチヘッドおよびマルチカーネルバリアント: このアーキテクチャは、独立した Toeplitz 行列が投影されたヘッド次元に適用されるマルチヘッド構成をサポートします。さらに、「カーネル」バリアントは、1 次元畳み込みと同様に、限られた数の隠れ次元要素に沿って混合を可能にします。

理論的基盤

本論文は、作用素指数理論の観点から TMM を分析します。具体的には、訓練された Toeplitz 層のフレドホルム指数を検討します。著者らは、TMM がすべての可能な入力に対して厳密に可逆ではないものの、訓練された因果的 TMM 層はしばしばシフト作用素に似た指数（具体的には -1）を示すと指摘しており、これは他のアーキテクチャに見られるバイアスとは異なる逐次的な情報シフトのメカニズムを示唆しています。

3. 主要な貢献と主張

本論文は、TMM の有効性について、Hyena、Mamba、ウィンドウ付き注意トランスフォーマーなどの他の二次未満のアーキテクチャと比較して、主に 3 つの主張を行っています。

1. 訓練効率: 一定の計算条件の下、TMM は他の二次未満のモデルと比較して同等または低い評価損失を達成します。Hyena や Mamba などの線形計算量のベースラインに対して、優れた収束率を示します。
2. 情報保持: TMM は、比較可能な二次未満のモデルよりもはるかに多くの入力情報を保持します。これは以下の点で実証されています。
   • 情報抽出: 最終隠れ状態から入力シーケンスを再生する際の忠実度が高いこと。
   • コピー能力: TMM は、Mamba や Hyena モデルが通常苦戦するタスクであるコピータスクにおいて、二次計算量のトランスフォーマーに匹敵またはそれ以上の訓練効率を達成します。
   • 情報容量: 情報保持を最大化するように明示的に訓練された場合、TMM は Hyena や Mamba モデルよりも高い情報容量を示します。これは、訓練に必要な計算量が大幅に少ない（1/2 から 1/8）場合でも同様です。
3. ベンチマーク性能: 優れた情報保持により、TMM は情報検索（IR）やコンテキスト内学習（ICL）など、正確な入力表現を必要とするベンチマークにおいて、他の低複雑度モデルを上回る性能を発揮し、一般的な言語理解タスクでは同等の性能を維持します。

4. 実験結果

著者らは FineWeb-edu の 100 億トークン subset 上でモデルを訓練し、トランスフォーマー、Mamba 2、Hyena モデルと比較評価しました。

• 訓練効率: 計算量とメモリを一致させた場合、TMM は Hyena や Mamba モデルよりも速く収束しました。著者らは、線形注意の実装が数値的不安定性と遅い収束に悩まされたため、TMM に焦点を当てたと述べています。
• コピータスク: 入力から 512 トークンを出力にコピーするタスクにおいて、TMM はトランスフォーマーに匹敵する訓練効率を達成しました。注目すべきは、Hyena モデルのモデル次元（$d_m$）の 1/32 の TMM が、ほぼ同一のコピー効率を達成したことです。研究により、学習可能な Toeplitz パラメータが重要であることが判明しました。これらの行列を固定するとコピーが完全に失敗する一方、学習可能なヘッド投影のみでは不十分でした。
• ベンチマーク:
  • IR/ICL: TMM は、SQuAD、SQuAD v2、LongBench、IFEval において Hyena を上回りました。
  • 一般タスク: TMM は、ARC-Easy、GLUE、HellaSwag などのタスクにおいて、Hyena やトランスフォーマーと同等の性能を発揮しました。
• アブレーション研究: 著者らは、非線形性で隔てられた 2 つの連続した Toeplitz 層を使用する（標準的な MLP ミキサースタイル）ことが、単一層よりも効率が低いことを発見しました。さらに、ヘッド数を増やすと漸近的な収束が向上する一方、カーネルサイズをある点を超えて増やすと効率が低下しました。

5. 意義と限界

意義

本論文は、TMM が競争力のある訓練効率と高速なデコーディングへの「単純で注意不要な道」を提供すると主張しています。その主な意義は、低複雑度モデルが、情報保持を劣化させるような厳格なアーキテクチャ的バイアス（例えば、最近性バイアスや状態空間の選択性）を本質的に必要としないことを実証している点にあります。明示的な Toeplitz パラメータ化を使用することで、モデルは Mamba や Hyena に見られる「コピーのギャップ」を回避し、Toeplitz という構造化されたアクセスパターンが、プレフィル段階における二次コストなしにトランスフォーマーの情報豊富な能力を保持できることを示唆しています。

著者らはまた、直感に反する理論的発見を強調しています。完全な入力空間に対して非可逆である訓練された因果的 TMM 層は、可逆性のために明示的に訓練されたモデル（オートエンコーダなど）よりも、フレドホルム指数がゼロに近い（または具体的にはシフト作用素として機能する -1）傾向があります。これは、因果モデルにおける層の「可逆性」が、数学的な可逆性というよりも、シフトを通じた情報フローの保持に関するものであることを示唆しています。

限界

著者らは、いくつかの限界を明確に認めています。

• スケーラビリティ: 実験は、2000 万から 3 億パラメータのモデルで、約 130 億から 260 億トークンで訓練されたものに対して行われました。これらの知見が、桁違いに多い計算量で訓練されたモデルにスケーリングする保証はありません。
• 自己回帰推論: TMM はプレフィル中は $O(dn \log n)$ ですが、最初のトークンが計算されると、入力行列の行をキャッシュして逐次処理する必要があるため、自己回帰生成においては$O(n^2)$ の複雑度に戻ります。したがって、TMM は長いプロンプトに対する最初のトークンまでの時間を短縮しますが、トランスフォーマーと比較して後続のトークンの生成を高速化するものではありません。
• 設定空間: 可能なハイパーパラメータの組み合わせが指数関数的であるため、本研究では限られた数のアーキテクチャ設定のみがテストされました。

結論として、TMM は、自己回帰生成の推論ボトルネックが管理される限り、複雑な推論や検索タスクに必要な高い情報保持と計算効率のバランスを取る、シーケンスモデリングにおける有望な方向性を表しています。