梯度平滑：耦合逐层更新以改进优化

# 基于层间更新耦合的优化改进 Source: https://arxiv.org/html/2606.30813 ## 梯度平滑：通过耦合层间更新改进优化 ###### 摘要 具有重复架构模块（如Transformer）的深度神经网络，在训练过程中往往会在各层之间涌现出结构化关系。受此观察启发，我们提出了*深度梯度增强*（Depth-wise Gradient Augmentation）——一种通用的优化范式，其中应用于每一层的更新是通过沿深度维度对逐块优化器更新集合进行变换得到的。在此框架内，我们研究了*梯度平滑*（Gradient Smoothing）这一深度方向平滑方法家族，并通过一个简单的局部*窗口平滑*（Window Smoothing）算子对其进行实例化。该方法直接作用于任意基础优化器（如SGD、Adam、Muon）生成的逐块更新，计算开销极小，且与现有优化流程兼容。我们在多种架构和训练场景下评估了梯度平滑，包括语言模型预训练、大语言模型推理的强化学习后训练、扩散建模以及视觉Transformer的图像分类。在这些设置中，梯度平滑在不修改模型架构或训练目标的情况下，持续改善了优化和泛化性能。我们进一步证明，它促进了表示在深度方向上的更结构化演化，这与将其解释为结构化深度预处理方法是一致的。这些结果共同确立了深度梯度增强作为利用跨深度结构优化的一个有前景的框架，并展示了梯度平滑作为一种简单且广泛适用的实例化方法。 Machine Learning, ICML ## 1 引言 许多现代深度神经网络由具有共享计算结构的重复架构模块组成，例如残差网络（ResNets）(He et al., 2015 (https://arxiv.org/html/2606.30813#bib.bib31))和Transformer (Vaswani et al., 2023 (https://arxiv.org/html/2606.30813#bib.bib33))。最近的研究表明，这种形式的网络在训练过程中往往会在深度方向上涌现出强烈的结构规律性。特别是，关于*Transformer模块耦合* (Aubry et al., 2025 (https://arxiv.org/html/2606.30813#bib.bib29))和*残差对齐* (Li and Papyan, 2024 (https://arxiv.org/html/2606.30813#bib.bib30))的研究表明，模块雅可比矩阵和残差表示的奇异向量在各层之间变得对齐，暗示着深度方向上的某种隐式协调。相关证据还出现在大语言模型中关于层相似性和剪枝的研究中 (Gromov et al., 2025 (https://arxiv.org/html/2606.30813#bib.bib15))，以及显示信息论量在深度上平滑演化的分析 (Skean et al., 2024 (https://arxiv.org/html/2606.30813#bib.bib186))。这些观察表明，深度网络中的重复模块并非独立演化，而是表现出跨深度的协调表示和动态结构。然而，标准优化方法通常独立地构造和应用逐块优化器更新（除通过前向和反向传播产生的依赖关系外），并未显式利用层间关系。这自然引出一个问题：*优化器更新本身能否被增强，以利用跨深度涌现的结构？*

在本文中，我们引入了一个通用范式——*深度梯度增强*。我们并非孤立地应用每个模块产生的优化器更新，而是通过对逐块优化器更新集合应用一个深度增强算子来获得应用于每一层的更新。这一视角涵盖了沿深度维度作用的一大类更新变换。作为该框架的第一个实例化，我们研究了*梯度平滑*，其动机来源于经验观察：相邻层通常学习相关的表示，并在训练过程中表现出越来越对齐的优化动态。我们聚焦于一个特别简单的实现——*窗口平滑*，它通过平均相邻层局部窗口内的优化器更新来增强每个模块的更新。窗口平滑实现简单，计算开销极小，并且与任意基础优化器兼容。此外，梯度平滑可以被自然地解释为一种作用于块结构参数空间的结构化预处理方法，同时为探索更丰富的深度更新增强形式提供了基础。

