Full-bandwidth transformer
Xi Wang, Ziyang Cai, Zheng Zhan, Harry Dong, Ying Fan, Gustavo de Rosa, Tim Pearce, John Langford
cs.AI
2026-08-10
微软把上一步的顶层隐状态融合回每一步解码的输入,推理开销不到 1%,1B 模型追平用 1.5 倍 token 训练的标准 Transformer,数学任务上接近 5 倍训练量基线。
自回归 Transformer 有两个计算轴:横向跨 token、纵向跨深度。横向是「全带宽」的,注意力让每个 token 能读到全部历史;纵向却很窄。每生成一个 token,模型把整个顶层隐状态压缩成一个离散词再喂回下一步,顶层那层 D 维向量本身直接被丢掉。更深层的中间状态虽然留在 KV cache 里,却是「深度冻结」的:新 token 的浅层只能看到过去被部分处理过的视图,已经算好的深层状态回不到栈底再加工。模型要么把中间状态写成文字(思维链),要么每个位置从头重算。高质量数据越来越稀缺的当下,问题变成:不动数据,能不能从每个 token 里压出更多学习信号。
核心是隐反馈(latent feedback)。每一步解码,把上一步的顶层隐状态和当前 token 的嵌入用一个门控线性单元融合,再喂回栈底:
et ⊗ h{t-1} = WU h{t-1} ⊙ σ(WG et)
设计刻意不对称:隐状态走数值通路,token 只做乘法门。对称的加法融合会留捷径,模型可以压制状态通路、退回普通 token 输入、白拿预训练 loss,于是宽带通道形同虚设;门控融合把状态变成必读项。额外开销只有两个 D×D 矩阵乘,每 token 不到 1%,架构、KV cache、serving 栈都不动,兼容 vLLM。
真正的难点在训练。预训练检查点从没在输入里见过隐状态,推理时不能直接打开;而且这个反馈在位置上是串行的,直接训会丢掉让 Transformer 高效的并行教师强制。作者用多趟(temporal parallelism)来近似:每一趟把上一趟的隐状态右移一位、融合、再并行重跑整个栈,串行性用几趟前向来付,而不是用整个序列。调度是渐进的:训练大部分时间用普通单趟,快结束时才引入反馈趟,典型配比 75% 单趟、22% 两趟、3% 三趟。一个反直觉的发现:只用单趟加两趟训练的模型,反馈深度一超出训练范围就发散;只要掺进 3% 的三趟批次,学到的映射就变成压缩向不动点的收敛,在 k=1000 趟反馈下仍然稳定。再加上前缀混入(prefix mixin)、深度缩放、权重绑定、噪声抖动(σ=0.02)几招稳定化。
1B 参数模型,Phi-4 数据配比,最多训 400B token,对照基线最多到 1T。
| 指标 | 标准解码 | 隐反馈 | 对照(标准模型) |
| Math500 Pass@1(200B) | 0.27 | 0.37 | 1T 基线被超过 |
| HumanEval Pass@3(200B) | 0.31 | 0.34 | |
| GSM8K Pass@1(200B,经指令微调) | 64.5 | 67.9(FUSED) | |
| GSM8K Pass@1(400B,经指令微调) | 67.9 | 71.8(FUSED) | 1T 标准 70.13 |
| MATH-500(400B,FUSED) | 46.0 | 48.4 | 1T 标准 47.4 |
| MBPP Pass@3(200B,FUSED) | 38.4 | 41.2 | 1T 标准 41.9 |
验证 loss 和 5-shot LM Eval 上,带两趟融合预填的 100B 全带宽模型追平 200B 标准、200B 追平 400B 标准,约等于 2 倍数据效率。在 base 模型上,隐反馈常常在同等或更高准确率下产出更短的推理轨迹(Math500 中位 token 数下降),这正是宽带通道的预言:原本要逐字「说出口」的中间计算,现在搭着隐状态走。作者还做了探针实验:在受控状态追踪任务上,一次反馈让栈底第 0 层的线性探针准确率从接近随机抬到 99.6%(完成追踪)和 100%(延迟记忆)。
不动架构、不动 serving,两行解码循环的改动就能拿收益,这是它最容易落地的地方。它给了一个新的扩展维度:用推理算力(或训练时的辅助目标)从同样的 token 里多挖信号,在独有数据见顶的当下尤其有用。即便推理时根本不开反馈,带反馈目标的训练也让模型表示更好,所以它同时是一笔训练效率红利。
只在 1B 规模做了,收益在大模型上是放大还是缩小并不知道。需要专门的训练配方,不能在现成检查点上一键开启。「更短推理」这个最实用的红利在指令微调后消失:SFT 数据是按逐字推理的啰嗦风格生产的,拟合它会重新把隐状态通道挤回去,作者把保留简洁性的 on-policy 后训练留给未来。最后,标题里的 1.5 倍是基准,最好看的数字往往要靠 FUSED(预填翻倍),并不总是免费的。