[論文レビュー] Model Fusion via Optimal Transport

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

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

ひとこと要約

本稿では、最適輸送（OT）を用いて、重みの平均化の前に異なるトレーニング履歴を持つモデル間のニューロンをアライメントすることで、1回のステップでモデルを融合する手法を提案する。層ごとの重み分布のワーサーテイン・バーチャルセンターを計算することにより、単純な平均化や蒸留よりも優れた性能を達成し、非同一分布（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.

研究の動機と目的

• リソース制約のある環境におけるアンサンブル手法の非効率性（メモリおよび計算コストの高さ）を解消すること。
• 同一の重みパラメータ化でないニューラルネットワークを統合する際の単純平均化の性能劣化を克服すること。
• 再トレーニングやデータ共有なしに、異種で非i.i.d.なデータでトレーニングされたモデルを1回のステップで統合できること。
• 特にモデル圧縮を目的とした、幅が異なるモデルの統合に原理的かつ整合的な手法を提供すること。
• フェデレーテッドおよび分散型学習において、蒸留やアンサンブル手法の効率的でデータ不要な代替手段としての役割を果たすこと。

提案手法

• 本手法は、活性化の類似性または入力重みの分布に基づいて、2つ以上の親モデルの対応する層内のニューロン間で柔らかく構造的なアライメントを最適輸送（OT）を用いて計算する。
• 層ごとの統合を、親モデルの重みベクトルが定める確率測度のワーサーテイン・バーチャルセンターを計算する問題として定式化する。
• 統合プロセスは層ごとに逐次適用され、再トレーニングなしに1回のステップで統合が可能となる。
• OTを用いたニューロンのアライメントにより、幅が異なるモデルの統合も可能となり、その後に平均化を行う。
• 構造的プルーニングのための事後処理ツールとして、および蒸留やファインチューニングの初期化戦略として利用可能である。
• 本手法は、VGG11、ResNet18、MLPを用いたCIFAR-10、CIFAR-100、MNISTの実験で評価され、ハイパーパramーターや初期化戦略に関するアブレーションも実施されている。

実験結果

リサーチクエスチョン

• RQ1 非i.i.d.データでトレーニングされたニューラルネットワークを統合する際、最適輸送に基づく層ごとの統合は、単純な平均化を上回る性能を示せるか？
• 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統合は、すべての設定で蒸留性能と同等またはそれを上回り、軽量な代替手法としての有効性を示した。
• 異なるラベル集合でトレーニングされたモデル間でも、知識の転送が成功し、両親モデルの独自の能力を統合しながら、個々のモデルを上回る性能を達成した。