有限宽度网络的精确梯度下降动力学

1. 引言

理解梯度下降在真实、有限宽度多层神经网络中的训练动力学仍是深度学习理论的核心挑战¹。现有理论如神经正切核（NTK）和均值场（Mean Field）在实际网络规模下存在局限性。

1.1 现有理论的局限

神经正切核（NTK）：

• 假设无限宽度（$n\to \infty$）
• 权重在训练中保持接近初始化
• 预测能力有限

均值场理论（Mean Field）：

• 假设宽度 $n\to \infty$
• 需要特殊初始化方案
• 仅适用于两层网络

实践中的Gap：

• 真实网络宽度有限（如512-4096）
• 训练远超”弱学习”区域
• 权重显著偏离初始化

1.2 本文贡献

本文提供首个精确的分布式刻画，涵盖：

1. 首层权重的高斯波动
2. 深层权重的确定性聚焦
3. 对非高斯特征的适用性
4. 明确的泛化误差估计

2. 理论框架

2.1 模型设定

考虑 $L$ 层全连接网络：

$$h_{l+1}=\sigma (W_l{h}_l),l=0,1,\dots ,L-1$$

其中：

• $h_0=x\in {\mathbb{R}}^{d_0}$ 为输入
• $W_l\in {\mathbb{R}}^{d_{l+1}\times d_l}$ 为权重矩阵
• $\sigma$ 为激活函数

单指标回归模型：

$$y=f(x)+\epsilon ,f(x)=⟨w^{\ast },\varphi (x)⟩$$

其中 $\varphi (x)=(h_1,\dots ,h_L)\in {\mathbb{R}}^p$ 为网络输出的特征向量。

2.2 有限宽度比例区域

关键设定：样本量 $n$ 和特征维度 $p$ 以相同速率增长：

$$n,p\to \infty ,\frac{p}{n}\to \gamma \in (0,\infty )$$

这被称为有限宽度比例区域，比 $p\ll n$ 或 $p\gg n$ 更具一般性。

2.3 状态演化方程

首层权重动力学：

对于 $W_1$，定义经验协方差矩阵：

$${\Sigma }_1^{(t)}=\frac{1}{p_1}\sum _{i=1}^{p_1}{w}_{1,i}^{(t)}(w_{1,i}^{(t)}{)}^{\top }$$

定理（首层状态演化）：

$${\Sigma }_1^{(t)}\stackrel{t}{\to }{\Sigma }_1^{\ast },\det ({\Sigma }_1^{\ast })\to 0\text{ 或 }{\Sigma }_1^{\ast }=I$$

即首层权重在训练后趋向低秩或单位阵。

深层权重动力学：

对于 $l\ge 2$，权重趋向确定性极限：

$$\frac{1}{p_l}{W}_l^{(t)}\stackrel{t}{\to }{\bar{W}}_l$$

其中 ${\bar{W}}_l$ 是确定性矩阵。

3. 训练动力学分析

3.1 首层：高斯波动

定理（首层高斯性）：在训练过程中，$W_1$ 的元素满足：

$$(W_1^{(t)}{)}_{ij}\sim \mathcal{N}\left({\mu }_{ij}^{(t)},({\sigma }_{ij}^{(t)}{)}^2\right)$$

其中均值和方差由以下 ODE 给出：

$$\frac{d{\mu }^{(t)}}{dt}=-\eta \cdot \mathbb{E}[{\nabla }_{\theta }\mathcal{L}],\frac{d({\sigma }^{(t)}{)}^2}{dt}={\eta }^2\cdot \text{Var}({\nabla }_{\theta }\mathcal{L})$$

关键发现：

• 首层权重在训练中保持随机波动
• 这是由于 $W_1$ 接收来自所有样本的梯度
• 深层权重由于共享梯度的平均效应趋于稳定

3.2 深层：确定性聚焦

定理（深层确定性）：对于 $l\ge 2$，

$$\Vert W_l^{(t)}-{\bar{W}}_l{\Vert }_F\stackrel{t}{\to }0$$

其中收敛速率：

$$\mathbb{E}[\Vert W_l^{(t)}-{\bar{W}}_l{\Vert }_F^2]\le \frac{C}{p_l}$$

物理解释：

• 深层权重接收来自首层的平均梯度
• 大数定律使深层权重稳定
• 首层的随机波动被深层平滑

3.3 特征演化

引理（特征集中）：

$$\frac{1}{n}\sum _{i=1}^n({\varphi }_i^{(t)}-{\bar{\varphi }}^{(t)})({\varphi }_i^{(t)}-{\bar{\varphi }}^{(t)}{)}^{\top }\to {\Sigma }_{\varphi }$$

特征协方差矩阵收敛到稳定值。

4. 泛化误差估计

4.1 主要结果

定理（泛化误差估计）：在有限宽度比例区域，泛化误差满足：

