MoMa: Efficient Early-Fusion Pre-training with Mixture of Modality-Aware Experts

## Abstract

Abstract We introduce MoMa, a novel modality-aware mixture-of-experts (MoE) architecture designed for pre-training mixed-modal, early-fusion language models. MoMa processes images and text in arbitrary sequences by dividing expert modules into modality-specific groups. These groups exclusively process designated tokens while employing learned routing within each group to maintain semantically informed adaptivity. Our empirical results reveal substantial pre-training efficiency gains through this modality-specific parameter allocation. Under a 1-trillion-token training budget, the MoMa 1.4B model, featuring 4 text experts and 4 image experts, achieves impressive FLOPs savings: 3.7 × \times × overall, with 2.6 × \times × for text and 5.2 × \times × for image processing compared to a compute-equivalent dense baseline, measured by pre-training loss. This outperforms the standard expert-choice MoE with 8 mixed-modal experts, which achieves 3× overall FLOPs savings (3 × \times × for text, 2.8 × \times × for image). Combining MoMa with mixture-of-depths (MoD) further improves pretraining FLOPs savings to 4.2 × \times × overall (text: 3.4 × \times × , image: 5.3 × \times × ),
although this combination hurts performance in causal inference due to increased sensitivity to router accuracy. These results demonstrate MoMa’s potential to significantly advance the efficiency of mixed-modal, early-fusion language model pre-training, paving the way for more resource-efficient and capable multimodal AI systems.

## 1 Introduction

Auto-regressive mixed-modal foundation models have shown significant promise in applications requiring the processing of mixed-modal inputs and the generation of mixed-modal outputs (Gemini et al., , ; OpenAI et al., ; Lu et al., ; Chameleon Team, ). These models have demonstrated remarkable capabilities in tasks ranging from visual question answering to multimodal content generation, pushing the boundaries of AI’s ability to understand and interact with our inherently multimodal world.

