跳转至

Spectral Geometry of Attention: From Information Routing to Uncertainty

会议: NeurIPS 2026
arXiv: 2610.06012
OpenReview: https://openreview.net/forum?id=TXNoJp3DRA
论文: https://arxiv.org/abs/2610.06012
领域: 其他
关键词: 谱几何, 注意力算子, Token差异算子, 注意力汇聚, 不确定性量化

一句话总结

本文将Transformer注意力头形式化为Token函数希尔伯特空间上的泛函映射,提出内在概率几何下的Token差异算子(TDO),在理论上解耦注意力汇聚与路由塌陷,并衍生出在上游精准捕捉路由容量衰减的无监督单遍不确定性评估器SAUQ。

研究背景与动机

注意力图一直被视为窥探Transformer内部计算机制的核心窗口。以往研究多将注意力视为词元两两间的相似度概率分布、离散时间马尔可夫链转移矩阵或图信号算子,亦或提取行级统计量(如注意力熵、前一Token关注度等)进行探测。然而,这些局部视角聚焦于词元配对交互,忽略了单步注意力变换对隐层表示维度的全局算子理论性质,难以从几何全局层面刻画注意力对下游信息路由的实质贡献。

这一局限在面对注意力汇聚(Attention Sinks)与表征塌陷时尤为突出。在标准的欧氏度量下,因果掩码和行随机性使早期词元天然汇聚大量注意力权重,导致常规欧氏奇异谱发生结构性扭曲——注意力在最初几个位置的质量堆积会人为抬高奇异谱,将纯粹的无信息汇聚误判为具有高信息路由能力。与此同时,基于模型输出概率(如困惑度、词元熵)的传统不确定性度量在长文本生成或复杂推理中极易失效,因为一旦关键语义在残差流中发生表征压缩或瓶颈挤压,输出Logits便无法反映真实的路由退化。

面对欧氏度量与注意力结构失配的核心矛盾,本文借鉴形状分析中的泛函映射与形状差异算子理论,为注意力机制赋予由其自身拉回诱导的内在概率几何。核心 idea:将单头注意力建模为加权希尔伯特空间间的泛函算子,通过由注意力转移诱导的内在概率测度构建Token差异算子(TDO),实现均值漂移与中心化路由的正交分解,从而在算子谱层面彻底解耦注意力汇聚与维度塌陷,并以此刻画路由容量以实现单遍不确定性量化。

方法详解

整体框架

本文将每个注意力头视为定义在Token空间 \(\mathcal{T}=\{1,\ldots,n\}\) 上的泛函映射 \(\mathcal{A}: L^2(\mu) \to L^2(\pi)\),其中输入与输出域分别配备正测度 \(\mu\) 与 \(\pi\)。该映射将值矩阵 \(V\) 变换为 \(AV\)。为了消除欧氏空间下注意力汇聚带来的虚假能量放大,本文将输出测度设为均匀概率测度 \(\pi = \frac{1}{n}\mathbf{1}\),并将输入测度设为注意力转置拉回测度 \(\mu = A^\top \pi\)。在此几何下,伴随算子 \(\mathcal{A}^* = \Omega_\mu^{-1} A^\top \Omega_\pi\),由此定义出对称自伴的Token差异算子 \(D_{\mu,\pi} = \mathcal{A}^* \mathcal{A}\)。

在推理生成阶段,模型基于中段Transformer层中路由头的TDO谱参与率,实时追踪序列生成过程中的有效路由维度,并通过递归置信度更新计算序列不确定性,整体流程如下图所示:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入前缀序列与注意力矩阵<br/>A ∈ ℝ^{n×n}"] --> TDO["Token差异算子构建<br/>拉回测度 μ = A^⊤ π,构建 D_{μ,π}"]
    TDO --> Decomp["均值-路由正交分解<br/>AV = 1m + AV_c (m=μ^⊤V, μ^⊤V_c=0)"]
    Decomp --> Screen["轻量级路由头筛选<br/>在中间层按生成块谱参与率选出 h*"]
    Screen --> SAUQRec["SAUQ递归置信度融合<br/>c_t = α p_t + (1-α) s_t c_{t-1}"]
    SAUQRec --> Out["输出序列不确定性评分<br/>U = 1 - (1/T) ∑ log c_t"]

