Sparse Weight Averaging with Multiple Particles for Iterative Magnitude Pruning

ICLR2024| Sparse Weight Averaging with Multiple Particles for Iterative Magnitude Pruning

• 论文链接：https://openreview.net/pdf?id=Y9t7MqZtCR
• 源码链接：https://openreview.net/attachment?id=Y9t7MqZtCR&name=supplementary_material

简介

深度神经网络通常高度参数化，并且可以在不牺牲模型性能的情况下对其大多数参数进行剪枝。彩票假说表明在初始化时存在一个稀疏自网络，可以对其进行训练，以实现与原始密集网络相同性能水平。这种匹配子网络可以通过带倒带的迭代幅度剪枝（IMP）得到，它包含以下三个步骤：1. 一定次数的迭代训练网络，2. 剪枝具有最小幅度的权重，3. 将权重倒回早期迭代，同时将剪枝后权重固定为0。该过程重复几轮，最终自网络可以实现整个网络性能的匹配彩票。尽管IMP简单，但它在寻找稀疏掩码方面提供了SOTA性能，尤其是在极端稀疏的情况下。
考虑到IMP的简单性，它的成功确实是违背直觉的。在这方面，. Pruning neural networks at initialization: Why are we missing the mark?揭示了彩票假说和线性模式连接性（Linear Mode Connectivity）之间的潜在联系，表明IMP的有效性取决于对随机优化的稳定性；IMP解决方案位于损失景观中统一吸引盆地中。Unmasking the lottery ticket hypothesis: What’s encoded in a winning ticket’s mask?深入研究了这一主题，发现具有不同稀疏性水平的连续IMP解也存在线性连通性。更准确地，他们得出结论是，如果连续轮次是不廉洁的，IMP就无法找到匹配的子网络，并进一步强调了简直率和倒带迭代对保持IMP解之间连通性的重要性。
受到IMP和线性模式连通性联系的启发，本文将理解扩展到损失景观的角度。分析深度神经网络损失景观是一种有效的工具，被广泛用于研究模式连通性，这促使本文找到位于损失景观平坦区域的解，以增强泛化能力。值得注意的是这两个领域都有一个共同目标，即识别以低损失值为特征的好子空间，这一目标与神经网络剪枝的最终目标一致–在给定的密集参数空间识别匹配的稀疏子空间。

本文方法

神经网络权重的线性连通性

考虑一个一维路径记为$P:[0,1]\to {\mathbb{R}}^D$，用于连接两个神经网络权重$w^{(0)},w^{(1)}$，开始和结束点分别是$P(0)=w^{(0)},P(1)=w^{(1)}$。在简化形式，如果条件 sup ⁡ λ ∈ [ 0 , 1 ] L ( P ( λ ) ) ≤ max ⁡ { L ( P ( 0 ) ) , L ( P ( 1 ) ) } + ϵ \sup_{\lambda\in [0,1]}\mathcal{L}(P(\lambda))\leq \max\{\mathcal{L}(P(0)),\mathcal{L}(P(1))\}+\epsilon supλ∈[0,1]​L(P(λ))≤max{L(P(0)),L(P(1))}+ϵ满足，可以说$w^{(0)},w^{(1)}$间存在连通性。尽管深度学习最近进展揭示了通过随机优化获得的局部极小值存在非线性路径，对于现代深度神经网络建立线性连通性（即与线性连接器$P(\lambda )=(1-\lambda )w^{(0)}+\lambda w^{(1)}$）仍然不是简单的。

IMP损失景观视角

. Pruning neural networks at initialization: Why are we missing the mark?证明了匹配彩票对神经网络对SGD噪声$\xi$的稳定性有显著影响。甚至当两个网络由相同随机初始化$w_0$训练，不同SGD噪声${\xi }^{(1)},{\xi }^{(2)}$破坏了通过SGD获得的解之间的线性连通性。例如不存在以下之间的线性连接：
$$w_T^{(1)}\leftarrow {\text{SGD}}_{0\leftarrow T}(w_0,{\xi }^{(1)},\mathcal{D}),w_T^{(2)}\leftarrow {\text{SGD}}_{0\leftarrow T}(w_0,{\xi }^{(2)},\mathcal{D})$$
因此优化对SGD噪声不稳定。他们进一步根据经验证实，通过IMP获得的稀疏解是匹配的，当且仅当他们对SGD噪声稳定时，并将这种不稳定性诊断为原始IMP算法失败情况。确保稳定性一种简单方法是共享优化轨迹的早期阶段。也就是说存在线性联通：
$$w_T^{(1)}\leftarrow {\text{SGD}}_{T_0\to T}(w_{T_0},{\xi }^{(1)},\mathcal{D}),w_T^{(2)}\leftarrow {\text{SGD}}_{T_0\to T}(w_{T_0},{\xi }^{(2)},\mathcal{D})$$
SGD从相同初始$w_{T_0}\leftarrow {\text{SGD}}_{0\to T_0}(w_0,\xi ,\mathcal{D})$开始。Unmasking the lottery ticket hypothesis证明了具有不同稀疏度水平的两个连续IMP解之间线性连通性，并将其作为IMP成功的关键因素。
然而尽管每对解之间存在线性连通性，但低损失子空间是否由三个或更多个解的凸包形成的问题仍然不确定。如果使用IMP解构建低损失子空间变得可行，则可能构建一个更有效的解，并在中点处改进该子空间泛化能力

SWAMP算法

SWAMP与普通IMP区别主要有两个方面：首先创建匹配彩票的多个副本，并使用不同随机种子同时训练它们，而IMP采用单个粒子。其次，本文用随机加权平均（Stochastic Weight Averaging,SWA）代替SGD训练，这是一种通过定期采样学习轨迹来构建参数移动平均值的方法。SWA使得我们能够在整个训练过程中积累更多粒子。然后，在进行剪枝步骤之前，对所有粒子进行平均。