Mixture of Modality-Experts (MoME) Integration

[githubshaurya]

I would like to add a Mixture of Modality Experts (MoME) block to nanoVLM in order to enable dynamic, modality‐specific routing between a small vision expert and a small text expert at each cross‐modal layer. The goal is to introduce a lightweight MoME architecture that allows it to learn richer modality‐specific representations.

Motivation

nanoVLM currently relies on a single shared transformer to process combined image and text inputs. Prior work shows that splitting modality processing into separate "experts" and routing between them can:

• Reduce cross-modal interference
• Enable modality-aware feature learning
• Provide interpretability via gating weights
  Integrating a lightweight Mixture-of-Modality-Experts (MoME) block into nanoVLM will let it learn to “soft switch” between a vision expert and a text expert per input, improving robustness without significantly increasing parameters.

Potential benefits

1. Reduced Modality Interference
   Separate routing of vision inputs through a Vision Expert and text inputs through a Text Expert prevents a single expert from juggling both modalities (MoME).

2. Modality-Aware Routing
   A small MLP router to decide how much weight to give to each expert (for every single input (every image+text pair)) (concatenating pooled vision/text embeddings → softmax) and multiply the output vector of each expert by its respective weight.

3. Interpretable Gating
   Logging "router weights" tell if an input was mostly handled by the Vision Expert or Text Expert. I²MoE suggests lightweight gating losses to encourage the router to use both experts in a meaningful way, rather than collapse to one .

Workflow

When you turn on MoME in nanoVLM, the model sends the picture through a small “vision expert” network and the words through a separate “text expert” network. Each expert distills its input into a representation of what it sees or reads. Then a tiny decision maker (the “router”) looks at both summaries and weighs each and blends them into one final representation. This representation goes into the usual nanoVLM head to make predictions. During training, the model learns both how to blend the representation and when to lean on vision versus text. Gating loss penalizes the model if it leans on one side.

Proposed Addition

1. In "models/config" :
   a. Add boolean flag 'use_mome': True.
   b. number of text and fusion experts.
   c. hidden dimension
   d. number layers

2. In "models/vision_transformer.py":
   a. number of vision experts
   b. return an output of each vision expert

3. In "models/vision_language_model.py":
   In class "VisionLanguageModel":
   a. Instantiation and implementation of Vision encoder, decoder
   b. Embedding text tokens while encoding
   c. Loss computation

4. New file: "models/mome_utils.py":
   a. Normalized weights for experts
   b. balancing loss for balancing the experts

References:

1. VLMo: Introduces Mixture-of-Modality-Experts(MOME) Transformer that can encode various modalities (images, text, and image-text pairs) within a Transformer block.
   

2. Uni-MoE: Introduces separate encoder for each modality and maps each modality into a unified representation space.

3. MoME: Introduces the idea of different experts for each modality.

4. I²MoE: Proposes to add a small regularization term (gating loss) on the router outputs to prevent one expert from dominating to ensure both experts stay useful. Introduces a multilayer perceptron (MLP) that assigns importance scores to each expert’s output.

[Naveen934]

In this code, we are using the Vision Transformer architecture and pre-trained model from 'google/siglip-base-patch16-224'. It doesn't have the MoME architecture, so if we added it, we would need to train the Vision Transformer from scratch. Do you know of any ViT pre-trained model with the MoME architecture?

[githubshaurya]

There isn’t a publicly available ViT checkpoint that already includes a MoME router. Instead, you can load a standard pretrained ViT (e.g., google/siglip-base-patch16-224) into your VisionExpert and only train the small routing and fusion layers on top. This way, you avoid retraining the entire ViT from scratch and leverage its existing pretrained weights.

[Naveen934]

There isn’t a publicly available ViT checkpoint that already includes a MoME router. Instead, you can load a standard pretrained ViT (e.g., google/siglip-base-patch16-224) into your VisionExpert and only train the small routing and fusion layers on top. This way, you avoid retraining the entire ViT from scratch and leverage its existing pretrained weights.

corrently after pixel_shuffle( ) they used self.proj = nn.Linear(self.input_dim, self.output_dim, bias=False) {in class ModalityProjector()} this line as fusion layer for converting img emb to text emb , do you want to replace this FFN to MoE architure ? or want to implement MoE in this line {class VisionLanguageModel() }
updated_token_embd = self._replace_img_tokens_with_embd(input_ids, token_embd, image_embd)

(#)The updated_token_embd is now the token_embd with image parts replaced.
(#)The attention_mask comes from the collator and should already cover the full sequence.
logits, _ = self.decoder(updated_token_embd, attention_mask=attention_mask)

[githubshaurya]

In the ModalityProjector class (where you currently have:
self.proj = nn.Linear(self.input_dim, self.output_dim, bias=False)

I want to replace that single Linear fusion layer with a small MoME subnetwork:

self.vis_expert = VisionExpert(hidden_dim=self.input_dim)
self.txt_expert = TextExpert(hidden_dim=self.input_dim)
self.router = MoMERouter(hidden_dim=self.input_dim)
self.fusion_proj = nn.Linear(self.input_dim, self.output_dim, bias=False)

A single FFN blindly projects the image representation into text space. That forces every example to use the same fusion, regardless of whether the image or text signal is more important.

MoME learns two separate “experts" and a router that decides trust in vision expert (when the image is most informative) vs. text experts (when the text is most informative). This reduces cross-modal interference and gives more flexibility than a single linear layer.

[Naveen934]

In the ModalityProjector class (where you currently have: self.proj = nn.Linear(self.input_dim, self.output_dim, bias=False)

I want to replace that single Linear fusion layer with a small MoME subnetwork:

self.vis_expert = VisionExpert(hidden_dim=self.input_dim) self.txt_expert = TextExpert(hidden_dim=self.input_dim) self.router = MoMERouter(hidden_dim=self.input_dim) self.fusion_proj = nn.Linear(self.input_dim, self.output_dim, bias=False)

A single FFN blindly projects the image representation into text space. That forces every example to use the same fusion, regardless of whether the image or text signal is more important.

MoME learns two separate “experts" and a router that decides trust in vision expert (when the image is most informative) vs. text experts (when the text is most informative). This reduces cross-modal interference and gives more flexibility than a single linear layer.

Hi @githubshaurya , Here they used ModalityProjector class only convert Image emp into text emp , your expectation is convert both input of text and image to Export block , ok will explain both model arch then you will understand

currently follow Arch

input (image , text) :
img -> ViT -> img emb ( this emb is unread by text decoder so we send through MP) -> MP -> img emb (able to read decoder)
text ->text encoder -> text emb

now input of img emb + text emb to flat --> docoder -> Answer

here MP not required Export Arch , current code is simple and enough i think