Microsoft的Full-bandwidth Transformers（26分钟阅读）

# 全带宽变换器  
来源：https://arxiv.org/html/2608.08888  
蔡子阳 詹正 董颖 范颖 古斯塔沃·德罗萨 蒂姆·皮尔斯 约翰·兰福德  
\[0.5em\]  
约翰斯·霍普金斯大学 普林斯顿大学 微软  
###### 摘要  
自回归变换器沿两个轴计算：水平方向跨生成词元，垂直方向跨模型深度。稠密注意力为每个词元提供了对过去的广泛水平访问，但解码步骤之间的垂直反馈通道仍然狭窄：只有采样词元返回堆栈底部，而顶层隐藏状态则被丢弃。我们引入了*全带宽变换器*，它通过*潜在反馈*来拓宽这一通道：在每个解码步骤中，前一个顶层隐藏状态通过门控线性单元与采样词元嵌入融合，并作为下一个输入反馈回去。潜在反馈允许未被言说的计算重新进入堆栈并获得新的深度预算，同时保留标准变换器架构、KV缓存和语言建模目标。为了在不丢失并行教师强制的情况下训练全带宽变换器，我们使用一种计划性多通道目标，该目标在预训练后期引入潜在反馈，并混合少量更深层的反馈通道以提高稳定性。我们训练了参数量达10亿、处理最多4000亿词元的全带宽变换器，发现潜在反馈改善了验证损失、五样本语言模型评估、数学和代码生成以及指令微调性能。在几乎不增加每词元解码开销的情况下，全带宽变换器匹配或接近用大约1.5×更多词元训练的标准变换器，并能够在相同或更好的准确性下产生更短的推理轨迹。  

††脚注文本：  
\*通讯作者：Xi Wang <[[email protected]](mailto:[email protected])>，John Langford <[[email protected]](mailto:[email protected])>。  
†\\dagger在微软AI前沿实习期间完成的工作。  

## 1 引言  
扩大语言模型规模在很大程度上意味着增加模型参数并在更多词元上训练\[22\]。然而，随着预训练规模的持续扩大，高质量独特数据的可用性成为一个日益突出的限制因素。这促使我们重新审视缩放轴本身：与其仅仅依赖更多数据，我们能否通过为每个词元分配更多计算来从中提取更有用的学习信号？循环、迭代和基于反馈的计算为实现这一方向提供了自然途径，但额外的浮点运算（FLOPs）只有在训练期间转化为更丰富的表示，或在推理期间转化为更有效的计算时才具有重要意义。自回归变换器暴露了一个特别未被充分利用的计算机会。它们已经包含一个反馈循环：在步骤\(t-1\)采样的词元成为步骤\(t\)的输入（图1，左）。正是这个循环使得思维链解码\[33\]能够执行深度随生成词元数量增长的计算\[24\]。但作为一个通信通道来衡量，这个循环极其狭窄：解码将模型的整个顶层状态（一个\(D\)维向量）压缩成最多携带\(\log_2 |V|\)位信息的单个符号。未被言说的计算并未被擦除——中间激活在KV缓存中持续存在并保持可访问——但它是*深度冻结*的：在层\(\ell\)产生的状态只能被高于\(\ell\)的层读取，因此它永远无法返回堆栈底部进行进一步处理，而最深的状态（顶层的输出）从不被缓存。因此，言说化是信息重新进入底层并接受新计算的唯一通道，代价是必须挤过单个词元。模型要么必须花费词元来叙述其中间状态，要么在每个位置从头重新计算该状态。  
在这项工作中，我们提出了*全带宽变换器*，我们将这个通道拓宽至其全部宽度。具体来说，我们引入了潜在反馈解码，它在解码过程中通过一个门控线性单元将前一个顶层隐藏状态与采样词元的嵌入融合（状态位于值通路上，词元充当门控），并将结果反馈作为下一个输入（图1右，第3.1节）。我们称能够以这种方式解码的变换器为全带宽变换器，因为其步骤间反馈现在携带的是整个隐藏状态而非单薄的词元。采样词元得以保留，因此模型仍然产生普通文本，并可以使用标准监督语言建模损失进行灵活训练；改变的是反馈不再局限于词元的身份。  
通过设计，这提供了标准解码所缺乏的两样东西：(i) 未被言说的状态——不确定性、部分结果、计划——可以重新进入堆栈底部，获得新的深度预算，并在步骤间进一步处理，而不是冻结在缓存中的生成层级；(ii) 每一层，包括最浅层，看到的过去都是由*整个*堆栈处理过的，而不仅仅是其下方层处理过的；  
关键的是，这些改进几乎没有带来架构变化和额外的部署成本：融合操作为每个生成的词元增加了两次矩阵乘法，注意力和KV缓存保持不变，预填充可以运行一次，或者可选地运行两次以获得更好性能。  
障碍在于训练。预训练模型从未在输入中见过隐藏状态，因此潜在反馈不能简单地在推理时切换开启；并且它定义的递归是位置序列化的，直接在其上训练将放弃使变换器高效训练的并行教师强制。我们通过*多通道*策略（第3.3节）解决这个问题：每个通道将前一个通道的隐藏状态向右移动一个位置，与词元嵌入融合，并在所有位置并行重新运行堆栈，因此序列化仅在少数几个通道上付出，而不是在整个序列上。两个要素使这变得可行。*渐进式计划*将训练的大部分时间花在普通单通道目标上，从而可以从标准预训练检查点开始运行，并在后期才引入额外的反馈通道；以及*前缀混合*在序列内随机化融合输入的开始位置，匹配推理时的提示后生成结构。  
经验上，我们发现计划的构成具有意想不到的重要性：仅使用两个反馈通道进行训练会产生一个在超出其训练深度后*发散*的递归，而混合哪怕只有3%的三通道批次则将学到的映射变为向一个*收缩*的不动点，该不动点在超过训练深度后保持稳定（图3）。  
经验上，全带宽变换器将微不足道的额外推理计算转化为相当于显著更多训练数据的增益。通过利用多次前向传递进行预填充，经过递归训练的模型在验证损失和多项选择准确性上匹配了在两倍词元上训练的无递归基线（图4）。在自由形式生成（图5）——GSM8K、Math500、HumanEval、MBPP——上，潜在反馈在所有任务上都优于*相同*权重的标准解码，匹配了2×词元基线，并在某些任务上接近了用多达5×词元训练的基线；这些增益通过长上下文扩展和指令微调得以保持（表1）。在基础模型上，潜在反馈通常在相同或更好的准确性下产生明显更短的推理轨迹（图6和图8）——这是拓宽通道所预测的行为，计算承载于隐藏状态之上，而非逐词元地言说。  

