SegMoE: Segmind Mixture of Diffusion Experts

SegMoE is an exciting framework for creating Mixture-of-Experts Diffusion models from scratch! SegMoE is comprehensively integrated within the Hugging Face ecosystem and comes supported with diffusers 🔥!

Among the features and integrations being released today:

• Models on the Hub, with their model cards and licenses (Apache 2.0)
• Github Repository to create your own MoE-style models.

Table of Contents

• What is SegMoE
  • About the name
• Inference
  • Samples
  • Using 🤗 Diffusers
  • Using a Local Model
• Comparison
• Creating your Own SegMoE
• Disclaimers and ongoing work
• Additional Resources
• Conclusion

What is SegMoE?

SegMoE models follow the same architecture as Stable Diffusion. Like Mixtral 8x7b, a SegMoE model comes with multiple models in one. The way this works is by replacing some Feed-Forward layers with a sparse MoE layer. A MoE layer contains a router network to select which experts process which tokens most efficiently. You can use the segmoe package to create your own MoE models! The process takes just a few minutes. For further information, please visit the Github Repository. We take inspiration from the popular library mergekit to design segmoe. We thank the contributors of mergekit for such a useful library.

For more details on MoEs, see the Hugging Face 🤗 post: hf.co/blog/moe.

SegMoE release TL;DR;

• Release of SegMoE-4x2, SegMoE-2x1 and SegMoE-SD4x2 versions
• Release of custom MoE-making code

About the name

The SegMoE MoEs are called SegMoE-AxB, where A refers to the number of expert models MoE-d together, while the second number refers to the number of experts involved in the generation of each image. Only some layers of the model (the feed-forward blocks, attentions, or all) are replicated depending on the configuration settings; the rest of the parameters are the same as in a Stable Diffusion model. For more details about how MoEs work, please refer to the "Mixture of Experts Explained" post.

Inference

We release 3 merges on the Hub:

Samples

Images generated using SegMoE 4x2