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.
翻译:将不同模型进行融合是机器学习应用中广泛使用的范式。虽然最常见的方法是形成模型集成并对其个体预测结果取平均,但这种方法常因给定的内存与计算资源限制而不可行——这些资源会随模型数量线性增长。我们提出一种面向神经网络的逐层模型融合算法,该算法利用最优传输在平均参数前对跨模型的神经元进行(软)对齐。实验表明,该方法能在异构非独立同分布数据上训练的神经网络之间成功实现"一次性"知识迁移(即无需任何再训练)。在独立同分布与非独立同分布场景下,我们证明该方法显著优于朴素平均方法;同时,通过适度微调,该方法可作为标准卷积网络(如VGG11)、残差网络(如ResNet18)及多层感知机在CIFAR10、CIFAR100和MNIST数据集上模型集成的高效替代方案。最后,本方法还为融合不同宽度的神经网络参数提供了理论依据,我们探索了其在模型压缩中的应用。代码已开源在:https://github.com/sidak/otfusion。