WhiteMatter: All-to-All Cross-Layer Connections via KV Mixing

arXiv:2608.184862026-08-20

让Transformer的每一层都能读取其他层的信息,像大脑白质那样跨层连线,从而在不加深模型的情况下提升语言建模效果

在普通的Transformer里,每一层回看之前的词元时只能用自己那一层当时生成的信息,哪怕更深的层已经算出了更丰富的表示也用不上。WhiteMatter给每个词元加了一个小型路由网络,把所有层的隐藏状态混合成少量共享通道,每一层按固定规则读取其中一个通道,这样浅层也能拿到深层算出的信息。在相同训练条件下,WhiteMatter的效果超过了层数多50%的普通Transformer,即使把缓存减半,大部分提升也依然保留。

METAL MEDIA 解读图

让Transformer的每一层都能读取其他层的信息,像大脑白质那样跨层连线,从而在不加深模型的情况下提升语言建模效果

01问题所在:Transformer回看历史词元时,每一层只能使用自己那一深度当时产生的信息,无法利用模型已经计算出的更深层表示。
02方法:在每个词元位置,一个路由网络读取所有层的隐藏状态,把它们混合成k个共享通道;每一层按固定规则只读取其中一个通道用于注意力计算。当通道数k小于层数L时,需要存储的缓存就会变小。
03训练难点:要生成某个词元的通道,必须先完成该词元所有层的计算,而这些计算又依赖之前词元的通道,形成循环依赖。作者用一种分组循环迭代的调度方式解决这个问题,在保持词元间并行计算的同时,经过几轮迭代逐步逼近正确结果。
04结果:在80亿词元的训练规模下,16层的WhiteMatter模型相比同样16层的普通Transformer,在验证集上的困惑度(衡量预测好坏的指标,越低越好)降低了8.2%,甚至优于24层的普通模型。把缓存减半后仍保留6.3%的改进,并比缓存大小相同的对比模型(LCKV)低5.0%。
05速度:这种分组循环迭代调度收敛速度最多比完全顺序计算快13.9倍,比最简单的并行迭代方式快11.2倍。

这是 METAL MEDIA 制作的解读图，并非论文作者提供的原图。

他们做了什么

问题所在:Transformer回看历史词元时,每一层只能使用自己那一深度当时产生的信息,无法利用模型已经计算出的更深层表示。
方法:在每个词元位置,一个路由网络读取所有层的隐藏状态,把它们混合成k个共享通道;每一层按固定规则只读取其中一个通道用于注意力计算。当通道数k小于层数L时,需要存储的缓存就会变小。
训练难点:要生成某个词元的通道,必须先完成该词元所有层的计算,而这些计算又依赖之前词元的通道,形成循环依赖。作者用一种分组循环迭代的调度方式解决这个问题,在保持词元间并行计算的同时,经过几轮迭代逐步逼近正确结果。
结果:在80亿词元的训练规模下,16层的WhiteMatter模型相比同样16层的普通Transformer,在验证集上的困惑度(衡量预测好坏的指标,越低越好)降低了8.2%,甚至优于24层的普通模型。把缓存减半后仍保留6.3%的改进,并比缓存大小相同的对比模型(LCKV)低5.0%。
速度:这种分组循环迭代调度收敛速度最多比完全顺序计算快13.9倍,比最简单的并行迭代方式快11.2倍。

Figure 4: Held-out language-modeling quality versus non-embedding parameter count at an 8B-token budget. The connected vanilla points form the depth-scaling reference; point labels report per-token KV-cache size relative to the L=16 vanilla model. WhiteMatter is shown in half- and full-cache configurations; the LCKV configurations have four and seven warmup layers.

Table 1: Downstream evaluation. LAMBADA and WikiText report perplexity; the remaining columns report normalized accuracy in percent. Bold denotes the best result among the 16-layer models.
Model	LAMBADA ↓	WikiText ↓	PIQA ↑	HellaSwag ↑	ARC-E ↑	OBQA ↑
Vanilla 16L	127.47	49.34	60.88	31.67	47.39	29.00
LCKV w=4	107.52	48.81	62.57	32.52	45.66	31.20
LCKV w=7	102.97	49.02	62.24	32.40	46.21	30.00
WhiteMatter k=8	71.58	44.40	62.35	33.61	45.71	29.60
WhiteMatter k=16	60.73	43.28	63.55	33.80	46.21	29.40
Vanilla 24L	97.40	44.71	62.73	33.21	47.94	31.80
Vanilla 32L	79.39	41.44	63.82	34.35	47.90	32.20

Figure 5: Prefill convergence wall time versus group count g. Jacobi (g=1) and autoregressive evaluation (g=T) form the two endpoints. The 4-layer model was trained with exact autoregressive execution at length 1024; evaluation used T=2048 and the same channel-read policy for pass selection and timing.

Table 2: Measured per-token FLOPs for training, prefill, and decoding. Each pair of columns reports GFLOPs per token and the ratio to the 16-layer vanilla model.
Model	Training	Prefill	Decode
	GFLOP/tok	×	GFLOP/tok	×	GFLOP/tok	×
