跳转至

Feature Learning Dynamics in Infinite-Depth Neural Networks

会议: NeurIPS 2026
arXiv: 2512.21075
领域: 学习理论
关键词: 特征学习动力学、无限深度极限、depth-μP、前向反向耦合、神经特征动力学 (NFD)

一句话总结

本文针对 depth-\(\mu\mathrm{P}\) 缩放下的单层 ResNet,利用条件高斯表示显式隔离反向传播中权重复用引起的前向-反向耦合项,证明该耦合在初始化时随宽度以 \(\mathcal{O}(n^{-1})\) 消失,而在训练时虽在无限宽度极限下存留,却随深度以更高阶 \(\mathcal{O}(L^{-2})\) 衰减,从而严格建立了无限深度特征学习的神经特征动力学(NFD)前向-反向 SDE 极限。

研究背景与动机

深度神经网络在现代人工智能中取得了巨大成功,但揭示训练过程中内部表征如何演化的机理理解仍不完备,在深层极限下尤其缺乏严格的理论分析。为了刻画大网络下的特征学习动力学,平均场理论(Mean-Field Theory)与张量程序(Tensor Programs / \(\mu\mathrm{P}\))在宽度维度上确立了原则性极限,确保特征更新在无限宽度极限下保持 \(\mathcal{O}(1)\),并成功实现了跨宽度的超参数迁移。然而,在深度维度上,特征学习动力学的严格数学刻画一直面临重大阻碍。

近年来,将单层残差网络分支缩放 \(1/\sqrt{L}\)(即 depth-\(\mu\mathrm{P}\))被证明是维持深层网络训练稳定性并实现跨深度超参迁移的关键手段。在此缩放下,离散层索引 \(\ell\) 可转化为连续时间变量 \(t_\ell = \ell / L \in [0, 1]\),初始化下的前向特征传播收敛至均值场随机微分方程(SDE)。然而,将该连续动力学推广到训练阶段时,现有研究(如基于物理启发的动态均值场理论 DMFT)大多依赖启发式响应函数,缺乏对反向传播核心机制的严格数学处理。深层特征动力学理论构建的真正瓶颈,在于标准反向传播中反向梯度传递必须复用前向权重矩阵的转置 \(\bm{W}_\ell^\top\),从而引发前向特征与反向梯度之间极度复杂的耦合依赖关系。

在懒惰训练(Lazy Training / NTK)体系下,梯度独立性假设(GIA)在无限宽度下成立且特征几乎保持不变;但在具有实质性表征演化的特征学习状态下,SGD 参数更新会将历史梯度与前向特征强力缠绕。本文的核心切入点是:在求取任何网络规模极限之前,通过条件高斯表示显式分解出权重复用带来的前向-反向耦合项,深入分析其在宽度和深度两个维度上的渐近行为。核心 idea:利用条件高斯表示显式解耦权重复用引起的前向-反向耦合,证明该耦合在训练时虽跨越无限宽度极限,却在 depth-\(\mu\mathrm{P}\) 深度缩放下表现为 \(\mathcal{O}(L^{-2})\) 的更高阶小量,从而建立以解耦反向权重驱动的神经特征动力学(NFD)极限 SDE 系统。

方法详解

整体框架

本文针对 depth-\(\mu\mathrm{P}\) 缩放下的单层 Pre-activation 残差网络展开研究。网络输入 \(\bm{x} \in \mathbb{R}^d\),残差流前向更新为 \(\bm{h}_\ell = \bm{h}_{\ell-1} + \frac{1}{\sqrt{Ln}} \bm{W}_\ell \phi(\bm{h}_{\ell-1})\),输出为 \(f(\bm{x}; \bm{\theta}) = \frac{1}{n} \bm{v}^\top \bm{h}_L\)。模型采用流式随机梯度下降(Stream SGD)优化,有效学习率设为 \(\eta = \eta_c n\)。整体理论推演分为四个有机串联的阶段:首先建立前向与反向增量的精确条件高斯表示;其次证明初始化时权重复用耦合随宽度 \(n \to \infty\) 均匀消失(收敛至独立前向-反向 SDE);接着证明训练更新引入的耦合项在无限宽度下不可忽略,但在深度维度具有 \(\mathcal{O}(L^{-2})\) 的更高阶压缩效应;最后由此定义并严格证明网络训练动力学向神经特征动力学(NFD)前向-反向 SDE 系统的收敛性。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入与初始状态:Pre-activation ResNet<br/>depth-μP 缩放与 Stream SGD"] --> B["条件高斯显式分解机制<br/>隔离权重复用耦合与独立高斯涨落"]
    B --> C["初始化 GIA 深度一致恢复<br/>有限宽度误差 O(1/n) 且随 L 均匀有界"]
    C --> D["训练步前向反向耦合与高阶深度压缩<br/>存留耦合具有单层 O(L⁻²) 与累积 O(L⁻¹) 衰减"]
    D --> E["神经特征动力学 (NFD) 极限收敛<br/>跨步双向 SDE 耦合核与 O(L⁻¹) 离散化极限"]

