技术摘要：SparseOpt——解决稀疏训练中由归一化引起的梯度偏斜

问题陈述

动态稀疏训练（DST）方法旨在通过在训练过程中动态调整网络拓扑（掩码），以高稀疏度训练神经网络。虽然 DST 有望降低计算成本和内存占用，但它存在一个关键的实践局限性：与稠密训练相比，其收敛速度显著更慢。现有的 DST 方法往往需要多达五倍的训练轮次（epochs）才能达到与稠密模型相当的泛化性能，从而抵消了潜在的效率增益。

作者指出，批归一化（Batch Normalization, BN） 是导致这种不稳定的主要且此前未被探索的原因。虽然 BN 在稠密网络中是稳定训练和加速收敛的标准组件，但本文证明，在稀疏网络中，BN 会引入梯度偏斜。具体而言，由于稀疏层中的神经元具有异质连接性（传入连接数量不同，即fan-in不同），BN 基于局部方差对输入进行归一化。这种方差直接取决于活跃连接的数量。因此，BN 会根据每个稀疏神经元各自的稀疏度水平，对梯度进行重缩放，导致层内梯度缩放不均匀。

在 DST 的语境下，掩码（从而单个神经元的稀疏度）在掩码更新步骤中会发生突变，这种由 BN 引起的重缩放会导致梯度的方向和幅度发生突发的、依赖于神经元的偏移。这破坏了优化轨迹，引入了噪声和不稳定性，从而减缓了收敛速度。

方法论：SparseOpt

为解决这一问题，作者提出了 SparseOpt，一种感知稀疏性的预条件梯度下降优化器。其核心洞察是：BN 引入的梯度失真纯粹是乘性的，并且取决于神经元级的稀疏度 $s_i$。

理论基础

作者推导得出，对于稀疏度为 $s_i$ 的神经元 $i$，在带有 BN 的稀疏层中，其梯度 $\frac{\partial L}{\partial x'_i}$ 相对于等效稠密层中的梯度，会被放大 $(1-s_i)^{-1/2}$ 倍。
$$\frac{\partial L}{\partial x'_i} = \frac{1}{\sqrt{1-s_i}} \frac{\partial L}{\partial x_i}$$
这种放大效应因神经元而异，改变了梯度分量的相对贡献，并在参数空间中旋转了更新方向。

SparseOpt 更新规则

SparseOpt 通过应用一个对角预条件器 $D$ 来校正这种偏斜，该预条件器对梯度进行重缩放以恢复均匀性。更新规则定义如下：
$$w_{t+1} = w_t - \frac{\eta}{\sqrt{1-s_{avg}}} D \nabla L(w_t)$$
其中：

• $\eta$ 是学习率。
• $s_{avg}$ 是神经元间的平均稀疏度。
• $D$ 是一个对角矩阵，其中对应神经元 $i$ 的条目为 $\sqrt{1-s_i}$。

通过将梯度乘以 $\sqrt{1-s_i}$，SparseOpt 有效地抵消了由 BN 引起的 $(1-s_i)^{-1/2}$ 放大效应。额外的 $\sqrt{1-s_{avg}}$ 归一化则保留了全局梯度范数。在稠密层的特殊情况（即所有 $i$ 的 $s_i=0$）下，预条件器退化为单位矩阵，使得 SparseOpt 等同于标准的随机梯度下降（SGD）。

主要贡献

1. 稀疏训练中 BN 的理论分析：本文提供了首个系统性的理论和实证分析，表明 BN 会因异质连接性而改变稀疏层中的梯度方向。它证明了 BN 通过 $(1-s_i)^{-1/2}$ 对梯度进行重缩放，从而在 DST 掩码更新期间导致不稳定性。
2. SparseOpt 算法：作者提出了一种新颖的预条件梯度下降方法，该方法显式地考虑神经元级稀疏度以校正梯度偏斜。该方法与现有的 DST 框架（如 RigL 和 SET）以及其他优化器（如 HAM）兼容。
3. 实证验证：在 CIFAR-100 和 ImageNet 数据集上的 ResNet 架构进行了广泛实验，结果表明 SparseOpt 一致地加速了收敛并提高了泛化能力，特别是在高稀疏度水平（0.90、0.95、0.97）下。

实验结果

作者使用 RigL（Rigging the Lottery）和 SET（Sparse Evolutionary Training）在 ResNet-20（CIFAR-100）和 ResNet-50（ImageNet）上评估了 SparseOpt。

• 收敛速度：与使用标准 SGD 的基线 DST 方法相比，SparseOpt 以更少的训练轮次实现了更高的准确率。例如，在 90% 稀疏度的 ImageNet 上，使用 RigL 的 SparseOpt 在 90 个轮次内达到了 75.06% 的 Top-1 准确率，而基线方法仅为 73.88%。
• 泛化能力：该方法在各种训练预算下始终产生更好的测试准确率。性能差距在高稀疏度水平下最为显著，因为那里的连接异质性最大。
• 掩码探索：该研究分析了"ITOP 率”（衡量掩码探索多样性的指标）。研究发现，虽然 SparseOpt 改善了训练动态，但在 RigL 的掩码再生步骤中使用原始（未校正的）梯度能带来更好的探索效果。这表明存在一种解耦：优化器受益于用于权重更新的校正梯度，而掩码探索机制则依赖于受 BN 影响的原始梯度幅度。
• 兼容性：研究表明，该方法与最先进的双曲感知最小化（HAM）优化器具有互补性，结合使用时可进一步改善收敛。

意义与主张

本文声称是首个工作，识别出批归一化是动态稀疏训练中不稳定的来源，并提出了一种系统性的解决方案。作者强调，当前的归一化层隐式地假设了均匀连接性，这一假设在稠密层中成立，但在稀疏网络中被破坏。

这项工作的意义在于：

• 弥合差距：通过解决根本的优化瓶颈，它迈出了使 DST 在实际中与稠密训练竞争的重要一步。
• 设计启示：它强调了在归一化机制中需要“感知稀疏性”的设计，而不是直接将为稠密模型开发的技术迁移过来。
• 实践效率：通过实现更快的收敛，SparseOpt 减少了训练稀疏模型所需的总训练时间和计算资源，使其更易于部署。

作者在范围上保持谦逊，指出虽然 SparseOpt 提高了收敛率，但 DST 的全部潜力还取决于优化动态与掩码探索策略之间的相互作用，这仍是未来研究的领域。他们还指出，其分析在初步研究中已扩展到层归一化（LN），表明其发现具有更广泛的适用性。