[论文解读] ZipIt! Merging Models from Different Tasks without Training

George Stoica, Daniel Bolya+4 more|arXiv (Cornell University)|May 4, 2023

Advanced Neural Network Applications被引用 6

一句话总结

ZipIt! 提出了一种新颖的方法，可在无需额外训练的情况下，将来自完全不同任务（如 CIFAR 和 ImageNet 上的图像分类）的预训练深度学习模型合并为单一多任务模型。通过将模型合并推广至模型内部及跨模型的相关特征‘压缩’（zipping），并支持最多至指定层的局部合并，ZipIt! 在不相交任务设置下，准确率相较先前方法最高提升 60%，显著优于基于排列的合并方法及现有基线方法。

ABSTRACT

Typical deep visual recognition models are capable of performing the one task they were trained on. In this paper, we tackle the extremely difficult problem of combining distinct models with different initializations, each solving a separate task, into one multi-task model without any additional training. Prior work in model merging permutes one model to the space of the other then averages them together. While this works for models trained on the same task, we find that this fails to account for the differences in models trained on disjoint tasks. Thus, we introduce "ZipIt!", a general method for merging two arbitrary models of the same architecture that incorporates two simple strategies. First, in order to account for features that aren't shared between models, we expand the model merging problem to allow for merging features within each model by defining a general "zip" operation. Second, we add support for partially zipping the models up until a specified layer, naturally creating a multi-head model. We find that these two changes combined account for 20-60% improvement over prior work, making it more feasible to merge models trained on disjoint tasks without retraining.

研究动机与目标

• 解决在不重新训练的情况下，将针对不相交任务训练的模型整合为单一多任务模型的挑战。
• 克服先前基于权重排列的模型合并方法在处理具有不同任务特异性特征分布的模型时所遭遇的失败。
• 通过支持模型内部及跨模型的特征级‘压缩’，实现有效合并，即使特征之间无直接相关性，也能捕捉共享表征。
• 开发一种通用、架构无关的方法，支持局部合并以构建多头模型，从而在复杂且异构的任务上提升性能。

提出的方法

• 提出一种广义的‘压缩’（zip）操作，允许不仅在模型之间，而且在每个模型内部合并相关特征，突破了传统的 1:1 权重排列限制。
• 应用基于图的算法，识别并合并跨模型的共享特征，实现灵活且可扩展的任意模型对合并。
• 通过仅合并至指定层的局部压缩支持部分合并，保留未合并的头部分支用于任务特定预测，从而构建多头架构。
• 基于大数定律不等式理论框架，对合并后特征幅度差异进行边界约束，确保方法的稳定性和鲁棒性。
• 采用对称合并策略，通过避免微调来保持模型完整性，防止灾难性遗忘。
• 通过共享合并协议对齐层与特征，使不同架构（只要结构相同）的模型能够合并。

Figure 1: Setting and ZipIt! (a) Prior work merges differently initialized models from the same dataset with the same label sets : e.g., merging two models both trained to classify dog breeds. (b) Our setting expands this to merging models from different datasets with different label sets : e.g., me

实验结果

研究问题

• RQ1 模型合并技术能否被推广至整合完全不相交任务（如犬种分类与鸟类物种识别）训练的模型？
• RQ2 当模型在无关任务上训练时，将一个模型的权重空间排列至另一模型是否失败？若失败，原因是什么？
• RQ3 在多任务设置中，内部特征扩展与模型内合并是否能优于跨模型权重平均？
• RQ4 局部合并（仅将合并限制在早期层，同时保留后续任务特定头部）在多大程度上能提升准确率与效率？
• RQ5 当应用于在不相交数据集上训练的模型时，ZipIt! 与最先进的合并方法（如 Git Re-Basin 和 Model Soups）相比表现如何？

主要发现

• 在 CIFAR-5+5 任务中，使用完整合并（19/20 层）的 ZipIt! 达到 87.4% 的联合准确率，较最近似基线（Git Re-Basin）高出超过 20 个百分点。
• 仅合并至第 19 层（ZipIt! 19/20）的局部合并方案达到 89.1% 的平均准确率，仅需一半的 FLOPs 即接近集成模型性能（94.1%）。
• 合并 13/20 层的 ZipIt! 达到 90.5% 的平均准确率，表明局部合并可在保持效率的同时显著提升性能。
• 在不相交任务基准上，该方法相较基于排列的合并技术，准确率最高提升 60%，验证了模型内与跨模型特征压缩的必要性。
• 即使未使用 REPAIR 或微调，ZipIt! 仍优于依赖这些技术的方法，证明其在零样本合并设置下的有效性。
• 消融研究证实，模型内压缩与局部合并均为关键组件，两者均对整体性能提升有显著贡献。

Figure 2: Task Loss Landscapes for models in Tab. LABEL:tab:cifar50+50 . Model A and Model B lie in low loss basins for their own tasks, but not for the other task . Thus, any interpolation between Model A and a permuted Model B (e.g., Git Re-basin) lies outside the minima for both tasks and thus pe