Multimodal Contrastive Learning with LIMoE: the Language-Image Mixture of Experts

Introduction

Sparsely activated mixture of expert (MoE) models have recently been used with great effect to scale up both vision and text models . The primary motivation for using MoEs is to scale model parameters while keeping compute costs under control. These models however have other benefits; for example, the sparsity protects against catastrophic forgetting in continual learning and can improve performance for multitask learning by offering a convenient inductive bias.

Given success in each individual domain, and the intuition that sparse models may better handle distinct tasks, we explore the application of MoEs to multimodal modelling. We take the first step in this direction, and study models that process both images and text. In particular, we train a single multimodal architecture that aligns image and text representations via contrastive learning .

When using a setup proposed in prior unimodal models , we find that feeding multiple modalities to a single architecture leads to new failure modes unique to MoEs. To overcome these, we present a set of entropy based regularisers which stabilise training and improve performance. We call the resulting model LIMoE (Language-Image MoE).

We train a range of LIMoE models which significantly outperform compute-matched dense baselines. We scale this up to a large 5.6B parameter LIMoE-H/14, which applies 675M parameters per token. When evaluated zero-shot on ImageNet-2012 it achieves an accuracy of 84.1%, competitive with two-tower models that make use of modality-specific pre-training and feature extractors, and apply 3-4x more parameters per token.

In summary, our contributions are as follows.

We propose LIMoE, the first large-scale multimodal mixture of experts models.

We demonstrate in detail how prior approaches to regularising mixture of experts models fall short for multimodal learning, and propose a new entropy-based regularisation scheme to stabilise training.

We show that LIMoE generalises across architecture scales, with relative improvements in zero-shot ImageNet accuracy ranging from 7% to 13% over equivalent dense models. Scaled further, LIMoE-H/14 achieves 84.1% zero-shot ImageNet accuracy, comparable to SOTA contrastive models with per-modality backbones and pre-training.

Lastly, we present ablations and analysis to understand the model’s behavior and our design decisions.

Multimodal Mixture of Experts

Multimodal contrastive learning typically works with independent per-modality encodings . That is, separate models $f_{m}$ are trained to provide a final representation for every input from the corresponding modality, $m$. In the case of some image and text inputs, ${\mathbf{i}}$ and ${\mathbf{t}}$, we have ${\mathbf{z}}_{{\mathbf{i}}}=f_{\text{image}}({\mathbf{i}})$ and ${\mathbf{z}}_{{\mathbf{t}}}=f_{\text{text}}({\mathbf{t}})$. For contrastive learning with images and text, this approach results in a “two-tower” architecture, one for each modality. We study a one-tower setup instead, where a single model is shared for all modalities, as shown in Figure 1. The one-tower design offers increased generality and scalability, and the potential for cross-modal and cross-task knowledge transfer. We next describe the LIMoE architecture and training routine.

2 The LIMoE Architecture

We use a single Transformer-based architecture for both image and text modalities. The model uses a linear layer per modality to project the intrinsic data dimension to the desired width: for text, a standard one-hot sentencepiece encoding and learned vocabulary , and for images, ViT-style patch-based embeddings . Then all tokens are processed by a shared transformer encoder, which is not explicitly conditioned on modality. The token representations from the final layer are average-pooled to produce a single representation vector ${\mathbf{z}}_{m}$ for each modality. To compute the training loss in (1), the paired image and text representations are then linearly projected using per-modality weight matrices ${\mathbf{W}}_{m}$’s and $\mathcal{L}_{j}$ is applied to $\{({\mathbf{W}}_{\text{image}}\ {\mathbf{z}}_{{\mathbf{i}}_{k}},{\mathbf{W}}_{\text{text}}\ {\mathbf{z}}_{{\mathbf{t}}_{k}})\}_{k=1}^{n}$.

This one-tower setup can be implemented with a standard dense Transformer (and we train many such models as baselines). Next, we describe how we introduce MoEs to this setup for LIMoE.

Note that, for computational efficiency and implementation constraints, experts have a fixed buffer capacity. The number of tokens each expert can process is fixed in advance, and typically assumes that tokens are roughly balanced across experts. If capacity is exceeded, some tokens are “dropped”; they are not processed by the expert, and the expert output is all zeros for those tokens. The rate at which tokens are successfully processed (that is, not dropped) is referred to as the “success rate”. It is an important indicator of healthy and balanced routing and often indicative of training stability.

We discovered that routing with tokens from multiple modalities introduces new failure modes; in the next sections we demonstrate this phenomenon, and describe our techniques to address it.

As mentioned, experts have a fixed buffer capacity. Without intervention, Top-$K$ MoEs tend to “collapse”, thus using only one expert. This causes most tokens to be dropped and leads to poor performance . Prior works therefore use auxiliary losses to encourage balanced routing .

In multimodal settings, new challenges arise; one is modality misbalance. In realistic setups, there will likely be more of one data type than another. Accordingly, we do not assume or enforce balanced data across modalities, and our experiments have $3-17\times$ more image tokens than text tokens.

Modality-specific experts tend to emerge naturally. In this imbalanced context, this leads to a scenario where all of the tokens from the minority modality get assigned to a single expert, which runs out of capacity. On a global level, routing still appears balanced: tokens from the majority modality are nicely distributed across experts, thereby satisfying modality-agnostic auxiliary losses. For example, in our standard B/16 setup, the router can optimize the importance loss to within 0.5% of its minimum value by perfectly balancing image tokens but dropping all text tokens. This however leads to unstable training and unperforming models.

