μ参数化与Rich Feature Learning

概述

本文深入解析 Greg Yang 等人在 arXiv:2503.09565 中提出的 μ参数化（μ-Parameterized Networks） 理论。这项工作首次在数学上证明了：深层神经网络可以同时实现全局收敛和有意义（non-trivial）的特征学习，突破了传统理论认为两者不可兼得的困境。

---

1. 背景：特征学习的理论困境

1.1 神经切核（NTK）理论与无限宽网络

神经切核理论描述了无限宽神经网络在梯度下降下的动态。当网络宽度 $n\to \infty$ 时：

• 网络输出近似线性：$f(\mathbf{x},\theta )\approx f_{\text{NTK}}(\mathbf{x})$
• 梯度下降等价于在 RKHS 中进行函数优化
• 训练动态可被精确预测

但 NTK 理论有一个致命缺陷：它描述的是”特征学习不存在”的状态。

1.2 有限宽度网络的困境

当宽度 $n$ 有限时，两种理论框架各执一词：

框架	预测	问题
NTK 理论	线性化动态，特征不学习	与实践不符（网络确实学习到有意义特征）
信号传播理论	保持方差，激活稳定	仅分析均值/方差，不涉及优化收敛

核心问题：能否有一个统一框架，同时解释特征学习和全局收敛？

1.3 特征学习的”非平凡性”定义

Yang 等人定义了什么是”有意义”的特征学习：

定义（非平凡特征学习）：设 $f(\mathbf{x},\theta )$ 是网络输出。如果存在 ${\theta }^{\ast }$ 使得：

1. $f(\mathbf{x},{\theta }^{\ast })=f^{\ast }(\mathbf{x})$ 对目标函数 $f^{\ast }$ 有好逼近
2. 且最优参数 ${\theta }^{\ast }$ 不在 NTK 极限路径上（即 $f(\mathbf{x},{\theta }^{\ast })\ne f_{\text{NTK}}(\mathbf{x})$）

则称网络实现了非平凡特征学习。

---

2. μ参数化的核心思想

2.1 从logits-space到mean-field

传统参数化在 logits 空间（输出层）定义参数。Yang 等人转向 均值场（mean-field） 参数化。

对于一个宽度为 $n$ 的层：

$${\mathbf{h}}^{(l)}=\frac{1}{\sqrt{n}}{W}^{(l)}{\mathbf{h}}^{(l-1)}$$

其中 $W^{(l)}\in {\mathbb{R}}^{n\times n}$ 是权重矩阵。

2.2 μ参数化的定义

核心思想：将网络参数分解为两部分：

1. NTK 组件：负责保持稳定的学习动态
2. 特征学习组件：负责捕获非平凡特征

$$\theta =\underset{\mu =0}{\underbrace{{\theta }_{\text{NTK}}}}+\underset{\mu \text{ 参数化}}{\underbrace{{\theta }_{\text{feat}}}}$$

2.3 μ参数的数学定义

对于第 $l$ 层的权重 $W^{(l)}$：

$$W^{(l)}=\underset{\text{NTK 缩放}}{\underbrace{\frac{1}{\sqrt{n}}{W}_{\text{ntk}}^{(l)}}}+\underset{\text{特征学习}}{\underbrace{\mu \cdot \frac{1}{\sqrt{n}}{W}_{\text{feat}}^{(l)}}}$$

其中 $\mu$ 是一个全局缩放参数：

• $\mu =0$：退化为标准 NTK 设置
• $\mu =1$：允许完整的特征学习
• $0<\mu \ll 1$：介于两者之间

---

3. 理论保证

3.1 全局收敛定理

定理（μ参数化全局收敛）：考虑深度为 $L$、宽度为 $n$ 的 ReLU 网络，使用 $\mu$ 参数化。假设：

1. 数据分布 $\mathcal{D}$ 满足 mild assumptions
2. 目标函数 $f^{\ast }$ 是 $L$-Lipschitz 的
3. 训练使用梯度下降，学习率 $\eta =O(1/n)$

则当 $n\to \infty$ 时，存在 ${\mu }^{\ast }(L)$ 使得：

$$\Vert f(\mathbf{x},\theta (t))-f^{\ast }(\mathbf{x})\Vert \to 0\text{当 }t\to \infty$$

关键：${\mu }^{\ast }(L)$ 只依赖于深度 $L$，不依赖于宽度 $n$。

3.2 特征学习的非平凡性

定理（非平凡特征学习）：在上述设置下，当 $\mu >0$ 足够大时：

1. 学习到的特征 $\varphi (\mathbf{x},{\theta }^{\ast })$ 与输入分布的语义结构相关
2. 特征维度之间存在有意义的交互（非对角协方差）
3. 性能超越纯 NTK 预测的基准

3.3 深度依赖性

关键发现：${\mu }^{\ast }(L)$ 随深度 $L$ 增长：

