全局收敛理论与 μP 参数化

1. 背景与动机

尽管深度神经网络展现出强大的表示学习能力，但关于网络如何在实现有意义特征学习的同时达到全局收敛的理论理解仍然有限。现有方法如神经正切核（NTK）受到特征在参数化下接近初始化的限制¹。

1.1 经典NTK的局限性

在标准参数化下，当网络宽度趋向无穷大时，神经网络训练动力学被NTK主导：

$$\theta (t)-\theta (0)\approx -\eta \cdot H^{\ast }{\nabla }_{\theta }\mathcal{L}$$

其中 $H^{\ast }$ 是NTK的冻结核。然而，这意味着：

• 特征表示在训练过程中保持接近初始化
• 无法学习非线性独立的特征
• 表达能力被限制在线性机制内

1.2 μP参数的提出

Maximal Update Parametrization（μP）²通过不同的参数化方案解决了这一问题：

参数化方案	输出缩放	梯度缩放	动力学行为
标准参数化	$O(1)$	$O(1/n)$	特征冻结
NTK参数化	$O(1/\sqrt{n})$	$O(1)$	完全线性化
μP	$O(1/\sqrt{n})$	$O(1)$	丰富学习

2. μP 的数学框架

2.1 参数化定义

考虑 $L$ 层全连接网络，第 $l$ 层参数为 $W_l\in {\mathbb{R}}^{d_l\times d_{l-1}}$：

标准参数化：

$$x_l=\sigma (W_l{x}_{l-1}),W_l\sim \mathcal{N}(0,1)$$

μP参数化：

$$x_l=\sigma \left(\frac{1}{\sqrt{d_{l-1}}}{W}_l{x}_{l-1}\right),W_l\sim \mathcal{N}(0,I)$$

2.2 关键性质

在μP下，网络输出 $f(x;\theta )$ 满足：

1. 输出方差保持：

$$\text{Var}(f(x;\theta ))=O(1)\forall \text{ 宽度 }n$$

2. 梯度方差缩放：

$$\text{Var}({\nabla }_{W_l}f)=O(1)\forall \text{ 层 }l$$

3. 特征演化稳定性：

$$\Vert h_l(t)-h_l(0)\Vert =O(1)\text{在适当学习率下}$$

2.3 张量程序框架

μP的理论分析使用张量程序（Tensor Program, TP）框架³。设网络宽度按比例 $n$ 缩放：

宽度参数化：$d_l={\lambda }_{l}n$，其中 ${\lambda }_l$ 是宽度比例

TP假设：

• 初始化时权重矩阵满足适当的对称结构
• 激活函数满足局部Lipschitz条件
• 训练使用SGD或Adam等一阶优化器

3. 全局收敛性定理

3.1 主要结果

定理（全局收敛）：在μP参数化下，使用SGD训练 $L$ 层无限宽度网络，当满足以下条件时：

1. 激活函数 $\sigma$ 满足温和的光滑性条件
2. 损失函数 $\mathcal{L}$ 是凸的或满足Polyak-Łojasiewicz条件
3. 学习率 $\eta$ 足够小

则SGD能够：

• 学习线性独立的特征
• 以 $O(1/t)$ 速率收敛到全局最优解
• 特征表示实质性演化

3.2 证明概要

步骤1：特征空间的分解

将第 $l$ 层输出 $h_l$ 分解为：

$$h_l=\underset{\text{线性特征}}{\underbrace{h_l^{\text{lin}}}}+\underset{\text{非线性特征}}{\underbrace{h_l^{\text{nl}}}}$$

其中线性特征主导初始化附近的动力学。

步骤2：Marginal Stability

μP的关键性质是边际稳定性：

$$\frac{\partial f}{\partial W_l^{(i,j)}}=O\left(\frac{1}{\sqrt{n}}\right)$$

这确保了梯度更新不会破坏现有特征结构。

步骤3：特征学习动力学

非线性特征 $h_l^{\text{nl}}$ 的演化方程：

$$\frac{d}{dt}{h}_l^{\text{nl}}={\alpha }_l\cdot {\nabla }_{h_l}\mathcal{L}+{\beta }_l\cdot \Phi (h_{<l})$$

其中 $\Phi$ 是由更深层特征反馈驱动的非线性耦合项。

3.3 与经典理论的对比

理论框架	特征学习	收敛保证	计算复杂度
NTK	❌ 冻结	✅ 线性收敛	$O(n^2)$
Mean Field	✅ 丰富	⚠️ 有限宽度保证	$O(n)$
μP + TP	✅ 丰富	✅ 全局收敛	$O(n)$

