Transformer动力学与平均场理论

概述

Transformer架构已经成为现代深度学习的基石，但其理论理解仍然相对有限。近年来，平均场理论（Mean-Field Theory）为理解Transformer的训练动态和表示学习提供了强有力的数学框架。¹

本文系统梳理Transformer动力学与平均场理论的核心内容，涵盖：

• 平均场近似的理论基础
• McKean-Vlasov动态建模
• Token作为粒子的物理图像
• 收敛性与稳定性分析

---

平均场理论基础

从粒子系统到神经网络

平均场理论起源于统计物理，用于描述大量相互作用粒子的集体行为。在深度学习语境下，我们可以将Transformer中的Token视为”粒子”，将注意力机制视为粒子间的”相互作用”。

设 $\mathcal{N}$ 为Token集合，每个Token $i$ 的状态由其表示向量 ${\mathbf{x}}_i\in {\mathbb{R}}^d$ 刻画。系统能量定义为：

$$H(\mathbf{X})=-\sum _{i,j}{\mathbf{x}}_i^{\top }{\mathbf{W}}_{ij}{\mathbf{x}}_j+\sum _{i}B({\mathbf{x}}_i)$$

其中 ${\mathbf{W}}_{ij}$ 是注意力权重矩阵，$B(\cdot )$ 是偏置项。

Kullback-Leibler平均场

平均场近似的核心思想是用”平均相互作用场”替代粒子间的复杂耦合。对于概率分布 $p(\mathbf{X})$，KL平均场变分推断给出：

$$q(\mathbf{X})=\prod _i{q}_i({\mathbf{x}}_i)$$

其中每个边际分布 $q_i$ 通过最小化 $\text{KL}(q\Vert p)$ 获得。平均场方程为：

$$\log q_i(\mathbf{x})={\mathbb{E}}_{q_{-i}}[\log p(\mathbf{X})]+\text{const}$$

自由能框架

定义平均场自由能：

$$\mathcal{F}[q]={\mathbb{E}}_q[\log q(\mathbf{X})]-{\mathbb{E}}_q[\log p(\mathbf{X})]$$

最小化自由能等价于最小化 $\text{KL}(q\Vert p)$。展开得：

$$\mathcal{F}[q]=-\sum _i{\mathbb{E}}_{q_i}[\log {\tilde{p}}_i({\mathbf{x}}_i)]+\text{const}$$

其中 ${\tilde{p}}_i$ 是有效势，包含来自其他粒子的平均相互作用。

---

Transformer的McKean-Vlasov动态

McKean-Vlasov方程

McKean-Vlasov方程描述了宏观粒子的概率分布随时间的演化。对于Transformer，设 $m_t$ 为时刻 $t$ 的Token分布，动态方程为：

$$\frac{\partial m_t}{\partial t}+\nabla \cdot (m_{t}v[m_t])=\frac{{\sigma }^2}{2}\Delta m_t$$

其中：

• $v[m_t]$ 是平均场速度，依赖于当前分布 $m_t$
• ${\sigma }^2$ 是噪声方差
• $\Delta$ 是拉普拉斯算子

注意力作为相互作用势

Transformer的注意力机制可以解释为一种软相互作用势。设 $Q,K,V$ 为查询、键、值投影，则注意力权重为：

$${\alpha }_{ij}=\frac{\exp (q_i^{\top }{k}_j/\sqrt{d})}{{\sum }_j\exp (q_i^{\top }{k}_j/\sqrt{d})}$$

这等价于 Boltzmann 分布：

$${\alpha }_{ij}\propto \exp \left(-\frac{E_{ij}}{\sqrt{d}}\right),E_{ij}=-q_i^{\top }{k}_j$$

因此，注意力机制在统计物理中对应平均场Ising模型或Potts模型的软版本。

连续动态系统视角

将Transformer层视为连续动力系统的离散化。设 ${\mathbf{h}}_i^{(l)}$ 为第 $l$ 层第 $i$ 个Token的表示，层间动态为：

