PaperlayerView exact revision on arXiv

Source-linked AI summary

MoEMB: Scaling Universal Multimodal Embeddings with Efficient Mixture-of-Experts Models

Xuanming Cui, Shlok Kumar Mishra, Wentao Bao, Aashu Singh, Zihao Wang, Xiangjun Fan, Jun Xiao, Ser-Nam Lim, Jianpeng Cheng

arXiv:2609.08663v1cs.LGcs.AI

TL;DR

UME scaling must increase encoder capacity without sacrificing contrastive-training efficiency or low-latency retrieval across tasks of varying complexity. MOEMB uses sparse expert scaling and adaptive computation, achieving state-of-the-art public-data results while reducing compute with minimal performance loss.

  • Problem

    UME encoders must handle diverse tasks while scaling remains constrained by retrieval latency, contrastive-training memory and batch size, and redundant computation on simpler tasks.

  • Method

    MOEMB scales UME encoders along the expert axis with sparse MoE and systematically studies adaptive computation across training-based and inference-only strategies.

  • Results

    MOEMB sets new state-of-the-art results on MMEB-V2 and MRMR among public-data models, while adaptive computation saves 50% compute with less than 1% performance degradation.

  • Takeaways & Limitations

    Expert scaling preserves single-vector, non-autoregressive embedding while improving capacity, and adaptive computation improves the efficiency of MoE-based UME.

  • Takeaways & Limitations

    MOEMB currently lags models trained with external data and does not support audio input.

Abstract

from arXiv · show

Universal multimodal embedding (UME) increasingly demands encoder's capacity for handling a broad range of tasks and modalities with increased complexity. Prior scaling methods either increase the representation size, retrieval effort, or scales the encoder into a heavy multimodal LLM. Recent works, such as Think-Then-Embed (TTE), explore scaling via reasoning tokens. However, embedding models are hard to scale up: increasing parameters directly tradeoffs for the large training batch size that contrastive learning needs, and retrieval has to be served under tight latency. Moreover, UME tasks are diverse in complexity, where scaling up embedders can bring significant redundant computation. In this work, we propose MOEMB, which instead scales UME along the expert axis through mixture-of-experts (MoE), growing encoder capacity while preserving single-vector, non-autoregressive encoding. Through a systematic study of the design space and training recipes for MoE-based UME, MoEMB sets a new state of the art on both MMEB-V2 and MRMR among models trained on public MMEB-family data: with only 3B active parameters, MoEMB surpasses TTE-based methods with >4x active parameters, using significantly less computes. To further improve the scalability and efficiency, we conduct the first comprehensive study of adaptive computation for MoE-based embedding, spanning diverse strategies across training-based and inference-only methods. Together, these results support expert scaling as an effective and efficient direction for UME, with adaptive computation further improving efficiency for MLLM-based embedding models towards large-scale retrieval and recommendation systems.

1 Introduction

UME tasks increasingly span diverse modalities and capabilities, making encoder capacity, efficiency, and generalist performance central scaling challenges. MOEMB addresses these challenges by scaling through sparse experts and adaptive computation, achieving strong public-data benchmark results with lower compute.

  • UME tasks now span classification, question answering, retrieval, grounding, and document search across heterogeneous inputs.
  • Representation-dimensionality and multi-vector methods expand storage or retrieval budgets without increasing the encoder’s understanding and reasoning capacity.
  • Scaling UME encoders is constrained by low-latency retrieval, doubled contrastive-training activation memory, batch-size tradeoffs, and redundant computation on simpler tasks.
  • MOEMB scales the encoder along the expert axis through sparse MoE rather than increasing embedding dimension, dense parameters, or reasoning tokens.
  • 70.4 overall accuracy on MMEB-V2 is achieved by MOEMB-A3B with approximately 3.1B active parameters, surpassing public-data UME models including higher-compute TTE encoders.
  • 50% compute savings are achieved with less than 1% performance degradation through adaptive computation, with token pruning and adaptive top-k providing the strongest trade-off.

2 Preliminaries

UME encoders map instruction-conditioned multimodal queries or targets to normalized single-vector representations. Training uses a one-directional contrastive objective based on cosine similarity.

  • An UME embedder maps an instruction-conditioned query or target containing optional visual and text inputs to a single vector.
  • The same pretrained multimodal encoder processes both queries and targets before pooling and normalization produce the final embedding.
  • For a batch of N pairs, training follows a standard one-directional InfoNCE objective.
  • The contrastive objective uses temperature τ and cosine similarity to compare query and target representations.

3 Building Sparse MoE UME Models

