论文

全带宽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的水平,并能在准确率相同或更高的前提下生成更短的推理轨迹。

全带宽Transformer:以潜在反馈拓宽解码垂直通道:论文配图
图1:标准解码与潜在反馈解码。左:在标准transformer中,当前状态只能访问较低层的过去状态(蓝色);更深的过去状态(白色)不可达,且唯一的跨步反馈是被采样的token嵌入(绿色)。右:全带宽transformer采用潜在反馈,通过一个保持维度的门控(⊗,式(4))将上一顶层隐状态与被采样的token嵌入融合,并将其作为下一输入反馈回去。这将完整的隐状态信息返回至堆栈底部,使得所有层处理过的过去状态都能被后续计算访问。