Alternating Gradient Flows（交替梯度流）

1. 引言

深度神经网络的一个核心问题是：网络如何学习有意义的特征表示？传统观点认为，梯度下降会同时优化所有参数，但最近的研究揭示了一种更精细的动力学模式——阶梯状学习（staircase learning）。

交替梯度流（Alternating Gradient Flows, AGF）理论提出了一种描述两层神经网络特征学习的数学框架，揭示了网络如何按优先级顺序学习傅里叶特征。¹

2. 问题设置

考虑一个两层神经网络：

$$f(x;\Theta )=\sum _{i=1}^H{a}_i\cdot \sigma (w_i^{\top }{g}_i(x))$$

其中：

• $x\in {\mathbb{R}}^d$：输入样本
• $w_i\in {\mathbb{R}}^{d\cdot m}$：第 $i$ 个隐藏神经元的输入权重
• $a_i\in {\mathbb{R}}^c$：第 $i$ 个隐藏神经元的输出权重
• $\sigma$：激活函数（满足 $\sigma (0)=0$，即原点传递性）
• $g_i:{\mathbb{R}}^d\to {\mathbb{R}}^m$：可选的输入变换

2.1 梯度流动力学

从随机初始化 ${\theta }_i(0)\sim \mathcal{N}(0,{\alpha }^{2}I)$ 开始，梯度流满足：

$$\frac{d\theta }{dt}=-{\nabla }_{\theta }\mathcal{L}(\Theta ),\Theta =\{w_i,a_i{\}}_{i=1}^H$$

其中 $\mathcal{L}$ 是经验风险。

3. Utility函数与Dormant神经元

3.1 Utility的定义

对于第 $i$ 个神经元，定义 Utility函数 $U_i$ 衡量其对减少当前损失的潜在贡献：

$$U_i(\theta ;r)={\mathbb{E}}_x\left[⟨f_i(x;\theta ),r(x)⟩\right]$$

其中 $r(x)=y(x)-f(x;{\Theta }_A)$ 是残差，由当前活跃神经元集合 ${\Theta }_A$ 决定。

3.2 Dormant神经元动力学

当神经元 $i$ 处于休眠状态（${\theta }_i$ 远离激活区域）时，其动力学由下式描述：

$$\frac{d{\bar{\theta }}_i}{dt}=\eta \Vert {\theta }_i{\Vert }^{\kappa -2}{P}_{{\bar{\theta }}_i}^{\perp }{\nabla }_{\theta }{U}_i({\bar{\theta }}_i;r)$$ $$\frac{d\Vert {\theta }_i\Vert }{dt}=\eta \kappa \Vert {\theta }_i{\Vert }^{\kappa -1}{U}_i({\bar{\theta }}_i;r)$$

其中：

• ${\bar{\theta }}_i={\theta }_i/\Vert {\theta }_i\Vert$：归一化方向
• $P_{{\bar{\theta }}_i}^{\perp }$：正交于 ${\bar{\theta }}_i$ 的投影
• $\kappa \in \{1,2\}$：参数化系数

3.3 Jump时间

定义累积Utility $S_i(t)$ 和跳跃时间 ${\tau }_i$：

$$S_i(t)={\int }_0^t\kappa U_i({\bar{\theta }}_i(s);r)ds$$ $${\tau }_i=\inf \{t>0:S_i(t)>c_i/\eta \}$$

当神经元 $i$ 在 ${\tau }_i$ 时刻跳跃时，其范数急剧增长，从休眠状态转变为活跃状态。¹

4. 阶梯状学习动力学

AGF揭示了梯度流的阶梯状行为，包含两种交替阶段：

4.1 Plateau阶段

• 活跃神经元 ${\Theta }_A$ 缓慢调整方向
• 休眠神经元 ${\Theta }_D$ 逐渐积累Utility
• 损失函数近似线性衰减

4.2 Sharp Drop阶段

• 某个休眠神经元 $i^{\ast }$ 跳跃（$\Vert {\theta }_{i^{\ast }}\Vert$ 急剧增长）
• 新特征被激活并开始影响预测
• 损失函数出现快速下降

$${\tau }_i=\inf \{t>0:S_i(t)>c_i/\eta \}$$