2.2 Auxiliary losses

We refer to auxiliary losses used in V-MoE as the classic auxiliary losses. We find that they do not yield stable and performant multimodal MoE models. Therefore, we introduce two new losses: the local entropy loss and the global entropy loss, which are applied on a per-modality basis. We combine these losses with the classic losses; see Appendix B for a summary of all auxiliary losses.

Intuitively, it is desirable for text tokens to use multiple experts, but not all of them. In order to allow flexibility, we threshold the global entropy loss as $\Omega^{\tau}_{\text{global}}({\mathbf{G}}_{m})=\max\{0,\tau+\Omega^{\text{global}}({\mathbf{G}}_{m})\}$, such that the model is encouraged to have a certain minimum entropy, but after exceeding that, the loss is not applied. This avoids distributional collapse but does not apply overly restrictive priors on the routing distribution, as there are many optimal solutions. This can be thought of as a “soft minimum” $S$. With $\tau=\log(S)$, the model must use at least $S$ experts to minimize the loss (either a uniform distribution across $S$ experts -with entropy $\log(S)$-, or a non-uniform distribution using more than $S$). Figure 3(b) shows the latter occurs; the empirical effect of these thresholds is analysed in Section 4.1.

Connection with mutual information. The sum $\Omega_{\text{local}}({\mathbf{G}}_{m})+\Omega_{\text{global}}({\mathbf{G}}_{m})$ corresponds to the (negative) mutual information between experts and tokens, conditioned on the modality $m$, which we write $-\text{MI}_{m}(\texttt{experts};{\mathbf{x}})$. For each modality taken separately, we are effectively encouraging the knowledge of the token representation to reduce the uncertainty about the experts selection. We also tried other variants of the losses which exploit this connection, such as the mutual information between the experts and modalities, $-\text{MI}(\texttt{experts};m)$, obtained by first marginalizing the tokens.

2.3 Priority routing

With Top-$K$ routing, some token dropping is virtually inevitable. Batch Priority Routing (BPR) actively decides which tokens to skip based on their routing weights. It assumes that tokens with a large routing weight are likely to be informative, and should be favored. BPR was mostly used at inference time in , allowing for smaller expert capacity buffers. In this setup, one must take care not to systematically favor one modality over the other, for instance, by determining which token to drop based on their rank in the batch, which are usually grouped according to the token modality. BPR provides an essential stabilisation effect during training (Figure 6); we show that it does not trivially rank one modality over another, and it cannot be replaced by other methods of re-ordering the batch. In the appendix we further show how routing priorities compare across text and images.

Experiments

We study LIMoE in the context of multimodal contrastive learning. We first perform a controlled comparison of LIMoE to an equivalent “standard” dense Transformer, across a range of model sizes. We then show that when scaled up LIMoE can reach a high level of performance. Finally, we ablate the various design decisions leading to LIMoE in Section 4.

Training data. By default, all models are trained on paired image-text data used in , consisting of 3.6B images and alt-texts scraped from the web. For large LIMoE-H/14 experiment, we also co-train with JFT-4B . We construct artificial text captions from JFT by comma-delimited concatenation of the class names . Appendix A contains full details of our training setup.

Evaluation. Our main evaluation is “zero-shot”: the model uses its text representations of the classes to make predictions on a new task without extra training data . We focus on image classification accuracy on ImageNet and cross-modal retrieval on MS-COCO , following the protocol in . We also evaluate LIMoE’s image representations via a linear adaptation protocol , and report 10-shot accuracy on ImageNet accuracy accordingly. Where ranges are given, they report 95% confidence intervals across three trials.

We train a range of LIMoE models at batch size $16$k for $781$k steps. This matches the number of training examples used for CLIP . Due to use of different training data and additional tricks, a direct comparison is difficult; we therefore train dense one-tower models as baselines. All models activate $k=1$ experts per token, similar to Switch Transformer .

Figure 4 shows the performance of each model (dense and sparse) against forward-pass FLOPs (for step times and further discussion on compute costs, see Appendix D.2.). The cost-performance Pareto frontier for LIMoE dominates the dense models by a wide margin, indicating that LIMoE offers strong improvements across all scales from S/32 , up to L/16. The effect is particularly large on zero-shot and 10-shot ImageNet classification, with absolute performance improvements of 10.1% and 12.2% on average. For text-to-image retrieval on COCO, LIMoE offers a strong boost at small scales, while at larger scales the gains are more modest but still significant.

2 Scaling up LIMoE

We increase the architecture size, training duration, and data size to assess the performance of LIMoE in the large-scale regime. In particular, we train a 32-layer LIMoE-H/14 with 12 expert layers; these are non-uniformly distributed, with 32 experts per layer, and $K=1$ activated per token. It was trained at a batch size of $21$k, introducing 25% JFT-4B images into each batch (with class names as texts). We average checkpoints towards the end of training ; refer to Appendix A.3 for details.

The model contains 5.6B parameters in total, but only applies 675M parameters per token. All routers combined account for less than 0.5M parameters. Table 1 shows its performance alongside current state-of-the-art contrastive models. LIMoE achieves 84.1% zero-shot ImageNet classification accuracy with a comparably modest architecture size and training counts. LIMoE is fully trained from scratch, without any pre-trained components, and is the first competitive model with a shared backbone.

