全帯域トランスフォーマー

要旨

自己回帰トランスフォーマーは、生成トークンにわたる水平方向と、モデルの深さにわたる垂直方向という2つの軸に沿って計算を行う。高密度アテンションは各トークンに過去への広い水平アクセスを与えるが、デコードステップ間の垂直フィードバック経路は狭いままである。すなわち、サンプリングされたトークンのみがスタックの底に戻り、最上層の隠れ状態は破棄される。我々は、この経路を潜在フィードバックによって拡張するフルバンド幅トランスフォーマーを導入する。各デコードステップにおいて、前の最上層の隠れ状態は、ゲート付き線形ユニットを介してサンプリングされたトークン埋め込みと融合され、次の入力としてフィードバックされる。潜在フィードバックにより、言語化されない計算が新たな深さの予算を得てスタックに再び入ることが可能になり、その一方で標準的なトランスフォーマーアーキテクチャ、KVキャッシュ、および言語モデリング目的関数は維持される。並列的なティーチャー・フォーシングを失うことなくフルバンド幅トランスフォーマーを訓練するために、我々はスケジュール化されたマルチパス目的関数を用いる。これは事前学習の後期に潜在フィードバックを導入し、安定性のために少量のより深いフィードバックパスを混合するものである。我々は1Bパラメータのフルバンド幅トランスフォーマーを最大400Bトークンで訓練し、潜在フィードバックが検証損失、5ショット言語モデル評価、数学およびコーディング生成、指示チューニング性能を改善することを見いだした。トークンあたりのデコードオーバーヘッドは無視できる程度でありながら、フルバンド幅トランスフォーマーは、およそ1.5倍多くのトークンで訓練された標準トランスフォーマーに匹敵するか、それに接近し、同等以上の精度でより短い推論トレースを生成することができる。

English

Autoregressive transformers compute along two axes: horizontally across generated tokens, and vertically through model depth. Dense attention gives each token broad horizontal access to the past, but the vertical feedback channel between decoding steps remains narrow: only the sampled token returns to the bottom of the stack, while the top-layer hidden state is discarded. We introduce the full-bandwidth transformer, which widens this channel with latent feedback: at each decoding step, the previous top-layer hidden state is fused with the sampled token embedding through a gated linear unit and fed back as the next input. Latent feedback lets non-verbalized computation re-enter the stack with a renewed depth budget, while preserving the standard transformer architecture, KV cache, and language-modeling objective. To train full-bandwidth transformers without losing parallel teacher forcing, we use a scheduled multi-pass objective that introduces latent feedback late in pretraining and mixes a small fraction of deeper feedback passes for stability. We train 1B-parameter full-bandwidth transformers up to 400B tokens and find that latent feedback improves validation loss, 5-shot language-model evaluation, math and coding generation, and instruction-tuned performance. With negligible per-token decoding overhead, full-bandwidth transformers match or approach standard transformers trained with roughly 1.5times more tokens, and manage to produce shorter reasoning traces at equal or better accuracy.