Mergekitを使ったMoE(Mixture of Experts)作成のテクニック

2024年2月4日 10:36

はじめに

以前行ったMergekitを使った日本語MoEの作成（以下記事）がそこそこ上手くいったものの、かなり手探りで実施した感があった。
そんな中、有志でベストプラクティスがまとめられているのを知ったのでその要約を自分のためにここにまとめようと思う。

忙しい方は4. まとめ だけ読めば雰囲気は確認できると思う。

1. 概要

章の構成は以下の通り。

Overview
What makes a perfect MoE: The secret formula
Using the same exact model together 4x or 8x or (etc) times is pointless
Why is a proper merge considered a base model, and how do we distinguish them from a FrankenMoE?
Why the community working together to improve as a whole is the only way we will get Mixtral right

最も重要なのはChapter2：What makes a perfect MoE
ここには、

base_model, expertsとして選択すべきモデル
positive_promptsの作成方法

が記載されている。詳細を見るべきはほぼここのみと言える。

2. What makes a perfect MoE

2-1. base_model, expertsとして選択すべきモデル

1. base_modelについて

base_modelには、expertsとして使用するモデルのbase modelを使用する。(例：Llama-2-base、Mistral-baseなど)

より具体的に述べると、私が以前作成してそこそこ上手くいったHachiML/youri-2x7b_devは、以下のようなconfigとなっている。これのbase_modelは、rinna/youri-7bとすべきだったようだ。

base_model: rinna/youri-7b-chat
gate_mode: hidden # one of "hidden", "cheap_embed", or "random"
dtype: bfloat16 # output dtype (float32, float16, or bfloat16)
experts:
  - source_model: rinna/youri-7b-chat
    positive_prompts: 
      - "質問と回答の選択肢を入力として受け取り、選択肢から回答を選択してください。"
      - "前提と仮説の関係を含意、矛盾、中立の中から回答してください。"
      - "以下のテキストを、ポジティブまたはネガティブの感情クラスのいずれかに分類してください。"
      - "以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。"
  - source_model: rinna/youri-7b-instruction
    positive_prompts: 
     - "質問に対する回答を題名と文章から一言で抽出してください。回答は名詞で答えてください。"
     - "与えられたニュース記事を要約してください。"
     - "与えられた文が文法的であるかを回答してください。"

2. expertsについて

expertsには上記で選択したbase_modelをFine-Tuningしたモデルを選択していくこととなる。

同じように考えることができない小さな脳を組み合わせると、大きな脳は正しく考えることができなくなり、精神衛生上の問題を抱えることになる。

と記されている。これを解釈すると、あまりに目的・用途が異なりすぎるモデルの組み合わせは良くないということだと思われる。
良くない組み合わせの例として、以下が挙げられている。

性的なロールプレイに特化して作られているモデル（Noromaid-7b）
多数の異なる書き言葉を理解し、翻訳し、教えるためのモデル（Openbuddy-zephyr-7b）

後者は、一見問題なさそうに見えるが『言語の壁』が生じるらしい。

逆に、以下の様な組み合わせが良いとしている

どれも何らかの形で数字やコーディングのトレーニングを受けている。（中略）全体的に、これらのモデルはコーディングだけでなく、コーディング以外の一般化されたタスクにも優れているので、お互いによく適しています。

つまり、特化領域を持ちつつもある程度一般的なタスクをこなせるモデルがexpertsとしては相応しいということだと思われる。営業とエンジニアの会話が噛み合わない、みたいなことがexperts同士でも起こるということかもしれない。

HachiML/youri-2x7b_devを振り返ると、rinna/youri-7b-chatとrinna/youri-7b-instructionは（組み合わせて効果が強いかどうかはおいておいて）ともに一般的なタスクに強いモデルでありexpertsとして適切だったと思われる。

2-2. positive_promptsの作成方法

これは、HachiML/youri-2x7b_devを作成した際にもどうしようか考えさせられた。私は、expertsとして選択したモデルのそれぞれが得意なベンチマークタスクのプロンプトを用いたが、ベストプラクティスとしては、Experts自身に自分の得意なプロンプトを作成させるというのが提案されていた。
プロンプト例を日本語訳すると以下の様な内容である。
プロンプト開始：

