论文
全带宽Transformer:以潜在反馈拓宽解码垂直通道
Full-bandwidth transformer
摘要
自回归Transformer沿两个轴计算:横向跨越生成的token,纵向贯穿模型深度。稠密注意力让每个token都能广泛地横向访问过去信息,但解码步骤之间的纵向反馈通道仍然狭窄:只有采样出的token回到栈底,顶层隐状态则被丢弃。我们提出全带宽Transformer(full-bandwidth transformer),通过潜在反馈(latent feedback)拓宽这一通道:在每个解码步骤,将上一步的顶层隐状态与采样token嵌入经门控线性单元融合,并作为下一步输入反馈回去。潜在反馈使未以文本表达的计算得以带着新的深度预算重新进入网络,同时保持标准Transformer架构、KV cache与语言建模目标不变。为在不牺牲并行教师强制的前提下训练全带宽Transformer,我们采用调度式多遍目标:在预训练后期才引入潜在反馈,并混入少量更深的反馈遍以保持稳定。我们训练了1B参数的全带宽Transformer,最多达400B token,发现潜在反馈改善了验证损失、5-shot语言模型评测、数学与代码生成以及指令微调后的性能。在每token解码开销可忽略的情况下,全带宽Transformer达到或接近多训约1.5倍token的标准Transformer的水平,并能在准确率相同或更高的前提下生成更短的推理轨迹。
