SegMoE: Segmind Mixture of Diffusion Experts

SegMoE is a framework designed to create Mixture-of-Experts (MoE) Diffusion models from scratch by combining pretrained models. It allows users to replace specific layers—such as Feed-Forward blocks, attention layers, or all of them—with sparse MoE layers that use a router network to efficiently assign tokens to the most suitable experts.

Technical Architecture and Naming Convention

SegMoE models maintain the same basic architecture as Stable Diffusion but incorporate multiple expert models into a single framework. The naming convention SegMoE-AxB defines the model's configuration:

  • A: The total number of expert models merged together.
  • B: The number of experts actively involved in generating each image.

Depending on the configuration, only specific layers (feed-forward blocks, attention, or both) are replicated; all other parameters remain identical to a standard Stable Diffusion model.

Available Model Releases

Three pre-merged models have been released on the Hugging Face Hub under the Apache 2.0 license:

  1. SegMoE 2x1: Utilizes two expert models.
  2. SegMoE 4x2: Utilizes four expert models.
  3. SegMoE SD 4x2: Utilizes four Stable Diffusion 1.5 expert models.

Comparison tests indicate that SegMoE models show improved prompt understanding compared to base models, such as RealVisXL_V3.0, specifically in complex scenarios like rendering multiple objects (e.g., "three green glass bottles") or specific spatial arrangements (e.g., "the statue of Liberty next to the Washington Monument").

Implementation and Customization

Users can create custom MoE models using the segmoe package, which was inspired by the mergekit library. The process involves defining a config.yaml file that specifies the base model path, the number of experts, the types of layers to mix (ff, attn, or all), and the source models for each expert along with their respective positive and negative prompts for computing gate weights.

Integration with Hugging Face Diffusers

SegMoE is integrated with the Hugging Face ecosystem. Inference can be performed using the SegMoEPipeline from the segmoe library:

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

Performance Trade-offs and Hardware Requirements

While SegMoE provides enhanced capabilities, it introduces specific computational overheads:

  • Inference Speed: When the number of experts per token is greater than one, the model must perform computations across multiple expert models, making it slower than a single SD 1.5 or SDXL model.
  • VRAM Usage: MoE models require significant VRAM. For example, SegMoE-4x2 requires 24GB of VRAM in half-precision, making it more suitable for multi-GPU deployments than local setups.

Tooling and Distribution

Custom models created via the segmoe tool can be saved locally or pushed to the Hugging Face Hub using the huggingface-cli or the upload_folder function from the huggingface_hub library.

Sources