MOEMB explores sparse MoE UME designs and their adaptation to contrastive embedding training. Native MoE routing performs best among the studied backbone strategies, while pooling, router adaptation, and training dynamics are also evaluated.

  • 3.1 MoE Backbone Designs: MOEMB compares native MoE, dense-to-sparse upcycling, and modality-level routing for sparse multimodal embedders.Native MoE reuses pretrained sparse experts and routers; upcycling splits dense feed-forward blocks into replicated experts; modality-level routing assigns each modality to a fixed expert.
  • 3.1 MoE Backbone Designs: Native MoE retains language-model-pretrained experts and router components, with each token routed to its top-k experts plus a gated shared expert.The router scores all experts, keeps the highest-scoring k, renormalizes their scores, and gates the shared expert separately.
  • 3.1 MoE Backbone Designs: Native MoE achieves the best empirical performance and becomes the default setting for subsequent experiments.Dense-to-MoE upcycling lags behind native MoE, while modality-level routing slightly improves over the dense baseline but remains behind native MoE.
  • 3.2 Training MoE-UME: The adaptation study examines pooling, router adaptation, auxiliary objectives, and contrastive-training dynamics for MoE-based UME.Training dynamics include global batch size, logit scaling, LoRA rank, learning rate, and training duration.
  • 3.2 Training MoE-UME: The designated <emb> token outperforms simple last-token pooling in the evaluated comparison.The study compares only last-token pooling with pooling from a designated <emb> token.

4 Adaptive Computation for MoE-UME

MOEMB studies six adaptive-computation actions for MoE-based UME across training-based and inference-only approaches. The methods target token, layer, expert, and routing computation while addressing the gap between logical compute reduction and actual inference acceleration.

  • 4.1 Training-Based Adaptive Computation Methods: The study evaluates six adaptive-computation actions spanning general methods and MoE-specific methods in both training-based and inference-only regimes.The actions are organized into six categories, with the top three general adaptive-computation actions and the bottom three specific to MoEs.
  • 4.1 Training-Based Adaptive Computation Methods: Token pruning removes tokens either once at an early layer or progressively across layers using per-token gates.Gradual pruning with STE and an AViT-style halting mechanism is unstable, so the paper realizes gradual pruning with GRPO policy optimization.
  • 4.1 Training-Based Adaptive Computation Methods: Learned expert skipping adds a skip gate that determines whether each token skips the entire experts computation at a layer.A model-wide skip budget permits different skip rates across layers, and soft fading trains the otherwise nondifferentiable hard gate.
  • 4.1 Training-Based Adaptive Computation Methods: Combining multiple adaptive actions requires training strategies that manage shared budgets, competing gates, and optimization instability.The study compares global or per-action budgets, alternating optimization, and sequential training with completed actions frozen.
  • 4.1 Training-Based Adaptive Computation Methods: GRPO trains discrete adaptive actions by sampling action groups, distilling toward a full-compute teacher embedding, and gating efficiency rewards by retrieval recall.This replaces a reference-policy KL term and reduces reward hacking by requiring efficiency rewards to depend on in-batch retrieval success.
  • 4.2 Inference-Based Efficiency: Gated DeltaNet’s writing-strength β provides a training-free token-importance signal when FlashAttention does not materialize attention maps.Averaged β scores rank visual tokens for pruning and drive subsequent inference-only token and expert decisions.
  • 4.2 Inference-Based Efficiency: Adaptive top-k routing dynamically selects fewer than k experts per token using calibrated thresholds and layer-specific scaling factors.The threshold retains experts whose scaled routing weights exceed τ, while scaling factors preserve more experts in critical layers.
  • 4.2 Inference-Based Efficiency: Total expert pruning removes experts permanently based on calibration-time saliency, reducing memory rather than compute.REAP generally performs better among the compared output-based saliency variants because it weights expert-output norms by routing weights.

5 Experiments