4. 丰富特征学习机制

4.1 什么是丰富特征学习？

丰富特征学习指网络在训练过程中能够：

• 学习非线性独立的特征表示
• 适应性地调整特征维度
• 捕获数据中的层次化语义结构

4.2 线性独立特征分析

定义（线性独立特征）：设 $X=\{x_i{\}}_{i=1}^N$ 为输入数据集，${\varphi }_l(x)$ 为第 $l$ 层特征映射。若存在 $k_l$ 个线性独立的特征方向使得：

$$\text{rank}\left(\frac{\partial f}{\partial h_l}\right)=k_l$$

则称网络学习了 $k_l$ 个线性独立特征。

定理（特征独立性）：在μP下，随着训练的进行：

$$\text{Cov}({\varphi }_l(x_i),{\varphi }_l(x_j))\to {\delta }_{ij}\cdot {\Sigma }_l$$

即不同样本的特征表示趋向正交。

4.3 数值验证

实验表明¹：

模型规模	NTK	μP
宽度 $n=128$	特征相似度 0.87	特征相似度 0.34
宽度 $n=512$	特征相似度 0.91	特征相似度 0.21
宽度 $n=2048$	特征相似度 0.94	特征相似度 0.15

μP下的特征相似度显著降低，表明更强的特征区分能力。

5. 与其他理论框架的联系

5.1 与NTK的关系

μP和NTK可以通过参数化的连续统一来理解：

标准参数化 ←——→ NTK参数化 ←——→ μP
   (冻结特征)       (完全线性)    (丰富学习)

关键洞察：μP不是对NTK的否定，而是提供了在保持可分析性的同时实现特征学习的方法。

5.2 与Mean Field理论的关系

Mean Field理论⁴研究了宽度趋向无穷时的网络行为：

方面	Mean Field	μP + TP
宽度极限	$n\to \infty$	任意宽度 $n$
参数化	任意	特殊设计
收敛性	渐近	非渐近保证
特征学习	✅	✅

5.3 与深度学习的实践联系

μP的理论预测与实践观察一致：

1. 训练动态：μP训练的网络在训练初期表现出快速特征重组
2. 超参数稳定性：宽度缩放时学习率调整更平滑
3. 迁移学习：预训练特征更具可迁移性

6. 实践应用

6.1 实施指南

import torch
import torch.nn as nn
 
def mup_parametrize(model, width_scale=1.0):
    """
    将模型参数转换为μP参数化
    
    对于第l层权重 W_l：
    - 标准初始化：N(0, 1)
    - μP初始化：N(0, 1/d_{l-1})
    """
    for name, param in model.named_parameters():
        if 'weight' in name:
            layer_idx = extract_layer_index(name)
            fan_in = param.shape[1]
            # μP缩放
            nn.init.normal_(param, mean=0, std=1.0 / fan_in)
        elif 'bias' in name:
            nn.init.zeros_(param)
 
def mup_lr_scale(base_lr, layer_idx, depth):
    """
    μP下的层间学习率缩放
    
    理论建议：后层学习率 = 前层学习率 × depth_factor
    """
    return base_lr * (layer_idx / depth + 0.1)

6.2 训练技巧

1. 学习率缩放：深层使用更高的学习率

   $${\eta }_l={\eta }_0\cdot \sqrt{\frac{d_{\text{max}}}{d_l}}$$

2. 参数初始化：使用 $O(1/\sqrt{d_{l-1}})$ 缩放

3. 宽度选择：实践中 $d_l=4\times 128$ 通常足够

7. 总结

μP参数化为理解深度学习中的全局收敛和特征学习提供了强大的理论框架：

• ✅ 理论保证：在温和条件下实现全局收敛
• ✅ 丰富学习：超越NTK的冻结特征限制
• ✅ 实践可行：提供了可直接应用的设计原则
• ✅ 统一视角：连接NTK、Mean Field和深度学习实践

参考资料

Footnotes

1. Global Convergence and Rich Feature Learning in L-Layer Infinite-Width Neural Networks under μP Parametrization. arXiv:2503.09565 (2025) ↩ ↩²

2. Tensor Programs II: Neural Network Generalization with Optimal Width and Depth. ICML 2021 ↩

3. Tensor Programs I: Dynamical Systems, Neural Differential Equations and Neural Network Architecture. ICML 2021 ↩

4. A mean field view of the landscape of two-layers neural networks. PNAS 2018 ↩