Vanilla 16L	0.444	1.00	0.142	1.00	0.179	1.00
Vanilla 24L	0.665	1.50	0.212	1.50	0.269	1.50
Vanilla 32L	0.887	2.00	0.283	2.00	0.359	2.00
LCKV w=4	1.421	3.20	0.935	6.61	0.173	0.97
LCKV w=7	1.174	2.65	0.738	5.21	0.175	0.97
WhiteMatter k=8	1.028	2.32	0.432	3.05	0.177	0.99
WhiteMatter k=16	1.111	2.50	0.467	3.30	0.184	1.03

Figure 6: Training-schedule and pool-rank ablations. (a) Model performance across training iteration schedules. Each cell reports one combination of iteration parameters, averaged over two seeds. The vertical axis nests the number of no-gradient passes nno-grad (inner labels) within the number of gradient-carrying passes ng (outer labels); the horizontal axis shows the training iteration schedule. (b) Test perplexity versus pool rank for WhiteMatter and two ablations.

Table 3: Complete zero-shot downstream results. LAMBADA and WikiText report perplexity; all other columns report accuracy in percent.
Model	LAMBADA	WikiText	PIQA	Wino.	BoolQ	Hella.	ARC-E	ARC-C	OBQA
Vanilla 16L	127.47	49.34	60.88	50.04	48.17	31.67	47.39	25.00	29.00
LCKV w=4	107.52	48.81	62.57	50.91	48.93	32.52	45.66	23.72	31.20
LCKV w=7	102.97	49.02	62.24	52.41	60.64	32.40	46.21	25.94	30.00
WhiteMatter k=8	71.58	44.40	62.35	51.38	60.40	33.61	45.71	25.26	29.60
WhiteMatter k=16	60.73	43.28	63.55	51.38	61.04	33.80	46.21	25.43	29.40
Vanilla 24L	97.40	44.71	62.73	50.99	56.67	33.21	47.94	26.37	31.80
Vanilla 32L	79.39	41.44	63.82	50.51	56.18	34.35	47.90	26.96	32.20

Figure 7: Whole-brain white-matter tractography. A population-averaged human structural connectome reconstructed from diffusion MRI, rendered as fiber tracts in sagittal (left) and coronal (right) views. The tracts span the brain and arc between distant regions in every direction; color encodes local fiber orientation (red: left–right, green: anterior–posterior, blue: superior–inferior). Rendered with DSI Studio (33) from its population-averaged human template (32), built from Human Connectome Project data (5).

为什么重要

模型生成时需要保存的缓存大小和计算量直接决定了实际部署的成本和速度。这项工作展示了一种不用加深模型就能提升效果、同时还能缩小缓存占用的方法,为用更少资源运行更强模型提供了参考思路。

Figure 8: Convergence timing for the larger cyclic-trained model. This 8-layer, D=1024, k=8 model is evaluated at T=4096.

本文术语

困惑度(perplexity) · 衡量语言模型预测下一个词好坏的指标,数值越低说明预测越准
KV缓存 · Transformer为避免重复计算而保存的历史词元的键(Key)和值(Value)数据
路由网络(router) · 一个小型神经网络,负责决定如何混合多个层的信息
自回归解码 · 逐个生成词元,每个新词元都基于之前已生成的所有词元
雅可比迭代/高斯-赛德尔迭代 · 用于求解循环依赖问题的数值方法,通过多轮近似计算逐步收敛,更新顺序和分组方式会影响并行度与收敛速度

论文原文摘要（英文）

In a Transformer, each layer attends to past tokens only through KV produced at its own depth, despite the presence of deeper representations during autoregressive decoding. Feedback architectures allow shallow consumer layers to attend to KV produced by deeper past-token representations, but give all consumer layers the same fixed connection patterns to source layers. We propose WhiteMatter, which connects every attention layer to the representations from all layers of each past token, with connection weights that can vary across consumer layers and adapt to the source token. For each token, a router implements these connections by mixing its $L$ layer states into $k$ KV channels that are cached for subsequent tokens; each consumer layer attends to one of the channels. The number of channels $k$ controls the KV-cache size. Setting $k<L$ reduces the cache's memory footprint. In our pretraining experiments, WhiteMatter outperforms a vanilla Transformer with 50% more layers and retains most of this gain with a 50% KV-cache compression.

作者 · Wenbo Zhang, Xiang Ren

在 arXiv 阅读

METAL MEDIA 最新报道

NCT 127发行时隔两年全新正规专辑《BLINGY》

图片来源： Wenbo Zhang et al., arXiv:2608.18486, CC BY 4.0

WhiteMatter: All-to-All Cross-Layer Connections via KV Mixing

让Transformer的每一层都能读取其他层的信息,像大脑白质那样跨层连线,从而在不加深模型的情况下提升语言建模效果

他们做了什么

为什么重要

本文术语

最新论文

METAL MEDIA 最新报道