#### 贡献。
1. 1. 我们提出了*深度梯度增强*，一个通用的优化范式，其中逐块优化器更新的集合在应用之前先被一个深度增强算子变换。在此框架内，我们研究了*梯度平滑*，一个结构化的深度方向平滑方法家族，并通过一个简单的、与优化器无关的*窗口平滑*算子对其进行实例化，为探索更一般的深度更新增强形式提供了基础。
2. 2. 我们证明，梯度平滑在多种架构和训练场景下持续改善了优化和泛化性能，包括语言模型预训练、大语言模型推理的强化学习微调、视觉Transformer的监督学习，以及扩散模型。
3. 3. 我们提供了经验性和理论性的证据，表明梯度平滑促进了表示在深度方向上的更结构化演化。经验上，我们观察到层间轨迹对齐和表示残差相似性增加；理论上，我们刻画了平滑如何影响表示残差结构。

## 2 背景与设置

### 2.1 块结构优化

我们考虑由L个重复模块通过跳跃连接组成的神经网络，每个模块具有相同（或相似）的函数形式，但拥有自己的参数。此外，模型可能包含一小部分与重复结构无关的参数，例如输入嵌入或输出头。形式化地，令F: R^d × R^p → R^d表示单个模块的残差映射。对于模块ℓ，残差函数为F(·; θ_ℓ)，其中θ_ℓ ∈ R^p是层特定参数。给定初始嵌入h_0，隐藏状态(h_ℓ)_{ℓ=1}^L ⊂ R^d的演化遵循：
h_{ℓ+1} = h_ℓ + F(h_ℓ; θ_ℓ), ℓ = 1, ..., L-1。 (1)
尽管我们为标准的残差连接给出了公式，但该框架自然扩展到具有更一般跳跃连接模式的架构。我们将完整参数向量写作 (θ, φ)，
θ = (θ_1, ..., θ_L), θ_l ∈ R^p，
其中 φ ∈ R^{D_φ} 收集所有非模块参数。这一设置涵盖了一大类现代架构，包括具有重复块结构的残差网络、视觉Transformer和基于Transformer的语言模型（通过令每个h_ℓ为一个向量序列）。我们用L(θ, φ)表示训练目标，并假设其可微。在优化步骤t，逐块梯度为：
g_l^{(t)} := ∇_{θ_l} L(θ^{(t)}, φ^{(t)}), g^{(t)} := (g_1^{(t)}, ..., g_L^{(t)})，
并用 g_φ^{(t)} := ∇_φ L(θ^{(t)}, φ^{(t)}) 表示关于φ的梯度。标准的一阶方法独立处理逐块梯度 {g_l^{(t)}}_{l=1}^L。然而，重复块架构在深度方向诱导出强烈的结构关系：相邻块通常学习相似的变换，并表现出对齐的表示，特别是在训练的初始瞬态阶段之后。这表明堆叠的梯度向量 g^{(t)} 在深度方向上具有有意义的结构。

f_L  f_3  f_2  f_1  ⋮  U(∇_{θ_L}L)  U(∇_{θ_3}L)  U(∇_{θ_2}L)  U(∇_{θ_1}L)  ⋮  滤波器 S  Ũ(∇_{θ_L}L)  Ũ(∇_{θ_3}L)  Ũ(∇_{θ_2}L)  Ũ(∇_{θ_1}L)  ⋮
图1: 梯度平滑。在深度方向对深度网络中反向传播后的更新应用梯度增强方案的示意图。对于一个具有L个相同架构模块（但参数不同）的深度网络，每个参数块θ^l中的梯度更新在深度方向重新加权，以稳定信息传播。

