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

In plain words

Most image-and-text models use two separate networks, one for pictures and one for words. A mixture-of-experts model is a single network made of many small specialist parts, and a router sends each piece of input to only a few of them. That keeps the cost low while the total size grows. This paper trains one such network on images and text together, so that a picture and its caption land close together in a shared space.

Mixing the two kinds of input in one expert pool causes trouble. Images are far more numerous than words, so the experts can end up handling mostly pictures and leave the words without room. The authors add extra training penalties that push the router to spread text over several experts and to keep its choices steady. With these in place, the models do better than dense networks of the same compute cost, and the largest one is reported as comparable to state-of-the-art models, though it was trained with extra data.

TL;DR

LIMoE is a one-tower, sparsely routed model trained on paired images and captions with a symmetric contrastive loss §2.1. Mixing modalities in the router causes capacity overflow and expert collapse, so the authors add per-modality local and global entropy losses and rely on batch priority routing §2.2.1 §2.2.3. Against dense one-tower baselines of equal compute, the paper reports gains of about 10 points zero-shot and 12 points 10-shot on ImageNet, on average across scales §3.1. LIMoE-H/14 reaches 84.1% zero-shot ImageNet, but it was trained with extra JFT-4B data and model soups §3.2 app. A.3.

Problem

Sparse mixture-of-experts models scale well in vision and in text, but they are usually trained on one modality at a time [abstract]. When images and text share one routed encoder, the router collapses. The paper reports two failure modes: a single expert that receives all tokens of the minority modality and drops them once full, and training instability that depends on the router's choices §2.2.1 §2.2.2. The classic V-MoE auxiliary losses, which balance total load across experts, do not prevent this, because the balance can hold globally while one modality is dropped §2.2.1.

The contrastive setting is assumed throughout: a batch of paired images and captions and a learned temperature §2.1. The results depend on the large paired web dataset used in the source paper §3 app. A.1, so the exact accuracy is tied to that data.

Method

Objective. Given n image-caption pairs, the model learns embeddings such that matched pairs are closer than unmatched ones. The loss is the average of an image-to-text and a text-to-image softmax cross-entropy over the batch, with a learned temperature T §2.1 eq. 1. The paper's appendix lists the temperature initialised at 10 app. A.

Architecture. One transformer encodes both modalities. Images are split into ViT-style patches and text is tokenised with a SentencePiece vocabulary; each modality has its own linear input projection §2.2. The shared encoder is not told which modality a token came from §2.2. The final token representations are average-pooled per modality §2.2. Each modality then gets its own linear projection W_m before the contrastive loss is applied §2.2.

Sparse MoE layers. MoE layers follow the V-MoE and GShard designs §2.2. Each token x is routed by a learned linear gate that is softmax-normalised over E experts, and the top K experts' MLP outputs are combined with their gate values §2.2. The router is one dense layer shared by all tokens, independent of modality, with no constraints on the gate app. E.

Capacity and dropping. Each expert has a fixed buffer. Tokens above capacity are dropped and their expert output is zero. The success rate, the share of tokens not dropped, is the paper's stability indicator §2.2.

Auxiliary losses. The paper uses the classic V-MoE losses, and adds two entropy losses applied per modality §2.2.2 app. B. The final combination, read from the bold entries of the auxiliary-loss sweep, is the load loss and z-loss, the global entropy loss on both modalities and the local entropy loss on text only §4.1. For modality m with n_m tokens and gate matrix G_m:

• Local entropy: the mean per-token entropy of the expert distribution. Minimising it makes each token's routing concentrated §2.2.2 eq. 2.
• Global entropy: the negative entropy of the expert distribution averaged over the modality's tokens. It is thresholded as max(0, tau + Omega_global), so the loss is zero once the marginal entropy reaches tau, which corresponds to a soft minimum of S = e^tau experts §2.2.2 app. B.

Their sum equals the negative mutual information between experts and tokens, conditioned on modality §2.2.2.