关键设计

1. 内在概率几何与Token差异算子(TDO):解耦汇聚能量与路由能力
传统欧氏度量采用计数测度(\(\mu = \pi = \mathbf{1}\)),导致Gram算子 \(A^\top A\) 将全1常数向量映射为列和向量 \(A^\top \mathbf{1}\)。因果掩码使前向Token(特别是第1个Token)被后续所有位置可见,汇聚效应导致列和极端不均衡,前向主奇异值暴增至 \(\sigma_1(A)^2 \approx n\),强行吃满Frobenius范数预算并掩盖其他方向的有效信息。本文通过为输入域赋予拉回概率测度 \(\mu = A^\top \pi\)(其中 \(\pi = \frac{1}{n}\mathbf{1}\) 为均匀概率测度),使得常数函数 \(\mathbf{1}\) 天然成为算子 \(D_{\mu,\pi} = \Omega_\mu^{-1} A^\top \Omega_\pi A\) 对应特征值 \(\lambda_1 = 1\) 的平凡特征向量,且不变测度满足 \(\mu^\top D_{\mu,\pi} = \mu^\top\)。由此,算子范数严格有界 \(\|\mathcal{A}\|_{L^2(\mu)\to L^2(\pi)} = 1\)。通过相似变换矩阵 \(S_{\mu,\pi} = \Omega_\mu^{-1/2} A^\top \Omega_\pi A \Omega_\mu^{-1/2}\),可以在对称半正定形式下以无求逆代价稳定获取算子谱。

2. 均值-路由正交分解与谱熵界:严格约束有效表示维度
在内在概率测度空间中,值矩阵 \(V \in \mathbb{R}^{n \times d}\) 被自然分解为加权均值项 \(m = \mu^\top V \in \mathbb{R}^d\) 与中心化偏差项 \(V_c = V - \mathbf{1}m\)(满足 \(\mu^\top V_c = 0\))。注意力输出精确呈现正交分解: $\(AV = \mathbf{1}m + AV_c, \quad \langle \mathbf{1}m, AV_c \rangle_{\pi,F} = 0\)$ 其中 \(\mathbf{1}m\) 代表与Token位置无关的纯平移均值分量(直接由汇聚权重与上下文均值主导),而 \(AV_c\) 才是真正承载词元间差异化路由的有效分量。定理证明中心化分量的Frobenius范数受到TDO第二特征值的严格收缩约束:\(\|AV_c\|_{\pi,F} \le \sqrt{\lambda_2(D_{\mu,\pi})} \|V_c\|_{\mu,F}\)。进一步地,定义有效谱熵 \(H(Z)\),中心化分量 \(Z_c = \Omega_\pi^{1/2} AV_c\) 的谱熵以二元熵和 \(\alpha = \frac{\|Z_c\|_F^2}{\|m\|_2^2 + \|Z_c\|_F^2}\) 比例给出双边界: $\(\alpha H(Z_c) \le H(Z) \le h_2(\alpha) + \alpha H(Z_c)\)$ 这表明,一旦TDO的中心化谱发生塌缩,注意力头输出的动态表征维度必发生硬性折损,从而为注意力头的维度衰减提供了首个坚实的算子谱理论刻画。

3. 基于谱参与率的不确定性评估器(SAUQ):捕获上游路由崩塌
基于TDO谱与输出表征维度的对应关系,本文提出谱注意力不确定性量化算法(SAUQ)。为了度量中心化活跃维度的丰富度,引入归一化谱参与率(Participation Ratio): $\(\mathrm{PR}(D_{\mu,\pi}) = \frac{1}{n} \frac{(\operatorname{tr}(D_{\mu,\pi}))^2}{\operatorname{tr}(D_{\mu,\pi}^2)} = \frac{1}{n} \frac{(\operatorname{tr}(S_{\mu,\pi}))^2}{\|S_{\mu,\pi}\|_F^2}\)$ 该指标仅需计算迹与Frobenius范数,规避了代价高昂的完整特征值分解。在自回归解码步 \(t\) 中,算法在中间层 \(\mathcal{L}_{\mathrm{mid}}\)(如LLaMA-3.1-8B的第10–22层)利用生成块代理筛选出最具路由活性的注意力头 \(h^*(l)\),提取该头的TDO谱参与率信号 \(s_t^{(l)}\)。随后利用递推式融合当前词元生成概率 \(p_t\) 与历史置信度:\(c_t^{(l)} = \alpha p_t + (1-\alpha) s_t^{(l)} c_{t-1}^{(l)}\)(超参数设为 \(\alpha=0.2\))。单层序列不确定性定义为 \(U^{(l)} = 1 - \frac{1}{T}\sum_{t=1}^T \log c_t^{(l)}\),最终取各中间层最大值 \(U_{\mathrm{SAUQ}} = \max_{l \in \mathcal{L}_{\mathrm{mid}}} U^{(l)}\) 作为序列判拒评分。

