Vision Transformer with Mixture of Experts (ViT-MoE)

[LawJarp-A]

Model description

Vision Transformer with Mixture of Experts (ViT-MoE) is an efficient scaling approach for vision transformers that replaces dense feedforward layers with sparse mixture of experts.

This architecture enables scaling model capacity while keeping per-token computation roughly constant, achieving better accuracy-efficiency trade offs compared to dense ViTs.

Key Benefits:

• 90.35% ImageNet accuracy with ~50% compute reduction vs dense models
• Expert specialization emerges naturally (deeper layers specialize in object classes)
• Scalable to large model sizes without proportional compute increase

The implementation would follow existing MoE patterns in transformers (Mixtral, Switch Transformers) but adapted for vision tasks, trying to filling the current gap of vision MoE models in the library.

Implementation Plan:

1. Base Implementation: Use Google V-MoE as foundation
2. Enhanced with ViMoE Research: Apply latest empirical insights (if required)

Evaluation: Start with ImageNet-1K classification, then expand to other vision tasks

Open source status

• The model implementation is available
• The model weights are available

Provide useful links for the implementation

Primary Implementation Reference:

• Google V-MoE Implementation: https://github.com/google-research/vmoe
• Google V-MoE Paper: https://arxiv.org/abs/2106.05974
• Google V-MoE Blog: https://research.google/blog/scaling-vision-with-sparse-mixture-of-experts/

Additional Research Foundation:

• ViMoE Paper (ICLR 2025): https://arxiv.org/abs/2410.15732
  • OpenReview: https://openreview.net/forum?id=KaYXsoCxV7
  • Authors: @Frank-ZY-Dou, @JoinWei-PKU (partial list)
• Mobile V-MoEs: https://arxiv.org/abs/2309.04354
• M³ViT: https://arxiv.org/abs/2210.14793

Existing MoE Patterns in Transformers:

• Mixtral: src/transformers/models/mixtral/modeling_mixtral.py
• Switch Transformers: src/transformers/models/switch_transformers/modeling_switch_transformers.py
• Qwen2-MoE: src/transformers/models/qwen2_moe/modeling_qwen2_moe.py

Pre-trained Weights:

• Google V-MoE checkpoints available in their repository for various model sizes
• ImageNet-21k and ILSVRC2012 trained models

[PamelaBha]

I would like to contribute if there is an opportunity

[Player256]

Hey is there an opportunity to contribute. I'd be happy to be involved in this

[Rocketknight1]

The Google V-MoE model is quite old and outdated at this point, right? Is there a reason people would use it instead of newer models?

[LawJarp-A]

Well I agree the original Google V-MoE is from 2021, but the sparse-expert concept remains relevant. My initial goal was to implement it based on the recent ViMoe paper, but I felt the goal here is toprovide a baseline MoE-ViT implementation in Transformers that can then be extended with more recent advances like ViMoE (ICLR 2025) or M³ViT. Starting from V-MoE makes sense because it’s well documented, open sourced, and already has pretrained checkpoints, making it a practical entry point to fill the current gap of vision oriented MoEs
@Rocketknight1 thoughts?

[Rocketknight1]

That kind of doesn't fit with our philosophy! We generally only add architectures when they are needed to support pre-trained models. We don't usually extend existing architectures with newer advances; we wait until a model has been released that uses those advances, and then add it as a separate architecture

[Rocketknight1]

The main reason we do things this way is that there are lots and lots of ML papers with new layers and methods and optimizers and so on, and a lot of them never make it all the way to production use in big models. If we try to implement them all before a pre-trained model exists then we'll increase our workload 10X and end up with a codebase full of unused, dead code

[LawJarp-A]

I get the concern. This is actually one of my first contributions to a big open-source project, and the idea came up because I needed ViMoE for my own work, ended up using my own patched up version i wrote, and thought it’d be neat to have it in transformers.
Plus it would be great experience to contribute as well

[LawJarp-A]

Would it help that the Google V-MoE model is the only openly available model for this domain, there are no newer model for OSS fold to use even if they want to @Rocketknight1. We would be implementing something that people would use

[Rocketknight1]

cc @qubvel for that question - I'm not certain about the SOTA for vision models right now!

[LawJarp-A]

@qubvel any thoughts on this?
Google V-MoE is the only fully OS implementation I could find so far, and a first step for anyone intending to experiment with vision moe LMs.

[qubvel]

Hi @LawJarp-A, thanks for raising the discussion, really appreciate you are interested in contributing 🤗

I agree with what was said above by @Rocketknight1, and I have doubts that it's a preferred area of research these days for vision models. Otherwise, we would see many more papers and open models available. But I might be mistaken, and we will see them in the coming months. That said, I would rather wait for the next SOTA vision MoE model to be released before adding it to the transformers

[LawJarp-A]

@qubvel
Hey, thanks for weighing in.
I understand that vision MOEs are not the hot topic right now and we don't have any latest SOTA open source models. I guess doesn't make sense to integrate it to transformers unless there is demand, not just a few people who want it. I will close this issue.
I have worked with creating hacked versions of vision MOEs with the latest training routines etc, in case any similar contributions come in the future, would be happy to help.
Thanks :)