While a popular architecture design for mixed-modal foundation models involves fusing modality-specific encoders or decoders (Gemini et al., , ; Lu et al., ; OpenAI et al., ), this approach can limit the model’s ability to integrate information across modalities and generate content with interleaved modalities. To address this limitation, Chameleon (Chameleon Team, ) recently introduced a single transformer architecture with a next-token prediction objective to model mixed-modal sequences consisting of discrete image and text tokens, allowing for seamless reasoning and generation across modalities.
Chameleon, pretrained on approximately 10 trillion mixed-modal tokens, has demonstrated broad vision and language capabilities across various downstream tasks. Notably, it surpasses commercial baselines such as Gemini 1.0 Pro(^1^11[https://ai.google.dev/](https://ai.google.dev/)) and GPT-4V(^2^22gpt-4-vision-preview: [https://openai.com/api/](https://openai.com/api/)) in generating mixed-modal long-form responses (Chameleon Team, ). However, scaling such mixed-modal early-fusion foundation models to greater capacities presents significant computational challenges.

To address these challenges, we investigate the application of routed sparse architectures (Lepikhin et al., ; Fedus et al., ; Clark et al., ; Jiang et al., ; Raposo et al., ). These architectures have previously shown effectiveness in scaling language and vision-specific foundation models, as well as multimodal contrastive learning (Mustafa et al., ). However, their application to mixed-modal early-fusion models presents unique opportunities and challenges.

The insight driving our approach is the inherent heterogeneity of modalities: text and image tokens possess distinct information densities and redundancy patterns (Liang et al., ). While we integrate these tokens into a unified early-fusion architecture, we propose further optimizing this framework by incorporating modality-specific modules. This concept, which we term *modality-aware sparsity* (MaS), enables models to better capture features specific to each modality while maintaining strong cross-modality integration through partial parameter sharing and attention mechanisms. Previous work such as VLMo (Bao et al., ), BEiT-3 (Wang et al., ) and VL-MoE (Shen et al., ) have applied mixture-of-modality-experts to train vision-language encoders and masked language models. We apply this type of sparse component to the mixed-modal, early-fusion language modeling setting.

Specifically, we adopt Chameleon as the base transformer architecture and introduce sparsity along two dimensions:

- 1.
Width: We apply *mixture-of-experts* (MoE) (Lepikhin et al., ) where tokens are routed across a set of feed-forward blocks (experts) at each layer. Crucially, we divide these experts into modality-specific groups, with each group processing only tokens of its designated modality while learned routing is done within each group for semantics-based adaptivity. We dub this approach *mixture of modality-aware experts* (MoMa).
- 2.
Depth: We apply *mixture-of-depths* (MoD) (Raposo et al., ) where tokens can selectively skip the attention and feed-forward computation at certain layers.

For both dimensions, we adopt expert-choice routing (Zhou et al., ) to ensure load balancing and maintain a static computation graph (Paszke et al., ), both important for high training throughput.

We conduct extensive FLOPs-controlled experiments comparing our proposed architecture with the dense baseline and multiple sparse variations. With a 1-trillion-token training budget, Chameleon-MoMa 1.4B employing 4 text experts and 4 image experts achieves a substantial 3.7$\times$ overall FLOPs savings (text: 2.6$\times$, image: 5.2$\times$) compared to the 1.4B compute-matched dense baseline as measured by pre-training loss,
while maintaining a relatively modest -17% throughput reduction.
In contrast, the standard expert-choice MoE with 8 mixed-modal experts achieves 3$\times$ FLOPs savings (text: 3$\times$, image: 2.8$\times$) under the same setting, with -9% throughput reduction.
Further combination with MoD (Chameleon-MoMaD) boosts the FLOPs savings to 4.2$\times$ overall (text: 3.4$\times$, image: 5.3$\times$) as measured by pre-training loss. However, the auto-regressive inference performance of the MoD model is compromised due to its sensitivity to routing accuracy, resulting in subpar performance compared to Chameleon-MoMa (§4.6).

We also demonstrate that it is possible to enhance the performance of Chameleon-MoMa using a modality-untied upcycling technique (Komatsuzaki et al., ). In particular, we can initialize the model with a seed sparse architecture consisting of 1 expert per modality, and effectively improve router learning for each modality. By training the seed model for only 10k steps, the resulting model achieves additional FLOP reduction while maintaining or improving performance.

Our main contributions can be summarized as follows:

- 1.
We introduce mixture of modality-aware experts (MoMa) for mixed-modal early-fusion models, allowing for more efficient parameter allocation and utilization.
- 2.
We investigated sparse scaling along both the width and depth dimensions of transformers, demonstrating substantial speedup of pre-training loss convergence can be realized in each direction. However, effectively combining these two approaches for better performance in an autor-regressive inference setup remains an open research challenge.
- 3.
We present a modality-untied upcycling technique that further enhances the efficiency of our sparse architecture.
- 4.
We provide extensive empirical results and analysis, offering insights into our approach’s scaling behavior and efficiency gains.

## 2 Model

Figure: Figure 1: Overview of our proposed multimodal early-fusion architecture. (a) The early-fusion architecture processes interleaved text and image data as a unified token sequence using a transformer model. (b) Detailed architecture of the mixture of modality-aware experts (MoMa) transformer block, where layer norms and the residual connection of attention layer are omitted for readability. The input sequence of tokens are first routed to different expert groups based on their modality, undergo learned routing within each group, and then have their output restored to a tensor of the original input shape of the FFN module.
Refer to caption: x1.png

### 2.1 Early Fusion

Our model builds upon the early fusion architecture introduced by Chameleon (Chameleon Team, ), which represents images and text as a series of discrete tokens within a unified transformer.
The core of Chameleon is a transformer-based model that applies self-attention mechanisms over the combined sequence of image and text tokens. This allows the model to capture complex relationships between and within modalities. The model is trained using a next-token prediction objective, learning to generate both text and image tokens autoregressively.

In Chameleon, images are tokenized using a learned image tokenizer that encodes a 512 $\times$ 512 image into 1024 discrete tokens from a codebook of size 8192. Text is tokenized using a BPE tokenizer with a vocabulary size of 65,536, which includes the 8192 image codebook tokens. This unified tokenization scheme enables the model to process arbitrary sequences of interleaved image and text tokens.

By adopting this early fusion approach, our model inherits several key advantages:

- 1.
*Unified representation:* The model learns a shared representation space for images and text, facilitating cross-modal reasoning and generation.
- 2.
*Flexibility:* The architecture can handle arbitrary sequences of images and text, enabling diverse multimodal tasks such as image captioning, visual dialogue, and mixed-modal document generation.
- 3.
*Scalability:* The token-based approach allows for uniform processing of both text and image data, enabling efficient scaling to larger model sizes and diverse datasets. This is evidenced by Chameleon’s successful training on approximately 10 trillion mixed-modal tokens.
- 4.
*End-to-end learning:* The entire auto-regressive model, is trained end-to-end, allowing for joint optimization of the representation and task-specific performance.

Building on this foundation (Figure 1a), our work introduces modality-aware sparsity techniques to further enhance the efficiency and performance of early fusion models, as detailed in the following sections. These techniques aim to address the computational challenges of scaling early fusion models while maintaining their powerful cross-modal reasoning capabilities.

### 2.2 Width Scaling: Mixture of Modality-Aware Experts

We propose a width scaling approach that incorporates modality-aware block sparsity in the feed-forward module, extending the standard mixture-of-experts (MoE) architecture (Lepikhin et al., ; Fedus et al., ; Wang et al., ). The insight driving this approach is that tokens of different modalities have distinct characteristics and information densities. By creating separate expert groups for each modality, we allow the model to develop specialized processing pathways while maintaining the ability to integrate information across modalities.

We describe the key components in the mixture of modality-aware experts (MoMa) formulation below (Figure 1b).

#### Modality-Specific Expert Groups.

We divide the experts in each MoE layer into distinct groups, each specialized in processing tokens from a specific modality $\mathcal{E}=\{E_{1}^{\mathcal{T}}\ldots,E_{m}^{\mathcal{T}},E_{1}^{\mathcal{I}
}\ldots,E_{n}^{\mathcal{T}}\}$: one group for processing text tokens and another for image tokens(^3^33While we focus on text and image modalities in this work, our formulation is designed to be easily extensible to accommodate an arbitrary number of modalities.). This separation allows each group to specialize in features relevant to its respective modality.

By implementing modality-aware block sparsity, we aim to achieve several benefits:

- •
Improved Efficiency: by routing tokens to modality-specific experts, we reduce the computational overhead of processing tokens with experts not specialized for their modality.
- •
Enhanced Specialization: modality-specific expert groups can develop more refined features relevant to their respective modalities.
- •
Maintained Cross-Modal Integration: despite the separation into modality-specific groups, the model still can integrate information across modalities through the shared self-attention mechanisms in non-MoE layers.

#### Hierarchical Routing.

We adopt a token-based routing mechanism (Lepikhin et al., ; Fedus et al., ; Jiang et al., ). For an input token $x$, our routing mechanism operates in two stages.

- 1.
Modality-aware routing: tokens are first routed to their corresponding modality-specific expert group based on their modality (text $\mathcal{T}$ or image $\mathcal{I}$).
- 2.
Intra-modality routing: within each modality-specific expert group $\mathcal{E}^{\mathcal{M}}$, tokens are then routed to specific experts using a learned routing function. Specifically, we use a projection matrix $W_{g}^{\mathcal{M}}\in\mathbb{R}^{d\times|\mathcal{E}^{\mathcal{M}}|},d=\text{
transformer hidden dimension}$, to compute the token-to-expert affinity scores, following other established MoE formulations (Lepikhin et al., ; Fedus et al., ; Jiang et al., ).

#### Expert choice.

Within each modality group, we implemented expert-choice (EC) routing (Zhou et al., ), where each expert has a fixed bucket size and processes the top-$k_{e}$ tokens in a batch. We set $k_{e}=b^{\mathcal{M}}\cdot c_{e}$,
where $c_{e}=\frac{1}{|\mathcal{E}^{\mathcal{M}}|}$ is a capacity factor and $b^{\mathcal{M}}$ is the total number of tokens of modality $\mathcal{M}$ in a batch.
As a result, each token can be routed to a variable number of experts.

EC routing ensures balanced expert utilization during training and eliminates needing a separate load-balancing loss term. Maintaining a single loss term promotes more stable optimization and a more straightforward training process.
However, EC routing compromises causality in auto-regressive language modeling, as each expert selects the top tokens to process in a batch by comparing token scores across all tokens. We use a combination of two techniques to address this issue and enable expert-choice training for auto-regressive LMs.

- 1.
We use Sigmoid as the non-linearity in the router scoring function, enabling independent calculation of token-to-expert affinity scores for each token.
- 2.
We introduce auxiliary routers, inspired by Raposo et al. (), which predict the likelihood of an expert selecting a token solely based on its hidden state representation. These routers are trained after the main model training is completed and employed during inference to ensure causality. We discuss the details of the auxiliary routers in §2.3.

In summary, the MoMa module for an input token $x$ can be formally defined as:

$$ $y=\begin{cases}\sum_{j=1}^{m}g^{\mathcal{T}}(x)_{j}\cdot{\text{FFN}^{\mathcal{ T}}_{\text{SwiGLU}}}_{j}(x)&\text{if }x\in\mathcal{T}\\ \sum_{j=1}^{n}g^{\mathcal{I}}(x)_{j}\cdot{\text{FFN}^{\mathcal{I}}_{\text{ SwiGLU}}}_{j}(x)&\text{if }x\in\mathcal{I}\end{cases}$ (1) $$

$$ $g^{\mathcal{M}}(x)_{j}=\begin{cases}\text{Sigmoid}(x\cdot W^{\mathcal{M}}_{g}) _{j}&\text{if }x\text{ is selected by }E_{j}\\ 0&\text{otherwise}\end{cases}$ $$

We further apply residual connection and the Swin Transformer normalization (Liu et al., ) post MoMa computation. Our experiments, detailed in later sections, demonstrate that MoMa significantly improves efficiency and performance compared to dense baselines and standard MoE architectures.

### 2.3 Mixture-of-Depths

Figure: Figure 2: Architecture of transformer layer consisting of MoMa combined with mixture-of-depths (MoD).
Refer to caption: extracted/5787609/figures/moma_arch_mod.png

We further investigate introducing sparsity in the depth dimension. Prior work explores sparisty in depth through either stochastic layer drop (Elhoushi et al., ) or through learnable routers (Raposo et al., ). We focus on learnable routers, and incorporate the recently proposed mixture-of-depths (MoD) technique (Raposo et al., ).
Specifically, in each MoD layer, we integrate MoD prior to any mixture-of-experts (MoE) routing, ensuring it is applied to the full batch before modality split (Figure 2).(^4^44This corresponding to the staged MoDE implementation presented in Raposo et al. (). We chose this implementation because it also introduces sparsity in attention, whereas the integrated MoDE implementation proposed in the same paper always computes full attention.)

Following Raposo et al. (), for each MoD layer, we use a projection matrix $\hat{W_{g}}\in\mathbb{R}^{d\times 1}$ to compute the token-to-layer affinity score, followed by a Sigmoid non-linearity

$$ $\hat{g}(x)=\text{Sigmoid}(x\cdot\hat{W_{g}}).$ (2) $$

Analogous to expert-choice routing in MoE, we set a fixed capacity of $k_{d}$ tokens, selected from the top-scoring tokens in a batch. We set $k_{d}=b\cdot c_{d}$, where $c_{d}\in(0,1]$ is a capacity factor chosen empirically. In practice, we first fixed the MoD layer interval and capacity factor $c_{d}$, then adjust the total number of transformer layers to ensure the resulting architecture has a comparable FLOPs per token to the base architecture.

### 2.4 Inference

We cannot directly apply the expert-choice routing for MoE and layer-choice routing for MoD during inference time, as the top-$k$ token selection within a batch breaks causality. Inspired by Raposo et al. (), to ensure causality during inference, we introduce auxiliary routers, to predict the likelihood of a token being selected by an expert or a layer solely based on its hidden representation.

Formally, for each MoE layer, we introduce an auxiliary router

$$ $g^{\mathcal{M}}_{\text{aux}}(x)=\text{Sigmoid}(\text{SiLU}(x\cdot W^{\mathcal{ M}}_{a1})\cdot W^{\mathcal{M}}_{a2}),$ (3) $$

where $W^{\mathcal{M}}_{a1}\in\mathbb{R}^{d\times(d/2)}$ and $W^{\mathcal{M}}_{a2}\in\mathbb{R}^{(d/2)\times|\mathcal{E}^{\mathcal{M}}|}$.

For each MoD layer, we introduce an auxiliary router

$$ $\hat{g}^{\mathcal{M}}_{\text{aux}}(x)=\text{Sigmoid}(\text{SiLU}(x\cdot\hat{W} _{a1})\cdot\hat{W}_{a2}),$ (4) $$

where $\hat{W}_{a1}\in\mathbb{R}^{d\times(d/2)}$ and $\hat{W}_{a2}\in\mathbb{R}^{(d/2)\times 1}$.

We employ a two-stage training approach, where the main model and auxiliary routers are trained separately. First, we train the main model to convergence. Then, we train the auxiliary routers using binary cross-entropy loss, supervised by the ground-truth top-$k$ routing assignments computed over an entire batch. At inference time, the main routers are only used to generate the weight values, and tokens are selectively routed to an expert or layer based on thresholding on the auxiliary router score (> 0.5).

### 2.5 Upcycling

Training MoE architectures with learnable routers from scratch presents unique challenges in optimizing both the representation space and routing mechanism (Xue et al., ).
The insight we identified is that MoE routers are responsible for partitioning the representation space for each expert. However, this representation space is sub-optimal in the early stages of model training, leading to a sub-optimally trained routing function.

To address this limitation of router training, we propose an upcycling approach, inspired by Komatsuzaki et al. ().
Specifically, we begin by training an architecture consisting of 1 FFN expert per modality. After a predetermined number of steps, we upcycle this model by converting each modality-specific FFN into an expert-choice MoE module, initializing each expert with the expert trained from the first stage.
We reset the learning rate scheduler while preserving the data loader state from the previous stage, thereby ensuring the second stage training is exposed to refreshed data.

To promote expert specialization, we augment the MoE routing function with Gumbel noise (Liu et al., ; Geng et al., ), allowing our router to differentiably sample experts. This is expressed in Equation 5:

$$ $\text{Gumbel-Sigmoid}(x)=\text{Sigmoid}(x+G^{{}^{\prime}}-G^{{}^{\prime\prime}})$ (5) $$

where $G^{{}^{\prime}}$ and $G^{{}^{\prime\prime}}$ are independent Gumbel noise samples.

The upcycling approach, combined with the Gumbel-Sigmoid technique, allows us to overcome the limitations of learned routers and achieve improved performance in our modality-aware sparse architecture.
In practice, we found that a short period of 10k-20k steps in the first-stage training significantly improves model training efficiency and stability, aligning with the findings of Xue et al. ().

## 3 Efficiency Optimization

To facilitate the distributed training of mixture of modality-aware experts (MoMa), we employ Fully Sharded Data Parallel (FSDP) (Zhao et al., ). However, this approach presents unique efficiency challenges compared to vanilla MoEs. This section discusses these challenges and our strategies to address them.

### 3.1 Load Balancing

Without constraints, load imbalance can occur in our system because the ratio of text to image tokens can vary significantly across different GPUs and iterations. Imbalances can create a cascading straggler effect, delaying weight prefetching for subsequent layers and gradient releases from previous layers. This effectively bounds the training latency by the maximum time required to process text and image experts across all GPUs within a batch.

To mitigate these issues, we developed a balanced data mix that aligns the text-to-image data ratio with the expert ratio on each GPU. This approach ensures load balancing in expectation. While alternative rebalancing algorithms, such as token redistribution at each FFN layer, are possible, they may introduce additional communication overhead.

### 3.2 Efficient Expert Execution

We explored several strategies to efficiently execute experts for different modalities. The first approach restricts to homogeneous experts across modalities and prohibits routing of text tokens to image experts and vice versa. This method allows processing of all tokens and modalities simultaneously, provided all experts share the same token count. Alternatively, we could enhance execution efficiency by employing block sparsity (Gale et al., ), which offers similar benefits to the first approach without requiring perfect expert token balance. Additionally, we considered running experts from different modalities sequentially when the number of modalities is limited. This approach allows better overlap of computation from previous modality experts with weight prefetching for subsequent modality experts, alleviating memory pressure. It also removes assumptions about expert load balance.

Given that our experiments involve a sufficiently large token count per GPU, hardware utilization is not a major concern even with multiple batched matrix multiplications. Thus, we find the sequential approach to be a clean and flexible choice for our experimental environment at the current scale.

### 3.3 Additional Optimizations

We implemented several optimizations to further enhance throughput. These include generic optimizations such as gradient communication quantization and automatic GPU kernel fusion, as well as graph optimizations via torch.compile (Ansel et al., ). Additionally, we developed MoMa-specific optimizations, including the reuse of modality token indices across different layers to minimize device synchronization between CPU and GPU(^5^55This optimization is incompatible with MoD in its current form and was not used in our final experiments to ensure fair cross-comparison across model variations. However, it can be modified to eliminate device synchronization with permutation while tracking live token counts in each MoD layer.). We also consolidated per-layer stats communication and moved these operations off the critical path of training.

## 4 Experiments

### 4.1 Setup

We use the same pre-training dataset and preprocessing as Chameleon Team (). To assess scaling performance, we train all models with over 1 trillion tokens. Unless specified otherwise, we employ a sequence length of 4096 tokens and a model parallel size of 1. Our training regimen includes a peak learning rate of $1e-4$, a 4000-step warm-up period, and linear annealing of the learning rate to 1% of its peak value.
For all MoE architectures, we implement MoE in every layer, setting each expert’s training capacity $k_{e}$ to $\frac{b^{\mathcal{M}}}{|\mathcal{E}^{\mathcal{M}}|}$ to maintain FLOPs per token comparable to the base dense model. In MoD architectures, we implement MoD in alternating layers, beginning with layer 0, using a layer capacity factor $c_{d}$ of 25%. To achieve FLOPs parity with the base dense model, we increase the total layer count while maintaining a constant hidden dimension. Table 1 provides detailed configurations for our dense and sparse models. Additional pre-training details are available in Appendix A.1.

For model comparison, we report training losses. Given that our 1-trillion-token training budget covers less than one epoch of our extensive pre-training data, we use training loss as a proxy for validation performance. Our use of expert-choice routing in both MoE and MoD modules introduces a caveat: the training loss calculation compromises causality, as token selection considers the top proportion from a batch, including future tokens. We address this in §4.6 by reporting test-time performances using auxiliary routers, demonstrating result generalization to validation settings and causal scenarios.

**Table 1: Specifications of dense and sparse architectures used in our experiments. Architectures listed in the same row are compute-matched with the same amount of active FLOPs per token. *params*: total number of parameters; $l$: number of layers; $d$: transformer hidden dimension; $ffn$: feed-forward module hidden dimension; $h$: number of attention heads; $e$: number of experts; $i$: MoE/D layer interval; $c$: MoE expert capacity/MoD layer capacity.**
| Dense |  | MoE |  | MoD |  | MoDE |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| *params* | $l$ | $d$ | $ffn$ | $h$ |  | *params* | $e$ | $i$ | $c_{e}$ |  | *params* | $l$ | $d$ | $h$ | $i$ | $c_{d}$ |  | *params* |  |  |
| 90M | 8 | 512 | 2048 | 8 |  | 210M | 8 | 1 | 0.125 |  | 110M | 14 | 512 | 8 | 2 | 0.25 |  | 317M |  |  |
| 435M | 24 | 1024 | 4096 | 16 |  | 1.9B | 8 | 1 | 0.125 |  | 635M | 40 | 1024 | 16 | 2 | 0.25 |  | 3B |  |  |
| 1.4B | 24 | 2048 | 8192 | 16 |  | 7.1B | 8 | 1 | 0.125 |  | 2.3B | 32 | 2304 | 18 | 2 | 0.25 |  | 12B |  |  |

Figure: Figure 3: Scaling of performance with compute. We consider three dense model sizes – 90M, 435M and 1.4B parameters – and their compute-matched sparse variations. We highlight three configurations of mixture-of-experts (MoE): 1) using 1 text expert and 1 image expert (moe_1t1i), 2) using 4 text experts and 4 image experts (moe_4t4i), and 3) using 8 mixed-modal experts (moe_8x) in addition to the dense baseline. All models were pre-trained and annealed to 1 trillion tokens. When controlling for FLOPs per token, all sparse architectures demonstrate better scaling properties compared to the dense baseline. In particular, architectures with modality-specific experts consistently outperforms architectures with only mixed-modal parameters, with particular large benefits shown on the image modality.
Refer to caption: extracted/5787609/figures/main_scaling_curves.png