Experiments evaluate MoE-based UME on MMEB-V2 and MRMR, compare expert-scaling designs and training choices, and analyze adaptive computation. Results show strong benchmark performance and substantial compute reductions with limited accuracy loss.

  • Experimental Setup: MMEB-V2 contains 78 multimodal retrieval tasks, while MRMR evaluates reasoning-intensive retrieval across complex multimodal contexts.The benchmarks cover image, video, visual-document, cross-modal knowledge, theorem-identification, and contradiction-detection tasks.
  • Main Results: 70.4 overall accuracy at 3.1B active parameters lets MOEMB-A3B surpass Embed-RL-4B and UME-R1-7B on MMEB-V2 despite their 4× larger active parameter counts.MOEMB-A10B reaches 72.4, while increasing total experts from E = 32 to E = 256 raises MOEMB-A3B accuracy by +10.4 points with negligible QPS decline.
  • Main Results: 51.6 overall score makes MOEMB-A3B the top MRMR system with only 3.1B active parameters, exceeding TTE-based methods by 3.4–11.6 points.The comparison includes UME-R1-7B, Embed-RL-4B, and PLUME, all using larger active parameter footprints and compute budgets.
  • MoE Design Ablations: Native pretrained sparse MoE fine-tuning raises the benchmark average by +6.8 points, whereas dense-to-MoE upcycling does not outperform the dense backbone and modality-wise MoE gains only +0.6.The authors conjecture that limited contrastive-training scale makes it difficult to learn a strong router for upcycled models.
  • Training Hyperparameter Ablations: Halving global batch size from 512 to 256 drops overall performance by 4.52 points, highlighting the importance of large in-batch negative pools for contrastive learning.Other training findings include a +1.02-point gain when reducing the learning rate from 10−4 to 5 × 10−5, followed by +0.55 points from a second epoch.
  • Adaptive Computation: Adaptive computation maintains performance within 1.2 points of the 70.3 unreduced baseline while reducing execution cost through expert, token, and sequence pruning.Token pruning and adaptive top-k provide the strongest trade-offs; inference-only methods are competitive at moderate budgets, while trained pruning is advantageous under extreme token removal.
  • Adaptive Computation: Token pruning acts mainly in early layers, with stronger pruning for video and image tokens than for visual documents and text queries.The β signal is smoother and more evenly distributed across modalities than attention scores, which concentrate mainly on text tokens.

6 Limitations and future work

MOEMB’s current scope is limited relative to models trained with external data and does not yet include audio input. Future work targets larger-scale training data and additional modalities.

  • MOEMB currently lags behind models trained with external data despite achieving state-of-the-art performance among models trained with public MMEB-V2 data.
  • Audio is not currently supported as an input modality.
  • Future work includes training MOEMB on larger-scale data and adding more input modalities.

7 Conclusion

MOEMB scales universal multimodal embedding capacity along the expert axis rather than through embedding dimension or test-time reasoning. The paper reports state-of-the-art public-data results and substantial compute savings from adaptive computation.

  • MOEMB scales universal multimodal embedder capacity along the expert axis rather than embedding dimension or test-time reasoning.
  • MOEMB adds over 10 accuracy points at near-constant active compute and establishes state-of-the-art performance among public-data models on MMEB-V2 and MRMR.
  • Adaptive computation eliminates up to half of logical compute for under a one-point accuracy drop in MoE embedders.
  • Gated DeltaNet write strength β serves as a zero-cost token-salience signal, while custom serving kernels convert theoretical FLOP savings into deployed throughput gains.

B.1 Model Architecture and Experimental Setup

The appendix describes MOEMB’s sparse-MoE backbones, multimodal training setup, routing objectives, and evaluation protocol. It also specifies the expert-scaling and baseline comparisons used for the experiments.

  • Model Architecture: MOEMB-A3B and MOEMB-A10B use 256 routed experts per layer, top-8 routing, and one shared expert, with 3.1B and 10.2B active parameters respectively.
  • Experimental Setup: Training uses 24 sampling streams spanning image, video, and visual-document retrieval, while excluding generated chain-of-thought sequences from targets.
  • Routing Objectives: The routing study evaluates load-balancing, z-loss, conflict-aware routing, task-ID priors, and instruction priors as alternatives for adapting expert assignments.
  • Routing Objectives: Conflict-aware routing uses task-pair gradient conflict to encourage distinct expert assignments through conflict-weighted Jensen–Shannon divergence.
  • Evaluation Protocol: Expert-scaling experiments prune the full 256-expert MOEMB-A3B model to evaluate expert counts of 32, 64, 128, 192, and 256.
  • Evaluation Protocol: Baselines include autoregressive think-then-embed systems with reasoning enabled and disabled, latent-reasoning scaling, and single-vector dense references.

C.3 Converting FLOP Savings into Wall-Clock Speedups

