SegMoE：Segmind 扩散专家混合模型

SegMoE：Segmind 扩散专家混合模型

SegMoE 是一个框架，通过在 Stable Diffusion 架构中用稀疏 MoE 层替换前馈层，以创建混合专家（MoE）扩散模型，从而提升对提示词的理解。

SegMoE：Segmind 扩散专家混合模型

SegMoE 是一个框架，旨在通过组合预训练模型从头创建混合专家（MoE）扩散模型。它允许用户将特定层——例如前馈块、注意力层或全部——替换为稀疏 MoE 层，这些层使用路由网络高效地将 token 分配给最合适的专家。

技术架构与命名约定

SegMoE 模型保持与 Stable Diffusion 相同的基本架构，但将多个专家模型整合到一个框架中。命名约定 SegMoE-AxB 定义模型的配置：

  • A：合并在一起的专家模型总数。
  • B：在生成每张图像时实际参与的专家数量。

根据配置，仅特定层（前馈块、注意力或两者）会被复制；所有其他参数保持与标准 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 包创建自定义 MoE 模型，该包受 mergekit 库启发。过程包括定义一个 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 模型更慢。
  • 显存使用：MoE 模型需要大量显存。例如，SegMoE-4x2 在半精度下需要 24GB 显存，更适合多 GPU 部署，而非本地单机设置。

工具与分发

通过 segmoe 工具创建的自定义模型可以本地保存，或使用 huggingface-cli 或 huggingface_hub 库中的 upload_folder 函数推送到 Hugging Face Hub。

Sources