### 4.2 Scaling of Performance with Compute

We present the scaling performance of various models across multiple compute levels, with FLOPs matched to three dense model sizes: 90M, 435M, and 1.4B parameters. We report two key metrics: (1) training loss and (2) pre-training speed-up factor $\eta$, which indicates that a sparse model can match the pre-training loss of the iso-FLOPs dense baseline using only $1/\eta$ of total FLOPs(^6^66Our definition of $\eta$ is analogous to the speed-up factor proposed by Artetxe et al. (), but is defined in terms of pre-training loss whereas the original definition uses validation perplexity.).

Figure: Figure 4: Effect of combining with mixture-of-depths (MoD). We consider the sparse model variations compute-matched to three dense model sizes – 90M, 435M and 1.4B parameters. Combining MoD with MoE in general lead to better pre-training loss convergence across model sizes and configurations.
Refer to caption: extracted/5787609/figures/combine_depth_and_width_scaling.png

#### Modality untying.

Introducing modality-specific expert groups improves pre-training efficiency across different scales, with particularly large benefits for the image modality. As shown in Figure 3, the moe_1t1i configuration, using 1 image expert and 1 text expert, significantly outperforms the dense baseline. The image loss of moe_1t1i nearly matches that of the regular expert-choice MoE model using 8 mix-modal experts (moe_8x), whereas the text loss remains substantially higher, indicating parameter untying has a disproportionate impact on the image modality. Scaling the number of experts within each modality group further improve the model performance. The moe_4t4i configuration, using 4 text experts and 4 image experts, consistently outperforms moe_8x across different scales, with a substantial margin in image loss. However, this comparison reveals a slight regression in the text modality, suggesting that processing text modality across more experts may be beneficial.