Wall-clock savings depend on fused kernel implementation, tile occupancy, sequence compaction, and workload size rather than FLOP reduction alone. The appendix quantifies when adaptive expert and token computation produces actual throughput gains.

  • Kernel Efficiency: 1189.5 ms versus 210.2 ms shows sequential expert looping is 5.7–8.0× slower than fused grouped GEMMs at B = 8 and S = 1024.
  • Tile-Bound Constraints: Grouped GEMM savings appear only when pruning empties a tile, because partially filled tiles cost nearly as much as full tiles.
  • Tile-Bound Constraints: At T = 4,096, a 52.9% skip policy causes a 0.91× slowdown, whereas at T = 8,192, a 69.6% skip policy reduces two tiles to one and yields 1.11× speedup.
  • Backend Customization: SonicMoE provides a 1.40× speedup over standard grouped GEMM at fixed k = 8, while further reducing active experts eventually encounters diminishing returns.
  • Token Pruning: Spatial masking produces a flat 1.02× speedup, whereas physical sequence compaction reaches 1.228× at a 25% visual-token keep ratio.
  • Dense-to-MoE Upcycling: Dense-to-MoE upcycling does not improve representation quality at equal active compute, and surpasses the dense baseline only after doubling active compute per token.

D.3 Expert Routing Behavior Across Modality and Task

Expert routing is organized more strongly by modality than by task category, with modality differences largest in early layers and diminishing with depth.

  • The routing analysis profiles top-8 assignments across 40 MoE layers using 10 MMEB-V2 categories and Jensen–Shannon distance over 256-expert distributions.
  • Routing divergence is lower within modalities than between modalities, while task-category grouping shows no significant separation.Mean JS distance is 0.114 within modalities versus 0.323 between modalities; task-category separation is not significant (z = −0.66, p = 0.73).
  • Image and video inputs form a unified visual routing cluster rather than sharply distinct modality groups.Their inter-modality JS distance is 0.117, lower than the intra-image distance of 0.136.
  • Modality divergence peaks in early decoder layers and steadily diminishes with depth.At layer 4, JS distance is 0.339 and 172 of 256 experts show strong modality preference.

E.1 Training-Based Methods

The training-based adaptive-computation methods combine learned gates, pruning, skipping, and policy optimization to reduce MoE embedding cost, but several joint strategies remain unstable.

  • Training setup: Training-based variants use large-scale BF16 contrastive training with LoRA adapters and a global batch size of 512.Training uses 64 NVIDIA H200 GPUs, expert parallelism of 64, AdamW, and sequence length up to 8,192.
  • Learned efficiency actions: Token pruning removes visual tokens after an early decoder block, while layer and expert skipping learn compute-reduction decisions.Layer skipping bypasses attention and MoE sub-layers; expert skipping penalizes routed-expert activity.
  • Adaptive top-k: Adaptive top-k routing selects physical or null experts jointly, allowing each token’s effective expert width to vary without manual thresholds.The router selects among 512 candidates, including 256 null experts and 256 physical experts.
  • Policy optimization: GRPO trains a per-layer policy that assigns each token one of ROUTE, SKIP, or DROP while preserving native greedy expert routing.The policy changes execution state but does not alter expert-selection logic.
  • Failure modes: Joint pruning methods can fail because aggressive regularization harms representations, whereas conservative settings produce near-zero gate gradients.The authors hypothesize that delayed or non-differentiable attention paths create high early-training gradient variance.
  • Inference comparison: Uniform top-k retains much of the pretrained router’s efficiency benefit: k = 4 loses 0.5 accuracy points, versus 6.6 for random selection at the same 11.6% FLOPs reduction.

F Additional Analysis on GDN Write Strength β

GDN write strength β provides a stable, spatially coherent signal across depth and video frames, supporting early single-shot token-pruning decisions distinct from attention.

  • Spatial behavior: β is semantically aligned, highlighting contiguous informative regions such as foreground objects and document text or plots.
  • Depth behavior: β remains more stable across depth than <emb> attention, with average cross-layer cosine similarity of 0.989 versus 0.752.β has a minimum similarity of 0.973, while attention ranges from 0.68 to 0.86.
  • Temporal behavior: Across video frames, β has cross-frame cosine similarity of 0.992 versus 0.744 for attention.
  • Pruning implication: The stability of β makes early single-shot pruning decisions remain useful through later layers.
  • Decision depth: Changing the pruning decision from layer 4 to layer 2 alters accuracy by at most +0.21 at tested keep rates, with modest additional FLOP savings.The reported FLOP savings are 0.35 and 1.03 percentage points at 75% and 25% keep rates, respectively.
  • Attention comparison: Visual comparisons distinguish β from <emb> attention: β is strong and smooth early, while attention is sparse and becomes more focused deeper in the network.

G Per-Task MMEB-V2 Results

This section presents per-task MMEB-V2 scores and includes a continuation marker for the associated table.

  • Table 14 reports MMEB-V2 scores separately for individual tasks.
  • The section includes a continuation of Table 14 from the previous page.
  • Single-pass TTE-based MOEMB is identified among the reported entries.
Loading 2609.08663v1…