关键设计

1. 条件高斯显式分解机制:隔离权重复用耦合与独立高斯涨落 在反向传播中,第 \(\ell\) 层的梯度增量 \(\bm{b}_\ell = \frac{1}{\sqrt{n}} \bm{W}_\ell^\top \bm{g}_\ell\) 复用了前向权重 \(\bm{W}_\ell\)。为了在有限宽度和有限深度下明确识别该耦合,作者在给定前向特征与反向梯度的条件下求解高斯线性回归。在前向特征方向定义投影算子 \(\bm{P}_{x, \ell-1} = \frac{\bm{x}_{\ell-1} \bm{x}_{\ell-1}^\top}{\|\bm{x}_{\ell-1}\|^2}\)(其中 \(\bm{x}_{\ell-1} = \phi(\bm{h}_{\ell-1})\)),将权重矩阵正交分解为其最小 Frobenius 范数特解与正交补空间高斯自由度。由此得到反向增量的精确条件高斯等价分布: $\(\bm{b}_\ell^{(0)} \overset{d}{=} \frac{\langle \bm{z}_\ell^{(0)}, \bm{g}_\ell^{(0)} \rangle}{\sqrt{n}} \frac{\bm{x}_{\ell-1}^{(0)}}{\|\bm{x}_{\ell-1}^{(0)}\|} + \frac{\|\bm{g}_\ell^{(0)}\|}{\sqrt{n}} (\bm{I} - \bm{P}_{x,\ell-1}^{(0)}) \widetilde{\bm{z}}_\ell^{(0)}\)$ 其中 \(\bm{z}_\ell^{(0)}, \widetilde{\bm{z}}_\ell^{(0)} \sim \mathcal{N}(0, \bm{I}_n)\) 为相互独立的标准高斯创新。该分解精巧地将增量拆成两部分:第一项是沿前向特征方向投影的耦合项,第二项则是正交补空间上的纯高斯涨落。通过这种表示,对梯度独立性假设(GIA)的偏离被精确数学具象化为第一项的存在。