这种Utility最大化→Cost最小化的交替过程，精确描述了特征学习的时间演化。

5. 傅里叶特征涌现

5.1 二次网络与模块加法

考虑一个简化的二次网络用于模块加法任务：

$$f(x)=\sum _{p=1}^P{x}_{a_p}\cdot x_{b_p}\cdot x_{c_p}$$

输入 $x\in \{0,1{\}}^n$，目标是学习表示 $x_a\cdot x_b\cdot x_c$ 的神经元。

5.2 主要结果

Theorem 5.1（傅里叶特征涌现）：对于二次网络，utility最大化的神经元方向与输入的主导谐波对齐：

$$u^{\ast }[a]=A_{\xi }\cos ({\omega }_{\xi }a+s_u),v^{\ast }[b]=A_{\xi }\cos ({\omega }_{\xi }b+s_v),w^{\ast }[c]=A_{\xi }\cos ({\omega }_{\xi }c+s_w)$$

其中 $A_{\xi }$ 是傅里叶系数，${\omega }_{\xi }=2\pi \xi /n$ 是频率，$s_u,s_v,s_w$ 是相位偏移。

Theorem 5.2（损失下界）：损失满足：

$$\mathcal{L}({\Theta }_A)\ge \frac{\Vert x{\Vert }^2}{2}-\frac{⟨x,1{⟩}^2}{2p}-\frac{|\hat{x}[\xi ]{|}^2}{p}$$

这个下界解释了为什么某些傅里叶特征比其他特征更难学习——它们对应的损失下界更高。

5.3 频率优先级

网络学习傅里叶特征的优先级顺序由傅里叶系数幅值 $|\hat{x}[\xi ]|$ 决定：

$${\xi }^{\ast }=\arg \underset{\xi }{\max }|\hat{x}[\xi ]{|}^3$$

即网络首先学习系数幅值最大的频率，这与频率原则（Frequency Principle）一致。²

6. 理论保证

6.1 AGF = Gradient Flow（对角网络）

Theorem 3.1：在vanishing initialization极限（$\alpha \to 0$）下，AGF收敛于对角线性网络的梯度流。

这建立了AGF与经典梯度流理论之间的联系。

6.2 收敛性分析

对于平方损失，AGF保证：

• 损失单调递减
• 每个神经元最多跳跃有限次
• 总跳跃次数 $\le H$（神经元数量）

7. 与其他理论的关系

7.1 与频率原则的联系

AGF提供了频率原则的机制性解释：网络按傅里叶系数幅值的优先级学习特征，这与经验观察到的”低频先于高频收敛”完全一致。²

7.2 与NTK的关系

在无限宽度极限下，AGF退化为Neural Tangent Kernel（NTK）理论，因为：

• 所有神经元同时活跃
• 没有跳跃发生
• 特征学习被抑制

7.3 与特征学习动力学的联系

AGF预测的阶梯状学习与以下现象一致：

• Grokking：长时间训练后突然泛化
• Catapult Phase：损失先升后降的异常阶段
• Progressive Sharpening：Hessian特征值的双峰化

8. 扩展到其他架构

8.1 全连接线性网络

对于全连接线性网络，AGF简化为学习矩阵的奇异向量：

$$\frac{dU}{dt}=-\eta U(U^{\top }U-{\Sigma }^2)$$

8.2 仅注意力Transformer

对于仅注意力的Transformer，AGF描述注意力头学习主成分的动态过程。

8.3 非线性激活

对于ReLU网络，AGF框架通过考虑激活区域的几何变化进行推广。

9. 实践意义

9.1 学习率调度

AGF揭示了阶梯状学习的存在，这意味着：

• 恒定学习率可能不是最优的
• 在跳跃时刻降低学习率可能有助于稳定训练
• 学习率衰减应该与特征学习节奏匹配

9.2 网络初始化

AGF预测初始化尺度 $\alpha$ 影响特征学习顺序：

• 较小的 $\alpha$ 使跳跃时间更可预测
• 较大的 $\alpha$ 可能导致多个神经元同时跳跃

9.3 课程学习

AGF为课程学习（Curriculum Learning）提供了理论支持：先训练简单特征，再引入复杂特征，符合utility积累的自然顺序。

10. 代码实现

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
 
