ZIPIT! 无需训练去合并来自不同任务的模型

大家好，欢迎来到IT知识分享网。

ICLR‘24乔治亚理工的论文“ZIPIT! Merging models from different tasks without training“。

典型的深度视觉识别模型能够执行它们所训练的一项任务。 本文解决了一个困难问题，即将具有不同初始化的不同模型（每个模型解决一个单独的任务）组合成一个多任务模型，而无需任何额外的训练。 模型合并中的先前工作将一个模型组合到另一个模型的空间，然后将它们平均在一起。 虽然这适用于在同一任务上训练的模型，但这无法解释在不相交任务上训练的模型差异。 因此，引入了“ZipIt！”，这是一种合并同一架构下两个任意模型的通用方法，该方法结合了两种简单的策略。 首先，考虑到模型之间不共享特征，扩展模型合并问题，允许通过定义通用的“压缩”操作来合并每个模型内特征。 其次，添加了直到指定层支持部分地压缩模型，自然地创建多头模型。 这两项变化结合后比之前的工作提高了 20-60%，使得合并在不相交任务上训练的模型而无需重新训练变得更加可行。

将多个模型组合成一个模型最近开始在视觉界引起关注。 Model Soups（Wortsman，2022a）可以添加根据相同预训练初始化进行微调的多个模型，提高准确性和鲁棒性。 Git Re-Basin（Ainsworth，2022）进一步推广到使用相同数据训练但初始化不同的模型，尽管精度显着下降。 REPAIR（Jordan，2022）通过添加新参数并在适用的情况下调整模型BN来改进 Git Re-Basin。 然而，所有这些方法都只结合了在同一任务上训练的模型。 本文将这一工作推向了一个逻辑端：合并在完全独立的任务上训练的不同初始化模型。 这对于之前的工作来说是一个极其困难的问题，并采用了两种简单的策略来使其可行。

首先，先前的工作重点是在合并模型时将一种模型置换为另一种模型。 这在两个模型之间创建了 1-1 映射，本质上假设它们之间的大多数特征都是相关的。 由于针对不同任务训练的模型不一定是这种情况，因此不能仅仅依赖组合。 相反，这里泛化模型合并以支持“压缩”每个模型内和模型间相关特征的任何组合。 在某些任务上，与基于组合的方法相比，仅此一项就可以将准确率提高高达 20%。 此外，在 （Entezari 2021）的理论背景下，模型内的合并可以产生更好的结果。

其次，现有方法合并整个网络。 虽然这可能适用于在相同设置中训练极其相似的模型，但在不相交任务上训练的模型特征在网络中变得不太相关（Kornblith，2019）。 为了解决这个问题，引入部分压缩，即仅“压缩”到指定的层。 然后，将合并模型的输出提供给原始网络剩余的未合并层，创建一个多头模型。 根据任务难度，这可以将准确度提高 15% 以上，同时仍然保持大多数层的合并。

如图所示，是设置和 ZipIt! 介绍。(a) 先前的工作将来自同一数据集的不同初始化模型与相同的标签集合并：例如，合并两个都经过训练以对狗品种进行分类的模型。 (b) 这里设置将其扩展为合并来自具有不同标不同数据集的模型：例如，将对狗品种进行分类的模型与对鸟类进行分类的模型进行合并。 © ZipIt！ 通过识别共享特征来合并这些模型，无需重新训练。

这项工作将模型合并视为将多个模型的检查点（即权重集合）组合成一个可以执行其组成单件的所有任务单个检查点。 为此，将模型的各层合并在一起。

在不同任务上训练的模型的特征可能不同，因为模型解决不同的问题。 强制组合这些不同的特征可能会产生在任一原始任务上都表现不佳的合并。 相反，这些特征可能与同一模型中的其他特征更兼容，组合后可以更好地保持性能。

事实上可以证明，在有限但普遍的设置中，允许在每个模型内（以及跨模型）合并的方法的性能等于或优于仅跨模型合并的方法（例如，依赖组合的方法）。

设输入特征xi，输出特征fi, 模型层L, 其参数为W和b，那么fi = Li(x) = Wi x + bi。合并层Li*，其输出特征为fi*。

与之前的工作一样，例如 Li （2015）和Jordan （2022），将特征相似性定义为一小组图像（无标签）神经元激活之间的成对相关性。 然而，与这些工作不同的是，本文计算完整串联空间中每个激活之间的相关性。

在将特征传递到下一层之前，需要去除合并操作。 因此，定义了一个“去合并”矩阵U。

