WhiteMatter: 通过KV混合实现全对全跨层连接

# WhiteMatter：通过KV混合实现全层全连接
来源：https://arxiv.org/html/2608.18486
Wenbo Zhang Xiang Ren

###### 摘要

在Transformer中，每一层仅通过自身深度产生的KV来关注过去的token，尽管在自回归解码过程中存在更深层的表示。反馈架构允许浅层的消费者层关注由更深层过去token表示产生的KV，但为所有消费者层提供了相同的固定源层连接模式。我们提出WhiteMatter，它将每个注意力层连接到每个过去token所有层的表示，连接权重可随消费者层变化并适应源token。对于每个token，路由器通过将其LL层状态混合到kk个KV通道中来实现这些连接，这些通道被缓存供后续token使用；每个消费者层关注其中一个通道。通道数量kk控制KV缓存的大小。设置k<Lk<L可减少缓存的内存占用。在我们的预训练实验中，WhiteMatter优于具有50%50\\%更多层的普通Transformer，并在KV缓存压缩50%50%时保留了大部分性能增益。

## 1 引言

在自回归Transformer的解码过程中\(26 (https://arxiv.org/html/2608.18486#bib.bib1)\)，模型在继续处理下一个token之前，为当前token产生一组隐藏状态。每个隐藏状态在对应的层生成，可能包含独特信息。然而，在处理下一个token时，每一层只能关注由相同深度的隐藏状态产生的KV。因此，模型无法利用其已产生的所有信息。特别是，过去token更深层状态的不可访问性被认为限制了计算深度和状态跟踪能力\(20 (https://arxiv.org/html/2608.18486#bib.bib19)\)。

有两条研究路线放松了这一限制的不同部分。反馈架构创建了跨token的从深到浅的路径。反馈Transformer\(6 (https://arxiv.org/html/2608.18486#bib.bib2)\)为每个消费者层提供了与每个过去token所有层状态的相同静态连接。LCKV\(27 (https://arxiv.org/html/2608.18486#bib.bib3)\)则使用顶层隐藏状态作为所有层的KV源，并引入雅可比迭代，使得具有反馈连接的模型在大规模上训练变得可行。这些架构对每个消费者层和输入使用相同的源层连接。因此，不同的消费者层无法选择不同的源。另一条研究路线提供从较早源层到较晚消费者层的前馈跨层连接。DenseFormer\(22 (https://arxiv.org/html/2608.18486#bib.bib6)\)、MUDDFormer\(29 (https://arxiv.org/html/2608.18486#bib.bib8)\)和相关方法\(37 (https://arxiv.org/html/2608.18486#bib.bib14);14 (https://arxiv.org/html/2608.18486#bib.bib16)\)为不同层提供了对当前token内较早层状态的不同连接。FusedKV\(16 (https://arxiv.org/html/2608.18486#bib.bib9)\)则为上层消费者层提供了对过去token底层和中层KV的静态、层特定连接。这些方法提供了消费者特定的连接性，其中一些是内容相关的。它们的连接仍然是前馈的：一个浅层的当前token层仍然无法访问过去token的更深层表示。

大脑将局部计算与长程通信相结合。灰质包含神经元细胞体，而白质包含连接远端区域的神经纤维（附录A (https://arxiv.org/html/2608.18486#A1)）。这些纤维在皮层区域之间形成密集、通常是双向的连接\(18 (https://arxiv.org/html/2608.18486#bib.bib39)\)。每个皮层区域都有独特的连接模式，且这些通路上的活动是动态调节的。

这种组织结构启发了四个架构特性：远距离层之间的直接连接、从深到浅的反馈连接、消费者特定的连接性以及连接的动态调制。我们提出WhiteMatter（图1 (https://arxiv.org/html/2608.18486#S1.F1)），它通过kk个共享的KV通道实现所有这些特性。在每个token位置，一个路由器将所有LL源层的隐藏状态混合到这些通道中。每个消费者层选择一个通道，因此不同的消费者可以接收连接到不同源深度的连接。由于路由器读取隐藏状态，连接权重会适应源token。

\(a\) 普通模型的同层KV过去token当前token \(b\) 反馈Transformer的单一静态共享混合过去token当前token \(c\) FusedKV的静态逐层KV融合过去token当前token \(d\) WhiteMatter的动态全深度通道过去token当前token block存储的KV隐藏状态源到KV KV到block

图1：跨层的KV生成与消费。灰色框表示解码器块，灰色箭头在深度方向传递隐藏状态。粉色箭头将源块连接到KV，蓝色箭头将KV连接到消费者块。多个箭头汇聚处，其源表示被组合。只有WhiteMatter的源到KV权重取决于token内容。\((a)\) 普通模型：每个块读取在相同深度产生的KV。\((b)\) 反馈Transformer：每个块对所有源深度具有相同的静态连接。LCKV具有类似的反馈路径，但仅使用顶层隐藏状态。\((c)\) FusedKV\(16 (https://arxiv.org/html/2608.18486#bib.bib9)\)：较低块存储KV，每个上层块读取底层和中层缓存的静态、块特定混合。\((d)\) WhiteMatter：一个路由器根据所有源深度形成kk个依赖于token的通道，一个固定分配将每个消费者块映射到一个通道（§3.1 (https://arxiv.org/html/2608.18486#S3.SS1)）。从深到浅的反馈在自回归解码过程中很直接，因为过去token的状态已经是最终的。然而，在并行训练和预填充过程中，每个token的KV是由其自身完成的隐藏状态构建的，而这些状态又依赖于更早token的KV；简单的从左到右解析这种循环依赖关系会在序列长度上顺序运行。我们通过使用循环高斯-赛德尔调度的迭代来解决这个问题，使计算保持token并行性。

我们在88B个FineWeb-Edu token上从头预训练了所有模型，使用相同的数据、token预算和优化器设置。在16层和完整的KV缓存（k=16）配置下，WhiteMatter达到了19.968的困惑度，比相同深度的普通模型（21.747）低8.2%，也略低于24层普通模型（20.181）。将缓存减半至k=8，困惑度为20.377，比具有相同缓存大小的LCKV基线低5.0%。两种配置在LAMBADA和WikiText上都优于所有其他16层模型。在使用精确自回归执行训练的受控模型中，g=16的循环高斯-赛德尔在4次内达到自回归困惑度的1%以内，并使收敛的预填充比精确自回归评估快13.9倍，比雅可比迭代快11.2倍。对于16层实验，循环训练仍然比普通模型贵约1.5倍。

我们总结我们的贡献如下：(1) WhiteMatter添加了每层内容相关的到过去表示的连接，连接来自所有源深度，通过从所有层状态的动态混合生成KV来实现。(2) 消费者层之间共享KV通道，当k<L时减少了KV缓存大小。(3) 我们应用了一种循环迭代调度，提高了训练和预填充的收敛速度，并系统地探索了迭代参数的选择。(4) 经验表明，完整缓存的WhiteMatter比相同深度的普通基线降低了8.2%的困惑度，并优于24层模型，而半缓存配置在保留大部分增益的同时实现了6.3%的困惑度降低。

## 2 相关工作

#### 从深到浅的反馈连接。

\(6 (https://arxiv.org/html/2608.18486#bib.bib2)\)用每个过去token的LL层状态的softmax混合池替换每层的KV，该混合池在所有消费者层之间共享。\(27 (https://arxiv.org/html/2008.18486#bib.bib3)\)仅将每个消费者层连接到顶层KV，并贡献了一种迭代训练过程，使得这种反馈架构在LLM规模上变得可行。\(2 (https://arxiv.org/html/2608.18486#bib.bib11)\)通过将缓存的中层状态注入较早层的残差流，在token之间传播单一的固定深层源。这三种方法使用跨消费者共享的单一连接模式，或来自单一深层源的固定连接。循环Transformer\(21 (https://arxiv.org/html/2608.18486#bib.bib5)\)则为每个消费者分配其自身层的输出作为KV。没有一种方法允许连接跨越所有源层，并随消费者层变化和适应每个过去token。

#### 前馈跨层连接。

在残差流中，DenseFormer和LAuReL-PA用早期层输出的混合替换每层的输入\(22 (https://arxiv.org/html/2608.18486#bib.bib6);19 (https://arxiv.org/html/2608.18486#bib.bib7)\)。MUDDFormer使混合权重内容相关，并为Q、K、V和残差流计算单独的聚合\(29 (https://arxiv.org/html/2608.18486#bib.bib8)\)。Hyper-Connections和mHC学习多个并行残差流之间的连接\(37 (https://arxiv.org/html/2608.18486#bib.bib14);30 (https://arxiv.org/html/2008.18486#bib.bib15)\)。DeepCrossAttention和Attention Residuals使用输入依赖的注意力关注早期层输出\(11 (https://arxiv.org/html/2008.18486#bib.bib20);14 (https://arxiv.org/html/2008.18486#bib.bib16)\)，而Delta Attention Residuals关注子层更新而非累积状态\(17 (https://arxiv.org/html/2008.18486#bib.bib17)\)。

相关方法通过键和值路径形成连接。值残差方法将第一层的值添加到后续层，使用逐层系数或逐token门控\(36 (https://arxiv.org/html/2008.18486#bib.bib12);9 (https://arxiv.org/html/2008.18486#bib.bib13)\)。其他方法使用分组模式（如CLA、MLKV和YOCO跨解码器）在层间共享KV\(1 (https://arxiv.org/html/2008.18486#bib.bib22);39 (https://arxiv.org/html/2008.18486#bib.bib23);25 (https://arxiv.org/html/2008.18486#bib.bib21)\)；这些是\(28 (https://arxiv.org/html/2008.18486#bib.bib4)\)路由框架的实例。FusedKV为每个上层提供来自底层和中层KV的静态混合。其Lite变体直接重用中层键和底层值\(16 (https://arxiv.org/html/2008.18486#bib.bib9)\)。\(7 (https://arxiv.org/html/2008.18486#bib.bib18)\)使用随机跨层注意力进行训练。这些方法可以通过跨层共享KV来减少KV缓存大小，但KV只能由相同层或更低层的隐藏状态产生。因此，它们没有将更深的过去token表示暴露给浅层消费者层。

#### 通过重复计算进行潜在推理。

Coconut\(10 (https://arxiv.org/html/2008.18486#bib.bib29)\)对语言模型进行微调，以将顶层隐藏状态反馈为连续潜在输入。PonderLM系列通过回收输入嵌入或插入潜在位置，将相关的重复计算引入预训练，其中一些变体使用自适应停止\(35 (https://arxiv.org/html/2008.18486#bib.bib24);24 (https://arxiv.org/html/2008.18486#bib.bib25);34 (https://arxiv.org/html/2008.18486#bib.bib26);15 (https://arxiv.org/html/2008.18486#bib.bib28)\)。插入位置变体通过将选定观察token的顶层隐藏状态反馈，在其后附加潜在输入。因此，从深到浅的反馈仅发生在其后有潜在思考token的token上。另一条研究路线对每个token应用权重共享的层堆栈进行多个循环步骤\(4 (https://arxiv.org/html/2008.18486#bib.bib31);8 (https://arxiv.org/html/2008.18486#bib.bib30);38 (https://arxiv.org/html/2008.18486#bib.bib27)\)。展开后，这些模型在深度上仍然是前馈的，注意力读取相同深度的状态。阶梯注意力\(13 (https://arxiv.org/html/2008.18486#bib.bib10)\)也在时间上递归，并扩展了\(6 (https://arxiv.org/html/2008.18486#bib.bib2)\)的反馈记忆。这些方法随着递归次数增加了每个token的计算量。WhiteMatter则相反，它将所有过去token状态暴露给每一层，不使用插入位置，并且解码成本与普通模型相似。

## 3 方法

我们修改了一个具有LL层和隐藏宽度DD的Transformer解码器。我们用TT表示序列长度，ii表示token位置，l\\ell表示层索引，jj表示通道索引。WhiteMatter保留了标准解码器块，但用*跨层KV池*替换了LL逐层的KV投影。在每个过去token处，一个数据依赖的路由器将所有LL层的隐藏状态混合到k≤L个共享通道中。一组k个共享投影对\{WjK,WjV\}j=0k−1将这些通道转换为键和值。因此，整个KV缓存的大小是标准LL层缓存的k/L。

每一层使用§3.1 (https://arxiv.org/html/2008.18486#S3.SS1)中描述的固定选择读取一个通道。键和值通道使用独立的带符号混合，权重为αK[i]和αV[i]（§3.1 (https://arxiv.org/html/2008.18486#S3.SS1)）。我们评估的配置学习αK和αV并使用固定的通道选择。

### 3.1 跨层KV池

源端KV生成每个过去缓存槽cc查询端KV读取新token的每一层l\\ellof新token步骤1：将LL状态混合成kk个通道步骤2：KV投影步骤3：逐层通道选择堆叠的隐藏状态h[c]h[c], L×DL\{\\times\}D路由器N×权重k×Lk\{\\times\}Lh~jK,h~jV\\tilde\{h\}^\{K\}\_\{j\},\\tilde\{h\}^\{V\}\_\{j\}k×Dk\{\\times\}D投影NKN\_\{K\}缓存K~j,Vj\\tilde\{K\}\_\{j\},V\_\{j\}\[⋅\]\[\\,\\cdot\\,\]选择一个通道attnQlQ\_\{\\ell\}输出NRMSNorm

图2：一个token位置ii的跨层KV池。虚线分隔线分隔了§3.1 (https://arxiv.org/html/2008.18486#S3.SS1)的三个步骤。在步骤1中，一个数据依赖的路由器将LL逐层状态混合成k个共享通道。在步骤2中，生成的通道进行KV投影；然后在缓存存储之前对键应用K归一化和RoPE。在步骤3中，每个查询侧层读取一个通道；虚线箭头表示缓存边界，因为存储的通道在处理后续token时被读取。键和值分支独立处理。#### 步骤1：将LL状态混合成kk个通道。

令hl[i]∈RDh\_\{\\ell\}\[i\]\\in\\mathbb\{R\}^\{D\}为进入层l\\ell在tokenii处的隐藏状态。在每个位置ii，池使用动态混合权重将LL源状态组合成k个通道，为键和值分支独立计算。我们描述键分支；值分支相同，具有自己的参数。

每个源状态首先进行RMS归一化，得到h^lK[i]\\hat\{h\}^\{K\}\_\{\\ell\}\[i\]。这种预混合归一化使LL层处于共同尺度，并防止它们的大小在通过反馈循环递归时增长。

混合权重αK[i]∈Rk×L由读取归一化状态的线性路由器生成。为了减少路由器的大小，它可能只读取每第p个源层，从上往下数