Predictive Geometry of Hidden Trajectories in Transformers¶

会议: NeurIPS 2026
arXiv: 2609.37717
OpenReview: https://openreview.net/forum?id=2abHWd81WV
领域: 其他
关键词: Transformer几何, Fisher信息, 拉回曲率, 模型压缩, 结构化剪枝

一句话总结¶

本文通过将终端预测损失沿固化前向计算拉回到中间层残差流，导出拉回Fisher曲率算子与因果Token敏感度，揭示出Transformer隐藏轨迹的高阶预测各向异性，并为非均匀秩分配、Token更新剪枝及低秩蒸馏提供了显式的预测几何指导。

研究背景与动机¶

大规模仅解码器（decoder-only）Transformer模型在训练时仅在输出层施加自回归的下一个Token预测损失，然而推断计算却在由多层残差流构成的长链隐状态轨迹中逐步展开。在可解释性与表征学习的研究中，研究者常常观察到中间表示呈现出某种低维流形特征（如内在维度沿网络深度的膨胀-收缩曲线），或者发现模型权重能够承受显著的低秩压缩（如SVD与LoRA类低秩适配）。然而，以往的工作大多只是将这些结构作为经验现象进行统计描述，未能将中间隐状态的几何形态与端到端训练时施加的终端预测目标直接在数学上联系起来，更缺乏一种由目标损失严格推导得出的、判定哪些隐空间方向与哪些Token位置在预测中至关重要的局部刻画工具。

这种割裂导致在模型压缩与分析时面临核心矛盾：既然最终的预测质量只取决于输出分布，那么隐状态空间中绝大多数自由度很可能并不影响预测结果，但通用的欧式距离或层内激活范数却无法区分“对输出致命的微扰”与“完全被后续层吸收的零扰动”。例如，在低秩权重分解与Token剪枝中，均匀分配秩或基于注意力权重幅值的剪枝策略，本质上都是在缺乏下游预测敏感度感知的情况下进行盲目裁剪，极易破坏极少数维度的关键几何信息。

针对这一问题，本文的切入角度是直接从固化模型的下游预测映射出发，建立自顶向下的拉回敏感度分析。通过将候选隐状态代入固定下游Transformer块并映射到终端预测KL散度，建立层级“待行损失（loss-to-go）”函数。核心 idea：在成功预测轨迹的低损失稳态附近，待行损失函数的二阶局部曲率由输出分布的拉回Fisher算子主导，该算子直接将残差流空间严格解耦为高敏感的可观测子空间与一阶预测零方向，并导出沿因果图传播的Token级曲率敏感度得分。

方法详解¶

整体框架¶

本文针对固化训练完成的仅解码器Pre-Norm Transformer模型，分析在给定目标位置 $t$ 处的预测几何。记中间第 $\ell$ 层的隐藏状态为 $X_\ell \in \mathbb{R}^{t \times d}$。从第 $\ell$ 层到终端Logits的映射表示为 $\Psi_\ell(X) = Z_t(\Phi_{\ell \to L}(X))$，其中 $\Phi_{\ell \to L} = B_{L-1} \circ \cdots \circ B_\ell$ 为后续Transformer块的复合前向计算，而 $Z_t$ 为终端归一化与LM Head投影。对于参考轨迹 $X_\ell^*$，其诱导的终端预测分布为 $p^* = \operatorname{softmax}(\Psi_\ell(X_\ell^*))$。当参考点达到低损失时，定义在中间层的待行损失函数 $\mathcal{J}_\ell(X)$ 的二阶局部几何由拉回Fisher算子 $K_\ell$ 主导。该算子不仅给出了残差流内部高曲率方向与零曲率方向的正交分解，还在因果祖先图上诱导了衡量各Token重要性的标量曲率得分 $\kappa_{\ell,s}$。

