全带宽变压器

摘要

自回归Transformer沿两个轴进行计算：水平方向跨越生成的token，垂直方向贯穿模型深度。密集注意力赋予每个token对过去内容的广泛水平访问能力，但解码步骤之间的垂直反馈通道仍然狭窄：只有采样得到的token返回栈底，而顶层的隐状态被丢弃。我们引入了全带宽Transformer，通过潜变量反馈拓宽这一通道：在每个解码步骤，先前顶层的隐状态通过门控线性单元与采样的token嵌入融合，并作为下一输入反馈回去。潜变量反馈使得未被言语化的计算能够以全新的深度预算重新进入堆栈，同时保持标准Transformer架构、KV缓存和语言建模目标不变。为了在不损失并行教师强制训练能力的前提下训练全带宽Transformer，我们采用了一种调度式多遍目标，在预训练后期引入潜变量反馈，并混合少量更深的反馈通道以确保稳定性。我们训练了参数量为10亿、数据量高达4000亿token的全带宽Transformer，发现潜变量反馈改善了验证损失、5-shot语言模型评估、数学与代码生成以及指令微调性能。在逐token解码开销可忽略不计的情况下，全带宽Transformer能够匹配或接近使用约1.5倍token训练的 standard 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.