(a) 深度24: 验证损失。 (b) 深度24: 验证BPB。 (c) 深度24: CORE指标。 (d) 深度30: 验证损失。 (e) 深度30: 验证BPB。 (f) 深度30: CORE指标。
图2: 使用梯度平滑的Nanochat预训练。验证损失、BPB和DCLM CORE指标，对比默认Adam + NorMuon优化器设置（表1 (https://arxiv.org/html/2606.30813#S4.T1)）的基线和使用相同超参数对Transformer模块应用深度更新平滑的运行。深度24（1.38B参数，上行）和深度30（2.40B参数，下行）的模型，α=0.1，针对seed ∈ {42,67,2026}重复基线。在两种深度下，平滑加速了验证损失的收敛，并在训练期间改善了CORE指标，在深度30时差距更明显。

### 2.2 通用一阶更新

我们不局限于特定优化器，而是考虑一个通用的一阶更新规则：
(θ^{(t+1)}, φ^{(t+1)}) = (θ^{(t)}, φ^{(t)}) − η (u^{(t)}, u_φ^{(t)})，
其中
u^{(t)} := U^{(t)}(g^{(t)}), u_φ^{(t)} := U_φ^{(t)}(g_φ^{(t)})
分别表示基础优化器对模块和非模块参数产生的更新。这一公式包括梯度下降（u^{(t)}=g^{(t)}）、动量方法、自适应优化器如Adam或AdamW，以及Muon和相关变体。重要的是，u^{(t)}可能依赖于当前梯度和优化器内部状态，并且可能涉及对g^{(t)}的非线性或时变变换。这种与优化器无关的视角允许我们直接在优化器更新层面上研究更新增强，而不依赖于构造这些更新的具体机制。因此，该框架可以作为对现有广泛优化器类的轻量级、模块化扩展而被集成。

## 3 深度梯度增强

### 3.1 通用框架

由重复架构模块构建的网络在训练过程中在其表示中发展出强烈的跨深度结构。受此观察启发，我们引入了*深度梯度增强*，一个通用的优化范式，其中应用于每一层的更新是通过沿深度轴对逐块优化器更新集合进行变换得到的。这一视角提供了一个统一框架，将跨深度信息融入优化器更新，同时保持与所用特定变换无关。

在迭代t，令
u_ℓ^{(t)} ∈ R^p, ℓ=1,...,L，
表示基础优化器（如SGD、Adam、Muon）对层ℓ参数产生的逐块更新，并记
u^{(t)} = (u_1^{(t)}, ..., u_L^{(t)}) ∈ (R^p)^L，
以及相应的非模块参数更新u_φ^{(t)}。在其最一般的形式中，深度梯度增强对每个块内坐标独立地沿深度方向应用一个算子。对于每个坐标 i = 1,...,p，令
C^{(i)}: R^L → R^L，
为一个深度增强算子。那么
(ũ_1^{(i)}, ..., ũ_L^{(i)}) = C^{(i)}(u_1^{(i)}, ..., u_L^{(i)})。
等价地，
ũ_ℓ^{(i)} = C_ℓ^{(i)}(u_1^{(i)}, ..., u_L^{(i)}), ℓ=1,...,L。
因此，层ℓ更新的每个坐标可能仅依赖于各层对应坐标，同时独立于所有其他块内坐标。在众多可能的深度增强算子选择中，我们专注于那些在深度方向上平滑优化器更新的算子。

### 3.2 梯度平滑

通过选择坐标级深度增强算子为一个共同的线性平滑算子，我们可以得到深度梯度增强方法的一个自然类别。我们将这个方法家族称为*梯度平滑*。具体地，我们取
C^{(i)} = S, i=1,...,p，
其中 S ∈ R^{L×L} 是一个作用于层索引的线性算子。等价地，梯度平滑由块结构算子表示：
P := S ⊗ I ∈ R^{D_θ × D_θ}，
其中 I 表示块内参数空间 R^p 上的恒等算子。变换后的逐块更新为：
ũ^{(t)} := P u^{(t)} = (S ⊗ I) u^{(t)}。

为了描述完整参数向量 (θ, φ) 上的更新，我们将 P 提升为块对角算子：
\(\bar{P} := \begin{pmatrix} P & 0 \\ 0 & I_\phi \end{pmatrix} \in \mathbb{R}^{(D_\theta + D_\phi) \times (D_\theta + D_\phi)}\)，
其中 I_φ 表示 φ 坐标上的恒等算子。给定完整优化器更新
u_full^{(t)} := (u^{(t)}, u_φ^{(t)})，
梯度平滑应用 \(\bar{P}\) 得到
ũ_full^{(t)} = \(\bar{P}\) u_full^{(t)} = (P u^{(t)}, u_φ^{(t)})。# 基于层间更新耦合的优化改进 Source: https://arxiv.org/html/2606.30813 ## 梯度平滑：通过耦合层间更新改进优化 ###### 摘要 具有重复架构模块（如Transformer）的深度神经网络，在训练过程中往往会在各层之间涌现出结构化关系。受此观察启发，我们提出了*深度梯度增强*（Depth-wise Gradient Augmentation）——一种通用的优化范式，其中应用于每一层的更新是通过沿深度维度对逐块优化器更新集合进行变换得到的。在此框架内，我们研究了*梯度平滑*（Gradient Smoothing）这一深度方向平滑方法家族，并通过一个简单的局部*窗口平滑*（Window Smoothing）算子对其进行实例化。该方法直接作用于任意基础优化器（如SGD、Adam、Muon）生成的逐块更新，计算开销极小，且与现有优化流程兼容。我们在多种架构和训练场景下评估了梯度平滑，包括语言模型预训练、大语言模型推理的强化学习后训练、扩散建模以及视觉Transformer的图像分类。在这些设置中，梯度平滑在不修改模型架构或训练目标的情况下，持续改善了优化和泛化性能。我们进一步证明，它促进了表示在深度方向上的更结构化演化，这与将其解释为结构化深度预处理方法是一致的。这些结果共同确立了深度梯度增强作为利用跨深度结构优化的一个有前景的框架，并展示了梯度平滑作为一种简单且广泛适用的实例化方法。 Machine Learning, ICML ## 1 引言 许多现代深度神经网络由具有共享计算结构的重复架构模块组成，例如残差网络（ResNets）(He et al., 2015 (https://arxiv.org/html/2606.30813#bib.bib31))和Transformer (Vaswani et al., 2023 (https://arxiv.org/html/2606.30813#bib.bib33))。最近的研究表明，这种形式的网络在训练过程中往往会在深度方向上涌现出强烈的结构规律性。特别是，关于*Transformer模块耦合* (Aubry et al., 2025 (https://arxiv.org/html/2606.30813#bib.bib29))和*残差对齐* (Li and Papyan, 2024 (https://arxiv.org/html/2606.30813#bib.bib30))的研究表明，模块雅可比矩阵和残差表示的奇异向量在各层之间变得对齐，暗示着深度方向上的某种隐式协调。相关证据还出现在大语言模型中关于层相似性和剪枝的研究中 (Gromov et al., 2025 (https://arxiv.org/html/2606.30813#bib.bib15))，以及显示信息论量在深度上平滑演化的分析 (Skean et al., 2024 (https://arxiv.org/html/2606.30813#bib.bib186))。这些观察表明，深度网络中的重复模块并非独立演化，而是表现出跨深度的协调表示和动态结构。然而，标准优化方法通常独立地构造和应用逐块优化器更新（除通过前向和反向传播产生的依赖关系外），并未显式利用层间关系。这自然引出一个问题：*优化器更新本身能否被增强，以利用跨深度涌现的结构？*

在本文中，我们引入了一个通用范式——*深度梯度增强*。我们并非孤立地应用每个模块产生的优化器更新，而是通过对逐块优化器更新集合应用一个深度增强算子来获得应用于每一层的更新。这一视角涵盖了沿深度维度作用的一大类更新变换。作为该框架的第一个实例化，我们研究了*梯度平滑*，其动机来源于经验观察：相邻层通常学习相关的表示，并在训练过程中表现出越来越对齐的优化动态。我们聚焦于一个特别简单的实现——*窗口平滑*，它通过平均相邻层局部窗口内的优化器更新来增强每个模块的更新。窗口平滑实现简单，计算开销极小，并且与任意基础优化器兼容。此外，梯度平滑可以被自然地解释为一种作用于块结构参数空间的结构化预处理方法，同时为探索更丰富的深度更新增强形式提供了基础。

#### 贡献。
1. 1. 我们提出了*深度梯度增强*，一个通用的优化范式，其中逐块优化器更新的集合在应用之前先被一个深度增强算子变换。在此框架内，我们研究了*梯度平滑*，一个结构化的深度方向平滑方法家族，并通过一个简单的、与优化器无关的*窗口平滑*算子对其进行实例化，为探索更一般的深度更新增强形式提供了基础。
2. 2. 我们证明，梯度平滑在多种架构和训练场景下持续改善了优化和泛化性能，包括语言模型预训练、大语言模型推理的强化学习微调、视觉Transformer的监督学习，以及扩散模型。
3. 3. 我们提供了经验性和理论性的证据，表明梯度平滑促进了表示在深度方向上的更结构化演化。经验上，我们观察到层间轨迹对齐和表示残差相似性增加；理论上，我们刻画了平滑如何影响表示残差结构。

## 2 背景与设置

### 2.1 块结构优化

我们考虑由L个重复模块通过跳跃连接组成的神经网络，每个模块具有相同（或相似）的函数形式，但拥有自己的参数。此外，模型可能包含一小部分与重复结构无关的参数，例如输入嵌入或输出头。形式化地，令F: R^d × R^p → R^d表示单个模块的残差映射。对于模块ℓ，残差函数为F(·; θ_ℓ)，其中θ_ℓ ∈ R^p是层特定参数。给定初始嵌入h_0，隐藏状态(h_ℓ)_{ℓ=1}^L ⊂ R^d的演化遵循：
h_{ℓ+1} = h_ℓ + F(h_ℓ; θ_ℓ), ℓ = 1, ..., L-1。 (1)
尽管我们为标准的残差连接给出了公式，但该框架自然扩展到具有更一般跳跃连接模式的架构。我们将完整参数向量写作 (θ, φ)，
θ = (θ_1, ..., θ_L), θ_l ∈ R^p，
其中 φ ∈ R^{D_φ} 收集所有非模块参数。这一设置涵盖了一大类现代架构，包括具有重复块结构的残差网络、视觉Transformer和基于Transformer的语言模型（通过令每个h_ℓ为一个向量序列）。我们用L(θ, φ)表示训练目标，并假设其可微。在优化步骤t，逐块梯度为：
g_l^{(t)} := ∇_{θ_l} L(θ^{(t)}, φ^{(t)}), g^{(t)} := (g_1^{(t)}, ..., g_L^{(t)})，
并用 g_φ^{(t)} := ∇_φ L(θ^{(t)}, φ^{(t)}) 表示关于φ的梯度。标准的一阶方法独立处理逐块梯度 {g_l^{(t)}}_{l=1}^L。然而，重复块架构在深度方向诱导出强烈的结构关系：相邻块通常学习相似的变换，并表现出对齐的表示，特别是在训练的初始瞬态阶段之后。这表明堆叠的梯度向量 g^{(t)} 在深度方向上具有有意义的结构。

f_L  f_3  f_2  f_1  ⋮  U(∇_{θ_L}L)  U(∇_{θ_3}L)  U(∇_{θ_2}L)  U(∇_{θ_1}L)  ⋮  滤波器 S  Ũ(∇_{θ_L}L)  Ũ(∇_{θ_3}L)  Ũ(∇_{θ_2}L)  Ũ(∇_{θ_1}L)  ⋮
图1: 梯度平滑。在深度方向对深度网络中反向传播后的更新应用梯度增强方案的示意图。对于一个具有L个相同架构模块（但参数不同）的深度网络，每个参数块θ^l中的梯度更新在深度方向重新加权，以稳定信息传播。

(a) 深度24: 验证损失。 (b) 深度24: 验证BPB。 (c) 深度24: CORE指标。 (d) 深度30: 验证损失。 (e) 深度30: 验证BPB。 (f) 深度30: CORE指标。
图2: 使用梯度平滑的Nanochat预训练。验证损失、BPB和DCLM CORE指标，对比默认Adam + NorMuon优化器设置（表1 (https://arxiv.org/html/2606.30813#S4.T1)）的基线和使用相同超参数对Transformer模块应用深度更新平滑的运行。深度24（1.38B参数，上行）和深度30（2.40B参数，下行）的模型，α=0.1，针对seed ∈ {42,67,2026}重复基线。在两种深度下，平滑加速了验证损失的收敛，并在训练期间改善了CORE指标，在深度30时差距更明显。

### 2.2 通用一阶更新

我们不局限于特定优化器，而是考虑一个通用的一阶更新规则：
(θ^{(t+1)}, φ^{(t+1)}) = (θ^{(t)}, φ^{(t)}) − η (u^{(t)}, u_φ^{(t)})，
其中
u^{(t)} := U^{(t)}(g^{(t)}), u_φ^{(t)} := U_φ^{(t)}(g_φ^{(t)})
分别表示基础优化器对模块和非模块参数产生的更新。这一公式包括梯度下降（u^{(t)}=g^{(t)}）、动量方法、自适应优化器如Adam或AdamW，以及Muon和相关变体。重要的是，u^{(t)}可能依赖于当前梯度和优化器内部状态，并且可能涉及对g^{(t)}的非线性或时变变换。这种与优化器无关的视角允许我们直接在优化器更新层面上研究更新增强，而不依赖于构造这些更新的具体机制。因此，该框架可以作为对现有广泛优化器类的轻量级、模块化扩展而被集成。

## 3 深度梯度增强

### 3.1 通用框架

由重复架构模块构建的网络在训练过程中在其表示中发展出强烈的跨深度结构。受此观察启发，我们引入了*深度梯度增强*，一个通用的优化范式，其中应用于每一层的更新是通过沿深度轴对逐块优化器更新集合进行变换得到的。这一视角提供了一个统一框架，将跨深度信息融入优化器更新，同时保持与所用特定变换无关。

在迭代t，令
u_ℓ^{(t)} ∈ R^p, ℓ=1,...,L，
表示基础优化器（如SGD、Adam、Muon）对层ℓ参数产生的逐块更新，并记
u^{(t)} = (u_1^{(t)}, ..., u_L^{(t)}) ∈ (R^p)^L，
以及相应的非模块参数更新u_φ^{(t)}。在其最一般的形式中，深度梯度增强对每个块内坐标独立地沿深度方向应用一个算子。对于每个坐标 i = 1,...,p，令
C^{(i)}: R^L → R^L，
为一个深度增强算子。那么
(ũ_1^{(i)}, ..., ũ_L^{(i)}) = C^{(i)}(u_1^{(i)}, ..., u_L^{(i)})。
等价地，
ũ_ℓ^{(i)} = C_ℓ^{(i)}(u_1^{(i)}, ..., u_L^{(i)}), ℓ=1,...,L。
因此，层ℓ更新的每个坐标可能仅依赖于各层对应坐标，同时独立于所有其他块内坐标。在众多可能的深度增强算子选择中，我们专注于那些在深度方向上平滑优化器更新的算子。

### 3.2 梯度平滑

通过选择坐标级深度增强算子为一个共同的线性平滑算子，我们可以得到深度梯度增强方法的一个自然类别。我们将这个方法家族称为*梯度平滑*。具体地，我们取
C^{(i)} = S, i=1,...,p，
其中 S ∈ R^{L×L} 是一个作用于层索引的线性算子。等价地，梯度平滑由块结构算子表示：
P := S ⊗ I ∈ R^{D_θ × D_θ}，
其中 I 表示块内参数空间 R^p 上的恒等算子。变换后的逐块更新为：
ũ^{(t)} := P u^{(t)} = (S ⊗ I) u^{(t)}。

为了描述完整参数向量 (θ, φ) 上的更新，我们将 P 提升为块对角算子：
\(\bar{P} := \begin{pmatrix} P & 0 \\ 0 & I_\phi \end{pmatrix} \in \mathbb{R}^{(D_\theta + D_\phi) \times (D_\theta + D_\phi)}\)，
其中 I_φ 表示 φ 坐标上的恒等算子。给定完整优化器更新
u_full^{(t)} := (u^{(t)}, u_φ^{(t)})，
梯度平滑应用 \(\bar{P}\) 得到
ũ_full^{(t)} = \(\bar{P}\) u_full^{(t)} = (P u^{(t)}, u_φ^{(t)})。