整个分析与应用流程如下所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入Token前缀与残差隐状态轨迹"] --> B["构建待行损失函数与下游Logit映射"]
    B --> C["拉回Fisher算子与局部二阶曲率分解"]
    C --> D["谱截断与非均匀层级低秩分配"]
    C --> E["因果Token纤维迹敏感度得分"]
    C --> F["低秩学生网络预测几何正则化蒸馏"]
    D --> G["输出：几何指导的模型降维与结构化更新剪枝"]
    E --> G
    F --> G

关键设计¶

1. 拉回Fisher算子与局部二阶曲率分解：将隐状态空间严谨解耦为可观测方向与预测零方向

针对隐空间微扰对预测影响不可知的问题，本文研究了待行损失函数 $\mathcal{J}_\ell(X) = D_{\mathrm{KL}}(q \parallel \operatorname{softmax}(\Psi_\ell(X)))$ 围绕参考隐状态 $X_\ell^*$ 的局部二阶泰勒展开。记中心化Logit投影为 $\bar{\Psi}_\ell = P \Psi_\ell$（其中 $P = I - \frac{1}{|\mathcal{V}|}\mathbf{1}\mathbf{1}^\top$ 为消除Logit平移不变性的中心化矩阵），对应雅可比矩阵为 $\bar{J}_\ell = D\bar{\Psi}_\ell(X_\ell^*)$。输出分布的Fisher信息矩阵定义为 $F^* = \operatorname{diag}(p^*) - p^* p^{*\top}$。本文证明了待行损失的Hessian矩阵满足： $$ \nabla^2 \mathcal{J}\ell(X\ell^) = \bar{J}_\ell^\top F^ \bar{J}\ell + R\ell = K_\ell + R_\ell $$ 其中残差项的算子范数由输出预测误差严格控制：$\|R_\ell\|_{\mathrm{op}} \leq M_\ell \|p^* - q\|_2 \leq M_\ell \sqrt{2\varepsilon_*}$（$M_\ell$ 为下游映射的二阶导数范数界）。因此，在模型预测较好的低损失轨迹邻域内，局部曲率完全由拉回Fisher算子 $K_\ell = \bar{J}_\ell^\top F^* \bar{J}_\ell$ 所主导，其诱导的二次型度量为 $g_\ell(\delta X, \delta X) = \langle \delta X, K_\ell \delta X \rangle$。这一算子的核空间 $\mathcal{N}_\ell := \ker \bar{J}_\ell$ 严格构成了首阶预测零方向（prediction-null directions），在此子空间内的微扰在首阶上完全不改变中心化目标Logits；而在其正交补空间中，高特征值对应的特征向量构成了高度敏感的局部可观测子空间（observable subspace）。

2. 因果Token纤维迹敏感度得分：基于因果雅可比耦合的损失感知Token重要性度量

为了解决传统注意力权重或激活范数在评估Token贡献时无法反映终端预测损失惩罚的问题，本文通过在Token纤维维度上取迹，导出了Token级的敏感度指标。设 $P_s: \mathbb{R}^{t \times d} \to \mathbb{R}^d$ 为将整个隐状态投影到第 $s$ 个Token纤维空间的投影算子，对应的下游Token到Logit的雅可比矩阵为 $H_{\ell,s} = D_{X_{\ell,s}} \bar{\Psi}_\ell(X_\ell^*) = \bar{J}_\ell P_s^\top \in \mathbb{R}^{|\mathcal{V}| \times d}$。定义Token级曲率得分为： $$ \kappa_{\ell,s} := \operatorname{tr}(P_s K_\ell P_s^\top) = \operatorname{tr}(H_{\ell,s}^\top F^* H_{\ell,s}) = |F^{* 1/2} H_{\ell,s}|_F^2 $$ 该公式表明 $\kappa_{\ell,s}$ 等价于在输出Fisher加权几何下，目标Logits对第 $\ell$ 层第 $s$ 个Token隐藏状态微扰的Frobenius范数平方敏感度。由于因果自注意力的因果掩码约束，如果第 $s$ 个Token不在目标Token $t$ 的下游因果祖先闭包中，下游雅可比必有 $H_{\ell,s} = 0$，进而严格导致 $\kappa_{\ell,s} = 0$。此外，本文证明了沿因果路径的乘积界：$\kappa_{\ell,s}$ 的大小由从 $(\ell,s)$ 沿所有有向计算路径到达 $(L,t)$ 的Token块雅可比乘积所控制，从数学上构成了比启发式注意力幅值更严密的损失感知剪枝信号。

