技術要約: TANGO および WANGO

問題提起

標準的なTransformerアーキテクチャは、トークン間の相互作用（自己注意機構によって処理）と、位置ごとの非線形変換（フィードフォワードネットワークによって処理）を分離している。この設計は、これら2つの異なるサブレイヤーをブロックごとに必要とする。これは効果的ではあるものの、効率性やパラメータ利用の観点からは最適ではない可能性がある。さらに、既存のリニア複雑度を持つシーケンスモデルは、トークン間の相互作用の豊かさを犠牲にする近似に依存しているか、あるいは、非線形変換を注意機構に完全には統合できていない複雑なゲーティングメカニズムを必要とする場合が多い。

手法

コアアーキテクチャ: トークン集約型非線形ゲーティング演算子 (TANGO)

本論文では、標準的な自己注意および位置ごとのフィードフォワード・サブレイヤーを、単一のクロス・トークン・ゲーテッド残差更新に置き換えるデコーダーブロックであるTANGOを導入する。

1. ソース条件付き SwiGLU ゲーティング:

   • ゲートと特徴量が同一のトークンから計算される従来の SwiGLU (Swish-Gated Linear Unit) とは異なり、TANGOは、ソース位置 ($x_j$) で計算されたゲートが、**デスティネーション（目的地）**位置 ($x_i$) の投影された特徴量を調節することを可能にする。
   • 各位置 $i$ は、ゲートベクトル $g_i = \text{SiLU}(W_g x_i)$ と、投影された特徴量ベクトル $v_i = W_v x_i$ を生成する。
   • デスティネーション $i$ に対して、モデルは、因果的に可視なすべてのソース ($j \le i$) からのゲートベクトルの、コンテンツ依存の加重平均 $\bar{g}_i$ を計算する。
   • 残差更新は次のように計算される：
     $$h^+_i = h_i + W_o(\bar{g}_i \odot v_i)$$
   • ここで、集約されたゲート $\bar{g}_i$ は、最終的な出力投影 $W_o$ の前段階において、デスティネーションの投影された特徴量 $v_i$ に対する対角スケーリング演算子として機能する。

2. 重み集約メカニズム:
   本論文では、ソースのゲートを集約するための重みの計算方法が異なる2つのバリアントを提案している。

   • TANGO (Full-Prefix): すべてのデスティネーションと、因果的に可視なすべてのソースとの間のスケーリングされたコサイン類似度に基づき、ソフトマックス重みを計算する。これにより、シーケンス長 $T$ に対して二次複雑度 ($O(T^2)$) となる。
   • WANGO (Windowed Aggregation): 線形複雑度を持つバリアント ($T$ または固定されたウィンドウ/特徴量次元に対して $O(T)$)。
     • 直近のソース (Window): TANGOと同様の指数関数的コサイン類似度スコアを、後方に続くウィンドウ $L_i$ 内のソースに対して使用する。
     • 古いソース (Prefix): スコアを計算するために決定論的な正の特性マップ $\phi(u)$ を使用する。これにより、モデルは変換されたキーとそのゲートベクトルとの外積の実行和（running sums）を維持することが可能になる。古いソースの寄与は、これらの実行和を介して計算されるため、過去のすべてのペアを再評価する必要がない。

実験設定

著者は、TANGOおよびWANGOを以下の4つのベースラインと比較している：

• Recurrent Transformer++: パラメータを共有して4回再利用される標準的なTransformerブロック。
• Untied Transformer++: 4つの独立したパラメータを持つTransformerブロック。
• Full-attention GAU (Gated Attention Unit): 完全なペアワイズ・アテンションを使用する4つの独立したブロック。
• FLASH: ローカルなペアワイズ・アテンションとグローバルな線形アテンション・サマリーを組み合わせた4つの独立したブロック。

制約事項:

• すべてのモデルは、非埋め込みパラメータ数（約44.3M個の異なるパラメータ）において一致している。
• すべてのモデルは、同一のデータ順序、初期化シード、およびトレーニング予算を用いて訓練されている。
• 実験は、FineWeb-Edu（自然言語）、Lean（形式的な証明およびソースコード）、および DeepMind Mathematics で行われている。
• 比較は、4ブロックおよび8ブロックの逐次適用を用いて行われている。