2. 初始化 GIA 深度一致恢复:建立跨越宽深两维度的交换极限 在初始化时刻,反向传播从与前向完全独立的读出向量 \(\bm{g}_L^{(0)} = \bm{v}^{(0)}\) 开始回传。利用随机矩阵在球面上的测度集中性质,前向高斯创新与反向梯度的归一化内积在坐标层面呈现弱依赖。作者证明了在假设激活函数导数 \(\phi'\) Lipschitz 连续的前提下,真实反向梯度过程与使用完全独立权重 \(\widetilde{\bm{W}}_\ell^{(0)}\) 构造的辅助 GIA 过程在坐标级 Wasserstein-2 距离上满足深度一致收敛: $\(\sup_{\ell \in [L]} W_2\left(\operatorname{Law}(\bm{g}_{\ell, i}^{(0)}), \operatorname{Law}(\bar{\bm{g}}_{\ell, i}^{(0)})\right)^2 \leq C n^{-1}\)$ 由于该界对深度 \(L\) 一致成立,结合离散 Euler-Maruyama 逼近分析与 McKean-Vlasov 过程的适定性,作者证明了当 \(\min(n, L) \to \infty\) 时,隐藏状态与梯度坐标联合收敛至相互独立布朗运动驱动的前向-反向 SDE 极限 \((H_t, G_t)\),总误差速率为 \(\mathcal{O}(n^{-1} + L^{-1})\)。这在数学上确立了 ResNet 在初始化阶段宽度极限与深度极限的可交换性。

3. 训练步前向反向耦合与高阶深度压缩:揭示深度抵消训练耦合的内在机制 网络经过一步 SGD 优化后,前向特征包含了沿上一步反向梯度方向的参数位移:\(\bm{h}_\ell^{(1)}\) 依赖于 \(\bm{W}_\ell^{(0)} \bm{x}_{\ell-1}^{(1)} - \frac{\eta_c}{L} \chi^{(0)} \frac{\langle \bm{x}_{\ell-1}^{(0)}, \bm{x}_{\ell-1}^{(1)} \rangle}{n} \bm{g}_\ell^{(0)}\)。由于第一步更新后的特征 \(\bm{x}_{\ell-1}^{(1)}\) 已嵌入了第一步反向梯度的信息,再次乘以含有 \(\bm{W}_\ell^{(0)}\) 的权重时,会在内部产生形如 \(\bm{W}_\ell^{(0)} (\bm{W}_\ell^{(0)})^\top \bm{g}_\ell^{(0)}\) 的非对角收缩。基于张量程序主定理(Tensor Programs Master Theorem)推导表明,当仅取宽度极限 \(n \to \infty\) 时,真实动力学与解耦反向权重的动力学不再一致,而是存留了一个确定性的训练耦合修正项: $\(\mathcal{C}_\ell^{(0 \to 1)} = -\frac{\eta_c}{L^2} \chi^{(0)} \mathbb{E}\left[\phi(H_{\ell-2}^{(0)}) \phi(H_{\ell-2}^{(1)})\right] \mathbb{E}\left[\phi'(H_{\ell-1}^{(0)}) \phi'(H_{\ell-1}^{(1)})\right] G_\ell^{(0)}\)$ 这项关键推导表明,单纯增加网络宽度无法消除训练时引入的权重复用耦合。然而,论文揭示了深度缩放带来的奇迹:该存留耦合项带有 \(1/L^2\) 的前置系数,而主 SGD 漂移项为 \(\mathcal{O}(L^{-1})\)。在全网络 \(L\) 层累积之后,总耦合误差为 \(\mathcal{O}(L^{-1})\),其平方二阶矩更是以 \(\mathcal{O}(L^{-2})\) 快速衰减。也就是说,在 depth-\(\mu\mathrm{P}\) 下,深度的增加自动抑制了权重复用耦合。

4. 神经特征动力学 (NFD) 极限收敛:无限深度特征演化的闭合 SDE 系统 基于深度的天然抑制效应,作者正式提出了神经特征动力学(Neural Feature Dynamics, NFD)。NFD 丢弃了高阶的权重复用耦合项(等价于将反向传播的权重矩阵替换为独立副本),但完整保留了跨训练步的特征内积核 \(\Sigma_t^{(i, k)} = \mathbb{E}[\phi(H_t^{(i)}) \phi(H_t^{(k)})]\) 与梯度协方差核 \(\Theta_t^{(i, k)} = \mathbb{E}[G_t^{(i)} G_t^{(k)}]\)。在假设协方差矩阵一致非退化(最小特征值 \(\lambda_{\min} \geq \varepsilon > 0\))的条件下,论文严格证明了对于任意固定训练步长 \(K < \infty\): $\(dH_t^{(k)} = -\eta_c \sum_{i=0}^{k-1} \mathring{\chi}^{(i)} \Sigma_t^{(i, k)} G_t^{(i)} dt + dW_t^{(k)}\)$ $\(dG_t^{(k)} = -\eta_c \sum_{i=0}^{k-1} \mathring{\chi}^{(i)} \Theta_t^{(i, k)} \phi(H_t^{(i)}) \phi'(H_t^{(k)}) dt + \phi'(H_t^{(k)}) d\widetilde{W}_t^{(k)}\)$ 有限网络训练动力学向该 NFD 系统的连续极限收敛,其主导误差仅为来自 SDE 离散化误差的 \(\mathcal{O}(L^{-1})\),而被忽略的权重复用累积耦合项的二阶矩衰减率达 \(\mathcal{O}(L^{-2})\),从而为深度学习特征演化提供了严谨的数学极限描述。

损失函数 / 训练策略

模型采用基于流式 SGD 的在线训练设定,每步迭代样本为 \((\bm{x}^{(k)}, y^{(k)})\),损失梯度标量为 \(\chi^{(k)} = \partial_f \mathcal{L}(f^{(k)}, y^{(k)})\)。优化学习率采用符合 depth-\(\mu\mathrm{P}\) 原则的 \(\eta = \eta_c n\)(其中 \(\eta_c = 0.01\) 为与网络宽高无关的常数)。网络结构采用 Pre-activation 设计,有效杜绝了 Post-activation 设计下隐藏状态范数随深度 \(\mathcal{O}((1 + c_1/\sqrt{Ln})^L)\) 发生指数级爆炸的发散问题(命题 3.1)。

实验关键数据

主实验

为了验证理论推导的 GIA 恢复现象(即真实耦合训练与解耦反向权重训练的一致性),作者在 CIFAR-10 上对三种模型族(Vanilla DNN、标准 \(\mu\mathrm{P}\) ResNet、depth-\(\mu\mathrm{P}\) ResNet)在宽度 \(n \in \{128, 256\}\) 和不同深度 \(L\) 下进行了系统对比评测。

架构类型 深度配置 (\(L\)) 宽度 (\(n\)) 标准训练测试损失 (对数尺度) 解耦训练测试损失 (对数尺度) 轨迹对齐与收敛表现
Vanilla DNN 2 \(\to\) 16 128 停滞在较高损失 (~2.30) 停滞在更高损失 (~2.30) 遭遇严重梯度消失,标准与解耦轨迹持续显著失配
Vanilla DNN 2 \(\to\) 16 256 损失略微改善 (~2.15) 严重失配偏离 增加宽度无法解决深层梯度消失与轨迹分裂问题
\(\mu\mathrm{P}\) ResNet 2 \(\to\) 4 128 稳步下降 (~1.65) 出现偏移 (~1.80) 浅层有一定进展但两线明显未对齐
\(\mu\mathrm{P}\) ResNet 8 \(\to\) 16 128 泛化崩溃/严重过拟合 显著失配且不稳定 缺少深度缩放,随着层数增加动力学极不稳定
\(\mu\mathrm{P}\) ResNet 16 256 不稳定震荡 依然明显分裂 扩大宽度无法修复缺乏深度缩放带来的结构性失调
depth-\(\mu\mathrm{P}\) (本文) 2 128 初始有微弱间隙 初始有微弱间隙 浅层时耦合项有轻微可见影响
depth-\(\mu\mathrm{P}\) (本文) 8 128 快速收敛下降 (~1.25) 高度重合贴合 (~1.26) 随深度提升训练与测试性能双升,轨迹极其接近
depth-\(\mu\mathrm{P}\) (本文) 16 128 最优收敛 (~1.08) 几乎完全重合 (~1.08) 验证推论 5.7:耦合项以 \(\mathcal{O}(L^{-2})\) 衰减,动力学几无二致
depth-\(\mu\mathrm{P}\) (本文) 16 256 最优收敛平滑曲线 完美对齐贴合 宽深结合表现出极佳的稳定性和泛化收益

消融实验

论文围绕核心理论假设(Assumption 5.4 协方差非退化性)以及 Pre-activation 对比 Post-activation 架构进行了严格的实证消融分析。

架构 / 宽度配置 评测维度 关键指标取值 / 表现 说明
Pre-activation (\(L=64\)) 隐藏状态范数稳定性 范数稳定保持在 \(\mathcal{O}(1)\) 区间 随层数递增特征范数保持平稳无爆炸
Post-activation (\(L=64\)) 隐藏状态范数稳定性 范数随深度快速爆炸式增长 (\(>10^3\)) 证实命题 3.1:正支配激活引发指数范数发散
Pre-activation (\(L=64\)) 训练与测试损失 训练损失快速降至 ~0.2,测试损失平稳方差极小 训练极其平稳,不同随机种子方差显著更低
Post-activation (\(L=64\)) 训练与测试损失 训练损失震荡发散,测试损失方差巨大 优化极其不稳定,深层训练性能显著劣化
ResNet (\(n=512\), \(L=4\)) 协方差最小特征值 \(\inf_t \lambda_{\min}(\bm{\Sigma}_t^{(k)}) \approx 0.002\) 宽度较窄时接近近退化边界,但仍保持严格正值
ResNet (\(n=1024\), \(L=4\)) 协方差最小特征值 \(\inf_t \lambda_{\min}(\bm{\Sigma}_t^{(k)}) \approx 0.015\) 特征值显著提升,支撑了非退化假设 5.4
ResNet (\(n=4096\), \(L=4\)) 协方差最小特征值 \(\inf_t \lambda_{\min}(\bm{\Sigma}_t^{(k)}) \approx 0.050\) 大宽度下特征矩阵远离奇异,SDE 解严格适定

关键发现

  • GIA 恢复现象在深度维度显著生效:在 CIFAR-10 实验中,depth-\(\mu\mathrm{P}\) ResNet 展现出了与其他架构截然不同的特性——随着深度 \(L\) 从 2 增加到 16,真实训练曲线与反向解耦权重曲线迅速靠拢贴合,充分证实了理论推导中耦合项以 \(\mathcal{O}(L^{-2})\) 快速衰减的预测。
  • 宽度与深度的误差收敛阶:在初始化与 30 轮训练后,真实网络与 NFD 极限的经验误差均严格遵从 \(\mathcal{O}(n^{-1} + L^{-1})\) 的衰减规律,表明在 depth-\(\mu\mathrm{P}\) 下宽度极限与深度极限具备优良的弱收敛一致性。
  • Pre-activation 的数学不可替代性:对于 ReLU 这类正支配激活函数,Post-activation 结构即使在 \(1/\sqrt{L}\) 缩放下也会产生指数级范数增长,而 Pre-activation 保证了残差流漂移项的有界性与平稳性。

亮点与洞察

  • 条件高斯表示显式隔离权重复用:以往理论多将反向权重复用视为黑箱或直接假设独立,本文首次利用条件正态分布将耦合精确分离为一个沿前向方向的有偏投影和一个正交高斯创新,使有限宽高下的误差可被定量测算。
  • 深度抵消训练耦合的反直觉理论洞察:直觉上层数越深,反向传播经过的转置权重越多,积累的相互依赖应越严重;但本文证明由于 depth-\(\mu\mathrm{P}\) 的残差分支带 \(1/\sqrt{L}\) 缩放,权重复用导致的单层收缩具有 \(L^{-2}\) 前置因子,最终使得深层网络中的耦合被深度本身所稀释。
  • NFD 奠定特征学习连续极限的新基石:提出了一套保留动态协方差核但解耦反向权重的闭合 SDE 系统(NFD),为后续研究无限深度网络中的特征演化、损失曲面与泛化误差提供了清晰的分析对象。

局限与展望

  • 模型与优化器范围局限:理论严格证明目前基于单层 Pre-activation 残差块和单样本 Stream SGD,尚未覆盖含有多层瓶颈残差块(Bottleneck)、多头自注意力(Transformer)、Mini-batch 批训练以及 AdamW 等自适应优化算法。
  • 有限时间步与顺次极限假设:收敛定理 5.6 依赖于固定训练迭代步数 \(K < \infty\) 和协方差矩阵一致非退化假设,且采用的是先取无限宽度再取无限深度的顺次极限,联合极限 \((n, L) \to \infty\) 的训练期动态证明仍留待未来探索。
  • 耦合消失对泛化的双刃剑效应尚待探明:论文证明了 depth-\(\mu\mathrm{P}\) 能有效抑制权重复用耦合,但完全消除该耦合是否必然有利于下游泛化与表征质量,抑或适当保留前向-反向自发对齐能带来潜在归纳偏置,仍是开放性问题。

相关工作与启发

  • vs Tensor Programs (\(\mu\mathrm{P}\)) 系列:Yang 等人的 TP 框架奠定了无限宽度特征学习的代数工具,而 TP VI 探讨了深度维度的特征学习;本文在此基础上,进一步深入到反向权重复用引起的显式耦合核结构,给出了具体的渐近阶数与闭合 SDE 表述。
  • vs 动态均值场理论 (DMFT):Bordelon 等人基于统计物理推导的 DMFT 引入了复杂的积分响应函数来隐含处理依赖;本文借助概率论中的条件高斯表示与 McKean-Vlasov SDE,直接给出了严谨的有限层截断误差界与连续极限。
  • vs 懒惰训练 / NTK 体系:NTK 在超大宽度极限下使特征冻结在初始化附近;而本文的 NFD 处于真实的非平凡特征学习体系,特征协方差核随 SGD 更新动态演进,真实刻画了表征演化。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次使用条件高斯表示显式隔离了反向传播中的权重复用耦合,并发现了深度抑制该耦合的精确数学阶。
  • 实验充分度: ⭐⭐⭐⭐ 理论推导完备扎实,针对 CIFAR-10 的模拟验证与特征值监测有力支撑了定理结论,但缺乏超大规模模型实证。
  • 写作质量: ⭐⭐⭐⭐⭐ 结构逻辑环环相扣,从初始化到单步训练再到连续 SDE 推进自然,数学表述极其精准严密。
  • 价值: ⭐⭐⭐⭐⭐ 攻克了深度残差网络特征学习动力学难以建立闭合严格极限的核心瓶颈,是深层学习理论的重要里程碑。