损失函数 / 训练策略

SAUQ是一种完全无监督且免训练的推断期评估方法,不需要对主干LLM进行任何参数微调或额外训练。在特征值提取与矩阵计算优化上,由于仅需获取矩阵迹与Frobenius范数,计算开销极低。在附录的有监督幻觉检测扩展实验中,仅提取所选注意力头的TDO前 \(k=16\) 个特征值,输入轻量级分类器(遵循LapEigval协议)进行二分类验证。

实验关键数据

主实验

在LM-Polygraph基准上评估无监督不确定性量化性能,测试模型采用LLaMA-3.1-8B,指标为归一化预测拒识率(Normalized Prediction Rejection Ratio, PRR,越高越好)。任务涵盖短文本问答(TriviaQA, TruthfulQA, MMLU, CoQA)与长文本摘要(SAMSum, XSum)。

评估方法 类型 / 采样数 TriviaQA TruthfulQA MMLU CoQA QA平均 SAMSum XSum 摘要平均
AttentionScore (Sriramanan et al.) 单遍 / 1 0.0796 -0.0505 0.2078 0.0933 0.0826 -0.3727 -0.9983 -0.6855
Perplexity 单遍 / 1 0.1088 0.2863 0.5158 0.2737 0.2962 0.5071 0.4053 0.4562
MeanTokenEntropy 单遍 / 1 0.5661 0.4199 0.2437 0.2568 0.3716 0.4999 0.3928 0.4464
SemanticEntropy (Kuhn et al.) 采样 / 10 0.2005 0.2113 0.4701 0.3662 0.3120 0.2525 -0.1886 0.0320
LUQ (Zhang et al.) 采样 / 10 0.5878 0.1820 0.4392 0.3588 0.3920 0.2900 0.3585 0.3243
RAUQ (Vazhentsev et al.) 单遍 / 1 0.2631 0.3369 0.5596 0.3110 0.3677 0.5269 0.3008 0.4139
LapEigValRAUQ (Binkowski et al.) 单遍 / 1 0.1190 0.3296 0.5578 0.3064 0.3282 0.5285 0.2907 0.4096
LogDetRAUQ (Sriramanan et al.) 单遍 / 1 0.1140 0.3313 0.5596 0.3087 0.3280 0.5288 0.3048 0.4168
SAUQ (仅生成块变体) 单遍 / 1 0.5013 0.3056 0.5375 0.3110 0.4139 0.4948 0.2578 0.3763
SAUQ (本文完整模型) 单遍 / 1 0.2914 0.3767 0.5588 0.3055 0.3831 0.5924 0.4214 0.5069

消融实验:测度空间选择对不确定性量化的影响

在LLaMA-3.1-8B上,比较SAUQ在不同几何测度定义下的表现。仅内在测度(Intrinsic)满足拉回守恒与算子非扩张性质,其余测度均破坏了平凡常数模与中心化子空间的正交性。

测度配置 (Measure) QA平均 PRR 摘要平均 PRR 说明
内在概率测度 (Intrinsic, 本文) 0.3831 0.5069 满足算子非扩张性,完全消除汇聚伪影
位置加权测度 (Positional) 0.4027 0.4771 对短序列局部依赖较好,但长文本退化
均匀计数测度 (Uniform / 欧氏) 0.3660 0.4204 欧氏度量严重受汇聚噪声干扰,表现最差
词频齐夫定律测度 (Zipf-Law) 0.3718 0.4223 引入词频先验但未能解决汇聚几何失配
解嵌入矩阵分解测度 (Unembedding) 0.3755 0.4202 语义先验对长程路由衰减缺乏灵敏度