あなたは（専門分野）を専門とするエキスパートAIモデルです。
あなたの名前は（AIモデルの名前）です。
あなたの仕事は、他のAIモデルがAIモデルのエキスパートとしての強みに合ったプロンプトを書いている以下の.ymlファイルを見て、あなたのAIモデルに使用されている.ymlファイルのセクションを、あなたの（専門分野）としての強みを最もよく表す説明で再現することです。（この太字のセクション以降の部分はオプションで、実際のプロンプトとモデル名をマージのモデルで使用する場合にのみ必要です：注：あなたの長所が他のモデルと似ている場合は、そのモデルとは異なる長所を書かなければなりません。）

```
base_model: mistralai_Mistral-7B-v0.1
gate_mode: hidden
dtype: float16
experts:
  - source_model: (Insert actual ai models name here)
    positive_prompts:
    - "(Insert some short description or keyword here)"
    - "(Insert some short description or keyword here)"
    - "(Insert some short description or keyword here)"
    - "(Insert some short description or keyword here)"
  - source_model: fblgit_UNA-TheBeagle-7b-v1
    positive_prompts:
    - "How do you"
    - "Explain the concept of"
    - "Give an overview of"
    - "Compare and contrast between"
    - "Provide information about"
    - "Help me understand"
    - "Summarize"
    - "Make a recommendation on"
    - "Answer this question"
  - source_model: LucciAI_openchat-3.5-0106-function-calling
    positive_prompts:
    - "Write a program to solve this problem"
    - "Modify this function to improve its performance"
    - "Refactor this code to enhance readability"
    - "Create a custom function for this specific use case"
    - "Optimize this algorithm to reduce computational complexity"
    - "Implement this feature by extending existing codebase"
    - "Integrate this API call into the application"
    - "Help me troubleshoot and fix this bug"
    - "Review and test this code snippet before deployment"
    - "Analyze this error log to identify potential issues"
    - "Generate a set of unit tests for this module"
    - "Evaluate different approaches to solving this problem"
    - "Do a web search for"
    - "Use the plugin to"
```

プロンプト終了

ここで気になるのが、上記プロンプトを試すモデルの順番だが、それについての記載はない。
また、ここで重要なテクニックとしてpositive_promptsは、

単一のキーワードとしない。（単語としない。）
詳細であること。（これは恐らくタスクが明確となっているプロンプトである必要があるということだと思われる。たとえば、『Pythonのコードを作成』や『要約する』など。）
モデル同士で重複させない。極力似たようなプロンプトも許さない。（ルーターモデルが推論中にどのエキスパートを使用するかを決定する際にエラーを引き起こさない程度に互いに異なるものであるべき。）

というものが挙げられている。

また、expertsの選択と内容は被るが、少なくとも上記のようなプロンプトに応えられるようなモデルでなければexpertsとして選択するに値しないと言えるかもしれない。

個人的には得意なベンチマークのプロンプトを使うことはそれほど悪くないと思っているが、網羅的にベンチマークテストが行われ、結果が公開されているFine-Tuningモデルはそれほど多くない。基本的にはモデルにプロンプトを作らせるのが良いと思う。（モデルが自身の強みをどれほど網羅的に理解しているのかというのには疑問の余地があるが、少なくともプラクティスとしてはある程度の成功を収めている様なので信頼できるやり方ではあるはず。）

3. その他のテクニックに関わる記載

あとは、

Chapter3：Using the same exact model together 4x or 8x or (etc) times is pointless（全く同じモデルを4倍、8倍、あるいは（その他）何回も使うことは無意味である。）

Chapter4：Why is a proper merge considered a base model, and how do we distinguish them from a FrankenMoE?（適切なマージはなぜベースモデルとみなされるのか？フランケンMoEとどう区別するのか？）

に若干プラクティス的なことが書いてあるがほぼ内容は以下の１文ずつで表せる。

同じモデルのマージは無意味である。（Chapter3）
最高品質のモデルと呼ぶには、ベンチマークのスコアだけでなくHands Testingも重要である。（Chapter4）
- Hands Testingの例としてユニークなリーダーボードが紹介されている。（🏆LLM Rubric Leaderboard）

4. まとめ

振り返りとして、Merge-kitを使ったより性能の高いMoEの作成のベストプラクティスは以下の通り。

同じモデルのマージは無意味である。
ベンチマークのスコアだけでなくHands Testingも重要である。
base_modelには、expertsとして使用するモデルのbase modelを使用する。(例：Llama-2-base、Mistral-baseなど)
特化領域を持ちつつもある程度一般的なタスクをこなせるモデルをexpertsとする。
Experts自身に自分の得意なプロンプトを作成させる。

今後は、上記を踏まえたMoEを作成したいと思う。