In light of its modality agnostic approach, this result is surprisingly strong. Large models handling dozens of distinct tasks are increasingly popular , but do not yet approach the state-of-the-art in these tasks. We believe the ability to build a generalist model with specialist components, which can decide how different modalities or tasks should interact, will be key to creating truly multimodal multitask models which excel at everything they do. LIMoE is a promising first step in that direction.

Ablations

We use a smaller setup to study various aspects of LIMoE. We train B/16 models at batch size 8096 for 100,000 steps (see Appendix A.2 for further details). Table 2 shows the average over three trials of this setting alongside dense one-tower and two-tower baselines. LIMoE greatly outperforms both dense models on ImageNet 0- and 10-shot, while confidence intervals overlap for retrieval with two towers. The two-tower model is twice as large and expensive, and still falls behind the sparse one.

Choice of auxiliary losses. With the introduction of the entropy based losses in addition to classic ones, there are 7 possible auxiliary losses. We aimed to find the simplest combination of these which obtains good performance. To study this, we performed a large sweep of auxiliary losses: for $N\in[2,\ldots,5]$, we considered all $\binom{7}{N}$ possible loss combinations. Table 3 shows, for each loss, the highest performing model with and without that loss. Some conclusions stand out: Both entropy losses are important for text, but for images, the global loss is not impactful and the local loss is harmful. The final combination of losses was chosen based on validation accuracy alongside qualitative observations around training stability and routing success rate.

Threshold for global entropy losses. In Section 2.2.2, we introduced a threshold $\tau$ to encourage balanced expert distributions without forcing all modalities to use all experts. To understand the importance of this threshold, we sweep over it for both the image and text global entropy losses. Appendix B.2 contains a full analysis; the most important conclusions are:

$\tau_{\text{image}}$ did not affect the number of experts used for images, as global entropy was always high. Aside from these threshold experiments with very high $\tau_{\text{image}}$, this loss is usually inactive. It was used in our main experiments, but can likely be removed in future work.

The threshold $\tau_{\text{text}}$ behaved exactly as a soft minimum for text experts: Sweeping $\tau_{\text{text}}$, we typically observed approximately $S=e^{\tau_{\text{text}}}$ text experts.

Performance is robust to different values of $\tau_{\text{text}}$, provided it is not too low. A low $\tau_{\text{text}}$ can be useful to limit the number of text experts, for later pruning, see Appendix E.4.

Mutual-information auxiliary loss. In Section 2.2.2, we discussed an alternative loss, namely $-\text{MI}(\texttt{experts};m)$, based on the mutual information between experts and modalities. While it has the advantage of merging the local and global entropy losses for both the text and image modalities into a single term, without threshold parameters, it leads to slightly worse results: in a comparable setup, it had 1.5% and 0.1% worse zero-shot and 10-shot performance compared to Table 2.

Our models use a text sequence length of 16, but image sequence lengths from 49 to 400 (for these ablations, 196).

Our ablations reveal that the entropy losses are most important when applied to the text tokens. This leads to a hypothesis that these are only necessary or useful in the imbalanced case. To test this, we vary the modality balance of LIMoE-B/16 by varying the patch size; this enables us to control the number of image tokens, and hence image:text balance, without changing the information content in the data. Figure 5 shows the results. First, we observe that, with entropy routing, a longer image sequence length is always better. This shows that entropy routing can effectively handle highly imbalanced setups, and mirrors the observation that for classical Vision Transformers: a longer sequence is better. Importantly, entropy routing is always far superior to the classical setup with growing gaps, even when the modalities are balanced 1:1 ($L_{\text{img}}=16$). This experiment also confirms the robustness of entropy routing to different setups.

Batch priority routing as a training stabilizer. Figure 6 shows the effect of BPR during training. BPR not only ameliorates against token dropping, but also improves training stability. Models with no dispatch order intervention (first-in-first-out) perform extremely poorly, whether we route images first or text first. These routers have low success rate. Randomly shuffling tokens (i.e. deciding which tokens to drop at random when an expert becomes full) partially ameliorates this, but its performance is still much worse than that of models trained with BPR. We further analyse BPR in Appendix F.5 and show that it does not simply rank one modality above another.

2 Other ablations

We summarize our other ablations here due to space constraints; details can be found in Appendix E.

Router structure (Appendix E.3). Our router is modality agnostic; we experiment with per-modality routers, and separate pools of per-modality experts. We find they all perform comparably to our generic, modality agnostic setup, but that separate pools of experts by design is more stable and does not require auxiliary losses for regularisation—while harder to scale to many modalities and tasks.

Increasing selected experts per token $K$ ( Appendix E.1). We propose modifications to BPR and the local auxiliary loss to generalise to $K>1$; by doing so we can steadily increase performance by increasing $K$, e.g. from 55.5% zero-shot accuracy with $K=1$ to 61.0% with $K=5$.

Total number experts (Appendix E.2). We show that increasing the pool of available experts at fixed $K$ improves performance (unlike what was observed for vision-only tasks ).

Expert pruning (Appendix E.4). We show using simple heuristics we can prune down to modality-specific experts for unimodal forward passes, thus avoiding expert collapse under unimodal batches.