$${\mu }^{\ast }(L)\approx \frac{C}{\sqrt{L}}$$

这意味着：

• 浅层网络（$L$ 小）：即使 $\mu$ 很小也能学习特征
• 深层网络（$L$ 大）：需要更大的 $\mu$ 来支持特征学习

μ* (最小特征学习参数)
    │
    │         ╭───────────── 理论上界
    │        ╱
    │       ╱
    │      ╱
    │     ╱
    │    ╱
    │   ╱  ─ ─ ─ ─ ─ ─ ─
    │  ╱           ╱
    │ ╱          ╱
    │╱─────────╱──────────→ 深度 L
    │
    0    L₁    L₂    L₃

---

4. 与 NTK 的关系

4.1 NTK 作为特殊 case

当 $\mu =0$ 时：

• $W^{(l)}=\frac{1}{\sqrt{n}}{W}_{\text{ntk}}^{(l)}$
• 网络退化为无限宽网络的线性化版本
• 学习动态完全由 NTK 决定

4.2 μ参数的物理意义

$\mu$ 参数控制了两个效应的相对强度：

效应	主导参数	特征
线性化（NTK）	${\theta }_{\text{ntk}}$	稳定但表达能力有限
非线性特征学习	${\theta }_{\text{feat}}$	表达力强但需要更大 $\mu$

4.3 临界 $\mu$ 值

存在一个临界值 ${\mu }_c$：

• $\mu <{\mu }_c$：网络行为接近 NTK
• $\mu >{\mu }_c$：网络进入特征学习 regime

实验发现 ${\mu }_c\approx 1$ 对大多数架构成立。

---

5. 与信号传播理论的联系

5.1 Edge of Chaos

信号传播理论告诉我们：存在一个”混沌边缘”（Edge of Chaos），在此处激活方差保持稳定。

μ参数化给出了这个现象的优化视角解释：

• 在 Edge of Chaos 附近，梯度方差稳定
• 这恰好是 $\mu$ 参数化的稳定训练区域
• 两者描述的是同一个现象的不同侧面

5.2 方差保持条件

设 ${\sigma }_w^2$ 和 ${\sigma }_b^2$ 是权重和偏置的方差。对于第 $l$ 层激活：

$$\text{Var}({\mathbf{h}}^{(l)})={\sigma }_w^2\cdot \mathbb{E}[{\mathbf{h}}^{(l-1)}\odot {\mathbf{h}}^{(l-1)}]+{\sigma }_b^2$$

在 Edge of Chaos 处：

$${\sigma }_w^2=\frac{1}{\mathbb{E}[{\varphi }^{\prime }(Z{)}^2]}$$

其中 $Z$ 是净输入，${\varphi }^{\prime }$ 是激活函数的导数。

---

6. 实验验证

6.1 设置

配置	值
架构	MLP, ResNet, Transformer
宽度	128, 256, 512, 1024
深度	1, 3, 6, 12, 24
数据集	CIFAR-10, ImageNet, 合成数据

6.2 特征学习度量

使用 神经切片方差分析（Neural Tangent Covariance） 来量化特征学习：

$${\text{NTC}}_{ij}=\frac{\partial f}{\partial {\theta }_i}\cdot \frac{\partial f}{\partial {\theta }_j}$$

如果特征在学习，则 NTC 矩阵会有非平凡的结构。

6.3 关键实验结果

实验1：深度 vs μ的关系

深度 $L$	$\mu =0$ 准确率	$\mu =1$ 准确率	提升
1	62.3%	63.1%	+0.8%
3	61.8%	64.7%	+2.9%
6	60.2%	66.2%	+6.0%
12	58.7%	65.8%	+7.1%
24	55.3%	64.1%	+8.8%

观察：深度越大，$\mu$ 参数化的优势越明显！

实验2：特征质量分析

模型	NTK 预测	$\mu$ 参数化	真实
隐藏层激活熵	0.42	0.78	0.85
特征相关性	0.12	0.45	0.52
下游任务转移	0.31	0.67	0.71

6.4 与其他特征学习理论的关系

方法	特征学习	全局收敛	可扩展性
NTK	❌	✅	✅
μ参数化	✅	✅	✅
贪心逐层预训练	✅	❌	✅
线性网络初始化	✅	✅	❌

---

7. 实际应用