$${\mathbf{h}}_i^{(l+1)}={\mathbf{h}}_i^{(l)}+f_{\theta }\left({\mathbf{h}}_i^{(l)},\sum _j{\alpha }_{ij}^{(l)}{\mathbf{h}}_j^{(l)}\right)$$

取连续极限 $l\to t$，定义 ODE：

$$\frac{d{\mathbf{h}}_i}{dt}=f_{\theta }\left({\mathbf{h}}_i,\int \alpha (s,i)\mathbf{h}(s)d\mu (s)\right)$$

其中 $\mu$ 是Token分布。

---

Token作为粒子系统

粒子的物理图像

将每个Token $i$ 视为 ${\mathbb{R}}^d$ 空间中的一个粒子，其位置为 ${\mathbf{x}}_i$，速度为 ${\dot{\mathbf{x}}}_i$。注意力机制定义了粒子间的成对力：

$${\mathbf{F}}_{ij}={\nabla }_{{\mathbf{x}}_i}\cdot \text{Attention}({\mathbf{x}}_i,{\mathbf{x}}_j)$$

典型的注意力核对应于势能函数：

$$\varphi ({\mathbf{x}}_i,{\mathbf{x}}_j)=-\frac{1}{\sqrt{d}}\cdot \frac{{\mathbf{x}}_i^{\top }{\mathbf{W}}_Q^{\top }{\mathbf{W}}_K{\mathbf{x}}_j}{\Vert {\mathbf{x}}_i\Vert \Vert {\mathbf{x}}_j\Vert }$$

吸引子与聚类

平均场动态会导致Token聚集形成吸引子。设 $m(\mathbf{x},t)$ 为时刻 $t$ 在位置 $\mathbf{x}$ 的粒子密度，聚类对应密度峰的演化。

不动点条件 $\frac{\partial m}{\partial t}=0$ 给出平衡方程：

$$m^{\ast }(\mathbf{x})=Z^{-1}\exp \left(-\int V(\mathbf{x},\mathbf{y})m^{\ast }(\mathbf{y})d\mathbf{y}\right)$$

其中 $V$ 是平均场势能。

Kuramoto振子模型

注意力机制与 Kuramoto 振子模型有深刻的类比。Kuramoto模型的同步动态为：

$${\dot{\theta }}_i={\omega }_i+\frac{K}{N}\sum _j\sin ({\theta }_j-{\theta }_i)$$

Transformer中的相位同步可以解释为”语义同步”——具有相似语义的Token趋于收敛到相似的表示。

---

收敛性分析

Wasserstein梯度流

平均场动态可以表述为Wasserstein空间中的梯度流。设 ${\mathcal{W}}_2$ 为Wasserstein-2距离，损失泛函 $\mathcal{L}[m]$，则：

$$\frac{\partial m_t}{\partial t}=-\nabla \mathcal{L}[m_t]\text{在 }{\mathcal{W}}_2\text{ 意义下}$$

对于Transformer，经验损失 $\mathcal{L}(\theta )$ 关于参数 $\theta$ 的梯度流对应Token分布的演化。

传输映射稳定性

最优传输理论提供了分析收敛性的新视角。Sinkhorn算法求解的熵正则化最优传输为：

$${\gamma }^{\ast }=\arg \underset{\gamma }{\min }⟨C,\gamma ⟩+ϵH(\gamma )\text{s.t.}\gamma 1=\mu ,{\gamma }^{\top }1=\nu$$

这与注意力计算的对应关系：

• $\mu ,\nu$ 对应源/目标Token分布
• $C_{ij}$ 对应 $q_i^{\top }{k}_j$
• $ϵ$ 对应温度参数

Lyapunov稳定性

定义Lyapunov函数：

$$\mathcal{V}(m)=\mathcal{F}[m]-\underset{m^{\prime }}{\min }\mathcal{F}[m^{\prime }]$$

其中 $\mathcal{F}$ 是平均场自由能。若 $\mathcal{V}$ 满足：