3. 无矩阵曲率估计与谱截断秩分配：无须显式构造海量矩阵的实用算子化实现

由于残差隐状态空间维度 $t \times d$ 与词表大小 $|\mathcal{V}|$ 均达数万，拉回矩阵 $K_\ell$ 绝无法显式具象化。本文利用JVP与VJP构成无矩阵乘法管线 $v \mapsto \bar{J}_\ell v \mapsto F^* (\bar{J}_\ell v) \mapsto \bar{J}_\ell^\top (F^* \bar{J}_\ell v)$，并结合Hutchinson迹估计器高效估计参与率有效秩 $r_{\mathrm{eff}}(K_\ell) := (\operatorname{tr}(K_\ell))^2 / \operatorname{tr}(K_\ell^2)$。在模型压缩时，根据设定的谱尾部质量阈值 $\tau$，自适应为各层分配最小截断秩 $k_\ell$ 使得 $\sum_{i > k_\ell} \lambda_i(K_\ell) \leq \tau \sum_i \lambda_i(K_\ell)$。由主干理论保证，丢弃尾部子空间 $\delta X_{\mathrm{tail}} = (I - \Pi_\ell^{(k_\ell)})\delta X$ 所导致的局部二次损失增量被严格限制在 $\frac{1}{2}\lambda_{k_\ell+1}(K_\ell)\|\delta X_{\mathrm{tail}}\|^2$ 以内。

损失函数 / 训练策略¶

在应用预测几何恢复压缩后的低秩学生模型时，针对注意力层与MLP层的低秩分解因子 $W \approx AB$（仅训练低秩因子，冻结其余权重），采用自回归知识蒸馏目标结合预测几何正则项： $$ \mathcal{L} = \mathcal{L}{\mathrm{KD}} + \lambda}} \mathcal{L{\mathrm{CE}} + \lambda}} \mathcal{L{\mathrm{geo}} $$ 其中几何正则化项的理论形态为惩罚隐状态匹配误差在教师拉回Fisher度量下的能量 $\frac{1}{2}(h_\ell^S - h_\ell^T)^\top K_\ell^T (h_\ell^S - h_\ell^T)$。为了保证跨模型体系的计算鲁棒性，实践中使用各层通过最终归一化与输出头投影后的读出Fisher代理： $$ \mathcal{L}}} = \frac{1}{|\mathcal{L{\mathrm{geo}}|} \sum{\mathrm{geo}}} \frac{1}{2} \Delta z\ell^\top F_T \Delta z_\ell, \quad \Delta z_\ell = \operatorname{LMHead}_T(\operatorname{FinalNorm}T(h\ell^S)) - \operatorname{LMHead}_T(\operatorname{FinalNorm}T(h\ell^T)) $$ 该项直接惩罚在输出分布敏感方向上的表征偏离，而容许在预测零方向上产生较大松弛，并配合Forward KL、Reverse KL（RKL）与Skew KL（$\alpha=0.1$）进行梯度反向传播。

实验关键数据¶

主实验¶

在低秩学生模型知识蒸馏恢复实验中，评估在相同参数规模、低秩因子维度、随机种子和初始化方式（SVD或ASVD）下，引入预测几何正则项（+Geo）相比基础KD目标的性能优势 $\Delta_{\mathrm{Geo}} = \mathrm{PPL}(\mathrm{KD}) - \mathrm{PPL}(\mathrm{KD+Geo})$（正值表示几何项使困惑度PPL降低，恢复效果更好）。评测覆盖SmolLM2-135M、SmolLM2-360M与Qwen2.5-0.5B等架构，成对测试结果如下：

