Your Transformer Can Hold Two Thoughts at Once: Evidence of Linear Superposition in LLMs
Pavel Tikhonov, Anton Korznikov, Matvey Mikhalchuk, Nikita Dragunov, Temurbek Rahmatullaev, Polina Druzhinina, Anton Razzhigaev, Ivan Oseledets, Elena Tutubalina
cs.CL, cs.AI
2026-09-24
把两段无关文本的 embedding 按位置取平均送进标准 LLM,输出仍接近两条独立 next-token 的算术平均。预训练会削弱这一性质,轻量蒸馏把 Pythia-2.8B 的 KL 从 1.86 降到 0.27。
Decoder-only Transformer 里全是非线性:softmax 注意力、MLP 激活。默认做法把一次前向当成一条语义流,要并行处理两段无关文本,就得开两次前向,或者改架构避免互相干扰。
残差流的层间映射经常能用仿射变换近似。这篇把观察推到端到端:把两段文本的 token embedding 按位置取平均,模型会不会输出两条独立 next-token 分布的叠加?命题叫 Superposition Linearity Hypothesis。
混合方式很简单。两段等长序列 x 和 y,位置 t 的输入是 (E(xt)+E(yt))/2,然后原样走因果注意力。不改 mask,不加 mux 层。
评测分三层。
注意力部分做了两个对照。Donor patching 把无关文本 C 的 post-softmax 注意力权重整层整头换进 A,Q/K/V 仍是 A 的。Permutation patching 打乱 A 自己每行因果前缀里的权重,保住行和与权重集合,毁掉位置结构。再按 token 类型切开:可预测位置约 65%,多是功能词和标点;内容位置约 35%。
要放大线性,用自蒸馏:冻结教师跑两条独立前向,学生吃混合 embedding,最小化 KL(0.5(Pteacher(A)+Pteacher(B)) ∥ Pstudent(z))。数据是 FineWeb 子集,约 20 万步,号称不到原预训练数据量的 0.025%。
解码另有一层麻烦。如果 logits 近似平均,概率会跟着几何平均走:P(t) ∝ √(PA(t) PB(t))。A 里很高、B 里很低的 token 会被压下去,直接从混合分布采样会在两条语义之间跳。Joint Contrastive 的做法是大模型吃混合输入,小模型分别看 A 和 B,用 ℓlarge(z)+α ℓsmall(A)−β ℓsmall(B) 把两条流拆开。
未改过的预训练模型已经远高于词频基线。Pythia-2.8B、Llama-3.2-3B、Qwen2.5-3B 上,单流预测 token 约 30–40% 落在混合分布 top-10,50–60% 在 top-50,60–65% 在 top-100。词表至少 5 万。对照里拿单流 A 的 logits 去找无关流 B 的真 token,top-10 只有 2.63%。
| 模型 | L=32 KL | RKL | 微调后 KL |
| Pythia-2.8B | 1.86 | 0.42 | 0.27 |
| Llama-3.1-8B | 2.16 | 0.37 | 未报告 |
| Pythia-160M | 0.91 | 0.35 | 未报告 |
L=512 时数字略好,没有随上下文变长而崩。N=3 条流等权混合时 RKL 只升 0.04 到 0.09,性质还在,定量变差。
训练轨迹指向架构先验。Pythia 各尺寸上,hidden-state 叠加误差在初始化最小,随预训练单调上升。层间线性分数呈 U 形:浅层高,中层掉到约 0.65,最后三分之一回到 0.95 以上。深层近线性,给叠加信号活到 logits 提供了几何上的出口。
注意力实验把「看起来能叠」拆开了。Qwen2.5-3B 上,可预测位置在 embedding mixing 下 median rank 是 6、top-1 24.8%;内容位置 median rank 284、top-1 只有 4.2%。Donor patching 在可预测位置更松(median 3),内容位置 111。Permutation 两边都塌:可预测 3079,内容 19246。LAMBADA 必须猜内容词,embedding mixing 的 argmax 准确率 2.25%、目标 median rank 339;donor patching 是 0.5%、2350。混合前向在难位置上保住的信号,比「换一套自然注意力形状」更多。
微调改的是内容位置。Qwen2.5-3B 微调后,可预测位置几乎不动(median 8),内容位置 median rank 从 284 掉到 5,top-1 到 22.8%。Pythia-2.8B 上混合 KL 从 1.86 降到 0.27,RKL 从 0.42 降到 0.06;true token 进 top-5 的概率从约 30% 升到超过 60%。
解码仍是短板。Joint Contrastive 把叠加前向的 LAMBADA 均值从预训练的 0.06–0.18 拉到 0.11–0.43:Llama-3.2-3B 到 0.430,对照小模型单流 0.540。Jaccard 重叠也降了。单流质量有代价:Pythia-2.8B LAMBADA 0.544→0.357,Qwen2.5-3B 0.602→0.460。纯蒸馏更狠,Qwen 单流 PPL 从 12.44 飙到 65.52;Joint Contrastive 把 PPL 压在 19.31,接近小模型的 18.96。吞吐上,双头 Separate 模式约等于 batch=2 的独立推理,大约是顺序跑两次的 2 倍;Pythia-2.8B Separate 109.6 tok/s,顺序基线 56.0。
这件事跟 DataMUX、RevMUX 那类外加 mux/demux 层不是同一类主张。标准预训练 LLM,只做 embedding 平均,就已经把两条流的信号压进同一条残差流。预训练在削弱它,轻量蒸馏能把它找回来。
对推理的含义很具体:一次前向、一份 KV cache,理论上吞吐翻倍、显存减半。附录里的 Separate 模式也显示,融合 embedding 的双头输出能摸到 batch=2 的速度。几何平均障碍还在,Joint Contrastive 仍要挂一个小模型,Llama 上 0.43 对不上 0.54 的单流小模型。把它当可部署的并行推理还早,更接近对「残差流能不能同时载两条语义」的存在性证明。
作者自己划了边界:分析实验上下文多在 L≤128,分布实验扩到 512,语料基本是单语;没测多模态,也没测把文本 token 和图像 patch 混进同一条残差流。
更大的裂缝在「拟合混合分布」和「真正拆开两条续写」之间。微调把 KL 压得很低,直接采样仍会串流。LAMBADA 上内容位置的存活,基座模型本来就差,微调救回来的同时单流能力掉一截。蒸馏目标用了不到 0.025% 的数据量,但 Pythia-2.8B 一轮仍要 2×A100 跑 114 小时,并不便宜。N=3 只做了分布距离,没有对应的解码方案。