1. $\mathcal{V}(m)\ge 0$，且 $\mathcal{V}(m)=0⟺m=m^{\ast }$
2. $\frac{d}{dt}\mathcal{V}(m_t)\le 0$

则系统收敛到平衡态 $m^{\ast }$。

---

深度Transformer的理论挑战

深度与表达能力的权衡

深层Transformer面临的核心问题：过度平滑（Over-smoothing）和过度压缩（Over-squashing）。

过度平滑：随着层数增加，Token表示趋于收敛到相同的向量。

设 ${\mathbf{H}}^{(l)}=[{\mathbf{h}}_1^{(l)},\dots ,{\mathbf{h}}_N^{(l)}]$，有：

$${\mathbf{H}}^{(l+1)}={\mathbf{A}}^{(l)}{\mathbf{H}}^{(l)}+\text{残差}$$

若注意力矩阵 $\mathbf{A}$ 的幂收玫到秩1矩阵，则表示趋于相同。

过度压缩：信息在长距离传播中衰减。

可使用消息传递的图论分析：

$$I_{ij}^{(l)}\approx {\beta }^l\cdot \cos ({\theta }_{ij})$$

其中 $\beta <1$ 是衰减因子，${\theta }_{ij}$ 是表示夹角。

注意力模式的相变

Transformer中的注意力模式会经历相变：

相	特征	物理类比
有序相	少数主导Token	铁磁有序
临界相	幂律衰减注意力	二级相变
无序相	均匀注意力	顺磁无序

相变点由温度参数 $T=\sqrt{d}$ 控制。临界温度 $T_c$ 对应初始化尺度 $1/\sqrt{d_{\text{model}}}$。

---

数值模拟与实验

平均场初始化

基于平均场理论，推荐的初始化策略：

import torch
import torch.nn as nn
import math
 
class MeanFieldInitializedTransformer(nn.Module):
    """
    基于平均场理论的Transformer层初始化
    
    理论依据：Geshkovski et al., "The mathematics of artificial intelligence"
    """
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.d_model = d_model
        
        # 查询/键/值投影：Xavier初始化修正
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
        # FFN层：平均场建议的初始化
        self.W_1 = nn.Linear(d_model, d_ff)
        self.W_2 = nn.Linear(d_ff, d_model)
        
        # 基于平均场的缩放因子
        self._mean_field_init()
    
    def _mean_field_init(self):
        """
        平均场初始化：根据谱分析和相变理论设置方差
        """
        d = self.d_model
        
        # 查询/键投影：保持注意力的临界性
        for W in [self.W_q, self.W_k]:
            nn.init.normal_(W.weight, mean=0, std=1.0 / math.sqrt(d))
            nn.init.zeros_(W.bias)
        
        # 值投影：单位方差
        nn.init.normal_(self.W_v.weight, mean=0, std=1.0 / math.sqrt(d))
        nn.init.zeros_(self.W_v.bias)
        
        # 输出投影：保持残差贡献
        scale = 1.0 / math.sqrt(2 * d)
        nn.init.normal_(self.W_o.weight, mean=0, std=scale)
        nn.init.zeros_(self.W_o.bias)
        
        # FFN：Grover初始化
        nn.init.normal_(self.W_1.weight, mean=0, std=2.0 / math.sqrt(d))
        nn.init.zeros_(self.W_1.bias)
        nn.init.normal_(self.W_2.weight, mean=0, std=1.0 / math.sqrt(d_ff))
        nn.init.zeros_(self.W_2.bias)

粒子动态模拟

模拟Token粒子的平均场动态：

import torch
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
 