Training on public data (Appendix E.6) The majority of LIMoE models were trained on proprietary data . We show that LIMoE works similarly well on publically available data, retaining performance improvements against a comparable dense model.

Model Analysis

In this section, we explore some of the internal workings of LIMoE. We use simple B/32 and B/16 models with 8 experts, and the large H/14 with 32. See Appendix F for further details and experiments.

Multimodal experts arise (Appendix F.1). Aside from encouraging diversity, we do not explicitly enforce experts to specialize. Nonetheless, we observe the emergence of both modality-specific experts, and multimodal experts which process both images and texts (per-expert distributions in F.1).

Qualitative analysis (Appendix F.2). We analyse some example data and show a clear emergence of semantically meaningful experts. With images for instance, some experts specialize on lower level features (colours, lines) while others on more complex features (faces and text), see Figure 2.

BPR ranking (Appendix F.5). The local loss encourages high max-routing weights for text, and BPR ranks according to this. We show however that this does not mean text is always prioritised first: Especially in later layers, the model often prioritises important image patches over text.

Related work

Unimodal, task-specific neural networks have long been researched, with increasing convergence towards Transformer-based architectures for both NLP and Computer Vision . Multimodal models aim to process multiple types of data using a single neural network.

Many approaches “fuse” modalities to tackle inherently multimodal tasks. LIMoE is more similar to approaches which do not do that, and still operate as unimodal feature extractors. Some co-train on distinct tasks without aligning or fusing representations—effectively sharing weights across tasks—whereas others include both unimodal aspects and fused multimodal aspects for functionality in both contexts .

We build on deep Sparse Mixture of Experts models, which have been studied independently in Computer Vision and NLP , typically in the context of transfer learning. These models use a learned gating mechanism whereby only a subset of $K$ experts out of $E\gg K$ are activated for a given input. Many works aim to improve the gating mechanism itself, by making it differentiable , reformulating as a linear assignment task or even swapping it out for a simple hashing algorithm . MoE models have also been studied for multitask learning , with per-task routers but a shared pool of experts. To our knowledge, sparse models have not been explored for multimodal learning.

A large body of research exists on contrastive learning, usually in self-supervised but also in supervised regimes . Multimodal contrastive learning trains on aligned data from multiple modalities. Originally studied for medical images and reports , it was recently scaled to noisy web data , where strong image-text alignments enabled performant image classification and cross-modal image-text retrieval without finetuning on downstream data. Follow up works improved upon this significantly by scaling up and using pretrained models and multitask training with generative modelling or other vision tasks . These works use unimodal models which separately process image and text data; we are not aware of previous research using a single model to process both images and texts for contrastive learning, neither with dense nor with sparse models.

Conclusions and Future Work

We have presented LIMoE, the first multimodal sparse mixture of experts model. We uncovered new failure modes specific to this setup and proposed entropy based auxiliary losses which stabilises training and results in highly performant models. It works across many model scales, with average improvements over FLOP-matched dense baselines of +10.2% zero-shot accuracy. When scaled to a large H/14 model, we achieve 84.1% accuracy, competitive with current SOTA approaches.

Societal impact and limitations: The potential harms of large scale models , contrastive models and web-scale multimodal data also carry over here, as LIMoE does not explicitly address them. On the other hand, it has been shown that pruning models tends to cause low-resource groups to be forgotten , causing performance to disproportionally drop for some subgroups. This would be worth considering for our expert-pruning experiments, but by analogue, the ability to scale models with experts that can specialize deeply may result in better performance on underrepresented groups.

Environmentally speaking, training large models is costly, though efforts are made to use efficient datacenters and offset emitted CO2. Prior works however show that most environmental impact occurs during model inference, and that MoEs are significantly more efficient in that regard ; LIMoE is naturally a good candidate for efficient, large-scale multimodal foundation models.

Future work: There are many interesting directions from here. The routing interference with multiple modalities still is not fully understood. In general, conclusions from applications of MoEs to NLP have not carried over perfectly to Vision, and vice-versa, and here we see again different behaviour between images and text. Naturally, extensions to more modalities should be explored; even with only two we see fascinating interactions between different data types and the routing algorithms, and that will only get more difficult, and interesting, with more modalities.

There are always more modalities to learn, and larger models to build: sparse models provide a very natural way to scale up while juggling very different tasks and data, and we look forward to seeing more research in this area.

Acknowledgements

We first thank Andreas Steiner, Xiao Wang and Xiaohua Zhai, who led early explorations into dense single-tower models for contrastive multimodal learning, and also were instrumental in providing data access. We also thank Andreas Steiner, and Douglas Eck, for early feedback on the paper. We thank André Susano Pinto, Maxim Neumann, Barret Zoph, Liam Fedus, Wei Han and Josip Djolonga for useful discussions, and Erica Moreira and Victor Gomes for help scaling up to LIMoE-H/14.

References

Appendix A Training details

All models were trained with adafactor, using the same modifications used for ViT-G . Unless otherwise specified, we use learning rate $1\text{\times}{10}^{-3}$ and decoupled weight decay of magnitude $1\text{\times}{10}^{-5}$. We use a cosine learning rate decay schedule, with a linear warmup (40k steps for longer scaling study models, 10k steps for ablations). Models were trained on a mixture of Cloud TPU-v2, v3 and v4 pods.