Batch priority routing (BPR). When an expert is full, tokens are kept in order of their largest routing weight rather than their position in the batch. This avoids favouring one modality through batch order §2.2.3. Ordering by batch position (first-in-first-out) makes training perform very poorly, and random shuffling is better but still weaker §4.2.

Training objective used in the experiments. The auxiliary losses are averaged and weighted by a single coefficient of 0.04 app. B. The text states this as a weighted average of all auxiliary losses; see Pitfalls.

# LIMoE forward pass with routing losses  [§2.2] [§2.2.2] [app. B]
# Input:  image_tokens (B, N_img, D); text_tokens (B, 16, D)
# Output: image_embedding (B, d), text_embedding (B, d), aux_loss (scalar)

def encode(image_tokens, text_tokens):
    # 1. Put both modalities into one sequence, with no modality flag  [§2.2]
    tokens = concat(image_tokens, text_tokens, axis=1)
    aux_losses = []
    for layer in encoder_layers:
        tokens = self_attention(tokens)
        # 2. Router: softmax over E experts for each token  [§2.2]
        gates = softmax(tokens @ W_gate)           # (B, N, E)
        # 3. Top-K assignment with buffer capacity; BPR chooses which
        #    tokens survive when an expert is full  [§2.2.3]
        tokens = moe_dispatch(tokens, gates, top_k=1)
        # 4. Load, z-loss and entropy losses on this layer's gates  [§4.1]
        aux_losses += routing_losses(gates, modality_of_tokens)
    image_embedding = pool(tokens[:, :N_img])
    text_embedding = pool(tokens[:, N_img:])
    return image_embedding, text_embedding, mean(aux_losses)


def training_loss(images, text_ids):
    # 1. Contrastive loss on projected, pooled embeddings  [§2.1] [eq. 1]
    image_emb, text_emb, aux = encode(embed_image(images), embed_text(text_ids))
    logits = temperature * (proj_image(image_emb) @ proj_text(text_emb).T)
    contrastive = symmetric_softmax_cross_entropy(logits)
    # 2. Add the routing regulariser with weight 0.04  [app. B]
    return contrastive + 0.04 * aux

Scaling recipe. The scaling study trains each model at batch size 16,384 for 781,250 steps at resolution 224, the same number of examples as CLIP §3.1 app. A.1. The largest model, LIMoE-H/14, uses batch size 21,502 at resolution 288, with JFT-4B images mixed 3:1 with the default data app. A.3. Full architecture definitions are in app. C.

Souping. For LIMoE-H/14 the authors average 8 checkpoints from one run, drawn from several reverse-square-root cooldowns app. A.3. The checkpoints were chosen greedily to maximise zero-shot accuracy on the ImageNet validation set, using a subset of the CLIP prompts app. A.3.

Hyperparameters

Results

Author-reported. Setup: web-scale image-text pairs (3.6B pairs from the LiT data, §3), dense one-tower and two-tower baselines at matched compute, and ImageNet and COCO evaluations §3 app. C.

Scaling study, zero-shot and 10-shot ImageNet top-1 and COCO recall@1 app. C:

Source: app. C, §3.1 (headline), §3.2 (H/14). The authors report average absolute gains of 10.1 points (zero-shot) and 12.2 points (10-shot) across scales §3.1.

Comparison with state-of-the-art contrastive models, zero-shot ImageNet top-1 §3.2:

Other numbers: LAION-400M, dense B/16 versus LIMoE B/16 at 16,384 batch and 1.4B examples, three trials each: zero-shot 56.0 vs 61.0, 10-shot 47.7 vs 54.3 app. E. The B/16 ablation (100k steps) gives dense one-tower 49.8 / 43.8, dense two-tower 54.7 / 47.1 and LIMoE 56.9 / 50.5 for zero-shot and 10-shot ImageNet §4 app. C.

Paper vs code

No official code.

I searched in three steps. (1) The paper, the arXiv comments and the LaTeX source contain no code link. The source's checklist, which is commented out of main.tex, answers "No" to releasing code checklist.tex. (2) A GitHub repository search for "limoe" returned only third-party reimplementations (for example kyegomez/LIMoE and YeonwooSung/LIMoE-pytorch), which are not used. (3) An organisation search for mixture-of-experts repositories under google-research returned none, and the google-research/big_vision README does not mention LIMoE. The paper does not describe a code release, and most of its models were trained on proprietary data §4.2.