class TokenParticleSimulation:
    """
    Token粒子系统的平均场动态模拟
    
    将每个Token视为d维空间中的粒子，跟踪其分布演化
    """
    def __init__(self, n_tokens, d_dim, device='cuda'):
        self.n = n_tokens
        self.d = d_dim
        self.device = device
        
        # 粒子位置：初始为球面均匀分布
        theta = torch.rand(n_tokens) * 2 * np.pi
        phi = torch.acos(2 * torch.rand(n_tokens) - 1)
        
        self.positions = torch.zeros(n_tokens, d_dim)
        self.positions[:, 0] = torch.sin(phi) * torch.cos(theta)
        self.positions[:, 1] = torch.sin(phi) * torch.sin(theta)
        self.positions[:, 2:] = torch.randn(n_tokens, d_dim - 2)
        
        # 归一化到单位球面
        self.positions = torch.nn.functional.normalize(self.positions, dim=1)
        self.positions = self.positions.to(device)
        
        # 历史记录
        self.history = [self.positions.cpu().clone()]
    
    def compute_attention(self, temperature=1.0):
        """
        计算注意力矩阵（作为粒子间软相互作用）
        """
        # 注意力分数
        scores = torch.mm(self.positions, self.positions.T) / temperature
        
        # 数值稳定softmax
        scores = scores - scores.max(dim=1, keepdim=True)[0]
        weights = torch.softmax(scores, dim=1)
        
        return weights
    
    def step(self, lr=0.01, temperature=1.0, noise_std=0.01):
        """
        一步平均场动态更新
        """
        # 计算注意力（平均场）
        A = self.compute_attention(temperature)
        
        # 平均场力：驱动Token向邻居聚集
        mean_field = torch.mm(A, self.positions)
        
        # 速度更新（带阻尼）
        velocity = mean_field - 0.1 * self.positions
        
        # 添加随机噪声（热波动）
        noise = torch.randn_like(self.positions) * noise_std
        velocity = velocity + noise
        
        # 位置更新
        self.positions = self.positions + lr * velocity
        
        # 归一化到单位球面（保持语义方向）
        self.positions = torch.nn.functional.normalize(self.positions, dim=1)
        
        # 记录历史
        self.history.append(self.positions.cpu().clone())
    
    def simulate(self, n_steps=100, lr=0.01, temperature=1.0):
        """
        运行完整模拟
        """
        for _ in range(n_steps):
            self.step(lr=lr, temperature=temperature)
    
    def compute_pairwise_distances(self):
        """
        计算成对距离（用于分析聚类）
        """
        with torch.no_grad():
            # 余弦距离
            cos_sim = torch.mm(self.positions, self.positions.T)
            return 1 - cos_sim  # 距离 = 1 - 相似度
    
    def compute_clustering_metrics(self):
        """
        计算聚类指标（Silhouette Score近似）
        """
        distances = self.compute_pairwise_distances()
        
        # 计算平均类内/类间距离
        # 简化：使用与自身的距离作为类内，与其他Token的距离作为类间
        intra_dist = distances.diagonal().mean()
        inter_dist = (distances.sum() - distances.diagonal().sum()) / (self.n * (self.n - 1))
        
        return {
            'mean_intra_distance': intra_dist.item(),
            'mean_inter_distance': inter_dist.item(),
            'clustering_coefficient': inter_dist.item() / (intra_dist.item() + 1e-8)
        }
    
    def visualize_3d(self, step_interval=10):
        """
        可视化3D轨迹
        """
        history = torch.stack(self.history[::step_interval])
        n_frames = history.shape[0]
        
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        def update(frame):
            ax.clear()
            positions = history[frame].numpy()
            ax.scatter(positions[:, 0], positions[:, 1], positions[:, 2], 
                      c=range(len(positions)), cmap='viridis', s=50)
            ax.set_xlim(-1.5, 1.5)
            ax.set_ylim(-1.5, 1.5)
            ax.set_zlim(-1.5, 1.5)
            ax.set_title(f'Step {frame * step_interval}')
            ax.set_xlabel('x')
            ax.set_ylabel('y')
            ax.set_zlabel('z')
            return []
        
        anim = FuncAnimation(fig, update, frames=n_frames, interval=100, blit=True)
        return fig, anim
 
 
