技術要約：フル帯域幅トランスフォーマー（Full-Bandwidth Transformer）

問題提起

自己回帰型トランスフォーマーは、生成されるトークンに沿った水平方向と、モデルの深さに沿った垂直方向の2つの軸に沿って計算を行います。高密度なアテンションは過去のトークンへの広範な水平方向のアクセスを提供しますが、デコーディングステップ間の垂直方向のフィードバックチャネルは依然として狭いままです。標準的なデコーディングでは、モデルは $D$ 次元の最上層隠れ状態全体を、単一の離散的なトークン（最大 $\log_2 |V|$ ビットの情報量）へと圧縮し、それを次のステップの入力として使用します。その結果、非言語化された計算（中間活性化、不確実性、部分的な計画など）は、KVキャッシュ内に「深さに凍結（depth-frozen）」されてしまいます。これらは生成ポイントより上のレイヤーにはアクセス可能ですが、スタックの底部に戻ってさらなる処理を受けることはできません。これにより、モデルは中間状態をトークンごとに言語化するか、あるいは最初から再計算することを余儀なくされ、トークンあたりの実効的な推論の深さが制限されます。

手法：潜在的フィードバック・デコーディング（Latent Feedback Decoding）

著者らは、前のステップの最上層隠れ状態を、サンプリングされたトークン埋め込みと共に次ステップの入力へとフィードバックすることで、ステップ間のフィードバックチャネルを広げる**フル帯域幅トランスフォーマー（Full-Bandwidth Transformer）**を提案しています。

1. 潜在的フィードバック・メカニズム

各デコーディングステップ $t$ において、入力 $u_t$ は、サンプリングされたトークン埋め込み $e_t$ と、前のステップの最上層隠れ状態 $h^L_{t-1}$ を融合させることで形成されます。
$$u_t = e_t \otimes h^L_{t-1}$$
融合演算子 $\otimes$ は**ゲート付き線形ユニット（GLU）**です。
$$e_t \otimes h^L_{t-1} = (W_U h^L_{t-1}) \odot \sigma(W_G e_t)$$

• 非対称性: 隠れ状態は値の経路（value pathway）を占め、トークンは乗算的ゲートとして機能します。この設計は、モデルが状態経路を無視することを学習してしまうこと（対称的な融合における潜在的なショートカット）を防ぎ、隠れ状態の全容を活用することを強制します。
• アーキテクチャ: これは、トークンあたりわずか2つの $D \times D$ 行列乗算を追加するのみです。標準的なトランスフォーマーのアーキテクチャ、KVキャッシュのレイアウト、およびサービング・スタックは変更なしでそのまま利用可能です。
• 効果: これにより、非言語化された状態が、新たな深さの予算を持ってスタックの底部へと再突入することが可能になります。すべてのレイヤー（最も浅いレイヤーを含む）は、自身の直下のレイヤーだけでなく、スタック全体によって処理された全履歴にアクセスできるようになります。

2. 並列学習戦略

この再帰構造を直接学習させると、ステップ $t$ の入力がステップ $t-1$ の出力に依存するため、教師強制による並列性が失われてしまいます。著者らはこれを解決するために**マルチパス目的関数（multi-pass objective）**を導入しました。

• 時間的並列性: シーケンスに対して逐次的にアンロールする代わりに、モデルは $K$ 回の並列フォワードパスを実行します。
  • パス 1: 通常のトークン埋め込みを用いたフォワードパス。
  • パス $k$ ($k > 1$): パス $k-1$ からの隠れ状態を右に1つシフトし、トークン埋め込みと融合させ、並列的にスタックへとフィードバックします。
• 損失関数: この目的関数は、第1パスからの標準的な次トークン予測（NTP）損失と、後続のパスからの補助的なNTP損失を組み合わせたものです。後のパスからの勾配は、より早い段階の潜在状態へとバックプロパゲーションされ、それらが再利用可能な入力となるよう促します。
• スケジューリング: 安定性を確保し、計算コストを管理するために、漸進的なスケジュールを使用します。
  • 学習は標準的なシングルパスの目的関数から開始します。
  • 潜在的フィードバックは、学習の中盤（例：2パスのバッチ）で導入されます。
  • より高いパス数（例：3パス）を持つバッチを、ごくわずかな割合（例：3%）で混合します。経験的に、この極めて小さな割合が、学習されたフィードバックマップを安定させ、推論時（数百パスのホライゾンへの外挿時）でも安定した不動点への収縮を実現します。
