Evolutionary Optimization of Model Merging Recipesについて

先日、sakana.aiが発表した論文になっています。大規模言語モデル(LLM)のパラメータと下流タスクでの性能の関係は、Task Arithmeticなどで研究されてきました。そこから、複数のLLMのレイヤーをパラメータを混ぜて一つにする方法が提案されてきました。この論文では、パラメータの混ぜ方だけでなく、レイヤーのマッチングまで一本化して最適化する手法を提案しています。
個人的に、この手法の発展が楽しみです。行く末は、ビッグテックが要するクローズドな大きなモデル(GPT4, Geminiなど)を、オープンソースの複数のモデルが超えていく未来が来るかもしれません(スイミーのように)。
それでは、この論文を理解して行きましょう。

この論文が扱う問題の理解

この論文が解く問題は、モデルの合体(model merging)の最適化・自動化です。先行研究から、複数のモデルを1つのモデルへ合体することは、下流タスクへの能力を向上させることにつながると知られています。ただ、問題なのが、合体後のモデルの構造の最適化がなされていないことです。ドメイン知識や人手による作業に頼っています。今回の論文は、ここに進化アルゴリズム(evolutionary algorithms)を用いて、複数のモデルから、それを超えるような最適な一つのモデル構造を自動で発見する手法を提案しています。

この論文が提案する問題解決方法の理解

解決方法の全体像

複数のモデルを合体させて、1つの強力なモデルを作るという問題を解くためには、以下の2つの課題を解く必要があります。

課題1: 複数のレイヤーが持つパラメータ配合の最適化 (論文では、Merging in the Parameter Space)
課題2: 複数のモデルの全レイヤーから、最適なレイヤーの配置の探索 (論文では、Merging in the Data Flow Space)

さて、この2つの課題を解決するアルゴリズムを具体的に見て行きましょう。

課題1 (Merging in the Parameter Space)

この課題解くために、パラメータを混ぜる・最適化することを行なっています。

複数のレイヤーのパラメータを混ぜる
論文では、TIES-MergingをベースとしたDAREというアルゴリズムを使っています。
パラメータの混ぜ方の最適化
進化アルゴリズム(CMA-ES)を活用して、パラメータの混ぜ方の最適化を行っています。最適化の軸は、ターゲットとなる下流タスクの指標となる評価指標です。進化アルゴリズムの言葉を使えば、それがfitnessです。

課題2 (Merging in the Data Flow Space)

まずは、理解しやすくするために、この課題を文章問題化して見ましょう。

モデルがN個あります。それぞれ、レイヤーがS個だとします。そうすると、全てモデルのレイヤーの総数は、$M = S\times N$で表せることができます。このM個のレイヤーから、ターゲットの下流タスクでの性能が一番高い、レイヤーの配置を探してください。

この文章問題を解いていくのですが、探索スペースは、$(M + 1)^T$パターンあります。論文でも述べられている通り、進化アルゴリズムでも、時間がかかります。論文では、進化アルゴリズムに工夫を加えて、探索スペースを減らしています。
その工夫の核となるのが、Indicator array $\mathcal{I}$です。それは、$T = M \times r$個の要素を持った配列です。$\mathcal{I}$は、使うレイヤーを表しています。まず、全てのレイヤーを順々に並べます。モデルiのj番目レイヤーを$L_{(i,j)}$だとすると、

$$
[L_{(1, 1)}, L_{(1, 2)}, L_{(1, 3)},...,L_{(1, S)}, L_{(2, 1)}, L_{(2, 2)}, L_{(2, 3)}, ..., L_{(2, S)}, ..., L_{(M, S)}]
$$
というふうに並べるイメージです。さらに、これを$r$回繰り返します。$\mathcal{I}$のk番目の要素$\mathcal{I}_k$は、この並べ方に対応しています。つまり、$\mathcal{I}_k > 0$の時に、並べた中のk番目のレイヤーを最終的な合体モデルに使うことを示します。逆に、そうでない場合は、使いません。こんな感じで、解き方を工夫することで、探索スペースを$2^T$に削減しています(二者択一の選択が$T$回あるので)。

進化アルゴリズムの中で、レイヤー配置の最適化だけでなく、レイヤー間のスケーリング調節の最適化も行なっています。なぜかというと、各モデルの学習で使われたデータの違いから、モデルAのレイヤーの出力を、モデルBのレイヤーへ入力すると分布違いから、最終的な生成結果がおかしくなってしまう可能性があるからです。そこで、スケーリングを$W \in \mathcal{R}^{M \times M }$を最適化することで対処しています。例えば、$W_{(1, 2)}$は、レイヤー1の出力へのスケーリング値を表しています。そのスケーリングされた出力が、レイヤー2への入力となります。

そして解く

これら2つのアルゴリズムを柔軟に適用しつつ、最適なモデル構造を探索します。特にどちらが最初で最後でなければならないという、処理順の制約はないです。論文で挙げられている例を拝借すると、Parameter Spaceでの探索でレイヤーを合体してから、それを探索対象に含めて、再度レイヤーの配置をData Flow Spaceで探索することもできたりします。レイヤーAとBからレイヤーCを生み出した後、レイヤーA,B,Cを候補として、その配置を最適化していくことも可能だということです。アメーバなどの無性生物の生殖を思い出しました。

関連文献・リンク

この論文のpreprint
- https://arxiv.org/abs/2403.13187
進化アルゴリズム系
- NSGA-Ⅱ
  - https://esa.github.io/pagmo2/docs/cpp/algorithms/nsga2.html
  - https://en.wikipedia.org/wiki/Multi-objective_optimization
- CMA-ES
  - https://en.wikipedia.org/wiki/CMA-ES
Model Merging系
- Task Arithmetic
  - https://arxiv.org/abs/2212.04089
- DARE
  - https://arxiv.org/abs/2311.03099
- TIES-Merging
  - https://proceedings.neurips.cc/paper_files/paper/2023/hash/1644c9af28ab7916874f6fd6228a9bcf-Abstract-Conference.html