关键发现

  • 长文本场景提升极其显著:在SAMSum与XSum摘要任务上,完整版SAUQ分别达到0.5924与0.4214的PRR,大幅超越RAUQ(0.5269/0.3008)和采样10次的高成本语义熵方法(0.2525/-0.1886)。这表明长上下文中输出Logits概率严重平滑,而上游TDO谱能够灵敏捕捉长程路由能力的丧失。
  • 与注意力汇聚彻底解耦:在头级行为诊断中,传统欧氏奇异谱和标量统计量(如注意力熵、前一Token得分)与汇聚得分(Sink Score)存在0.5至1.0的强相关性;而TDO谱参与率与汇聚相关性大幅降低,同时与实际输出谱熵 \(H(AV)\) 的Spearman相关性显著提升,证明其真正反映了有效表征维度。
  • 测度几何的必要性:消融实验清晰证明,简单使用均匀测度(欧氏空间)在摘要任务上比内在测度低0.0865个PRR点,证实了拉回概率测度对压制因果掩码累积偏差的决定性理论作用。

亮点与洞察

  • 算子理论破局注意力汇聚:将注意力机制抽象为泛函映射,用拉回概率测度 \(\mu = A^\top \pi\) 抹平因果注意力早期Token的列和不均,数学推导极其优雅且具启发性。
  • 均值与路由的严格正交解耦:证明了 \(AV = \mathbf{1}m + AV_c\) 的正交性与收缩界,澄清了学界对“注意力汇聚是表示崩溃还是特征平移”的长期争论,确立了中心化谱作为表征维度的理论度量。
  • 低成本上游诊断赋能置信度估计:无须特征值分解,仅通过矩阵迹与Frobenius范数便可在单遍前向中提取高阶几何信号,为大模型实时幻觉拦截与不确定性拒识提供了通用范式。

局限与展望

  • 单步与单头隔离分析的局限:理论分析聚焦于单层单头的注意力矩阵 \(A\),未将残差连接、跨头投影矩阵 \(W^O\) 以及后续MLP块的非线性动力学完整纳入全局演化方程中。
  • 多头协同效应未被显式建模:当前SAUQ采用启发式规则挑选单个活跃路由头,忽略了多头之间的联合子空间互补性与信息重叠。
  • 长序列矩阵乘法开销:尽管无需SVD,但在长文本推理中频繁计算 \(n_t \times n_t\) 注意力矩阵的迹与Frobenius范数,仍存在显存与计算延迟负担,未来需探索基于稀疏注意力或分块近似的快速实现。

相关工作与启发

  • vs 标量/行统计量方法 (RAUQ / Zhai et al.):以往方法依赖上一Token注意力强度或注意力熵等启发式指标,无法区分局部偏置与全局路由;本文从全局算子谱切入,不仅理论自洽,且在长序列摘要上较RAUQ提升近0.1 PRR。
  • vs 图拉普拉斯算子方法 (LapEigval / Noël):图拉普拉斯方法虽做了一定度归一化,但未利用注意力的概率转移特性构建输入-输出双重测度空间;本文TDO与输出谱熵相关性更强,且免受无向图对称化假设带来的失真。
  • vs 输出概率不确定性方法 (Semantic Entropy / LUQ):输出概率方法在表征挤压(Over-squashing)时无法反映错误;本文SAUQ直接监控上游路由容量的塌陷,且作为单遍方法免去了10次采样的巨大推理开销。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从算子理论与内在概率谱几何切入注意力机制,理论推导极具原创性与严谨性。
  • 实验充分度: ⭐⭐⭐⭐☆ 覆盖理论验证、代表头可视化、多模型跨架构验证及LM-Polygraph多基准对比,若补充更大规模模型(如70B)则更完备。
  • 写作质量: ⭐⭐⭐⭐⭐ 概念清晰、数学表述精准流畅,从算子泛函到工程落地的逻辑链条紧凑完整。
  • 价值: ⭐⭐⭐⭐⭐ 兼具深刻的理论解释力与极高实用价值,为注意力机制分析和LLM不确定性量化开辟了全新方向。