[论文解读] Model Fusion via Optimal Transport

Sidak Pal Singh, Martin Jaggi|arXiv (Cornell University)|Oct 12, 2019

Advanced Neural Network Applications参考文献 32被引用 7

一句话总结

本文提出一种基于最优传输（OT）的单次模型融合方法，通过在平均权重前对不同训练模型中的神经元进行对齐，实现神经网络的模型融合。通过计算层间权重分布的Wasserstein中位数，该方法在性能上优于简单的平均方法，并能在异构、非独立同分布（non-i.i.d.）数据上实现有效的知识迁移，尤其在模型压缩和作为集成模型的替代方案方面表现优异。

ABSTRACT

Combining different models is a widely used paradigm in machine learning applications. While the most common approach is to form an ensemble of models and average their individual predictions, this approach is often rendered infeasible by given resource constraints in terms of memory and computation, which grow linearly with the number of models. We present a layer-wise model fusion algorithm for neural networks that utilizes optimal transport to (soft-) align neurons across the models before averaging their associated parameters. We show that this can successfully yield "one-shot" knowledge transfer (i.e, without requiring any retraining) between neural networks trained on heterogeneous non-i.i.d. data. In both i.i.d. and non-i.i.d. settings , we illustrate that our approach significantly outperforms vanilla averaging, as well as how it can serve as an efficient replacement for the ensemble with moderate fine-tuning, for standard convolutional networks (like VGG11), residual networks (like ResNet18), and multi-layer perceptrons on CIFAR10, CIFAR100, and MNIST. Finally, our approach also provides a principled way to combine the parameters of neural networks with different widths, and we explore its application for model compression. The code is available at the following link, https://github.com/sidak/otfusion.

研究动机与目标

• 为解决在资源受限环境下，集成方法因高内存和计算成本导致的效率低下问题。
• 克服在合并具有非相同权重参数化结构的神经网络时，简单权重平均方法性能不佳的问题。
• 实现在未重新训练或共享数据的前提下，对在异构、非独立同分布数据上训练的模型进行一次性融合。
• 为不同宽度的模型融合提供一种系统化的方法，尤其适用于模型压缩任务。
• 在联邦学习与去中心化学习中，作为蒸馏与集成方法的高效、无需数据的替代方案。

提出的方法

• 该方法利用最优传输计算两个或多个父模型对应层中神经元之间的软性、结构化对齐，依据激活相似性或输入权重分布。
• 将层间融合建模为基于父模型权重向量定义的概率测度的Wasserstein中位数计算问题。
• 融合过程按层逐层进行，实现无需微调的一次性合并。
• 通过在平均前利用OT对齐神经元，该方法支持不同宽度模型的融合。
• 可作为结构化剪枝的后处理工具，也可作为蒸馏或微调的初始化策略。
• 在CIFAR-10、CIFAR-100和MNIST数据集上，基于VGG11、ResNet18和MLP模型进行了评估，并对超参数与初始化策略进行了消融研究。

实验结果

研究问题

• RQ1 基于最优传输的层间融合是否能在非独立同分布数据上训练的神经网络融合中优于简单平均方法？
• RQ2 该方法是否能成功实现跨不同任务或标签集训练的模型之间的知识迁移？
• RQ3 OT融合是否可作为无需微调或数据访问的集成平均的有效替代方案？
• RQ4 在蒸馏不可行的情况下，OT融合在模型压缩方面效果如何？
• RQ5 与随机初始化或基于模型的初始化相比，基于OT的初始化是否能提升蒸馏性能？

主要发现

• 在CIFAR-10上，使用VGG11时，OT融合达到98.30%的测试准确率，显著优于简单平均（98.11%）和蒸馏（平均98.26%）。
• 在CIFAR-100上使用ResNet18时，OT融合达到94.45%的准确率，超过简单平均（93.89%），并匹配或超过蒸馏结果。
• 当将更大的教师模型压缩至原大小的一半（ρ=2）时，OT融合达到98.26%的准确率，优于随机初始化（98.13%）和模型$M_B$初始化（98.17%）的蒸馏结果。
• 在极端压缩设置下（ρ=10），OT融合在蒸馏中达到97.59%的准确率，高于随机初始化的97.21%和$M_B$初始化的97.55%，表现出一致的优势。
• 结合微调后，OT融合在所有设置中均匹配或超过蒸馏性能，证明其作为轻量级替代方案的可行性。
• 该方法成功实现了在不同标签集上训练的模型之间的知识迁移，同时继承了各父模型的独特能力，并优于任一模型单独表现。