The Geometry of Inference in Transformer Residual Streams
Timur Mudarisov, Mikhail Burtsev, Radu State
cs.CL, cs.LG
2026-09-29
以其它上下文的终点为参照,六个模型实测残差流:首层起平均偏好为正,数百竞争者到深层才清场,轨迹被证明不是直线收敛。
Tuned lens 一类工作早就发现,Transformer 中间层的残差里已经能解出接近最终答案的预测;对最终残差态的研究也找到了输出相关的几何结构。但「中间态里有信息」和「中间态分得清自己的去向」是两回事:一个第 10 层的状态离自己的终点有多近、比离别的上下文的终点近多少、这个偏好从哪层开始、随深度怎么收紧,此前没有系统测量。
拿「到终点的距离」当尺子也不行。欧氏距离把范数和方向混在一起,曲线走平可能只是范数在涨,方向其实一直在对齐。论文据此提出「几何推断假设」:中间残差态表达的是对可能终点集合的部分区分,残差更新靠改变相对几何来细化这些区分。
测量框架很朴素。对每个上下文取最后一个 token 位置、每个 block 之后的残差态,最后一层的态记为该上下文的「终点」;其它 1,023 个上下文的终点构成固定的参照银行。第 ℓ 层上,凡是比自己的终点更近的终点都算「竞争者」,竞争者比例就是几何特异性的度量。欧氏与余弦两种距离并行,余弦下的竞争区域是球面上一顶帽子,欧氏下是一个球。
实验对象是 Gemma-2B/7B、Qwen2.5-1.5B/7B、Mistral-7B、Llama-3-8B,每个 1,024 条 256-token 的 FineWeb 上下文。另做一组拟合实验:只用终点数据拟合五种方向分布(均匀球、球冠、单个 vMF、vMF 混合、投影正态),检验它们能否预测 held-out 轨迹上的竞争者曲线。
| 观察 | 指标 | 结果 |
| 六个模型 | 首个测量层起平均偏好为正 | 双度量均成立,文档级 p<10⁻³ |
| 其中五个模型 | 欧氏竞争者数量 | 深度大半区间维持数百个 |
| 六个模型 | token 排名与终点余弦距离的 Spearman | 0.726(Gemma-2B)至 0.982(Gemma-7B) |
| 端点分布拟合 | held-out 诊断相对均匀球面 | 提升 4.6–7.5 倍 |
| Qwen2.5-1.5B 首层 | 同/异输入 token 的竞争者比例 | 0.364 vs 0.058 |
三层结论。平均意义上的「认亲」出现极早,选择性却来得很晚:自己的终点从第一个测量层起就比平均替代者近,但数百个个体竞争者还在,到深层才骤降几个数量级(Gemma-7B 收缩得最早)。收缩也不是从固定名单里逐个划掉:余弦度量下竞争者集合有进有出,Qwen2.5-1.5B 的换手贯穿大部分深度,结合直线定理,这直接排除了直线收敛。欧氏距离的平台期还能和方向特异性的持续改善并存,Mistral-7B 和 Llama-3-8B 的中层最明显。
终点几何确实连着输出:查询上下文里概率越低的 token,其对应终点在余弦距离上越远,六个模型全部显著,欧氏版本只有四个显著。终点间距离与输出分布 KL 散度的 Pearson 相关在 0.25 到 0.63 之间。
控制实验排除了平凡解释。去掉同文档替代者,曲线几乎不变;换成同输入 token 的银行,竞争者比例升高(0.364 vs 0.058)但偏好仍在;把输入 token 的嵌入方向投影掉,六个模型首层余弦偏好依然为正(margin 0.0159–0.1124),只是三个模型的欧氏偏好推迟到第 2、5、13 层。早期偏好由 token 身份和更广的上下文共同贡献。
拟合实验给出外推证据:投影正态模型在六个模型中的五个上预测误差最低,Gemma-7B 上 vMF 混合最好,两者都压过球冠与单个 vMF,均匀球面在五个模型上垫底。
对做可解释性的人,「终点银行加竞争者比例」把三件常被混谈的事拆开了:到终点的距离、竞争者数量、幸存终点的集中度。用 logit lens 评估中间层时,「能解出信息」和「几何上已经区分」不是同一个属性,这篇给了度量后者的工具。
直线定理是最可复用的部分:相邻层观察到一次竞争者进入,就能证明该轨迹不是直线收敛,只需要逐层激活,不碰权重。做早退和层剪枝的人也可以直接引用「平均偏好不等于选择性」这一条。
定位要诚实:这是描述性的几何分析,没有干预实验,不改进能力,交付的是测量框架和否证基线,不是机制解释。
作者自述:拟合的分布能沿已观测轨迹预测竞争,但不解释残差更新的来源与时机;几何竞争的功能角色未知;竞争者集合依赖银行和度量的选择,且在终点处按构造为零;终点几何只解释输出变化的一部分;换手现象不能证明模型内部存在显式搜索。在 margin 过零附近做干预的实验只是被提出,没有做。
读下来另需留意:每个模型只有 1,024 条 256-token 上下文,且只看每个上下文的最后一个位置,对终点分布的刻画有限;全部是不超过 8B 的 base 模型;「竞争者」的定义依赖参照银行的组成,换一批上下文数值就会变。推理用低精度执行,作者也提醒接近打平处的进入与退出要打数值折扣。