class TwoLayerNetwork(nn.Module):
    """两层神经网络用于AGF验证"""
    def __init__(self, d, m, H, c, init_scale=1e-3):
        super().__init__()
        self.w = nn.Parameter(torch.randn(H, d * m) * init_scale)
        self.a = nn.Parameter(torch.randn(H, c) * init_scale)
        self.g = nn.Linear(d, m)  # 可选的输入变换
        
    def forward(self, x):
        x_transformed = self.g(x)
        # 展平权重用于矩阵乘法
        w_matrix = self.w.view(self.w.shape[0], -1)
        h_pre = x_transformed @ w_matrix.T  # (batch, H, m)
        h = torch.relu(h_pre)  # origin-passing activation
        return (self.a.unsqueeze(0) * h).sum(dim=1)
 
def compute_utility(model, x, y, active_idx, device):
    """计算休眠神经元的utility"""
    model.eval()
    x, y = x.to(device), y.to(device)
    
    with torch.no_grad():
        residual = y - model(x)
    
    gradients = []
    for idx in active_idx:
        model.zero_grad()
        # 这里需要手动计算f_i的梯度
        pass
    return utility
 
def train_with_agf_dynamics(model, train_loader, epochs=100, device='cuda'):
    """训练并记录AGF动力学"""
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    norms_history = []
    loss_history = []
    
    for epoch in range(epochs):
        epoch_loss = 0
        for x, y in train_loader:
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            loss = nn.functional.mse_loss(model(x), y)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        
        # 记录神经元范数
        with torch.no_grad():
            norms = model.w.norm(dim=1).cpu().numpy()
        
        norms_history.append(norms)
        loss_history.append(epoch_loss / len(train_loader))
    
    return norms_history, loss_history
 
# 实验示例
def visualize_staircase_learning():
    """可视化阶梯状学习"""
    np.random.seed(42)
    torch.manual_seed(42)
    
    # 生成多频率数据
    n_samples, d = 1000, 10
    x = torch.randn(n_samples, d)
    
    # 目标函数：多个频率的组合
    y = (x[:, 0] * x[:, 1] + 
         0.5 * x[:, 2] * x[:, 3] * x[:, 4] + 
         0.1 * torch.randn(n_samples))
    
    # 训练模型
    model = TwoLayerNetwork(d=d, m=4, H=20, c=1, init_scale=1e-3).cuda()
    train_loader = torch.utils.data.DataLoader(
        torch.utils.data.TensorDataset(x, y), batch_size=64, shuffle=True
    )
    
    norms, losses = train_with_agf_dynamics(model, train_loader, epochs=500)
    
    # 可视化
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))
    
    # 损失曲线
    ax1.plot(losses)
    ax1.set_xlabel('Epoch')
    ax1.set_ylabel('Loss')
    ax1.set_title('Staircase Learning in Loss Curve')
    ax1.grid(True, alpha=0.3)
    
    # 神经元范数演化
    norms_arr = np.array(norms)
    for i in range(norms_arr.shape[1]):
        ax2.plot(norms_arr[:, i], alpha=0.7)
    ax2.set_xlabel('Epoch')
    ax2.set_ylabel('Neuron Norm')
    ax2.set_title('Neuron Norm Evolution (Staircase Jumps)')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('agf_staircase_learning.png', dpi=150)
    plt.show()
 
if __name__ == '__main__':
    visualize_staircase_learning()

11. 总结

交替梯度流（AGF）理论提供了一种描述神经网络特征学习的统一框架：

1. 核心机制：Utility最大化与Cost最小化的交替驱动特征涌现
2. 预测现象：阶梯状学习曲线、傅里叶特征按优先级学习
3. 理论保证：AGF收敛于vanishing initialization极限下的梯度流
4. 实践指导：学习率调度、初始化策略、课程学习设计

AGF桥接了微观神经元动力学与宏观学习现象，为理解深度学习提供了新的视角。

---

Footnotes

1. Kunin, D., et al. “Alternating Gradient Flows: A Theory of Feature Learning in Two-layer Neural Networks.” NeurIPS 2025. https://arxiv.org/abs/2506.06489 ↩ ↩²

2. 本理论与频率原则密切相关，AGF提供了频率原则的机制性解释。 ↩ ↩²