[論文レビュー] 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 における画像分類といったまったく異なるタスクで事前学習された深層学習モデルを、追加の訓練を一切行わずに統合する、革新的な手法を導入する。特徴量の相関関係をモデル内およびモデル間で「ジッピング」できる一般化されたモデル統合手法を導入し、指定されたレイヤーまで部分的統合を可能にすることで、別々のタスクに分かれた設定において、従来手法よりも最大 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.

研究の動機と目的

• 別々のタスクで学習されたモデルを再訓練せずに、1つのマルチタスクモデルに統合する課題に対処すること。
• 異なるタスク固有の特徴分布を持つモデルに適用した場合、重みのパーミュテーションに基づく従来のモデル統合手法が失敗する理由を克服すること。
• 特徴レベルでの「ジッピング」を可能にすることで、モデル内およびモデル間で共有される表現を捉えること。特徴量が直接相関がなくても有効に機能する。
• アーキテクチャに依存しない一般化された手法を構築し、部分的統合を可能にすることでマルチヘッドアーキテクチャを実現し、複雑で多様なタスクにおける性能を向上させること。

提案手法

• 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 別々のデータセットで学習されたモデルに適用した場合、Git Re-Basin や Model Soups といった最先端の統合手法と比べて ZipIt! はどのように差をつけるか？

主な発見

• ZipIt! は、全統合（19/20 レイヤー）を用いた CIFAR-5+5 タスクで 87.4% の統合精度を達成し、最も近いベースライン（Git Re-Basin）を 20 パcentage point 以上上回った。
• 19 番目のレイヤーまで部分統合した ZipIt! 19/20 では、平均 89.1% の精度を達成。FLOPs を半分に抑えたにもかかわらず、アンサンブル性能（94.1%）にほぼ並ぶ。
• 20 レイヤー中 13 レイヤーを統合した ZipIt! では、平均 90.5% の精度を達成。部分的統合が効率を保ちつつも顕著な性能向上をもたらすことを示した。
• 別々のタスク設定において、パーミュテーションベースの統合手法よりも最大 60% の精度向上を達成。モデル内およびモデル間の特徴量ジッピングの必要性を裏付けた。
• REPAIR や微調整を一切使用しない状態でも、これらの技術に依存する手法を上回る性能を発揮。ゼロショット統合設定でも有効であることを示した。
• アブレーションスタディにより、モデル内ジッピングと部分的統合の両方が、全体の性能向上に顕著な寄与をしていることが確認された。

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