def run_particle_simulation():
    """
    运行完整的Token粒子模拟实验
    """
    # 参数设置
    n_tokens = 50
    d_dim = 16
    n_steps = 200
    
    # 初始化模拟器
    sim = TokenParticleSimulation(n_tokens, d_dim)
    
    # 运行不同温度的模拟
    results = {}
    
    for temp in [0.5, 1.0, 2.0]:
        sim_temp = TokenParticleSimulation(n_tokens, d_dim)
        sim_temp.simulate(n_steps, lr=0.05, temperature=temp)
        
        # 计算聚类指标
        metrics = sim_temp.compute_clustering_metrics()
        metrics['final_positions'] = sim_temp.positions.cpu()
        
        results[f'temp_{temp}'] = metrics
    
    return results

---

与Transformer训练动态的联系

连续时间梯度流

Transformer的训练可以被视为参数空间中的梯度流：

$$\frac{d\theta }{dt}=-{\nabla }_{\theta }\mathcal{L}(\theta )$$

在无限宽极限（平均场极限）下，网络输出关于参数的导数可以用Neural Tangent Kernel (NTK)近似：

$$f(\mathbf{x};\theta )=f_0(\mathbf{x})+ϵ\sum _{i}K(\mathbf{x},{\mathbf{x}}_i){\nabla }_{f({\mathbf{x}}_i)}\mathcal{L}$$

经验NTK与Transformer

对于Transformer，经验NTK定义为：

$$K_{\text{emp}}^{(l)}(i,j)=\frac{\partial {\mathbf{h}}_i^{(l)}}{\partial \theta }{\left(\frac{\partial {\mathbf{h}}_j^{(l)}}{\partial \theta }\right)}^{\top }$$

随层演化的NTK满足Dyson方程：

$$K^{(l+1)}=K^{(l)}+H^{(l)}\cdot K^{(l)}\cdot (H^{(l)}{)}^{\top }+\cdots$$

初始化的临界性

平均场理论解释了为什么 $1/\sqrt{d_{\text{model}}}$ 是Transformer的标准初始化尺度：

• 太小：梯度消失，训练缓慢
• 太大：熵主导，表示失去区分性
• 临界：系统在有序/无序相变边缘，最优可塑性

---

数学附录

定理：平均场收敛性

定理（平均场收敛）：设 $\{{\mathbf{X}}_i^N{\}}_{i=1}^N$ 为 $N$ 个相互作用粒子，其经验测度为 ${\mu }_N=\frac{1}{N}{\sum }_i{\delta }_{{\mathbf{X}}_i^N}$。若初始经验测度 ${\mu }_N(0)\Rightarrow {\mu }_0$ 弱收敛，则当 $N\to \infty$ 时，${\mu }_N(t)\Rightarrow m(t)$，其中 $m(t)$ 满足McKean-Vlasov方程。

证明概要：

1. 构造粒子系统的martingale表示
2. 应用McKean耦合方法证明相对熵有界
3. 使用 Gronwall不等式完成收敛性证明

公式速查

概念	公式	说明
注意力权重	${\alpha }_{ij}=\frac{e^{q_i^{\top }{k}_j/T}}{{\sum }_j{e}^{q_i^{\top }{k}_j/T}}$	Boltzmann分布
平均场能量	$E_{\text{MF}}(\mathbf{x})=\int V(\mathbf{x},\mathbf{y})m(\mathbf{y})d\mathbf{y}$	积分势能
自由能	$\mathcal{F}[m]=\int m\log m-mV\ast m$	变分自由能
Wasserstein梯度流	${\partial }_{t}m=\nabla \cdot (m\nabla \delta \mathcal{F}[m])$	传输方程

---

参考文献

---

相关主题

• Nexusformer非线性注意力
• 注意力与最优传输理论
• Transformer粒子系统模型
• Post-LayerNorm深度模型
• 深度诅咒理论

Footnotes

1. Geshkovski, B., Zuazua, E., &蔗. (2025). The mathematics of artificial intelligence. AMS Bulletin. arXiv:2505.XXXXX ↩