参见标题  
图1：标准解码 vs. 潜在反馈解码。左图：在标准变换器中，当前状态只能访问较低层的过去状态（蓝色）；更深层的过去状态（白色）不可达，唯一的步骤间反馈是采样词元嵌入（绿色）。右图：*全带宽变换器*使用*潜在反馈*，通过一个维度保持门控（\(\otimes\)，公式4）将前一个顶层隐藏状态与采样词元嵌入融合，并将其反馈作为下一个输入。这将完整的隐藏状态信息返回堆栈底部，使得所有层处理过的过去信息可供后续计算访问。  

## 2 背景  
给定一个大小为\(|V|\)的词汇表和一个\(D\)维残差流，一个仅解码器的LLM将包含\(T\)个词元的输入序列，其嵌入为\(\{\bm{e}_t\}_{t=1}^T \in \mathbb{R}^{T \times D}\)，通过\(L\)个注意力-MLP块映射。最终层的隐藏状态\(\{\bm{h}_t^L\}_{t=1}^T\)由语言模型头\(W_{\text{head}} \in \mathbb{R}^{|V| \times D}\)投影为下一个词元分布：  
\[
\bm{h}_t^L = f_\theta(\bm{e}_t;\, C), \qquad \bm{e}_t \leftarrow \mathrm{Decode}\!\left(\bm{h}_{t-1}^L\right), \quad C=\bm{e}_0, \bm{e}_1, \ldots, \bm{e}_{t-1}.
\]  
(1)  

#### KV缓存。  
在使用变换器进行自回归解码时，先前计算的键和值被缓存并重用，避免了对前缀的重复计算。与将历史压缩到固定大小循环状态的RNN和状态空间模型不同，稠密注意力变换器保留了所有过去词元的显式表示，因此每个新的隐藏状态都可以直接关注完整的缓存历史。  

