ICLR‘23华盛顿大学论文“Git Re-basin: merging models modulo permutation symmetries“。
深度学习的成功在很大程度上归功于相对轻松地解决某些大规模非凸优化问题。尽管非凸优化是 NP 困难的,但简单的算法(通常是随机梯度下降的变型)在实际拟合大型神经网络方面表现出令人惊讶的有效性。考虑到隐藏单元的所有可能的排列对称性,神经网络损失域通常包含(几乎)一个盆地,就像 (Entezari 2021)一样。引入三种算法来排列模型的单元,使它们与参考模型对齐,以便在权重 空间中 合并 两个模型。这种变换产生了一组功能等效的权重,这些权重位于参考模型附近近似地凸盆中。
对于MLP任何经过训练的权重 θ,存在功能等效权重 分配的整个等价类,而不仅仅是一个这样的 θ,并且与任何其他等价类相反,收敛到该等价类的任何一个特定元素仅由随机种子确定。一个功能保留的 权重排列 表示为 π(θ)。
考虑将两个独立训练的模型 A 和 B 的权重 θA 和 θB 进行协调的任务,以便可以在它们之间进行线性插值。假设模型 A 和 B 使用相同的架构进行训练,但随机初始化、数据顺序以及可能超参或数据集也不同。中心问题是:给定 θA 和 θB ,能否识别某个π ,以便在 θA 和 π(θB) 之间进行线性插值时,所有中间模型都享有与 θA 和 θB 类似的性能?
对损失域凸性的任何声明,都基于多维每凸性的通常定义,即一维凸性。函数的任意一维切片都是凸的,以便推理复杂高维函数的凸性。在实践中,很少观察到完美的凸性,而是希望尽可能接近它。
同样采用“屏障”来测量凸性的近似值。两个模型A和B的损失L近似,那么损失屏障是非负的,定义如下:
其中零值表示平坦或正曲率的插值。
本文给出三个方法讨论模型A和模型B的单元匹配。
经典的Hebbian咒语,“[神经网络单元]一起激发,连接在一起”(Hebb,2005),考虑在它们激活之间回归来关联两个模型之间的单元。模型之间的匹配激活非常引人注目,因为它捕捉到了两个模型必须学习相似的特征才能完成相同任务的直观想法(Li et al., 2016)。为每个模型提供激活z,目标是将 A 的每个单元与 B 的单元关联起来。按理说,两个模型的激活之间可能存在线性关系。
将普通最小二乘法 (OLS) 限制为置换矩阵 Sd 集的解,则这个问题可以拟合到回归框架中。这样定义MLP第 l 层的排列矩阵
其中Z为d-维激活。 (1)公式给出,所求解的是一个线性分配问题(LAP)。
每层的激活匹配独立于其他层的匹配,从而产生可分离且简单的优化问题。
另外,可以选择检查模型本身的权重,而不是通过激活来关联单元。问题变成如下优化目标:Wi为i层的权重
其构成一个“双线性分配问题之和”(SOBLAP)。 作为一个NP-hard问题,其近似求解,如下算法总结:
受到直通估计器(STE)在其他离散优化问题中成功的启发(Bengio et al., 2013; Kusupati et al., 2021; Rastegari et al., 2016; Courbariaux & Bengio, 2016),这里“学习”权重 π(θB) 的理想排列。 具体来说,目标优化如下
然而,等式(3) 涉及不方便的不可微分投影操作 proj(·),使优化变得复杂。 通过“直通”估计器克服了这个问题: 用一组权重 θ ̃B ≈ π(θB) 来参数化问题。 在前向传播中,将 θB 投影到最接近可实现的 π(θB )。 在向后传播中,切换回不受限制的权重 θB。 通过这种方式,可以保证在评估损失时保持真实的投影约束,但仍然可以在当前参数 θB下计算可用的梯度。
方便的是,可以重新利用上述算法求解 proj(θB)。此外,初始化 θ ̃B = θA 比随机初始化表现更好。这是在初始化时立刻期待的,因为初始匹配将等同于权重匹配方法。然而,目前尚不清楚为什么这些解决方案继续渐近地优于随机初始化。
与上述方法不同,算法 2 尝试使用传统的训练循环显式“学习”最佳排列 π。通过初始化权重匹配解决方案并利用匹配激活的数据分布,它试图提供一个两全其美的解决方案。然而,相对于其他两种方法,这种方法的计算成本非常高。
如图所示是Git Re-Basin将解决方案传送到损失域的单个盆地中来合并模型。θB 被置换为功能相等的 π(θB),以便与 θA 位于同一盆地中。