SegMoE：Segmind 擴散專家混合模型

SegMoE 是一個框架，旨在透過結合預訓練模型，從頭開始建立 Mixture-of-Experts（MoE）擴散模型。它允許使用者將特定層——例如 Feed-Forward 區塊、注意力層或全部——替換為稀疏 MoE 層，該層使用路由網路有效地將 token 分配給最適合的專家。

技術架構與命名慣例

SegMoE 模型保留與 Stable Diffusion 相同的基本架構，但將多個專家模型整合到單一框架中。命名慣例 SegMoE-AxB 定義模型的配置：

  • A：合併在一起的專家模型總數。
  • B：在生成每張圖像時實際參與的專家數量。

根據配置，僅會複製特定層（feed-forward 區塊、注意力層或兩者皆有）；所有其他參數與標準的 Stable Diffusion 模型保持一致。

可用模型發布

已在 Hugging Face Hub 上以 Apache 2.0 授權釋出三個預先合併的模型：

  1. SegMoE 2x1：使用兩個專家模型。
  2. SegMoE 4x2：使用四個專家模型。
  3. SegMoE SD 4x2：使用四個 Stable Diffusion 1.5 專家模型。

比較測試顯示，SegMoE 模型在提示詞理解方面較基礎模型（如 RealVisXL_V3.0）有提升，尤其在複雜情境下，如渲染多個物件（例如「三個綠色玻璃瓶」）或特定空間排列（例如「自由女神像位於華盛頓紀念碑旁」）。

實作與自訂

使用者可以使用 segmoe 套件（受 mergekit 函式庫啟發）建立自訂 MoE 模型。此過程包括定義 config.yaml 檔案，該檔案指定基礎模型路徑、專家數量、要混合的層類型（ff、attn 或 all），以及每個專家的來源模型，並提供其正向與負向提示詞以計算門控權重。

與 Hugging Face Diffusers 的整合

SegMoE 與 Hugging Face 生態系統整合。可使用 segmoe 函式庫中的 SegMoEPipeline 進行推論：

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")

效能取捨與硬體需求

雖然 SegMoE 提供了增強的功能，但也帶來了特定的計算開銷：

  • 推論速度：當每個 token 的專家數量大於一時，模型必須在多個專家模型上執行計算，因而比單一的 SD 1.5 或 SDXL 模型更慢。
  • VRAM 使用量：MoE 模型需要大量 VRAM。例如，SegMoE-4x2 在半精度下需要 24GB VRAM，較適合多 GPU 部署，而非本機單卡環境。

工具與發佈

透過 segmoe 工具建立的自訂模型可本地保存，或使用 huggingface-cli 或 huggingface_hub 函式庫的 upload_folder 功能推送至 Hugging Face Hub。

Sources