Models were trained with 32 experts, with experts placed every 2 layers – except where explicitly stated. Otherwise, architecture parameters (e.g. hidden size, number of layers) follow those of ViT . All models except for LIMoE-H/14 use dimensionality 512 for the final output representation; this final representation is cast to bfloat16 precision for reduced all-to-all costs and increased memory efficiency. The learned contrastive temperature parameter is initialised at 10. Text sequences are tokenized to a sequence length of 16 using the T5 SentencePiece vocabulary . Images were linearly renormalized to a value range of .

We train models at batch size 16,384 for 781,250 steps at resolution 224. This trains for the same number of examples as CLIP ; they however use a larger batch size (32768), increase resolution in the final epoch, and use a larger dimensionality for the final contrastive feature representation, all of which improve performance.

A.2 Ablations

These are B/16 models trained for 100,000 steps at batch size 8192. The threshold used for the text global entropy loss is $\tau_{T}=\log(9)$ – that is, we incentivize the use of at least 9 experts (uniformly) or more (not necessarily in a uniform way). For images, $\tau_{T}=\log(20)$, but with this threshold, the loss is not applied at all and it can be ignored.

A.3 LIMoE-H/14

The largest scale model is trained at batch size 21502, with resolution 288 and text sequence length 16. The global entropy loss thresholds are $\tau_{\text{text}}=\log(4)$ and $\tau_{\text{text}}=\log(25)$ for text and image respectively. There are MoE layers in 12 encoder blocks, namely, in 3, 7, 11, 15, 18, 21, 24, 26, 28, 30, 31, 32. The default training data is mixed with data from JFT-4B with a ratio of 3:1. Text strings are generated from JFT-4B by simply concatenating the class names. JFT-4B was also deduplicated using the same method as previous works .

Checkpoint souping. We adapt the methodology developed for finetuning , but instead combine checkpoints from the same run. We used a reverse-sqrt schedule , which has a linear cooldown at the end. To generate diversity for the model soup, we launched multiple cooldowns, and greedily selected checkpoints to maximize zero-shot accuracy on the ImageNet validation set, using the smaller subset of prompts from CLIP . Checkpoints could be reused multiple times.

The model was trained for 700k steps pre-cooldown. There was one cooldown of length 125k steps from the final step, and 3 of length 40k steps starting from step 650k. Two of the cooldowns had no changes to the original setup described above. To generate diversity for the soup, we also trained one 40k cooldown with only JFT data, and one with no JFT data at all.

Figure 7 shows the zero-shot accuracy evaluated at 12.5k step intervals during training, for all the different cooldowns, and the end of training. The final model soup consisted of 8 checkpoints in total.

Appendix B Auxiliary losses

Importance loss. We consider the definition from , inspired by the original proposal of . The importance loss $\Omega_{\text{imp}}$ enforces a balanced profile of the gating weights across the experts. More formally, for any expert $e\in\{1,\dots,E\}$, we consider

and define the loss $\Omega_{\text{imp}}$ via the squared coefficient of variation for $\text{imp}({\mathbf{X}})=\{\text{imp}_{e}({\mathbf{X}})\}_{e=1}^{E}$, namely

Load loss. Like previously, we follow whose definition is inspired by the original proposal of . We assume throughout that paragraph that the gating weights $g_{\text{noisy}}({\mathbf{x}})$ are obtained by a noisy version of the routing, i.e., $g_{\text{noisy}}({\mathbf{x}})=\texttt{softmax}({\mathbf{W}}{\mathbf{x}}+\varepsilon)$ with $\varepsilon\sim{\mathcal{N}}({\mathbf{0}},\sigma^{2}{\mathbf{I}})$ and $\sigma=1/E$ (see details in ). We introduce $\eta_{K}$ the $K$-th largest entry of ${\mathbf{W}}{\mathbf{x}}+\varepsilon$.

The load loss $\Omega_{\text{load}}$ complements the importance loss $\Omega_{\text{imp}}$ by trying to balance the number of assignments across the experts. To circumvent the fact that the assignments are discrete, $\Omega_{\text{imp}}$ focuses instead on the probability of selecting the expert. For any $e\in\{1,\dots,E\}$, the probability is understood as the probability of having the expert $e$ still being among the Top-$K$ while resampling only the noise of that expert. More formally, this corresponds to

with $\Phi$ the cumulative distribution function of a Gaussian distribution.

The load loss $\Omega_{\text{load}}$ is eventually defined by

The mutual-information loss $-\text{MI}(\texttt{experts};m)$. In Section 2.2.2, we allude to a variant of the local and global entropy losses in the form of the mutual information between the experts and the modalities (as a reminder, the sum of the local and global entropy losses corresponds instead to the (negative) mutual information between the experts and tokens, conditioned on the modality). Let us assume we have a total of $M$ modalities. Formally, and reusing the notation from Section 2.2.2, we define $-\text{MI}(\texttt{experts};m)$ as

where, for each modality $m^{\prime}$, we have computed the approximate marginal probability over the $n_{m^{\prime}}$ tokens of that modality

Final aggregated auxiliary loss. When considering the combination of several auxiliary losses, the final auxiliary loss is computed as the average over all the losses. The average is weighted by a single regularization parameter that is a hyperparameter of our approach. After some preliminary tuning phase, we have set its value to $0.04$ in all our experiments and found this choice to be robust.

