SegMoE: Segmind Mixture of Diffusion Experts の概要

2024年2月4日 05:45

以下の記事が面白かったので、かるくまとめました。

・SegMoE: Segmind Mixture of Diffusion Experts

1. SegMoE

「SegMoE」は、MoE Diffusionモデルをゼロから作成するためのフレームワークです。「SegMoE」はHuggingFaceエコシステムに統合されており、diffusersでサポートされています。

「SegMoE」は、「Stable Diffusion」と同じアーキテクチャに従います。「Mixtral 8x7b」と同様、「SegMoE」には複数のモデルが1つにまとめられています。これが機能する方法は、一部のフィードフォワードレイヤーを疎なMoEレイヤーに置き換えることです。MoEレイヤーには、どのエキスパートがどのトークンを最も効率的に処理するかを選択するためのルーターネットワークが含まれています。「segmoe」パッケージを使用して独自のMoEモデルを作成できます。このプロセスには数分しかかかりません。詳細については、Githubリポジトリを参照してください。人気のあるライブラリ「mergekit」からインスピレーションを得て「segmoe」を設計しています。

「MoE」の詳細については、HuggingFaceの投稿を参照してください。

2. モデルとカスタムMoE作成コード

2-1. モデル

本日 (2月3日) リリースされるモデルは、次のとおりです。

・SegMoE 4x2 : 4つのエキスパートモデルで構成
・SegMoE 2x1 : 2つのエキスパートモデルで構成
・SegMoE SD 4x2 : 4つのStable Diffusion 1.5 エキスパートモデルで構成

「SegMoE MoE」は「SegMoE-AxB」と呼ばれます。ここで、Aはエキスパートモデル「MoE-d」の数を指し、2番目の数字は各画像の生成に関与するエキスパートの数を指します。構成設定に応じて、モデルの一部のレイヤー (フィードフォワードブロック、アテンション、またはすべて) のみが複製されます。残りのパラメータはStableDiffusionモデルと同じです。

2-2. カスタムMoe作成コード

カスタムMoE作成コードは以下のリポジトリで公開しています。

3. サンプル

「SegMoe 4x2」のサンプルは、次のとおりです。

「SegMoe 2x1」のサンプルは、次のとおりです。

「SegMoe SD 4x2」のサンプルは、次のとおりです。

4. diffusersの使用

(1) パッケージのインストール。

pip install -U segmoe diffusers transformers

(2) SegMoE-4x2 による画像生成。

from segmoe import SegMoEPipeline

pipeline = SegMoEPipeline("segmind/SegMoE-4x2-v0", device="cuda")

prompt = "cosmic canvas, orange city background, painting of a chubby cat"
negative_prompt = "nsfw, bad quality, worse quality"
img = pipeline(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=1024,
    width=1024,
    num_inference_steps=25,
    guidance_scale=7.5,
).images[0]
img.save("image.png")

ローカルモデルを使用することもできます。

from segmoe import SegMoEPipeline

pipeline = SegMoEPipeline("segmoe_v0", device="cuda")

prompt = "cosmic canvas, orange city background, painting of a chubby cat"
negative_prompt = "nsfw, bad quality, worse quality"
img = pipeline(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=1024,
    width=1024,
    num_inference_steps=25,
    guidance_scale=7.5,
).images[0]
img.save("image.png")

独自の SegMoE を作成する方法は、「Creating your Own SegMoE」を参照してください。

5. 比較

以下の画像に示すように、プロンプトの理解が向上しているようです。各画像は、左からSegMoE-2x1-v0、SegMoE-4x2-v0、ベースモデル (RealVisXL_V3.0) です。

・three green glass bottles (緑色のガラス瓶3本)

・panda bear with aviator glasses on its head (頭に飛行用の眼鏡をかけたパンダのクマ)

・the statue of Liberty next to the Washington Monument (ワシントン記念塔の隣にある自由の女神像)

