全頻寬變壓器

摘要

自回歸Transformer沿兩個軸向進行計算：水平方向跨越生成出的token，垂直方向貫穿模型深度。密集注意力讓每個token都能對過去進行廣泛的水平訪問，但解碼步驟之間的垂直反饋通道仍然狹窄：只有採樣出的token會返回堆疊底部，而頂層隱藏狀態則被丟棄。我們引入了全頻寬Transformer（full-bandwidth transformer），透過潛在反饋（latent feedback）拓寬此通道：在每個解碼步驟中，先前的頂層隱藏狀態會透過門控線性單元與採樣出的token嵌入融合，並作為下一個輸入反饋回去。潛在反饋讓未言語化的計算能以重新獲得的深度預算進入堆疊，同時保留標準Transformer架構、KV快取和語言建模目標。為了在不失去並行教師強制（teacher forcing）的情況下訓練全頻寬Transformer，我們使用一種排程多遍次目標：在預訓練後期才引入潛在反饋，並混合少量更深層的反饋遍次以維持穩定性。我們訓練了多達4000億token的10億參數全頻寬Transformer，發現潛在反饋改善了驗證損失、5-shot語言模型評估、數學與程式碼生成，以及指令微調的效能。在每token解碼開銷可忽略不計的情況下，全頻寬Transformer能夠匹配或接近以約1.5倍token訓練的標準Transformer，並且在相同或更好的準確度下產生更短的推理軌跡。

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.