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

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) 这里设置将其扩展为合并来自具有不同标不同数据集的模型：例如，将对狗品种进行分类的模型与对鸟类进行分类的模型进行合并。(c) ZipIt！通过识别共享特征来合并这些模型，无需重新训练。

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

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

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

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

接下来，如果两个特征相关性很好，可以对它们进行平均而不会丢失太多信息。因此，可以通过在完整串联空间中找到 ni 对相似特征并将它们平均来构造 fi*。默认情况下，贪婪地执行此操作：即迭代地匹配具有最高相关性的特征而不进行替换，最后用这些匹配构成 fi*。

定义合并矩阵M*，其平均匹配的特征，其每个匹配对应一个输出特征fi*。

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

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

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

因此，通过这些层“传播”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 应用于所有参数（权重、偏差）。继续传播。由于 层归一化会动态计算平均值和标准差，因此不需要做任何特殊的事情。