・Taj Mahal with its reflection. detailed charcoal sketch. (タージ・マハルとその反射。詳細な木炭スケッチ)

6. 独自の SegMoE の作成

「config.yaml」を準備するだけです。

base_model: 基本モデルのパス、モデル カード、CivitAIダウンロード リンク
num_experts: 使用するエキスパートの数
moe_layers: 混合するレイヤーのタイプ (ff, attn, all : defaultはattn)
num_experts_per_tok: 使用するエキスパートの数
experts:
  - source_model: Expert 1 パス、モデル カード、CivitAIダウンロード リンク
    positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
    negative_prompt: ゲート重みを計算するためのネガティブ プロンプト
  - source_model: Expert 2 パス、モデル カード、CivitAIダウンロード リンク
    positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
    negative_prompt: ゲート重みを計算するためのネガティブ プロンプト
  - source_model: Expert 3 パス、モデル カード、CivitAIダウンロード リンク
    positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
    negative_prompt: ゲート重みを計算するためのネガティブ プロンプト
  - source_model: Expert 4 パス、モデル カード、CivitAIダウンロード リンク
    positive_prompt: ゲート重みを計算するためのポジティブ プロンプト
    negative_prompt: ゲート重みを計算するためのネガティブ プロンプト

任意の数のモデルを組み合わせることができます。「config.yaml」の詳細については、githubリポジトリを参照してください。

HuggingFace モデルとCivitAI モデルの両方がサポートされています。CivitAIモデルの場合は、モデルのダウンロードリンクを貼り付けます。
(例: "https://civitai.com/api/download/models/239306")

次に、コマンドを実行します。

segmoe config.yaml segmoe_v0

これにより、次の構造を持つ segmoe_v0 というフォルダが作成されます。

├── model_index.json
├── scheduler
│   └── scheduler_config.json
├── text_encoder
│   ├── config.json
│   └── model.safetensors
├── text_encoder_2
│   ├── config.json
│   └── model.safetensors
├── tokenizer
│   ├── merges.txt
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   └── vocab.json
├── tokenizer_2
│   ├── merges.txt
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   └── vocab.json
├── unet
│   ├── config.json
│   └── diffusion_pytorch_model.safetensors
└──vae
    ├── config.json
    └── diffusion_pytorch_model.safetensors

あるいは、Python APIを使用して、MoEを作成することもできます。

from segmoe import SegMoEPipeline

pipeline = SegMoEPipeline("config.yaml", device="cuda")

pipeline.save_pretrained("segmoe_v0")

7. HubへのPush

モデルは、huggingface-cli 経由でHubにPushできます。

huggingface-cli upload segmind/segmoe_v0 ./segmoe_v0

モデルは Python から直接HubにPushすることもできます。

from huggingface_hub import create_repo, upload_folder
 
model_id = "segmind/SegMoE-v0"

repo_id = create_repo(repo_id=model_id, exist_ok=True).repo_id

upload_folder(
    repo_id=repo_id,
    folder_path="segmoe_v0",
    commit_message="Inital Commit",
    ignore_patterns=["step_*", "epoch_*"],
)

8. 免責事項と進行中の作業

・速度が遅い : トークンあたりのエキスパートの数が1より大きい場合、MoE は複数のエキスパートモデルにわたって計算を実行します。これにより、単一のSD 1.5またはSDXLよりも遅くなります。

・VRAMの使用量が多い : MoE は推論を非常に高速に実行しますが、それでも大量の VRAM (したがって高価な GPU) が必要です。このため、ローカルで使用するのは困難ですが、複数のGPUを使用する展開には最適です。参考として、SegMoE-4x2には半精度で24GBのVRAMが必要です。

9. おわりに

推論時間を低く抑えながら、事前学習されたモデルを組み合わせるだけで、SOTA diffuserモデルを簡単に作成できる可能性のある新しいツールをコミュニティに提供するため、「SegMoE」は作成されました。これを使って何を構築できるか楽しみです。

この記事が気に入ったらサポートをしてみませんか？