KD基线	配对组数 (Pairs)	几何项胜出组数 (Wins)	平均 $\Delta_{\mathrm{Geo}} \uparrow$	中位数 $\Delta_{\mathrm{Geo}} \uparrow$	SVD初始化均值	ASVD初始化均值
FKL (Forward KL)	14	8	-38.4	-12.6	-35.0	-41.7
RKL (Reverse KL)	14	9	+21.3	+12.3	+44.5	-1.9
SkewKL ($\alpha=0.1$)	14	10	+59.3	+8.6	+66.7	+51.9

消融实验¶

以Qwen2.5-0.5B模型在低秩压缩（压缩倍率与分解秩 $r$ 对应，训练温度 $\tau=2$）下的最终困惑度（Final PPL）消融对比为主，考察不同蒸馏准则与几何项的交互影响：

压缩配置 / 秩	蒸馏算法	初始PPL (Init PPL) $\downarrow$	最终PPL (Final PPL) $\downarrow$	相对初始PPL改善量 $\Delta_{\mathrm{init}} \uparrow$	说明
ASVD / Rank 16 (40.5x)	asvd + rkl	$2.44 \times 10^5$	653.0	$2.43 \times 10^5$	传统RKL基线
ASVD / Rank 16 (40.5x)	asvd + skewkl	$2.44 \times 10^5$	551.2	$2.43 \times 10^5$	强自回归SkewKL基线
SVD / Rank 16 (40.5x)	svd + rkl	$2.15 \times 10^6$	761.5	$2.15 \times 10^6$	SVD初始化的RKL基线
SVD / Rank 16 (40.5x)	svd + rkl_geom	$2.15 \times 10^6$	756.7	$2.15 \times 10^6$	引入几何正则项，PPL下降4.8
SVD / Rank 16 (40.5x)	svd + skewkl	$2.15 \times 10^6$	690.9	$2.15 \times 10^6$	SVD初始化的SkewKL基线
SVD / Rank 16 (40.5x)	svd + skewkl_geom	$2.15 \times 10^6$	682.3	$2.15 \times 10^6$	引入几何正则项，PPL下降8.6
ASVD / Rank 32 (20.3x)	asvd + skewkl	$1.63 \times 10^6$	503.3	$1.63 \times 10^6$	ASVD下SkewKL基线
ASVD / Rank 32 (20.3x)	asvd + skewkl_geom	$1.63 \times 10^6$	495.5	$1.63 \times 10^6$	引入几何正则项，PPL下降7.8
SVD / Rank 32 (20.3x)	svd + rkl	77866	648.9	77217	SVD下RKL基线
SVD / Rank 32 (20.3x)	svd + rkl_geom	77866	617.1	77249	引入几何正则项，PPL显著下降31.8

关键发现¶

残差流极度各向异性：在所有测试模型中，可观测子空间内的等模长微扰所引发的KL散度变化，比零空间扰动高出数个数量级，且该差距随网络深度加深而急剧扩大，表明中间隐藏状态仅有极薄的一组各向异性方向与最终预测直接挂钩。
非均匀秩分配优势：基于拉回Fisher谱截断进行的非均匀秩分配在适中阈值（$\tau=0.05, 0.1$）下对中小模型（如Gemma-2B、Qwen2.5-3B、SmolLM2-1.7B）表现出显著超越均匀SVD/LoRA的PPL保留能力，证实了不同深度层对于终端预测的关键容量需求高度异构。
与自回归蒸馏互补：在学生蒸馏恢复中，几何项对前向KL（FKL）未能带来增益，但在偏向模式寻求（mode-seeking）的强自回归目标（RKL与SkewKL）中展现了显著且稳定的改进（胜率分别达 9/14 与 10/14），证明几何约束与输出层模式对齐形成了有效的多层梯度互补。

亮点与洞察¶