#### Mixture-of-Depths-and-Experts.

We observe accelerated convergence of training loss when adopting both MoE and MoD, as well as their combinations. As shown in Figure 4, adding MoD to the moe_1t1i architecture (mod_moe_1t1i) significantly improves model performance across various model sizes. Moreover, mod_moe_1t1i performs on par with or better than moe_4t4i across model sizes and modalities, indicating that introducing sparsity over the depth dimension can also effectively improve training efficiency. On the other hand, we observe diminishing returns when stacking MoD and MoE.
Adding MoD to the moe_4t4i architecture results in only a mild performance boost compared to mod_moe_1t1i and moe_4t4i. The improvement is also more visible for the text modality, while the image improvement is less significant. These findings suggest that future research may explore the combination of width and depth scaling to further enhance text modality performance. In contrast, improving image modality performance will require additional exploration of alternative approaches.

### 4.3 Scaling Number of Experts

Figure: (a) The same number of experts are allocated to each modality
Refer to caption: extracted/5787609/figures/scaling_num_experts_balanced.png

We conducted further ablation experiments to investigate the impact of scaling the number of experts. We explored two scenarios: allocating an equal number of experts per modality (balanced expert allocation) and allocating different numbers of experts per modality. We included the dense model and expert-choice MoE with 8 mixed-modal experts (moe_8x) as baselines.

#### Balanced expert allocation.

Figure 5(a) demonstrates that training loss consistently improves as the number of experts increases. Text and image losses exhibit distinct scaling patterns. While text loss improves steadily with each doubling of expert count, image loss shows diminishing returns as the number of experts increases from 2 to 4. This suggests that the intrinsic characteristics of each modality lead to different behaviors in sparse modeling. This finding aligns with our previous observation in §4.2, where untying modalities over 8 experts significantly enhanced image performance but not text performance, indicating that the text modality may benefit more from a larger number of experts.

#### Imbalanced expert allocation.