B.2 In-depth analysis of global entropy threshold

Note again that we can view a threshold $\tau$ as a soft minimum, as the minimum number of experts which must be used by a modality to satisfy the loss is $S=e^{\tau}$. We find it more intuitive to think in terms of this soft minimum threshold $S_{T}$.

Performance. Figure 8 shows the effect of the threshold on performance.

When the text threshold is too low, models are unstable and performance is poor.

Past some limit however, performance of models w.r.t. text threshold is fairly consistent.

Outside (and probably inside) the unstable region, the image threshold makes no systematic difference.

Actual global entropies. Looking at the actual entropies of model routing helps at least explain why the image threshold is unimportant. Figure 9 shows the empirical entropy. The image entropy is always large; note that when it is higher than the threshold $\tau$, the loss is not applied; ergo, for most of the settings, the global entropy loss is not applied to images. This also applies to almost all models trained for this paper. On the other hand, analysing text entropies, it is clear that the model closely tracks the threshold $\tau_{\text{text}}$. As a side effect, image entropy tends to reduce as $\tau_{\text{text}}$ increases.

Expert specialization. As discussed, the threshold can be viewed as setting an implicit soft minimum $S_{T}=e^{\tau}$. The number of experts actually used for each modality is shown in Figure 10. The text threshold exactly behaves as a soft minimum; as it is increased, the model has more text experts and less image experts.

Overall. For text, the entropy loss behaves as expected; as it is increased, there are more text experts. A few questions remain: why does it not impact performance? Why does text behave differently than images - is it due to the imbalance between them during training, or is it simply a fundamental difference in routing behavior for the two modalities?

Appendix C Tabular results

All results and parameters from Figure 4 are shown in Table 4, alongside the results of LIMoE-H/14.

C.2 All tabular results

Appendix D Computational costs of LIMoE

Recall that each expert has a capacity $C$ - it can process at most $C$ tokens, and if it is assigned more, those above $C$ will not be processed. This capacity is usually set relative to some ‘ideal’. If there are $N$ tokens and $E$ experts, we usually assume each expert can handle at most $C_{R}\times\frac{N}{E}$ tokens, where $C_{R}\geq 1$ is a slack factor. This way we try to reach a balanced setup where most expert process a similar number of tokens.

Multimodal routing presents a unique issue here. During training, the model learns to balance tokens when it has both images and text available to it. When there is only one modality, it will not use all the experts due to natural emergence of modality-specific experts - but the expert capacity size will be set assuming all experts are used. This results in high rates of token dropping, depending on the ratio of modality-specific experts.

In this effort, we encounter this during zero-shot classification and retrieval; models first compute representations for all text tokens, and then separately for all image tokens. In order to get around this token dropping, we simply evaluate with a high slack factor $C_{R}=16$.

There are however other natural solutions; for many circumstances, one could trivially restructure evaluation such that image and text inputs are processed simultaneously. A more interesting, MoE specific solution is pruning modality specific experts, which is explored and shown to work in E.4. LIMoE models could have been evaluated at a ‘normal’ capacity, with pruned experts.

D.2 Understanding the compute costs of LIMoE

Zero-shot evaluation on ImageNet with 6 prompts requires 6000 text forward passes and 50000 image forward passes. With 80 prompts, a la CLIP , it is 80000 for text instead. How does one compare compute cost vs. performance? The costs of LIMoE, its dense baselines, and other two-tower models, were computed assuming a full batch of images and texts, as this is the approach which makes the least assumptions about the downstream setup. This does not generalise perfectly: if, for example, a particular use case processed very large numbers of texts but only few images, models with smaller/cheaper text towers would be clearly advantaged.

