Quantifying Concentration Phenomena of Mean-Field Transformers in the Low-Temperature Regime¶
会议: NeurIPS 2026
arXiv: 2605.10931
OpenReview: https://openreview.net/forum?id=tLcRuEazgY
领域: 优化/理论
关键词: 平均场极限、自注意力动力学、Wasserstein距离、低温度机制、Laplace原理
一句话总结¶
本文通过连续时间平均场连续性方程严谨刻画了深层 Transformer 的低温度推理演化,给出了 token 分布在 \(W_2\) 距离下向 \(VB^\top\) 主特征子空间投影前推分布聚集的非渐近误差界,并揭示了由 \(V\) 主特征子空间主导的终态动力学。
研究背景与动机¶
基于自注意力机制的 Transformer 架构已成为现代大型语言模型与基础模型的核心底座。然而,深层网络内部 token 表示随着推理深度的推进到底如何演化,其内部机制在数学上仍面临巨大的分析挑战。为了探究深层 Transformer 的内部表征流动,近年来学界引入了粒子系统与偏微分方程(PDE)的动力学视角:将各 token 视作超球面上的粒子,在连续深度极限与平均场视角下,自注意力加 LayerNorm(或 RMSNorm)的推理前向传播可以被精准刻画为超球面 \(\mathbb{S}^{d-1}\) 上的连续性方程。
在注意力机制中,逆温度参数 \(\beta\) 控制着 Softmax 分布的锐度。已有理论定性证明了当 \(\beta \to \infty\)(即零温度极限)时,平均场连续性方程的解会收敛到确定性向量场驱动的动力系统,且在长时间演化下 token 会塌缩到投影矩阵 \(VB^\top\) 的主特征子空间上。然而,现有的平均场理论存在一个致命断层:它们大多局限于定性极限或者理想的零温度情况,缺乏在有限低温度(\(\beta^{-1} \to 0\) 且 \(\beta < \infty\))条件下的定量误差阶与时间收敛率分析;更无法解释有限温度与推理时间 \(t\) 之间的非平凡耦合关系,导致理论难以直接指导有限深度的真实模型。
为了解决这一理论痛点,本文借助统计物理多粒子系统分析工具(特别是基于共识的全局优化 CBO 中的定量 Laplace 原理与 Lyapunov 稳定性分析),在 Wasserstein 空间下建立了首个定量收敛与聚集估计。核心 idea:利用定量 Laplace 原理建立有限低温方程与零温方程在 \(W_2\) 空间下的稳定性耦合,并构造基于特征子空间距离比值的非负 Lyapunov 泛函证明零温稳态为初始分布的前推投影 \(\Pi_\sharp \rho_0\),由此在 \(t \lesssim \mathcal{O}(\log \beta)\) 早期亚稳态时间尺度内严格量化 token 分布的指数级聚集界。
方法详解¶
整体框架¶
本文考虑带有单头自注意力和 RMSNorm 简化的深层编码器 Transformer 在大 token 极限(\(n \to \infty\))下的平均场连续性方程。设 token 在超球面 \(\mathbb{S}^{d-1}\) 上的空间经验测度流为 \(\rho_t^\beta \in \mathcal{P}(\mathbb{S}^{d-1})\),其满足球面流形连续性方程: $$ \partial_t \rho_t^\beta = -\operatorname{div}\left(\rho_t^\beta \, \mathrm{P}x\left(\int\rho_t^\beta(y)\right)\right), $$ 其中 }^{d-1}}\frac{\exp(\beta\langle x, By\rangle)}{\int_{\mathbb{S}^{d-1}}\exp(\beta\langle x, Bz\rangle)\,\mathrm{d}\rho_t^\beta(z)}Vy\,\mathrm{d\(B = Q^\top K \in \mathbb{R}^{d \times d}\) 为 Query-Key 相互作用矩阵,\(V \in \mathbb{R}^{d \times d}\) 为 Value 矩阵,\(\mathrm{P}_x(y) = y - \langle x, y\rangle x\) 为球面切空间投影算子。
整个理论分析由两大支柱和三角不等式串联而成:第一阶段,运用球面上的定量 Laplace 原理,将有限低温方程 \(\rho_t^\beta\) 与零温度极限方程 \(\rho_t\) 的局部漂移场差异严格量化,得到 \(W_2(\rho_t^\beta, \rho_t)\) 随时间指数增长的稳定性上界;第二阶段,针对零温度连续性方程,识别其平稳分布为初始测度在 \(VB^\top\) 主特征子空间 \(E\) 上的前推投影测度 \(\Pi_\sharp \rho_0\),并通过构造专用的 Lyapunov 能量泛函导出以速率 \(\exp\left(-\frac{p\gamma}{\sigma_{\max}(B)}t\right)\) 衰减的指数收敛界;最终三角不等式合成给出两阶段动力学行为。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["初始分布 ρ₀ ∈ 𝒫(𝕊ᵈ⁻¹)"] --> B["基于定量 Laplace 原理的 Wasserstein 稳定性估计<br/>W₂(ρₜᵝ, ρₜ) ≤ 2√(log(β+1)/β) (e^{C₁t} - e^{C₀t})"]
A --> C["零温方程不变投影与 Lyapunov 泛函收敛分析<br/>W₂(ρₜ, Π_♯ ρ₀) ≤ √(2𝒱_p(ρ₀)) exp(-pγ/σ_max(B) · t)"]
B --> D["三角不等式耦合非渐近全界<br/>W₂(ρₜᵝ, Π_♯ ρ₀) ≤ 温度误差项 + 动力学收敛项"]
C --> D
D -->|时间 t ≲ 𝒪(log β)| E["早期亚稳态相<br/>Token 分布指数聚集于 VBᵀ 的主特征子空间 E"]
D -->|时间 t ≳ 𝒪(log β)| F["终态相变动力学<br/>Token 聚集漂移向 Value 矩阵 V 的主特征子空间 F"]
关键设计¶
1. 基于定量 Laplace 原理的 Wasserstein 稳定性估计:控制有限低温与零温流场偏差
在逆温度 \(\beta \to \infty\) 时,加权注意力向量场会收敛至点态最大化方向 \(y^*(x) = \frac{B^\top x}{\|B^\top x\|}\)。然而,为了得到有限 \(\beta\) 下的非渐近定量界,必须精确控制积分平均 \(m_{\beta, \varrho}(x)\) 与极值点 \(y^*(x)\) 之间的距离。本文将共识优化(CBO)中的定量 Laplace 原理推广到球面几何: $$ |m_{\beta, \varrho}(x) - y^(x)| \leq \sqrt{r^2 + \frac{2q}{\sigma_{\min}(B)}} + \frac{2e^{-\beta q}}{\varrho(\mathcal{B}_r(y^(x)))}, \quad \forall r, q > 0. $$ 结合初始分布在球面上具有下界 \(\ell_0 > 0\) 的非退化条件,证明了有限时间内测度密度下界保持 \(\min f_t^\beta(x) \geq \frac{1}{2}\ell_0 e^{-C_0 t}\)。进而通过最优输运理论中的稳定性估计,成功建立了两个解在 2-Wasserstein 距离下的增长控制: $$ W_2(\rho_t^\beta, \rho_t) \leq 2\sqrt{\frac{\log(\beta+1)}{\beta}}\left(e^{C_1 t} - e^{C_0 t}\right). $$
2. 零温方程平稳态的前推几何刻画:揭示初始分布的投影记忆
以往研究仅定性说明零温方程解会塌缩到 \(VB^\top\) 的主特征子空间 \(E\),但从未识别出具体的极限定量分布。本文在 \(VB^\top\) 对称性假设(特征值差 \(\gamma = \mu_1 - \mu_2 > 0\))下,引入非线性径向投影映射 \(\Pi(x) = \frac{\mathrm{P}_E x}{\|\mathrm{P}_E x\|}\)。通过微分几何计算,本文证明了沿零温流场的速度场与 \(\Pi\) 的导数处处正交,即: $$ D\Pi(x) \mathrm{P}_x\left(\frac{VB^\top x}{|B^\top x|}\right) = 0, \quad \forall x \in \mathbb{S}^{d-1} \setminus E^\perp. $$ 该微分不变量直接诱导了测度守恒律:对所有时间 \(t \geq 0\),均严格有 \(\Pi_\sharp \rho_t = \Pi_\sharp \rho_0\)。这一结论从根本上证实了极限稳态并不是单一的狄拉克点,而是完全保留了输入初始数据在主子空间几何轮廓的前推测度 \(\Pi_\sharp \rho_0\)。
3. 奇异能量泛函的 Lyapunov 耗散分析:推导指数级聚集速率
为精确度量测度向特征子空间 \(E\) 的聚集速度,本文构建了定义在非正交余空间能量比值上的 Lyapunov 泛函: $$ \mathcal{V}p(\varrho) \coloneqq \int\varrho(x), \quad p \in (0, 1]. $$ 通过精细的截断函数正则化 }^{d-1}} \frac{|(\mathrm{Id} - \mathrm{P}_E)x|^{2p}}{|\mathrm{P}_E x|^{2p}} \,\mathrm{d\(R_{p, K}(x) = \min\{R_p(x), K\}\) 与 Rademacher 几乎处处可微分析,结合 \(E\) 与 \(E^\perp\) 上的谱间隙 \(\gamma\),导出了该泛函沿零温动力学轨迹的严格微分耗散不等式: $$ \frac{\mathrm{d}}{\mathrm{d}t} \mathcal{V}p(\rho_t) \leq -\frac{2p\gamma}{\sigma}(B)} \mathcal{V}_p(\rho_t) \implies \mathcal{V}_p(\rho_t) \leq \mathcal{Vp(\rho_0) \exp\left(-\frac{2p\gamma}{\sigma t\right). $$ 再结合几何不等式 }(B)\(\|x - \Pi(x)\|^2 \leq 2 R_p(x)\),将测度在 Wasserstein 距离下的收敛紧致耦合于能量衰减,得到 \(W_2(\rho_t, \Pi_\sharp \rho_0) \leq \sqrt{2\mathcal{V}_p(\rho_0)} \exp\left(-\frac{p\gamma}{\sigma_{\max}(B)} t\right)\)。
损失函数 / 训练策略¶
本文研究深层 Transformer 在测试推理阶段(Inference Time)的表示演化,模型权重 \(Q, K, V\) 在推理期间冻结保持不变。为了模拟连续时间平均场连续性方程,数值求解时采用显式 Euler 离散化方案(时间步长 \(\Delta t = 0.01\)),利用共识优化库 CBX 进行粒子系统模拟: $$ w_{ij} \leftarrow \frac{\exp(\beta \langle x_i, B x_j\rangle)}{\sum_{k=1}^n \exp(\beta \langle x_i, B x_k\rangle)}, \quad m_i \leftarrow \sum_{j=1}^n w_{ij} x_j, \quad x_i \leftarrow \frac{x_i + \Delta t \, \mathrm{P}{x_i} V m_i}{|x_i + \Delta t \, \mathrm{P}. $$} V m_i|
实验关键数据¶
主实验¶
为了验证理论推导的早期聚集相与定理 1 的界限,作者在合成高斯/均匀分布及 von Mises-Fisher 分布初始化的粒子系(\(n=200 \sim 2000\) 个 token)上进行了系统仿真,观测不同温度 \(\beta\) 下粒子到特征子空间 \(E\)(\(VB^\top\) 的主特征空间)以及 \(F\)(\(V\) 的主特征空间)的对齐度(Alignment,满分为 1.0)与 \(W_2\) 距离演化。
| 实验配置 (\(\beta\) 设定) | 理论有效时间窗口 \(t \lesssim \mathcal{O}(\log \beta)\) | 主子空间 \(E\) 最大对齐度 | 是否出现向 \(F\) 的终态相变 | \(W_2(\rho_t^\beta, \Pi_\sharp \rho_0)\) 演化趋势 |
|---|---|---|---|---|
| 低逆温 \(\beta = 10\) | \(t \lesssim 2.30\) | ~0.72 | 是 (在 \(t \approx 5\) 时迅速主导) | 先降至最低点后因正温漂移迅速回升 |
| 中逆温 \(\beta = 30\) | \(t \lesssim 3.40\) | ~0.89 | 是 (在 \(t \approx 8\) 时开始漂移) | 呈现典型 U 型曲线,验证误差项竞争 |
| 高逆温 \(\beta = 10^3\) | \(t \lesssim 6.91\) | > 0.98 | 是 (相变显著推迟至 \(t > 15\)) | 长时间保持极低 \(W_2\) 距离,吻合亚稳态预测 |
| 零温度 \(\beta = \infty\) | \(t \in [0, \infty)\) 全程有效 | 1.00 (渐近完全对齐) | 否 (终态不发生相变,永久停留在 \(E\)) | 单调指数衰减至 0,完全契合推论 2 |
消融实验¶
针对定理核心假设条件(对称性假设 \(VB^\top\)、变分梯度流假设 \(B = \pm V\) 及非对称扰动),作者在附录进行了系统性矩阵结构消融分析。
| 消融配置 | 矩阵设定 (\(B, V\)) | 特征值性质 | 早期聚集行为 (\(E\)) | 终态演化行为 |
|---|---|---|---|---|
| 对称基准 (\(VB^\top\) 对称) | \(V=\mathrm{diag}(-1,1,-2), B=\mathrm{diag}(-1,-1,1)\) | 谱间隙 \(\gamma > 0\), 全实数 | 迅速指数聚集在 \(e_1\) 方向 | 稳定相变迁移至 \(V\) 主特征向 \(e_2\) |
| 非对称但实特征值 | \(V_2 = \mathrm{diag}(-1,1,-2), B_2 \text{ 非对称}\) | 特征值全为实数 (\(1\pm\sqrt{2}, -2\)) | 依旧保持两阶段聚集特性 | 同样在后期迁移至 \(V\) 的主特征子空间 |
| 非对称且复特征值 | \(V_1 \text{ 伴随旋转块}, B_1 = \mathrm{diag}(-1,-1,1)\) | 出现共轭虚特征值 \(\pm i\) | 粒子轨迹在球面上剧烈旋转 | 无法收敛,完全破坏平稳聚类结构 |
| 梯度流极大化 (\(B = V\)) | \(B = V\) (正定矩阵) | 变分能量极大化 \(\mathcal{E}_B\) | 大 \(\beta\) 下偏向单一或对称 Dirac | 能量达到理论全局极值,吻合定理 1 |
| 梯度流极小化 (\(B = -V\)) | \(B = -V\) (负定矩阵) | 变分能量极小化 \(\mathcal{E}_B\) | 中间 \(\beta\) 下先对齐 \(E\) 后对齐 \(F\) | 揭示第二阶段相变本质源于 \(B\) 与 \(V\) 谱冲突 |
关键发现¶
- 双时间尺度竞争机制:理论证明的界 \(\sqrt{\frac{\log(\beta+1)}{\beta}} e^{C_1 t} + \exp(-ct)\) 极其精准地解释了数值模拟中的 U 型曲线——在 \(t \lesssim \log \beta\) 前,负指数衰减项主导,距离快速缩小;而在 \(t \gtrsim \log \beta\) 后,正指数膨胀项暴涨,正温误差接管动力学。
- 终态相变的发现:长期演化并不像过去文献假设的那样永久停留在 \(VB^\top\) 的主子空间;有限逆温度 \(\beta\) 下,长时间动力学脱离 \(VB^\top\) 约束,不可逆地汇聚向 Value 矩阵 \(V\)(或 \(|V|\))的主特征子空间 \(F\)。
- 谱对称性是稳态驻留的关键:一旦 \(VB^\top\) 出现复数谱,系统丧失 Lyapunov 单调势能,粒子呈现周期性球面旋转,无法形成稳定的语义聚类。
亮点与洞察¶
- 定量 Laplace 原理的跨学科迁移:将非凸全局优化领域的 CBO 工具创新性地引入球面流形上的 Transformer 连续性方程,打通了统计物理、偏微分方程与深度学习动力学分析的理论鸿沟。
- 打破「表征完全塌缩」的粗糙假设:明确给出了零温极限平稳分布 \(\Pi_\sharp \rho_0\) 的解析结构,证明了初始 token 的空间几何信息被精确保留在前推测度中,而不是平凡地凝聚成无信息量的单一狄拉克脉冲。
- 解释实际 Transformer 深度与温度设计的理论参考:指出有效表征聚类必须发生在 \(t \sim \mathcal{O}(\log \beta)\) 早期窗口内,揭示了过深的网络若没有非线性层(MLP)介入可能会跌入 \(V\) 矩阵主导的平凡吸引子陷阱。
局限与展望¶
- 依赖 \(VB^\top\) 对称性假设:为了构造精确解析的不变测度投影 \(\Pi\),理论要求 \(VB^\top\) 为对称矩阵,这在非对称一般训练权重下需要进一步放宽。
- 未覆盖多头注意力和 MLP 层:分析基于单头自注意力,多头自注意力的凸组合平均以及多层感知机(FFN/MLP)的非线性激活尚未纳入平均场流场中。
- 终态相变仍为猜想性结论:对于 \(t \gg \log \beta\) 阶段聚集到 \(V\) 主特征子空间 \(F\) 的现象,文中仅给出了启发式猜想与仿真证据,尚缺乏完整的严谨非渐近收敛证明。
相关工作与启发¶
- vs Bruno et al. (2025a):Bruno 等人首次建立了零温连续性方程并指出了三阶段演化定性概念,但其分析是渐近非定量的(无 \(\beta\) 依赖阶与收敛率)。本文不仅给出了首个定量收敛上界,还纠正了其未识别稳态分布 \(\Pi_\sharp \rho_0\) 且忽略有限温终态相变的不足。
- vs Burger et al. (2025):Burger 等人仅在 \(B = \pm V\) 的对称变分能量梯度流下分析稳态。本文放宽了 \(B = \pm V\) 限制,适用于更一般的 \(B \neq V\) 动力系统。
- vs Altafini (2026):Altafini 从 Oja 流视角讨论了自注意力的多稳态性。本文实验证实的 \(t > \log \beta\) 终态聚集与 Altafini 的谱主导结论形成了深刻的呼应。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立平均场 Transformer 在有限低温下的定量收敛与 Wasserstein 稳定性理论。
- 实验充分度: ⭐⭐⭐⭐☆ 理论证明极其严密完整,数值仿真充分验证了亚稳态相与终态相变,附录探讨详实。
- 写作质量: ⭐⭐⭐⭐⭐ 数学推导严密、逻辑链条清晰,PDE 与几何测度论工具运用炉火纯青。
- 价值: ⭐⭐⭐⭐⭐ 为理解深层自注意力机制的内部表示演化与几何浓度现象提供了里程碑式的数学基石。