Given the diminishing returns observed when allocating more than 2 experts to the image modality, we experimented with configurations allocating fewer experts to the image modality and more to the text modality. Figure 5(b) compares three configurations with the same total number of experts (8), varying the allocation between text and image modalities (moe_7t1i, moe_6t2i and moe_4t4i). We observed that allocating more experts to a modality generally enhances its performance. However, assigning more than 4 experts to the text modality also yields diminishing returns. The overall training losses of the three configurations converge to similar levels. Considering the better load balancing of balanced expert assignment with our pre-training data mixture, we selected moe_4t4i as our optimal MoE configuration, despite moe_6t2i having a slightly lower average loss. We leave the design of sparse architectures that can effectively leverage the intrinsic characteristics of different modalities to future work.

### 4.4 Upcycling

We investigate further the influence of upcycling as described in §2.5. Specifically, we examine the 2.3B MoD model, comparing the training dynamics of mod_moe_4t4i when trained from scratch versus when initialized from an mod_moe_1t1i checkpoint. To ensure a fair comparison, we adjust the data loader and training steps to account for the number of training steps already completed by mod_moe_1t1i, thereby maintaining equivalent training FLOPs. We ablate initializing the model from 10k and 20k steps.

Figure 6 presents a comparison of the training curves for the various model variations. The training curves for the upcycled models have been adjusted to account for the computational cost of the first stage. We experimented with two seed checkpoints: mod_moe_1t1i trained for 10k and 20k steps, respectively. Our results show that upcycling further enhances model training, yielding a 1.2$\times$ FLOPs gain with 10k steps in the first stage and a 1.16$\times$ FLOPs gain with 20k steps. We observe the performance gap between upcycled models and the from-scratch model widens throughout training.

#### Optimal Upcycling Period.

As shown in Figure 6, both initializations outperform the base MoE model, with 10k steps providing a 1.2$\times$ speedup and 20k steps providing a 1.16$\times$ speedup compared to the model trained from scratch. This suggests that a sweet spot likely exists for upcycling, where undertraining the seed model offers minimal benefits for router convergence, while overtraining it impedes future specialization. Based on these findings, we recommend upcycling from 10k steps. However, we hypothesize that the optimal number of upcycle steps may change when training beyond 1T tokens, and we leave this exploration for future research on in-depth upcycling of MoMa.

Figure: Figure 6: Upcycling experiments. We compare the training curves of 2.3B mod_moe_4t4i starting from three different initialization points: (1) scratch, (2) 2.3B mod_moe_1t1i trained for 10k steps, and (3) 2.3B mod_moe_1t1i trained for 20k steps. All curves are FLOP adjusted to be equivalent, by shifting the upcycled models by the number of steps the dense model has been trained on.
Refer to caption: extracted/5787609/figures/upcycling_23b_finished.png

### 4.5 Throughput Analysis

Sparse models often cannot immediately lead to performance gains due to their added dynamism and associated data balancing issues (§3). To quantify the impact of our proposals on training efficiency, we conducted a controlled experiment comparing the training throughput of various architectures, including a combination of MoMa, mixture-of-experts (MoE), and mixture-of-depths (MoD).
We compare architectures FLOPs-matched to the 435M dense baseline. We performed our experiments on 256 A100 GPUs with a sequence length of 4096 and a batch size of 6 per GPU and summarized the model throughput in Table 2.

**Table 2: Training throughput (wps/GPU) of various architectures compared with the same flops as a 435M dense model on 256 A100 GPUs. Modality-conditioned sparsity offers a good quality-throughput tradeoff compared to the dense model, and exhibited reasonable scalability with more experts. On the other hand, while MoD variants achieves the best absolute loss, they generally incurs larger overheads due to additional dynamism and imbalance.**
| Configuration | WPS | % diff |
| --- | --- | --- |
| Dense | 31,970 | – |
| MoMa 8X | 28,990 | -9% |
| MoMa 1t1i | 30,196 | -6% |
| MoMa 4t4i | 26,424 | -17% |
| MoD MoMa 1t1i | 25,367 | -21% |
| MoD MoMa 4t4i | 22,114 | -31% |

Comparing expert-choice MoE (moe_8x) and dense models, introducing sparsity does incur overheads (-9%). This loss in throughput likely comes from the need to compute routing decisions and synchronize gradients for all experts despite being FLOPs-equivalent. On the other hand, as discussed in §3, running experts sequentially by modality (moe_1t1i) does not suffer from large execution overheads, and we can attribute most of the loss (-6%) to computing token indices for each modality, which can be amortized by precomputing the indices and share with each transformer layer.

When combining modality conditioned feed-forward with learned routing (moe_4t4i), we observe a smooth degradation in throughput as the number of experts increases, incurring additional 11% overheads with 8 experts, which is on par with the throughput loss of 9% when transitioning from the dense model to MoE with 8 experts.

While combining MoD and MoE achieves the best training loss in our experiments, introducing MoD triggers an estimated throughput loss of 15% . This is because MoD architectures require an additional router on the depth dimension, introducing complexity and potential bottlenecks. Furthermore, when combined with MoMa, MoD may exacerbate system imbalance due to the varying active tokens for each modality at each layer, deviating from the predefined token mix ratio in the dataset. This deviation undermines the assumption of efficient MoMa execution, particularly with larger world sizes.
To mitigate this issue, we can force the MoD routers to admit tokens based on the predefined token mix ratio. However, the impact on model quality requires further investigation, which we leave for future work.

### 4.6 Inference-time Performance