• プレフィックス混合（Prefix Mix-in）: 学習時（完全に融合された入力）と推論時（プレフィックス＋融合された生成）の間の分布の不一致を処理するため、学習プロセスにおいて融合された入力が始まるプレフィックスの長さをランダム化します。

3. 安定化技術

長期ホライゾンの推論における発散を防ぐため、以下の手法を採用しています。

• 深さのスケーリング（Depth Scaling）: 最上層の状態のノルムが、深さに応じて増大することなく $O(1)$ に保たれるようにします。
• 重みの共有（Weight Tying）: 埋め込み層とリードアウト層の重みを共有し、プレーンなトークン入力と融合された状態入力の間の互換性を確保します。
• ノイズ正則化: 学習中に、運搬される隠れ状態に小さなジッターノイズを加えることで、堅牢性を向上させます。

主な結果

著者らは、10億パラメータのモデルを最大4000億トークンまで学習させ、1000億〜1兆トークンで学習された標準的なトランスフォーマーと比較評価を行いました。

• データ効率: フル帯域幅トランスフォーマーは、標準的なトランスフォーマーが約1.5倍から2倍のトークンを用いて達成する性能に匹敵、あるいは接近します。例えば、2000億トークンのフル帯域幅モデルは、4000億トークンの標準的なベースラインの性能に近づきます。
• 性能向上:
  • 検証損失: あらゆるスケールにおいて改善が見られました。
  • 推論およびコーディング: 自由形式の生成タスク（GSM8K, Math500, HumanEval, MBPP）において、潜在的フィードバック・デコーディング（SOFTおよびFUSEDレジーム）は、同じ重みを用いた標準的なデコーディングを一貫して上回りました。
    • SOFT (潜在的フィードバック): 数学タスクに最適であり、隠れ状態が言語化されたトークンよりも推論をより良く運んでいることを示唆しています。
    • FUSED (マルチパス・プリフィル): コーディングタスクに最適であり、洗練されたプロンプト表現が重要であることを示唆しています。
  • 指示チューニング: 指示チューニング後も性能向上は維持されます。
• 簡潔性: ベースモデル（指示チューニング前）において、潜在的フィードバックは、同等の精度を維持しながら著しく短い推論トレースを生成します。モデルは、トークンとして語るのではなく、隠れ状態の中で「静かに」計算を行っています。この効果は、指示チューニング後に減少しますが、これはチューニングデータが冗長な標準的推論を模倣しているためと考えられます。
• 状態のアクセシビリティ: プロービング実験により、再帰的なプリフィルが、標準的なプリフィルよりもはるかに効果的にグローバルな状態情報を浅いレイヤー（レイヤー0）に露出させていることが確認されました。これは「フル帯域幅」という主張を裏付けるものです。

意義と主張

本論文は、フル帯域幅トランスフォーマーが、大規模言語モデルにおける新しいスケーリング軸、すなわち**「トークンあたりの計算量を割り当てる」**ことを提供すると主張しています。

• スケーリングの再考: 高品質でユニークなデータが制約となる中で、著者らは、各トークンからより多くの学習信号を引き出すことを提唱しています。潜在的フィードバックは、言語化というボトルネックを経由せずに内部状態を反復させることで、これを実現します。
• 推論効率: この手法は、無視できる程度のデコーディング・オーバーヘッド（わずか2つの行列乗算）で性能向上をもたらします。追加の計算は、主にオプションのマルチパス・プリフィルまたは学習時に発生し、逐次的なトークン生成時には発生しません。
• アーキテクチャの最小限性: 本手法は標準的なトランスフォーマーのアーキテクチャとKVキャッシュを維持しており、入力の構築方法と特定の学習スケジュールを変更するだけで済みます。
• 安定性: 本研究は、並列マルチパス目的関数を用いることで、再帰的なフィードバックループをトランスフォーマー内で効果的に学習できることを示し、再帰ネットワークの学習に伴う不安定性を克服しています。

著者らは、フル帯域幅トランスフォーマーが、単なるデータの規模に頼るのではなく、再帰的かつフィードバックに基づいた計算を活用して、表現の質と推論の効率を向上させるというパラダイムシフトを提示していると結論付けています。