#### 模型水平轴与垂直轴的带宽。  
将水平轴（跨位置）与垂直轴（跨深度）分开是有用的，因为两者以不同的速率携带信息。*水平上*，稠密注意力实际上是全带宽的：当生成词元\(t\)时，层\(\ell\)的状态\(\bm{h}_t^\ell\)可以读取每个更早位置的缓存表示。*垂直上*，访问受限：\(\bm{h}_t^\ell\)不能读取任何更深的过去状态\(\bm{h}_{t'}^{\ell'}\)，其中\(t'<t\)且\(\ell' \geq \ell\)（图1，左）。形式上，当在层\(\ell\)计算位置\(t\)时可到达的状态为  
\[
\mathcal{R}_{\text{std}}(t,\ell) = \big\{(t',\ell'): t' < t,\; \ell' < \ell \big\}, \qquad \bigl|\mathcal{R}_{\mathrm{std}}\bigr| = \Theta(T\ell),
\]  
(2)  
因此，新词元的浅层只能看到过去的*部分处理*视图，即使这些相同位置的更深、更充分处理的状态已经计算完毕并位于缓存中。过去的计算因此持续存在，但*深度冻结*，因为在层\(\ell\)产生的表示只能被高于\(\ell\)的层读取，且永远无法被路由回底层进行进一步处理。这就是第3.1节所拓宽的狭窄垂直通道。  
重要的是，这种深度方向的依赖性约束也是变换器能够跨位置并行训练的原因：仅需要跨层的序列化计算，而非跨词元的。然而，在解码时，生成已经是跨词元序列化的，因此该约束没有带来任何好处——这就为更丰富地依赖过去隐藏状态打开了大门，我们接下来将展开这一点。  

## 3 通过潜在反馈解码拓宽带宽  
### 3.1 潜在反馈解码  
全带宽变换器的核心创新是潜在反馈解码，它将前一个顶层隐藏状态反馈回输入。在步骤\(t\)，  
\[
\bm{h}_t^L = f_\theta(\bm{e}_t \otimes \bm{h}_{t-1}^L;\; C), \qquad \textrm{其中}~\bm{e}_t \leftarrow \mathrm{Decode}\!\left(\bm{W}^{\textrm{head}}\bm{h}_{t-1}^L\right),\; C=\bm{e}_0, \bm{e}_1\otimes\bm{h}_0^L, \ldots, \bm{e}_{t-1}\otimes\bm{h}_{t-2}^L
\]  
(3)  
其中\(f_\theta\)是\(L\)层变换器堆栈，\(\cdot\otimes\cdot\)将采样词元的嵌入与前一个潜在状态融合，\(C\)是过去上下文（所有更早位置的KV缓存）。标准解码（公式1）是仅采样词元在步骤间传递的特殊情况。  
融合\(\otimes\)是一个门控线性单元：  
\[
\bm{e}_t \otimes \bm{h}_{t-1} = \bm{W}^U \bm{h}_{t-1} \odot \sigma(\bm{W}^G \bm{e}_t),
\]  
(4)  
其中\(\bm{W}^U, \bm{W}^G \in \mathbb{R}^{D \times D}\)。这种不对称是刻意为之：隐藏状态占据值通路，而词元嵌入仅作为乘法门控进入。像\(\bm{e}_t + \bm{W}\bm{h}_{t-1}\)这样的对称融合会留下一条捷径：模型可以抑制状态通路，恢复普通词元输入，并达到普通预训练损失，从而使宽通道闲置。当训练从标准检查点开始时（其低损失可由加法路径重现），这条捷径尤其诱人。公式4关闭了它，因为丢弃\(\bm{h}_{t-1}\)就丢弃了输入本身，词元的身份仅存活于它对状态施加的\(D\)维门控模式中。读取状态因此成为强制性而非可选的。  

#### 潜在反馈的推理成本很低。  
增加的推理成本与上下文长度和模型深度无关，且每词元低于1%。状态\(\bm{h}_{t-1}^L\)在标准解码期间已经计算过，因此唯一额外的工作是融合：两次\(D \times D\)矩阵乘法，相对于通过\(L\)个块的前向传递来说微不足道。因为融合保持了输入维度\(D\)，架构、KV缓存布局和部署堆栈保持不变，解码循环仅改变两行（图2，右）。该方案也与vLLM兼容：我们将顶层状态存储在专用缓冲区中，借鉴了多词元预测实现所使用的机制（附录D）。  

### 3.2 潜在反馈解码 vs. 标准CoT  
标准CoT通过单个反馈通道执行串行计算：每个词元被附加到上下文并成为下一个输入。状态是词元序列：  
\[
s_{t+1} = s_t \| a_t, \qquad a_t \sim \pi_\theta(\cdot \mid s_t) \in \mathcal{V}, \qquad s_t = x_{1:t},
\]  
(5)  
因此，跨越步骤的唯一东西是离散的动作序列。潜在的问题解决状态原则上可能是过去动作的确定性函数，但从词元历史中恢复它本身就是一个状态跟踪问题，而固定深度变换器只有有界的串行计算能力。