将下游计算视为测量算子：突破了以往仅从自编码器或层内表征相似度看待隐藏表示的局限，首次从严格的二次型分析将下游所有Transformer层作为“拉回测量映射”，用数学语言精确定义了中间隐空间“对预测有用”与“对预测无形”的边界。
损失感知的Token剪枝指标：导出的 $\kappa_{\ell,s}$ 具有严格的因果祖先支撑，不仅在形式上统一了解码器因果掩码与敏感度传播，而且相比注意力分数具有更强的损失敏感度理论依据。
通用矩阵无矩阵计算范式：通过巧妙编排JVP与VJP，绕开了在庞大残差维度与词表空间显式物化Fisher矩阵的物理不可能，使微分几何工具能够实际部署于数十亿参数规模的开源LLM上。

局限与展望¶

计算开销依赖反向图展开：虽然采用了JVP/VJP无矩阵实现，但评估单层曲率仍需要重跑后续所有Transformer层的微分计算图，在长上下文下无法做到零开销实时推断调度。
蒸馏正则项采用了读出近似：出于跨模型架构的数值稳定性与训练效率考量，蒸馏实验中采用了输出头直连的读出Fisher代理（readout-Fisher proxy），并未在每一步完全拉回完整的深层复合雅可比。
有监督交叉熵的一阶主导性：在监督单热微扰实验中，交叉熵损失由一阶梯度项完全主导，纯二阶Fisher二次型只有在Teacher-KL场景下才表现出决定性相关，因此对于监督微调场景仍需更完善的“一阶+二阶”联合几何理论。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从严格的待行损失与拉回Fisher流形视角刻画Transformer中间隐状态轨迹几何，理论推导严密。
实验充分度: ⭐⭐⭐⭐☆ 涵盖十款开源主流大模型与多个测试基准，从微扰验证、秩分配、Token剪枝到低秩蒸馏全流程闭环验证。
写作质量: ⭐⭐⭐⭐⭐ 理论框架清晰，符号系统严谨自洽，附录提供了详尽的数学证明与算法伪代码。
价值: ⭐⭐⭐⭐☆ 为大模型低秩压缩、结构化剪枝和高效表征对齐提供了全新且坚实的几何理论基础与实用算子。

压缩配置 / 秩	蒸馏算法	初始PPL (Init PPL) \(\downarrow\)	最终PPL (Final PPL) \(\downarrow\)	相对初始PPL改善量 \(\Delta_{\mathrm{init}} \uparrow\)	说明
ASVD / Rank 16 (40.5x)	asvd + rkl	\(2.44 \times 10^5\)	653.0	\(2.43 \times 10^5\)	传统RKL基线
ASVD / Rank 16 (40.5x)	asvd + skewkl	\(2.44 \times 10^5\)	551.2	\(2.43 \times 10^5\)	强自回归SkewKL基线
SVD / Rank 16 (40.5x)	svd + rkl	\(2.15 \times 10^6\)	761.5	\(2.15 \times 10^6\)	SVD初始化的RKL基线
SVD / Rank 16 (40.5x)	svd + rkl_geom	\(2.15 \times 10^6\)	756.7	\(2.15 \times 10^6\)	引入几何正则项，PPL下降4.8
SVD / Rank 16 (40.5x)	svd + skewkl	\(2.15 \times 10^6\)	690.9	\(2.15 \times 10^6\)	SVD初始化的SkewKL基线
SVD / Rank 16 (40.5x)	svd + skewkl_geom	\(2.15 \times 10^6\)	682.3	\(2.15 \times 10^6\)	引入几何正则项，PPL下降8.6
ASVD / Rank 32 (20.3x)	asvd + skewkl	\(1.63 \times 10^6\)	503.3	\(1.63 \times 10^6\)	ASVD下SkewKL基线
ASVD / Rank 32 (20.3x)	asvd + skewkl_geom	\(1.63 \times 10^6\)	495.5	\(1.63 \times 10^6\)	引入几何正则项，PPL下降7.8
SVD / Rank 32 (20.3x)	svd + rkl	77866	648.9	77217	SVD下RKL基线
SVD / Rank 32 (20.3x)	svd + rkl_geom	77866	617.1	77249	引入几何正则项，PPL显著下降31.8