We evaluate our models (1.4B dense, 1.4B moe_1t1i, 1.4B moe_8x, 1.4B moe_4t4i, 2.3B mod_moe_4t4i) on held-out language modeling data and downstream tasks. We report perplexity on subsets of OBELICS (Laurençon et al., ) and Shutterstock(^7^77[https://www.shuitterstock.com/](https://www.shuitterstock.com/)), which are held-out from our pre-training dataset. In addition, we report 0-shot performance on a set of commonsense reasoning of tasks commonly used for benchmarking pre-trained language models. We also selected several vision-language task datasets and report the perplexity of the ground truth output in these datasets for cross model comparison, where we format the examples in a zero-shot manner.(^8^88We do not report the generation evaluation metrics for these tasks, as our experiment checkpoints are under-trained and have not achieved sufficient image understanding and generation capabilities to produce meaningful results on downstream tasks.)

- •
Commonsense Reasoning: We evaluated on a set of text-only tasks commonly used for benchmarking the commonsense reasoning capabilities of language models: PIQA (Bisk et al., ), SIQA (Sap et al., ), HellaSwag (Zellers et al., ), WinoGrande (Sakaguchi et al., ), ARC-Easy (Clark et al., ), ARC-Challenge (Clark et al., ), OpenBookQA (Mihaylov et al., ), and BoolQ (Clark et al., ). We score the prompt with each candidate’s answer and compute accuracy using the candidate with the highest score.
- •
Image Captioning: We take the Karpathy test split of MS-COCO (Lin et al., ) and the Karpathy test split of Flickr30k (Plummer et al., ), and report text-to-image and image-to-text conditional perplexity using these two datasets.
- •
Visual Question Answering: We report the perplexity of ground truth answers on the test-dev split of VQA-v2 (Goyal et al., ).

**Table 3: Interleaved data modeling performance.**
|  | Image to text (PPL $\downarrow$) | Text to image (PPL $\downarrow$) | Interleaved (PPL $\downarrow$) |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | COCO | Flickr | VQAV2 | COCO | Flickr | Obelics<br>Overall | Obelics<br>Text | Obelics<br>Image | SSTK<br>Overall | SSTK<br>Text | SSTK<br>Image |
| 1.4B Dense | 21.7 | 28.6 | 20.0 | 458.3 | 559.4 | 44.6 | 14.3 | 64.7 | 238.5 | 48.7 | 245.8 |
| 1.4B MoMa 1t1i | 21.9 | 29.3 | 19.8 | 416.7 | 508.6 | 41.5 | 13.8 | 59.3 | 208.1 | 7.6 | 221.5 |
| 1.4B MoE 8x | 18.9 | 24.2 | 18.9 | 426.5 | 518.5 | 40.4 | 12.5 | 59.2 | 218.8 | 33.7 | 226.6 |
| 1.4B MoMa 4t4i | 20.9 | 26.6 | 19.8 | 392.5 | 479.6 | 39.0 | 13.0 | 55.8 | 194.7 | 10.5 | 205.7 |
| \hdashline2.3B MoD MoMa 4t4i | 22.9 | 28.0 | 20.5 | 428.1 | 518.0 | 41.4 | 13.2 | 60.0 | 215.1 | 15.0 | 226.1 |

**Table 4: Text-to-text commonsense reasoning task performance.**
| Model | BoolQ | PiQA | SiQA | Winogrande1.1 | OBQA | Hellaswag | Arc-E | Arc-C | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 1.4B Dense | 61.8 | 72.1 | 44.3 | 54.2 | 42.0 | 51.0 | 59.8 | 34.3 | 52.4 |
| 1.4B MoMa 1t1i | 61.5 | 71.3 | 44.4 | 54.7 | 46.0 | 52.8 | 61.1 | 34.9 | 53.3 |
| 1.4B MoE 8x | 62.2 | 73.5 | 44.4 | 58.3 | 42.2 | 58.2 | 62.6 | 36.9 | 54.8 |
| 1.4B MoMa 4t4i | 62.3 | 72.5 | 44.6 | 55.9 | 47.0 | 55.2 | 60.7 | 35.7 | 54.2 |

For all sparse models, we perform the second-stage training of auxiliary routers and use the auxiliary router for causal inference. The details of auxiliary router training can be found in Appendix A.2.

#### Interleaved data modeling.

The relative performance of the dense model and various MoE configurations is consistent with the trends observed in pre-training loss (§4.2). As shown in Table 3, by using 1 more image expert, the 1.4B MoMa 1t1i model significantly outperforms the dense baseline on most metrics, except for image-to-text conditional perplexity on COCO and Flickr. Similar to the trend shown in §4.2, adding the image expert leads to substantial performance gains on the image modality. Further scaling up the number of experts also improves the performance, with 1.4B MoE 8x achieves the best image-to-text performance. Additionally, the model excels on text-to-text tasks (Table 4). 1.4B MoMa 4t4i performs the best on all conditional image perplexity metrics, with the text perplexity closely matches that of 1.4B MoE 8x on most of the benchmarks. Overall, the 1.4B MoMa 4t4i model achieves the best interleaved data modeling results.

#### MoD performance regression.

According to Table 3, we observe significant performance regression of the 2.3B MoD MoMa 4t4i model across all metrics, leaving it underperforming the 1.4B MoMa 4t4i model despite having better pre-training loss (§4.2). Our analysis identifies the auxiliary router’s performance as the primary cause. We used a smaller model, 635M MoD MoMa 4t4i, in the analysis. Initially, we validated that using the training router and selecting top $k_{d}$ tokens within a batch for inference enables the 635M MoD MoMa 4t4i model to outperform the 435M MoMa 4t4i model. However, when introducing random noise to the training router selection with a noise ratio of $\sigma$, the 635M MoD MoMa 4t4i model starts underperforming the 435M MoMa 4t4i model when $\sigma$ exceeds 0.5%. Achieving an error rate below this threshold is impractical for auxiliary router training. Therefore, further research is necessary to develop a robust approach for practical MoD applications.

On the other hand, MoE models do not exhibit similar regression during inference, suggesting they are less sensitive to router errors. This is likely due to the MoE model’s multi-expert architecture, which allows mis-routed tokens to be processed by other experts, preserving their information. In contrast, the MoD model’s single selection switch per layer means routing mistakes cannot be remedied, leading to greater performance degradation.

## 5 Related Work

#### Early-fusion vision-language models.

Early fusion techniques have gained traction in multi-modal learning due to their ability to capture cross-modal interactions from the onset of processing. Several notable works have explored this paradigm. PerceiverIO (Jaegle et al., ) introduced a fully attentional read-process-write architecture that operates over a modality-agnostic latent space for processing diverse inputs, including text and images. NÜWA (Wu et al., ) presented a 3D-attention transformer capable of both understanding and generating text, image, and video in various combinations. CM3 (Aghajanyan et al., ) adopted a causally masked transformer to learn from mixed-modal documents on the Internet, showcasing the scalability of this paradigm to large-scale pretraining. We adopt Chameleon as the base transformer architecture and demonstrate that modality-aware sparsity can effectively improve its scaling performance further.

#### Multi-modal representation learning.

BEIT-3 (Wang et al., ) introduced a general-purpose multimodal foundation model that achieves state-of-the-art performance on both vision and vision-language tasks. It employs a Multiway Transformer architecture with modality-specific expert modules and uses a unified masked data modeling objective for pretraining on images, texts, and image-text pairs.

#### Sparse neural networks.

Sparse neural networks have emerged as a promising approach to improve the efficiency and scalability of deep learning models. One notable architecture is the Mixture-of-Experts (MoE) model, which dynamically selects a subset of experts to process each input, reducing computational costs and promoting specialization (Shazeer et al., ; Lepikhin et al., ; Fedus et al., ). These advances in sparse neural networks have paved the way for the development of increasingly large and powerful models, enabling state-of-the-art results in various domains, including natural language processing (Jiang et al., ) and computer vision (Riquelme et al., ). Recent work have also proposed approaches that trains mixture-of-experts network using dense model as the seed to improve training stability (Komatsuzaki et al., ) and leverage efficiency gains from asynchronous computation (Sukhbaatar et al., ).

#### Sparse multi-modal language models.

Recent advancements in sparse modeling techniques have also shown promising results in efficient scaling of multimodal language models. VL-MoE (Shen et al., ) proposed a vision-language model that uses a sparse mixture-of-experts (MoE) to scale efficiently. It applies modality-aware MoE (Bao et al., ; Wang et al., ; Shen et al., ; Chen et al., ) to the vision and language feed-forward networks in a unified framework, using masked data modeling objectives for pretraining. VL-MoE achieves strong performance on vision-language tasks while using fewer parameters per token than dense models, providing additional evidence for MoE training for multi-modal settings.

## 6 Limitations

Our current MoMa implementation relies on matching the token mix ratio in the dataset with the expert mix ratio in the model to maintain load balance across GPUs. Even so, minor imbalance may still occur because there is no hard limit for a batch to deviate from that ratio at per-GPI per-iteration level. We leave further improvements in this area as future work.

Expert-choice routing alleviates the expert load balancing issue during training but presents additional challenges for auto-regressive Language Models (LMs) during inference (Zhou et al., ). Although auxiliary routers comprise only a small fraction of the network’s parameters, their role is crucial. In our study, we trained the auxiliary router after completing the whole network training and limited this process to a few thousand steps, while previous work have demonstrated the possibility to jointly train such modules with the full network (Raposo et al., ). Future research should explore the architecture and training techniques for auxiliary routers to prevent them from becoming a performance bottleneck and ensure generalizability across diverse data distributions. Especially, further investigation is necessary for training mixture-of-depths architectures, including both the auxiliary routers and the original model, to ensure effective performance in causal inference scenarios.

In our work, we experimented only the vanilla formulation of MoD and its staged integration with MoE. We leave the investigation of other MoD variations, including modality-aware MoD to future work. In addition, batched sequence generation with mixture-of-depths (MoD) is non-trivial as unlike standard sequence generation there are dynamics shapes and dynamic updates to the KV cache for each layer as certain sequences and layers may skip different tokens. There remains further room to optimize the inference implementations for MoD models.

## 7 Conclusion

In this work, we have introduced a series of modality-aware sparse architectures for early fusion, mixed-modality foundation models. Our approach leverages domain specificity while preserving cross-modality knowledge sharing and feature interaction. We have developed highly efficient and effective model architectures by incorporating sparsity in both the width dimension (through modality-aware mixture-of-experts) and the depth dimension (via mixture-of-depths).

Our best architecture, Chameleon-MoMa, demonstrate significant improvements over state-of-the-art baselines. In complexity-controlled experiments, we reduce total FLOPs requirements by up to 3.7$\times$. Importantly, our experimental findings reveal that our modality-aware sparse architectures maintain an empirical scaling law. This characteristic suggests that our approach provides immediate performance benefits and a scalable framework for future developments in mixed-modal foundation model training.

The promising potential of Chameleon-MoMa opens up several promising directions for future research. These include exploring more sophisticated routing mechanisms, investigating the impact of different sparsity patterns across modalities, and extending our approach to a broader range of modalities and tasks.

## Acknowledgments

We thank Baptiste Rozière, Alexei Baevski, Naman Goyal, and Horace He for helpful discussions on mixture-of-experts model training. We thank Ram Pasunuru and Andrew Cohen for helpful feedback on model evaluation. We thank the members of the Chameleon Team for insightful discussions and feedback.

## Appendix A Additional Implementation Details

### A.1 Pre-training

Table 5 reports the hyperparameters of our pre-training aproach.

**Table 5: Hyperparameters used for pre-training our models. We used the same pre-training hyperparameters for models compute-matched to the same dense size.**
| model size | peak lr | end lr | lr scheduler | warm-up | # steps | batch size (tokens) | model parallel | seq len |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 60M | 1e-4 | 1e-6 | linear | 4000 | 160000 | 6291456 | 1 | 4096 |
| 435M | 1e-4 | 1e-6 | linear | 4000 | 160000 | 6291456 | 1 | 4096 |
| 1.4B | 1e-4 | 1e-6 | linear | 4000 | 120000 | 8388608 | 1 | 4096 |

### A.2 Auxiliary Router Training

Table 6 reports the hyperparameters for training the auxillary routers for both Mixture of Depth and Mixture of Experts.

**Table 6: Hyperparameters for aux router training.**
| model size | peak lr | end lr | lr scheduler | warm-up | # steps | batch size (tokens) | seq len |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 2.3B MoD MoE 4t4i | 1e-4 | 1e-6 | linear | 2000 | 5000 | 4194304 | 4096 |
| 1.4B MoE 4t4i | 1e-4 | 1e-6 | linear | 2000 | 10000 | 4194304 | 4096 |
| 1.4B MoE EC8 | 1e-4 | 1e-6 | linear | 2000 | 10000 | 4194304 | 4096 |
| 635M MoD MoE 4t4i | 1e-4 | 1e-6 | linear | 1000 | 5000 | 4194304 | 4096 |
| 435M MoE 4t4i | 1e-4 | 1e-6 | linear | 1000 | 5000 | 4194304 | 4096 |
| 435M MoE EC8 | 1e-4 | 1e-6 | linear | 1000 | 5000 | 4194304 | 4096 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: see above weaknesses
soundness: 3
strengths: The paper addresses the challenge in mixed-modal learning with an architecture that significantly reduces computational overhead through modality-specific expert allocation. 
It presents a thorough experimental setup, including detailed ablation studies that underline the architecture's strengths in different settings.
summary: The manuscript presents MoMa, a novel architecture for pre-training mixed-modal early-fusion language models utilizing a mixture-of-modality-aware experts strategy. By dividing experts into modality-specific groups and using learned routing, MoMa achieves significant efficiency gains in pre-training, highlighted by a comprehensive suite of experiments. This approach addresses the challenges associated with computational costs in scaling mixed-modal models and offers a modular framework that could be extended to various multimodal tasks.
weaknesses: 1. The paper could be improved by a more detailed comparison with existing methods that utilize modality-specific experts. The current discussion does not fully delineate how MoMa advances beyond prior art in terms of both design and performance metrics.

2. For performance evaluation, the results in Table 3 to assess the multimodal capabilities of the pre-trained models through perplexity are not enough. Common benchmarks (e.g., SEED, MMBench, VQAv2) are encouraged to introduced to compare it with other state-of-the-art models under the same setting.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: - The clarification of the technical novelty.
- The training/test time comparison.
- The performance comparison with the MoE model.
soundness: 2
strengths: - This paper shows the effectiveness of a modality-aware processing structure on a large image-language model.
- The evaluation performance shows that this method achieves a clear reduction in FLOPs.
- The ablation and analysis are comprehensive, considering the significant training cost of the pre-training experiment.
summary: This paper proposes a modality-aware mixture-of-experts architecture for pre-training mixed-modal, early-fusion language models. This structure processes image and text tokens with two groups of expert modules in the FFN layer. Two types of routers are trained to control expert selection: the main routers generate the weight values, and tokens are selectively routed to an expert based on thresholding the auxiliary router score. Under a 1-trillion-token training budget, the performance of a model with 4 text experts and 4 image experts achieves significant FLOPs savings.
weaknesses: - The technical novelty of the proposed method is limited, since grouping has been used in FFN layers, and previous works have separately processed image and text tokens, e.g., using image encoders to process the image tokens. It requires further justification of the novelty of the proposed method.

- Is there an ablation study to see if using both image/text and image+text groups of MoE would be better than using image/text groups of MoE?

- The FLOPs saving for training is given. While the FLOPs during training do not strictly align with training time, and MoE structure would make the training more complex, the full training time comparison should be given.

- Another concern is that using two groups of experts will lead to more MoE layers. The inference cost, e.g., inference time, should also be compared with the standard MoE model.

- From Tables 3 and 4, the performance improvement is inferior to the MoE model in some cases. Is it possible that the FLOPs saving comes at the expense of performance drop?

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: The questions listed above are essential for readers to gain a clearer understanding of the proposed method. A detailed response to these points is expected to help clarify potential confusion.
soundness: 3
strengths: The introduction of modal-specific MoE into the early-fusion arch foundation models is interesting. Many scaling pre-training experiments are conducted to verify the effectiveness of the proposed method.
summary: This paper mainly focuses on the early-fusion of multi-modal model pretraining. The authors proposed a Mixture of Modality-aware Experts (MoMa) Module to replace the standard dense connection layer. To ensure a balanced loading for each expert, they predefined a bucket size to select topk tokens among the batch data. An auxiliary router is trained to secure a causality for model inference. Experimental results show the effectiveness of the proposed method.
weaknesses: 1. In the Expert-choice routing, loading balance is mainly considered where each expert will select topk tokens from the current batch data. However, how to ensure that all input tokens can be selected by at least one expert? How can the diversity of the experts be maintained?
2. In the abstract, it is claimed that MoMa 4t4i achieves impressive FLOPs savings: 3.7× overall, with 2.6× for text and 5.2× for image compared to a compute-equivalent dense baseline. How these numbers are calculated? In Table1, MoMa has much more parameters than the dense one for the same amount of active FLOPs per token. Additionally, in Table2, MoMa shows lower training throughput than the dense one. Hoth of these two tables may not clearly verify the efficiency of the proposed method.
3. For the modality-aware expert, the motivation of this work argues that different modalities have distinct characteristics and information densities (L131) and proposes the MoE for the dense layer. How about the modal-specific attention layers? Or the interleaved mix-modal and modal-specific attention layers?
4. In Table 3, MoE 8x outperforms MoMa 4t4i in 4 out of 11 settings, and in Table 4, it shows better performance in 5 out of 8 benchmarks. Given that MoE 8x uses a mixture-of-modality design, does this suggest that a mixed-expert approach itself is more beneficial, without necessarily requiring a unique expert for each modality?
5. In L168-171, “However, EC routing compromises causality in auto-regressive language modeling, as each expert selects the top tokens to process in a batch by comparing token scores across all tokens.” I am a little bit confused about why causality is an issue here. Although selecting the top-k tokens involves looking at future tokens, this process only affects the selection and doesn't incorporate future tokens' embeddings into the calculation of each token itself.
6. Some minor questions:
- Lines 254-255 state, “In MoD architectures, we implement MoD in alternating layers, beginning with layer 0, using a layer capacity factor cdc_dcd​ of 25%.” The term MoD is unclear as it lacks a prior introduction in the paper. Although supplementary section B presents MoD-related experiments, none are included in the main paper, making this mention unnecessary.
- L178, details of auxiliary routers should be in Supplementary section A.2.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please refer to the weaknesses section.
soundness: 3
strengths: The paper introduces MoMa, a modality-aware mixture-of-experts (MoE) architecture tailored for early-fusion language models. This approach represents a creative evolution in multimodal language processing by combining modality-specific experts that process either text or image tokens separately. 

The methodology is rigorous, with a strong experimental design that includes detailed comparisons between MoMa and dense baseline models. 

The paper is well-structured, and the key concepts are presented clearly, such as modality-specific expert grouping, hierarchical routing, and load balancing strategies.
summary: The paper introduces MoMa (modality-aware mixture-of-experts), a novel architecture designed for pre-training mixed-modal, early-fusion language models that process images and text in arbitrary sequences. MoMa divides expert modules into modality-specific groups, which exclusively process designated tokens while maintaining semantics-based adaptivity through learned routing within each group. 

The key contributions and findings of the paper are:
- MoMa achieves significant pre-training efficiency gains by allocating parameters modality-specifically.
- MoMa outperforms standard expert-choice MoE with 8 mixed-modal experts, which achieves 3× overall FLOPs savings.
- MoMa maintains a relatively modest throughput reduction while significantly improving efficiency and performance. The paper also evaluates MoMa's inference-time performance on language modeling data and downstream tasks, showing consistent improvements over the dense baseline.
weaknesses: - The innovation in the methodology presented in this paper is limited. The proposed decoupled modality approach lacks originality, as it has already been introduced in [1]. Additionally, the use of mixture of experts (MoE) is not a novel contribution, as multimodal MoE has been previously established[2].

- The paper lacks extensive validation and comparison across a broad range of benchmarks. While the authors provide comparisons for purely language downstream tasks in Table 4, there is a notable absence of direct performance comparisons on visual-language tasks. I appreciate the approach taken in Table 3 to assess the multimodal capabilities of the pre-trained models through perplexity, but I would encourage the authors to evaluate the model's performance on general visual-language tasks and compare it with some state-of-the-art understanding models, such as Flamingo[3], MM1[4], etc. Given that the primary innovation of this paper focuses on addressing multimodal scenarios, such comparisons are essential.

[1] Wang, Weihan, et al. "Cogvlm: Visual expert for pretrained language models." arXiv preprint arXiv:2311.03079 (2023).

[2] Lin, Bin, et al. "Moe-llava: Mixture of experts for large vision-language models." arXiv preprint arXiv:2401.15947 (2024).

[3] Alayrac, Jean-Baptiste, et al. "Flamingo: a visual language model for few-shot learning." Advances in neural information processing systems 35 (2022): 23716-23736.

[4] McKinzie, Brandon, et al. "Mm1: Methods, analysis & insights from multimodal llm pre-training." arXiv preprint arXiv:2403.09611 (2024).

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: Please refer to Weaknesses.
soundness: 3
strengths: * Considering the current trends in multimodal pre-training, which increasingly rely on large-scale language models and extensive data, pursuing research in efficient architectures, exemplified by MoE, to optimize computational efficiency is a promising research direction.
* Generally, this paper is well-written and very easy to understand.
* The performance experiments and ablation studies are comprehensive and rich in detail.
summary: This paper introduces MoMa, a novel architecture for early-fusion pre-training of mixed-modal language models that utilizes modality-specific mixture-of-experts to enhance computational efficiency and model scalability across various multimodal tasks.
weaknesses: * I have some concerns regarding the novelty of the paper. To my knowledge, there have been numerous attempts to apply MoE architecture in MLLMs to achieve a balanced performance and efficiency. These attempts have explored two paradigms: **shared modal experts** (*e.g.*, MoE-LLaVA [1]) and **decoupled modal experts** (*e.g.*, CogVLM [2]).  They correspond exactly to the MoMa 8x and MoMa 1t1i baselines. Given that these studies share a similar research motivation with MoMa, a comprehensive comparison with related works would better highlight the novelty of this paper.

* Regarding the experimental results, it can be observed from Fig 3 that the modality decoupling strategy (*i.e.*, MoMa 4t4i) appears to exhibit better convergence during pre-training compared to the modality sharing baseline (*i.e.*, MoMa 8x). However, the results presented in Tab 3 and Tab 4 suggest that MoMa 8x actually performs better across various benchmarks than MoMa 4t4i. It does not sufficiently demonstrate the superiority of expert decoupling.

* I wonder whether the authors have considered simultaneously modeling both modality sharing and decoupling.  Could it lead to improved outcomes?  Specifically,  I am curious if there has been any attempt to assign certain experts as modality-shared while designating others as modality-specific. It may potentially balance the commonalities and specificities between modalities more effectively. 

[1] Lin B, Tang Z, Ye Y, et al. Moe-llava: Mixture of experts for large vision-language models[J]. arXiv preprint arXiv:2401.15947, 2024.

[2] Wang W, Lv Q, Yu W, et al. Cogvlm: Visual expert for pretrained language models[J]. arXiv preprint arXiv:2311.03079, 2023.