主な貢献

1. TANGOモデル: 自己注意とフィードフォワードネットワークを、ソース・トークンがデスティネーションの特徴量を調節する単一のゲーテッド更新へと統合する新しいブロックを導入した。
2. WANGOモデル: ローカルウィンドウ内での正確なペアワイズ相互作用を保持しつつ、古いソースに対しては特性マップのプレフィックス和を使用することで、TANGOと同じパラメータ化を維持したまま、線形複雑度を実現するバリアントを開発した。
3. 厳密な比較: パラメータ数、トレーニング目的、およびデータの順序を厳格に一致させることで、アーキテクチャの設計による影響を、訓練の分散から分離した制御された比較を提供した。

結果

ベンチマークにおける性能

• FineWeb-Edu:
  • TANGOは、全体で最も低い平均検証負の対数尤度（NLL）（3.293）を達成した。
  • WANGOは、線形のシーケンス長スケーリングを持つ全アーキテクチャの中で最も低いNLL（3.355）を達成し、FLASH（3.465）およびRecurrent Transformer++（3.392）を上回った。
  • 特筆すべき点として、8,192トークンのシーケンスにおいて、WANGOは、解析的なフォワードパスの積和演算（MAC）数がほぼ同一であるにもかかわらず、Recurrent Transformer++よりも低いNLL（1.90T vs 1.94T）を達成した。
• Lean (ソースコードおよび証明完了):
  • TANGOは、最小の結合NLL（1.719）および、ソースコードモデリングと証明完了それぞれの個別NLLにおいて最小値を獲得した。
  • WANGO（2.271）はFLASH（2.377）を上回ったが、4ブロック実験においてはRecurrent Transformer++（2.530）を上回らなかった。しかし、8ブロック実験では、WANGO（2.231）はRecurrent Transformer++（2.364）を上回った。
• DeepMind Mathematics:
  • TANGOは、最も低い検証NLL（0.473）を達成した。
  • WANGO（0.487）は、Recurrent Transformer++（0.537）およびUntied Transformer++（0.556）よりも優れた性能を示した。

計算複雑度

• TANGO: シーケンス長に対して二次的である ($O(T^2)$)。8,192トークンのシーケンスの場合、約9.44TのMACを必要とする。
• WANGO: 固定次元に対してシーケンス長に対して線形である ($O(T)$)。同じシーケンスに対して、約1.90TのMACを必要とする。
• Recurrent Transformer++: 二次複雑度であるが、パラメータ共有および完全なクロス・トークン・ゲーティング集約の欠如により、TANGOよりも低いMAC数（~1.94T）となる。

意義および主張

本論文は、評価されたパラメータスケールにおいて、自然言語、形式コード、および数学的推論のタスク全体で、TANGOがテストされたアーキテクチャの中で最先端の検証尤度を達成したと主張している。

著者は、WANGOが、解析的な演算数およびパラメータサイズが一致している場合に、線形複雑度のモデルが二次複雑度のベースライン（Recurrent Transformer++など）を上回り得ることを示している（特にFineWeb-Eduにおいて）。具体的には、SwiGLUゲートを実行和を介して集約できるWANGOの能力により、高い性能を維持しながら線形スケーリングが可能となっている。

著者は、自らの主張の限界を明示している：

• 結果は教師強制下での検証NLLに基づいており、生成された回答の正確性、証明の受理率、または実時間のトレーニング速度を測定したものではない。
• 実験は、テストされたパラメータ数（約7,000万パラメータ）を超えたスケーリング則を確立するものではない。
• TANGOの低いNLLは、速度やメモリの利点として解釈されるべきではない。なぜなら、TANGOはWANGOやRecurrent Transformer++よりも大幅に高い計算コストを課すからである。

要約すると、本論文は、クロス・トークン・ゲーティングを通じて注意とフィードフォワードを統合する新しいアーキテクチャブロックを提示しており、このアプローチが優れた予測尤度をもたらすこと、そして線形複雑度のバリアント（WANGO）が、長文脈アプリケーションのための競争力のある代替案を提供することを実証している。