7.1 PyTorch 实现

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class MuParameterizedLinear(nn.Module):
    """μ参数化的线性层"""
    def __init__(self, in_features, out_features, mu=1.0, bias=True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.mu = mu
        
        # NTK 组件（固定为小初始化）
        self.weight_ntk = nn.Parameter(
            torch.randn(out_features, in_features) * 1e-3
        )
        
        # 特征学习组件
        self.weight_feat = nn.Parameter(
            torch.randn(out_features, in_features) * 0.01
        )
        
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter('bias', None)
    
    def forward(self, x):
        # 组合权重
        weight = self.weight_ntk / self.in_features**0.5 + \
                 self.mu * self.weight_feat / self.in_features**0.5
        
        return F.linear(x, weight, self.bias)
 
 
class MuParameterizedMLP(nn.Module):
    """μ参数化的MLP"""
    def __init__(self, input_dim, hidden_dim, output_dim, depth, mu=1.0):
        super().__init__()
        self.layers = nn.ModuleList()
        self.depth = depth
        
        # 输入层
        self.layers.append(MuParameterizedLinear(input_dim, hidden_dim, mu))
        
        # 隐藏层
        for _ in range(depth - 2):
            self.layers.append(MuParameterizedLinear(hidden_dim, hidden_dim, mu))
            self.layers.append(nn.LayerNorm(hidden_dim))
            self.layers.append(nn.ReLU())
        
        # 输出层
        self.layers.append(MuParameterizedLinear(hidden_dim, output_dim, mu))
    
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

7.2 训练脚本

def train_with_mu_parametrization():
    # 超参数
    depth = 12
    mu = 1.0
    lr = 0.001
    
    model = MuParameterizedMLP(
        input_dim=784,
        hidden_dim=512,
        output_dim=10,
        depth=depth,
        mu=mu
    )
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    
    # 训练循环
    for epoch in range(100):
        for x, y in dataloader:
            optimizer.zero_grad()
            output = model(x)
            loss = criterion(output, y)
            loss.backward()
            
            # 可选：梯度裁剪
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            
            optimizer.step()
    
    return model

7.3 超参数选择指南

深度 $L$	推荐 $\mu$	学习率调整
1-3	0.1-0.5	标准
4-8	0.5-1.0	略低
9-24	1.0-2.0	较低
> 24	1.5-3.0	最低

---

8. 与 AGF 框架的联系

8.1 AGF 框架简介

AGF（Aligned Gradient Flow）框架提出了一个统一视角：特征学习的关键是输入-输出雅可比的对齐。

8.2 μ参数化与 AGF 的等价性

可以证明：当 $\mu >0$ 时，μ参数化网络恰好满足 AGF 的对齐条件：

$$\frac{\partial f}{\partial \mathbf{x}}\approx \mathbf{U}\cdot \varphi (\mathbf{x})$$

其中 $\mathbf{U}$ 是固定的方向矩阵，$\varphi (\mathbf{x})$ 是学习到的特征。

8.3 统一理论框架

                    ┌─────────────────────────────┐
                    │     深度网络特征学习统一理论    │
                    └─────────────────────────────┘
                                    │
            ┌───────────────────────┼───────────────────────┐
            │                       │                       │
            ▼                       ▼                       ▼
    ┌───────────────┐       ┌───────────────┐       ┌───────────────┐
    │  信号传播理论  │       │  μ参数化理论   │       │  AGF框架     │
    │ (方差保持)    │       │ (全局收敛)     │       │ (梯度对齐)   │
    └───────────────┘       └───────────────┘       └───────────────┘
            │                       │                       │
            └───────────────────────┼───────────────────────┘
                                    │
                                    ▼
                        ┌───────────────────────┐
                        │   实际训练中的特征学习   │
                        │   (非平凡 + 可收敛)    │
                        └───────────────────────┘

---

9. 局限性与未来方向

9.1 当前局限性

1. 理论仅针对 ReLU：其他激活函数（如 GELU、SiLU）的分析尚未完成
2. 假设数据分布：理论基于一定的数据假设，与实际数据可能有差距
3. 计算开销：μ参数化增加了参数数量（约 2x）

9.2 开放问题

1. 最优 μ 的自动调整：能否在训练中自适应调整 μ？
2. 与其他技术的组合：与 BatchNorm、残差连接的相互作用？
3. Transformer 的扩展：能否将 μ 参数化推广到注意力机制？

9.3 后续工作

Yang 等人在后续工作中（arXiv:2506.06489）进一步提出了 AGF 框架，将 μ 参数化的洞察推广到更一般的设置。

---

10. 总结

10.1 核心贡献

1. 提出 μ 参数化：一种新的神经网络参数化方法
2. 证明全局收敛 + 特征学习可兼得：打破了两者不可兼得的传统观点
3. 揭示深度依赖性：${\mu }^{\ast }(L)\approx C/\sqrt{L}$
4. 统一多个理论框架：将 NTK、信号传播、AGF 等联系起来

10.2 实践意义

• 训练更深网络：μ 参数化使得训练超深网络更稳定
• 改善特征质量：学习到的特征更有意义
• 理论指导实践：提供了选择超参数的数学依据

10.3 理论意义

• 解释”为什么深度学习有效”：提供了特征学习的理论基础
• 指导新架构设计：μ 参数化的洞察可应用于架构创新
• 连接优化与表示学习：桥接了两个传统上分离的领域

---

参考资料