Full profiling data for Section 3. For training and evaluation, we used a variety of TPU versions. For consistency, we profiled computation times on a TPUv3 (v3-32 to be more precisehttps://cloud.google.com/tpu/docs/types-topologies). Figure 11 shows performance with respect to different proxies for compute cost. As discussed in Section 3, LIMoE is clearly pareto optimal with respect to total FLOPs. However, this does not fully account for certain costs related to MoE models, such as cross-device communication. Figures 11(b) and 11(d) show the performance with respect to step time. With respect to zeroshot and 10-shot classification accuracy, the performance improvements of LIMoE are significant enough that it is still clearly pareto optimal; for retrieval metrics on COCO, LIMoE’s gains exactly justify the costs, and it is not significantly more efficient than dense baselines. The story is similar whether looking at train or evaluation cost.

Appendix E Further experiments

In this section, we present further ablations not included in the main text due to space constraints.

All models in this paper select $K=1$ expert per token to match the cost of a dense backbone.

There are two main challenges with increasing $K$:

Modifications to auxiliary losses. The local entropy loss effectively encourages that router choices are one-hot. When increasing $K$, the model is still incentivized to only use 1 expert, assigning other experts weights near 0, thereby effectively behaving as $K=1$.

We try two modifications to the local loss to ameliorate this:

Target entropy: Encourage the local entropy to be $\log{K}$ – at least a uniform distribution over $K$ experts – instead of 0: we minimize $\Omega_{\text{target}}^{\text{K}}=(\log(K)-\Omega_{\text{local}}({\mathbf{G}}_{m}))^{2}$.

Merged entropy: We sum the top $K$ and the bottom $N-K$ routing probabilities to give a binomial distribution, and optimize this to have entropy 0. This encourages the routing weight to all be in the top $K$ experts, but does not care exactly how it is distributed among them.

BPR modifications With these losses, the router uses $K>1$ experts per token. However, BPR prioritises tokens according to their max routing probability, which decreases when probabilities are distributed over $K$ choices. The stabilisation effect BPR provides training is consequently lost. We alter it to prioritise tokens by the sum of top $K$ probabilities. In vision tasks, the two approaches perform identically , but here the latter stabilises training and unlocks $K>1$.

Table 6 shows the final results. Without changes to the local entropy loss (BPR score = max, local entropy method = default), there are some improvements which stem from using $K>1$ for image tokens - the local loss on text means it is effectively using $K=1$ for text anyway. Without the modifications to the BPR score, the modifications to the local loss can result in fairly unstable models. Once the BPR score is modified, we see consistent improvements with increasing $K$, particularly with the ’Merged’ variant.

E.2 Increasing the total number of experts

There is thus far no consensus on the optimal number of experts in MoEs; early NLP research scaled to 1000s of experts , before reducing to 32 or 64 , which is the standard setup for vision . In Figure 12, we vary this for LIMoE, and show that larger expert pools yield consistent performance improvements.

E.3 Router design choice

Recall the router is simply a dense layer; by default we have a joint router for all tokens, independent of modality, with no constraints on gating. We consider two other options:

Per-modality router. We consider modality-dependent routers which can leverage knowledge of token modality to improve performance (that is, one router for image tokens, and a different one for text tokens). They both output routing distributions over a shared pool of experts, similar to prior works have per-task routers for multitask learning .

Disjoint experts and routers. We define separate pools of image and text experts. This way, image tokens can only go to a set of experts $\mathcal{E}_{\text{img}}$, and text tokens can only be assigned to another set of experts $\mathcal{E}_{\text{txt}}$. In principle, these sets may or may not intersect. In Table 7, we report results when the sets are indeed disjoint.

The results in Table 7 show the three approaches lead to comparable performance. In general, the disjoint setup was more stable, and did not need entropy regularisation as per-modality balance/independence is enforced by design. While convenient and well-behaved here, this approach may not be as general for the case with dozens of tasks and modalities.

E.4 Pruning Multimodal Experts

During training we track what fraction of each modality’s tokens went to each expert. It is therefore trivial to identify which experts are processing predominately text and which are processing predominately images. We show here that this information can be trivially used to prune experts for single-modality forward passes, demonstrating on two 32-expert LIMoE-S/16 models: one trained with global text entropy threshold $\tau_{\text{text}}=\log(4)$, and one with $\tau_{\text{text}}=\log(9)$.

Choosing what to prune. Note that we separately choose what experts to prune per-modality; we use text as an illustrative example. Pruning is simple: For each MoE layer, we rank experts according to the fraction of text tokens they processed during training (we average over the last 2500 steps with measurements sampled every 50 steps). We then start pruning according to the one that processed the least tokens, and so on. Figure 13 shows how the coverage of different modalities changes as experts are pruned. Following the relationship between the global text entropy threshold and the idea of the ‘soft minimum’, we see that around $e^{\tau_{\text{text}}}$ text experts are needed to process the majority of text tokens; e.g. with $\tau_{\text{text}}=4$ for a single-modality forward pass, 28 experts could be comfortably pruned. Image experts are more distributed; almost all the experts are needed to process all image tokens, as expected.

How to run LIMoE inference with fewer experts. While some experts are pruned, the model is not further trained to adapt to this new situation. One must therefore think carefully on the best way to apply models with a subset of experts. The router predicts $p(\texttt{expert}|{\mathbf{x}})$. The top-$K$ experts are activated, and the output of the expert layer is the weighted average of the expert outputs. The weighting used for expert $i$ is the unnormalized $p(\texttt{expert}_{i}|{\mathbf{x}})$. This is important, as removing some of the experts and their logits modifies the concentration of $p(\texttt{expert}|{\mathbf{x}})$, and could result in expert weights higher than those used at training time.

When removing some of the experts, there are therefore two natural options:

router-drop: Completely remove the experts from the router. The softmax for $p(\texttt{expert}|{\mathbf{x}})$ will be computed over a subset of experts, thereby adjusting the weights as discussed above.

router-pred: The router still predicts probabilities for pruned experts. However, it is unable to actually use the pruned experts; the top-$K$ operation will ignore those that are unavailable. This preserves the original scaling the model was trained with.

The two approaches are naturally very similar if very few experts are removed. Illustrating with ImageNet-10shot (linear few-shot evaluation), Figure 14 compares the two options. When a large number of experts are pruned, the router-pred is significantly better, but if only a few experts are pruned, they both perform similarly.

Figure 15 shows the impact of pruning image and text experts on zero-shot ImageNet accuracy. Recall that image and text inputs are processed independently for this evaluation, and so the experts used for each modality can be independently pruned.

As expected, we can prune down to only 4 experts during text evaluation without significantly harming performance. On the other hand, the less pruning of image experts, the better.

E.5 Grouped routing

Splitting batches into groups before dispatching can reduce routing cost significantly, which depending on implementation can scale $\sim\mathcal{O}(\texttt{num tokens}^{2})$. There are two sources of potential issues though: in our implementation, auxiliary losses are computed in each group then averaged. The necessary batch-wise statistics become less reliable with more numerous, smaller groups. Secondly, with smaller groups, it is more likely to get an almost homogenous batch, which makes distributing across experts harder. To study this, we sweep the group size in a parallel setup with 128 examples per device. Group size 1 means processing and dispatching $128\times(196+16)=27136$ tokens at once, whereas e.g. group size 8 involves splitting into 8 groups of 3392 tokens. Figure 16 shows the effect of this; up to 4 groups, performance is good, but any more than that and training becomes unstable, harming performance. This is more fragile than image-only routing, where group sizes as small as 400 are stable (equivalent to $\sim 68$ groups here). Nonetheless, with 4 groups, step time is reduced by 30%, capturing 75% of potential efficiency gains from grouped routing.

E.6 Experiments on public data

In order to ascertain LIMoE’s efficacy on public data, and reproducibility, we train B/16 models on LAION-400M . We train for 5 epochs at batch size 16,384. Table 8 shows the outcome of three trials, compared against a dense baseline. Once again, we see significant improvements performance, especially in ImageNet zero-shot (+5.0% absolute, +8.9% relative) and 10-shot (+6.6% absolute, +13.8% relative) performance.

Appendix F Model Analysis

In this section, we explore how routing is distributed across different layers, experts, and modalities. In particular, we focus on which tokens are dropped. We analyze two models, B/32 and B/16, each with 8 experts. This way we can appreciate the impact of having a significantly different ratio of text:image tokens. Moreover, the global entropy targets $S=e^{\tau}$ for (text, image) tokens are (3, 25) and (6, 6) for the B/32 and B/16 models, respectively.

We first show the routing distributions under the training distribution in Figures 17 and 18. In both cases –as expected– routing works very well. Moreover, most experts handle both image and text tokens.

F.2 Routing Examples

In this section, we share practical examples of image and text token routing on the B/32 and B/16 models introduced at the beginning of the section. All evaluations are on ImageNet (that is, not on the training data). While the number of experts is clearly smaller than the number of different semantic concepts in images and text, we still highlight some cool patterns in most experts – especially in the context of images, as text tokens tend to use a reduced number of experts. We show some of the patches with the highest routing confidence, as analyzing all the thousands of patches that are assigned to each expert is difficult. However, we expect many other semantic concepts present in the training data to be almost exclusively served by individual experts.

F.3 Routing for Individual Inputs

In this subsection, we show the expert split for a specific given input – image and text. Recall tokens from different modalities do not interact in the forward pass (other than via sharing expert capacity).

F.4 Routing Trajectories

In this section, we try to have a look at the overall trajectories followed by both image and text tokens across the network. While definitely a complex endeavor, we show in Figure 30 for B/32 and Figure 31 for B/16 the main trajectories followed by such tokens. Interestingly enough, it seems that for both models and image tokens, the first two/three MoE layers are fairly interconnected – in other words, given the expert selected for some token in one layer, it may be hard to predict the next steps. Text tokens (probably given that very few experts are indeed often used for text) have more predictable trajectories.

F.5 BPR rankings

The local entropy loss encourages concentrated routing predictions with high $p_{\text{max}}$ for text. At the same time, BPR prioritises tokens with high $p_{\text{max}}$. One might assume that this combination is effectively just ranking all text tokens first. The following plots give us some insight into how the buffers end up sorting tokens from both modalities. Figures 32 and 33 show the priority distribution on the training data for the B/32 and B/16 models, respectively. Under a data shift, Figures 34 and 35 show the same statistics for COCO data, and Figures 36 and 37 for ImageNet. In these cases, no extra training was performed (i.e., it is zero-shot). Overall, we see that while text tokens enjoy by default a much higher priority, this is not always the case, and some (important?) image patches are sometimes processed before other text tokens.

Appendix G LIMoE-H/14 Analysis

In this section, we share some details and analysis regarding our largest model, the LIMoE-H/14. Figure 38 shows the development of the max routing probability across different MoE layers. Figure 2 shows qualitatively the specialization of image experts. Experts naturally specializing on semantic concepts such as body parts (hands, eyes), textures, fauna, food and doors. In Figure 39, we show the distribution of tokens per type and expert for every layer. Note that we set the entropy loss to approximately require at least 4 text experts, something that seems to agree well with the plot (in this case the ratio text:image tokens was close to 1:27).

We analyse the routing distributions of text tokens for LIMoE-H/14, using NLTK to distinguish between verbs, nouns, adjectives, prepositions and determiners. Note that the SentencePiece tokenizer breaks words into smaller units, which are not necessarily always handled by the same expert, so it is not possible to perfectly parse every token processed by every expert.

The majority of tokens are from images, so only 3-4 experts handle text in this scenario. Figure 40 contains preliminary analysis, showing for each expert the breakdown of tokens it handles. Though some experts process a bit of everything (e.g. experts 0 and 1 in layer 6 and 31), there are signs of some semantic specialization. There are often experts which process mostly padding tokens. In Layer 14, expert 1 processes no prepositions, determiners or verbs, focussing on nouns and adjectives (and some padding); similarly expert 1 processes very few nouns or adjectives, instead handling padding tokens.