这样一个合并操作包括：首先，匹配特征 fiA 和 fiB 来计算两个矩阵 Mi 和 Ui。 然后将 Ui 传递到下一层并从上一层接收 Ui−1。 使用 Mi 和 Ui−1，将合并和去合并操作“融合”到各层参数中。

如图所示：ZipIt！ 利用模型特征中的冗余来逐层合并模型。 (a) 来自两个不相交层的输出特征 fA 和 fB 与其他特征基于其激活的相似性进行配对。 © 生成一个合并矩阵 M 将这些匹配对组合成一个共享输出特征空间，并生成一个相应的去合并矩阵 U 来撤消此操作。 (d) 然后，将 U 网络向上传播以对齐下一层的输入空间，并同时接收前一层的 U 来对齐输入空间。 (e) 用 M 作为输出、U 作为输入将各层“压缩”在一起，生成单个层 (f)。 然后在下一层重复 (a)。

然而，大多数现代神经网络并不是简单地相互堆叠的线性层集合。 在实践中，不能将合并和去合并矩阵组合到网络的每一层中，因为局部的压缩操作期望该层具有权重矩阵，即该层必须具有单独的输入和输出空间，以便去合并输入空间，以及合并输出空间。 其他层（例如 BatchNorm、ReLU）没有这样的权重矩阵。

因此，通过这些层“传播”Mi 和 Ui。 例如，如图展示典型 ConvNet 中常见的层堆栈。 跟随 (Jordan2022)方法，用网络的激活（即在每个 ReLU 之后）计算 Mi 和 Ui。 不去将 Mi 与 ReLU 层融合，因为它没有任何参数。 类似地，可以合并前面 BatchNorm 层的参数（即与偏差相同）。 但它没有权重矩阵，所以也无法将Mi融合到其中。 只有到达 Conv 层后，才能使用方程式将 Mi 和 Ui 融合到其中（在这种情况下，每个内核单元独立）。

压缩传播的一些网络层规则：

• 线性。 应用 Mi 和 Ui。 停止传播。
• 卷积。将 Mi 和 Ui 应用于每个内核位置（即将 k × k 内核维度移动到批量维度）。 停止传播。
• 批归一化（BN）。 将 Mi 应用于所有参数（权重、偏差、均值、方差），将其平方作为方差项。 继续传播。 正如（Jordan2022）指出，如果不知道两个模型之间的协方差（无法访问），就无法计算正确的方差。因此，在合并后重置BN以正确评估方差。
• 层归一化（LN）。 将 Mi 应用于所有参数（权重、偏差）。 继续传播。 由于 层归一化会动态计算平均值和标准差，因此不需要做任何特殊的事情。
• ReLU。 没有什么可以合并的。 继续传播。 请注意，将合并矩阵不变地传递到 ReLU 层是一个近似值，因为非线性特征使用线性合并操作。解决这个问题可能是未来工作的一个有趣话题，因为即使是先前工作的组合和添加方法也存在这个问题（ReLU 对于组合是不变的，但肯定不是添加这样）。
• 平均/最大池化。 没有什么可合并的。 继续传播。
• 跳连接。 继续传播到跳连接的每个输入（每个输入使用相同的 Mi 和 Ui）。

其实并不希望压缩两个网络的每一层，特别当它们的输出空间不兼容情况，或者这样做会损失太多的准确性。 相反，可以执行部分压缩。 也就是说，将大部分层压缩在一起，但后面的层解压缩，如图所示。在框架中实现此操作很简单：正常压缩直到指定层 i，然后剩余的解压缩层将通过压缩传播接收 Ui。 如果UiA 应用于层 LAi+1，将 UiB 应用于层 LBi+1，则剩余的解压缩层将形成“头”，期望合并的特征作为输入。 然后，可以集成头或选择一个在运行时进行评估。

重复匹配（α）。 在某些情况下，希望将两个以上的模型合并在一起。 为此，允许“重复匹配”。 也就是说，当两个特征在贪婪算法中匹配时，它们将被删除并替换为合并的特征。 为了确保一个特征不会无休止地合并，将新特征的相关性设置为旧特征相似度的最小值，并按 α ε (0, 1] 加权。较小的 α 值通常效果最好 。

相同-模型预算（β）。 为了证明同一模型合并的有效性，引入了一个“预算”参数 β ∈ [0, 1]，它表示总合并特征百分比可以来自模型自身的合并，每个模型接收该预算的相等部分 。 请注意，预算为 0 意味着模型内不能合并任何特征。