$$R({\hat{f}}_t)=\underset{\text{贝叶斯最优}}{\underbrace{R^{\ast }}}+\underset{\text{方差项}}{\underbrace{\frac{\text{tr}({\Sigma }_{\varphi }{\Sigma }_w)}{n}}}+\underset{\text{偏差项}}{\underbrace{O\left(\frac{1}{\sqrt{n}}\right)}}$$

其中 ${\Sigma }_w$ 是权重后验协方差。

4.2 在线学习视角

将训练过程视为在线学习：

累积风险：

$${\bar{R}}_T=\frac{1}{T}\sum _{t=1}^{T}R(f_t)$$

遗憾界（Regret Bound）：

$${\bar{R}}_T-R^{\ast }\le \frac{B}{\sqrt{T}}+\frac{\gamma \cdot \log T}{T}$$

4.3 早停与泛化

定理（早停估计）：定义最优停止时间 $t^{\ast }$：

$$t^{\ast }\approx \arg \underset{t}{\min }\left\{\text{bias}(t)+\frac{\text{variance}}{t}\right\}$$

这提供了早停策略的理论基础：

• 早期：偏差主导，应继续训练
• 后期：方差主导，应提前停止

5. 与现有理论的关系

5.1 与NTK的关系

方面	NTK	本文理论
宽度假设	$p\to \infty$	有限 $p$
权重演化	完全冻结	首层波动，深层聚焦
表达能力	线性	非线性丰富
泛化分析	仅训练误差	泛化误差

5.2 与均值场的关系

方面	Mean Field	本文理论
初始化	特殊方案	任意
动力学	概率分布	点估计+波动
适用深度	两层	任意深度
宽度要求	$p\to \infty$	有限宽度

5.3 统一视角

有限宽度比例区域 (n,p → ∞, p/n → γ)

                    ↓
         ┌─────────┴─────────┐
         ↓                   ↓
    γ << 1 (欠参数化)    γ >> 1 (过参数化)
         ↓                   ↓
    接近Mean Field      接近NTK
         ↓                   ↓
    首层主导           深层主导

6. 实践应用

6.1 泛化误差预测

import numpy as np
import torch
 
def estimate_generalization_error(model, train_loader, device):
    """
    基于状态演化的泛化误差估计
    
    使用首层权重的经验协方差估计方差项
    """
    model.eval()
    
    # 收集首层权重统计
    w1 = model.fc1.weight.data.flatten().numpy()
    
    # 计算经验方差
    sigma_w = np.var(w1)
    
    # 估计特征协方差trace
    # (实际中需要更复杂的计算)
    n_samples = len(train_loader.dataset)
    
    # 泛化误差估计
    variance_term = sigma_w / n_samples
    bias_term = 1 / np.sqrt(n_samples)
    
    return variance_term + bias_term
 
def predict_optimal_stopping(model, val_loader, thresholds):
    """
    预测最优早停时间
    
    基于方差累积的启发式方法
    """
    w1_history = []
    val_losses = []
    
    for threshold in thresholds:
        # 训练到指定阈值
        train_with_threshold(model, val_loader, threshold)
        
        # 记录首层方差
        w1 = model.fc1.weight.data.flatten()
        w1_var = torch.var(w1).item()
        w1_history.append(w1_var)
        
        # 验证损失
        val_loss = evaluate(model, val_loader)
        val_losses.append(val_loss)
    
    # 找到方差-损失权衡点
    best_idx = np.argmin(val_losses)
    return thresholds[best_idx]

6.2 网络设计建议

基于理论的建议：

1. 首层宽度选择：

   • 较窄的首层减少高斯波动
   • 建议：$d_1\ge 2d_0$

2. 深层宽度设计：

   • 深层宽度影响特征表达能力
   • 建议：$d_l$ 随深度逐渐增加或保持

3. 正则化时机：

   • 早期可依赖隐式偏差
   • 后期添加正则化抑制方差

7. 扩展与展望

7.1 当前理论限制

1. 激活函数：需要有界或Lipschitz连续
2. 损失函数：主要针对平方损失
3. 网络结构：全连接网络

7.2 未来研究方向

理论扩展：

• 卷积层和注意力层
• 更一般的损失函数
• 批量梯度下降的随机性

实践应用：

• 自适应学习率设计
• 早停策略优化
• 网络架构搜索

8. 总结

本文提出了首个精确的有限宽度多层神经网络梯度下降动力学理论：

• ✅ 首层高斯波动：揭示了首层权重的随机演化
• ✅ 深层确定性聚焦：深层权重趋向确定极限
• ✅ 非渐近分析：适用于实际网络规模
• ✅ 泛化误差估计：提供了可计算的估计方法

这一理论与NTK和Mean Field互补，为理解深度学习训练提供了新的视角。

参考资料

Footnotes

1. Precise Gradient Descent Training Dynamics for Finite-Width Multi-Layer Neural Networks. arXiv:2505.04898 (2025) ↩