Pitfalls

• No code, so every implementation detail is from the paper. Some details that matter for matching the reported numbers are not stated: the exact form of the combined auxiliary loss, the gate initialisation, the capacity factor used in training, and the H/14 learning rate.
• Auxiliary loss weighting is ambiguous. The text says the final auxiliary loss is an average of all losses, "weighted by a single regularization parameter" set to 0.04 app. B. It is not stated whether all auxiliary losses (classic plus entropy) are averaged together, or whether the two entropy losses are averaged and then combined with the classic ones. Read the statement as an implementer's choice and note it in your code.
• Entropy losses on images are mostly inactive. The global image threshold of log 20 is never reached, and the paper says the global image loss is not applied in most settings app. B. The local entropy loss on images is harmful in the sweep §4.1. Enable entropy losses on text tokens.
• Text expert count is a soft minimum, not a hard count. The threshold gives roughly e^tau experts for text, but the model may use more §2.2.2 §4.1.
• Entropy losses are not only a fix for imbalance. With text length 16 and image length 16 (roughly balanced), the entropy losses still beat the classic setup, so do not drop them when the data is balanced §4.1 app. C.
• Unimodal evaluation drops tokens. Expert capacities are set for mixed batches. Evaluating text alone or images alone with normal capacity causes heavy dropping, so the reported zero-shot numbers use capacity factor 16 app. D. Without this, results will differ. Pruning modality-specific experts is an alternative the paper explores app. E.
• Batch priority routing is needed. Batch-order dispatch performs very poorly, and random shuffling is still worse than BPR §4.2 app. C.
• Grouped routing trades stability for speed. Up to 4 groups is stable; 8 or more groups degrade training. With 4 groups, step time falls by about 30% app. E.
• Batch size and data quantities disagree across sections. Ablations use 8,096 in the text §4 but 8,192 in the appendix app. A.2 app. C. H/14 uses 21k in the text §3.2 and 21,502 in the appendix app. A.3. The text describes 3.6B pairs from LiT §3, which the paper does not reproduce.
• Results depend on proprietary data. Most models use internal web data. The LAION-400M experiments show gains of similar sign but a smaller absolute size (zero-shot +5.0 points) app. E.
• H/14 accuracy depends on JFT-4B co-training and model soups. The paper's own headline pairs the 84.1% number with JFT-4B co-training and souping. The souping selection uses ImageNet validation accuracy, and the reported result is zero-shot on ImageNet app. A.3. Treat this as an optimistic number for a clean comparison.
• Cited gains differ. The introduction says 10.1 points zero-shot, while the conclusions say 10.2% §1 §7. The gains are averages over scales and are sensitive to which scales are counted.
• Compute claims depend on the measure. LIMoE is Pareto-optimal in FLOPs, but for COCO retrieval its step-time gain over dense baselines is not significant app. D.
• Expert count matters. More total experts at fixed K=1 improve results, unlike in vision-only MoEs app. E.
• Mutual-information loss variant. Replacing the two entropy losses with the single mutual-information loss costs about 1.5 points zero-shot and 0.1 points 10-shot §4.1.
• Modality projection weights are per modality. The W_m projections before the loss are per modality; the shared transformer and router are not §2.2.

Lineage

• Builds on: V-MoE [riquelme2021vmoe], which supplies the sparse MoE layer design and the classic importance and load losses, and BPR [riquelme2021vmoe] §2.2.3. GShard [lepikhin2021gshard] and Switch Transformer [fedus2022switch] are cited for MoE routing and capacity §2.2.
• Compared against: CLIP [radford2021clip] and ALIGN [jia2021align] as two-tower contrastive baselines; LiT [zhai2021lit], CoCa [yu2022coca] and BASIC [pham2022basic] as state-of-the-art contrastive models §3.2.
• Superseded by: not known from content in this wiki.