UEmbed: Unified Sparse and Dense Multimodal Embeddings
Tingyu Song, Mingxin Li, Yanzhao Zhang, Dingkun Long, Pengjun Xie, Zhijie Nie, Yilun Zhao, Shu Wu
UEmbed unifies dense and sparse multimodal retrieval within a single decoder-only backbone.
How can a single multimodal model generate both dense and sparse embeddings to support both semantic and lexical retrieval?
Learned sparse retrieval methods typically rely on bidirectional encoders, which prevents them from using the high-throughput, causal architectures standard in modern multimodal models. UEmbed solves this by appending learnable special tokens to a decoder-only backbone, where each token acts as a "soft topic specialist" responsible for predicting sparse weights over a specific subset of the vocabulary. This unified approach achieves state-of-the-art performance in sparse multimodal retrieval while maintaining competitive dense retrieval scores on standard benchmarks.
Paper Primer
The core challenge is the information bottleneck: a single token cannot effectively project into a massive vocabulary space. UEmbed circumvents this by partitioning the vocabulary into $N$ disjoint subsets via k-means clustering, assigning each subset to a dedicated learnable special token that summarizes the input sequence.
The model is trained with a unified objective that combines dense and sparse InfoNCE losses, alongside FLOPS regularizers to enforce sparsity. This design allows the model to produce both dense and sparse representations in a single causal forward pass, ensuring compatibility with standard inverted indices and high-throughput serving stacks like vLLM.
UEmbed-9B establishes a new state of the art for sparse multimodal retrieval on the MMEB-v2 benchmark.
The model achieves a sparse score of 71.0, outperforming established dense models and closing the performance gap between dense and sparse modes to within 1.0 point.
The sparse mode provides significant cost advantages in agentic search workflows.
On the BrowseComp-Plus benchmark, UEmbed's sparse mode consistently requires fewer tool-call search rounds than its dense counterpart while maintaining comparable recall.
Why is a decoder-only architecture preferred over the bidirectional encoders used in previous sparse retrieval methods?
Decoder-only models are compatible with high-throughput serving frameworks like vLLM and allow for a unified backbone that supports both dense and sparse retrieval without needing auxiliary cross-modal modules.
Does the joint training of dense and sparse modes lead to negative transfer?
No; ablation studies show that the jointly trained UEmbed model closely matches the performance of single-mode specialists, indicating that the dual-mode capability incurs a negligible performance penalty.
Motivation and Problem Framing
Unifying dense and sparse retrieval in a single decoder‑only model eliminates duplicated systems.
Information retrieval powers web search and question answering, but sparse lexical methods like BM25 trade semantic depth for speed and interpretability. Learned Sparse Retrieval (LSR) improves semantics yet remains tied to encoder‑style, bidirectional models and struggles to handle multimodal inputs. UEmbed solves this by appending $N$ learnable special tokens to a decoder‑only MLLM, partitioning the vocabulary into $N$ disjoint clusters so each token predicts sparse weights, while the $EOS$ hidden state yields a dense embedding—all in a single causal forward pass.
Running separate dense‑embedding and sparse‑lexical models doubles compute, memory, and engineering effort.
LSR trains neural networks to emit sparse lexical vectors that capture semantics beyond exact term overlap.
Separate dense and sparse retrieval models double system complexity.
The UEmbed Architecture
UEmbed adds learnable tokens to a decoder‑only MLLM to enable joint dense and sparse retrieval.
Decoder‑only models cannot attend to future tokens, so max‑pooling over hidden states fails for sparse retrieval. UEmbed solves this by appending learnable special tokens that can summarize the entire input. The method also retains the ability to produce a dense embedding from the EOS token.
InfoNCE pushes the query toward its positive document while pulling it away from all other in‑batch negatives.
How does InfoNCE differ from the standard cross‑entropy loss?
InfoNCE normalizes over all in‑batch negatives, turning each query‑document pair into a separate softmax problem, whereas cross‑entropy typically normalizes over a fixed set of class labels.
Before partitioning the vocabulary, UEmbed compresses it by stripping accents, lowercasing, and collapsing whitespace, merging variants such as “Hello”, “HELLO”, and “héllò” into a single entry. This reduces the token count from 248,320 to 184,016, simplifying the later partitioning step.
The special tokens act like mail sorters that read the entire letter and then drop the relevant words into their assigned bins, each bin covering a disjoint part of the vocabulary.
Append tokens $\langle s_{1}\rangle$–$\langle s_{4}\rangle$ to the end of a 5‑token input.
Each token attends to the five preceding tokens and produces a hidden state $\mathbf{h}_{s_{k}}$.
For head 1, compute weights for its subset $\{t_{1},t_{2},t_{3}\}$ using $W^{(1)}$ and $b^{(1)}$, yielding $w_{t_{1}},w_{t_{2}},w_{t_{3}}$.
Repeat for heads 2–4, obtaining three weights per head.
Concatenate the four 3‑dimensional vectors to form $\mathbf{w}\in\mathbb{R}^{12}$.
Partitioning guarantees that each special token focuses on a distinct lexical region, preventing overlap and making the sparse representation more interpretable.
When only dense retrieval is required, UEmbed simply takes the hidden state of the EOS token that precedes the special tokens as the dense embedding $\mathbf{d}$, incurring no extra forward computation.
UEmbed jointly optimizes dense and sparse retrieval by summing their InfoNCE losses and adding regularizers that push term weights toward sparsity.
Why combine dense and sparse InfoNCE losses instead of training separate models?
Joint training forces a single model to produce both a semantic embedding and a lexical activation map, saving parameters and inference time while allowing the two signals to complement each other during retrieval.
UEmbed’s training data merges three public corpora, yielding 3.94 M query‑document pairs that span textual and multimodal domains.
The sources are (1) Echo‑embedding data for broad domain coverage, (2) MLDR for long‑document training, and (3) MMEB for multimodal examples.
Because multimodal datasets lack hard negatives, UEmbed uses the teacher model Qwen3‑VL‑Embedding‑8B to retrieve the top‑k non‑relevant documents for each query, treating those as hard negatives during training.
Multimodal Retrieval Performance
UEmbed matches or exceeds specialized multimodal models across image, video, and document retrieval.
UEmbed‑9B (dense) achieves the highest aggregate score among publicly‑available open models.
Table 1 shows an aggregate score of 71.8, outperforming RzenEmbed‑V2‑7B (71.1) and Ops‑MM‑Embed‑7B (67.1).
MMEB‑v2 measures multimodal embedding quality on image, video, and visually‑rich document tasks, reporting per‑task scores and an overall aggregate.
**Table 1.** Results on MMEB-v2. Bold marks the best score in each column among models (i.e., excluding Qwen3-VL-Embedding, which is trained with large-scale proprietary data).
UEmbed matches or exceeds specialized models across image, video, and document retrieval.
Text Retrieval and Ablations
Component ablations reveal each design choice’s impact on retrieval performance.
UEmbed appends learnable special tokens to a decoder‑only MLLM, enabling a single model to emit dense embeddings for semantic search and sparse lexical activations for exact matching.
BEIR aggregates a diverse set of text‑only retrieval tasks, providing a common nDCG@10 evaluation across domains.
UEmbed‑2B outperforms the bidirectional SPLADE baseline by 3.2 nDCG@10 points in dense mode.
Table 3 shows dense scores rising from 61.3 to 64.5 when using UEmbed‑2B.
UEmbed‑2B also beats the SPLADE baseline by 2.1 nDCG@10 points in sparse mode.
Table 3 reports sparse scores improving from 61.3 to 63.4 with UEmbed‑2B.
Hybrid scoring adds the largest gain on VisDoc (+0.5 nDCG@10) and a modest boost on Text (+0.3 nDCG@10).
Table 5 lists hybrid scores that exceed dense‑only results by 0.5 on VisDoc and 0.3 on Text.
**Figure 2.** Ablation studies of UEMBED: (a) joint training matches single-mode specialists; (b) sweeping the sparse softmax temperature $\tau_s$; and (c) sweeping the number of appended special tokens $N$.
**Table 3.** Comparison with the bidirectional SPLADE baseline on the image subset of MMEB-v1.
**Figure 3.** Case study of UEMBED on different modalities. We visualize the top-10 activated sparse tokens.
**Table 5.** Modality-wise performance of UEMBED-2B under dense, sparse, and hybrid modes.
**Table 6** Results on BrowseComp-Plus. “Avg. Search” denotes the average number of search rounds per question.
Overall, the ablations confirm that each component—joint training, semantic vocabulary clustering, and carefully tuned hyper‑parameters—contributes positively, while the unified model retains strong performance across dense, sparse, and hybrid retrieval settings.
Implementation and Data Details
Implementation specifics for training data, hyperparameters, and analysis.
We assemble a multi‑source training mixture covering text, image, video, and visually‑rich document retrieval. Figure 4 shows the proportion of each data source, while Tables 7 and 8 list the per‑dataset retrieval instructions.
**Figure 4.** Distribution of training data sources.
For text training we draw pairs from Echo‑Embedding and M3‑Embedding, covering fifteen datasets (e.g., NLI, DuReader, FEVER). Standard text tasks use a maximum sequence length of $1{,}500$ tokens, extended to $1{,}800$ tokens for long‑document retrieval (MLDR).
Multimodal training includes video (VideoCaption300k, VideoQA240k) sampled at eight frames per video with up to $200$ tokens per frame, visual‑document corpora (ColPali, VisRAG‑Synthetic/IID), and image‑text datasets (MMEB training set). Non‑video multimodal tasks use a $1{,}500$‑token limit, reserving up to $1{,}000$ tokens for the image component.
**Figure 5.** Training dynamics comparison. We plot the loss, max negative gap, and mean negative gap during training.
We initialize from a pretrained multimodal backbone and prepend $N=16$ learnable special tokens to the tokenizer and embedding matrix. LoRA adapters are applied to attention and MLP projections, while the visual encoder remains frozen.
Training runs in bf16 mixed precision with DeepSpeed ZeRO and gradient checkpointing. The cosine learning‑rate schedule peaks at $3\times10^{-5}$ with a warm‑up ratio of $0.1$. Dense retrieval temperature is $\tau=0.03$, sparse retrieval temperature $\tau_{s}=32$, FLOPS regularizer weights $\alpha_{q}=\alpha_{d}=1\times10^{-4}$, and sparse loss factor $\lambda=1.0$.
Hardware configurations: UEmbed‑2B and UEmbed‑4B train on $16\times$ A100 GPUs (per‑device batch size $16$), while UEmbed‑9B uses $32\times$ A100 GPUs (per‑device batch size $8$), maintaining a total batch size of $256$ across all variants.
Because sparse scores are raw inner products, their dynamic range can cause instability. Incorporating text data markedly stabilizes training, as evidenced by the rapid reduction of negative gaps in Figure 5.
Table 9 details the baselines compared against UEmbed on the multimodal benchmark MMEB‑v2 and the text benchmark BEIR.
Evaluation uses nDCG@10 for BEIR and the task‑specific metrics defined by MMEB‑v2. Dense scores are cosine similarities of the EOS‑token hidden states; sparse scores are inner products over the partitioned vocabulary representation.
Hybrid scoring combines dense cosine similarity $s_{\text{dense}}(q,d)$ and sparse inner product $s_{\text{sparse}}(q,d)$ as $s_{\text{hybrid}}(q,d)=\alpha\,s_{\text{dense}}(q,d)+\beta\,s_{\text{sparse}}(q,d)$. We fix $\alpha=1.0$ and set modality‑specific $\beta$ values: $\beta=5\times10^{-8}$ (Text), $5\times10^{-4}$ (Image), $1\times10^{-4}$ (Video), $7\times10^{-4}$ (VisDoc).
Sparse representations integrate directly with inverted indices, enabling large‑scale lexical retrieval. A benchmark comparing a Faiss dense index to a Lucene sparse index shows that, while dense search attains slightly higher NDCG@5 on the current corpus, sparse retrieval scales more favorably as corpus size grows.
Table 10 reports that joint dense‑sparse training preserves dense performance while delivering competitive sparse results.
Table 11 demonstrates that increasing the number of special tokens $N$ generally maintains performance, but $N=32$ causes a noticeable drop in VisDoc VDRv2 and overall VisDoc averages.
Extended Performance Metrics
Additional ablations and qualitative analysis of UEmbed on MMEB‑v2.
Table 10 reports the MMEB‑v2 scores for four representation modes: dense‑only, sparse‑only, and the two UEmbed‑2B variants (dense and sparse).
Table 11 explores how the number of special tokens $N$ (2, 4, 8, 16, 32) influences performance; increasing $N$ generally improves the average score, with $N\!=\!16$ yielding the highest average of 81.2.
Table 12 varies the temperature hyper‑parameter from 0.5 to 1.0; moderate temperatures (0.8–0.9) give the best average results (≈81.8), while very low or high values degrade performance.
C.1 Qualitative Analysis identifies modality‑specific gaps: for images, UEmbed struggles with fine‑grained categories and misleading lexical cues; for video, the lack of dedicated temporal supervision hurts action‑recognition tasks; for document vision, the rank‑KL objective of the baseline gives it an edge on multilingual and table‑reasoning data.
C.2 Case Study visualizes the top‑10 activated sparse tokens per query; blue tokens mark query activations, green tokens corpus activations, and red tokens shared activations. The analysis shows strong relevance for QA, Retrieval, and Grounding, but over‑association with tangential concepts in Classification tasks.
Questions & answers
What is UEmbed's main contribution?
UEmbed introduces a unified sparse and dense multimodal embedding framework that appends N learnable special tokens to a decoder-only backbone, where each token is responsible for predicting sparse weights over a disjoint subset of the vocabulary, enabling both dense and sparse retrieval from a single causal forward pass without separate models or auxiliary cross-modal modules.
What problem does UEmbed address?
UEmbed addresses the incompatibility of learned sparse retrieval (LSR) methods with modern decoder-only multimodal architectures: prior LSR methods rely on bidirectional encoders, which cannot leverage high-throughput causal serving stacks like vLLM and require separate models for dense and sparse retrieval, doubling system complexity.
Why is a decoder-only architecture preferred over bidirectional encoders for sparse retrieval?
Decoder-only models are compatible with high-throughput serving frameworks like vLLM and allow a unified backbone that supports both dense and sparse retrieval without needing auxiliary cross-modal modules, whereas bidirectional encoders are incompatible with these modern multimodal serving stacks.
How does UEmbed solve the information bottleneck of projecting into a large vocabulary?
UEmbed partitions the vocabulary into N disjoint subsets via k-means clustering and assigns each subset to a dedicated learnable special token; each token summarizes the input sequence and predicts sparse weights only over its assigned vocabulary partition, distributing the projection load across multiple tokens.
How does UEmbed produce dense embeddings?
For dense retrieval, UEmbed takes the hidden state of the EOS token that precedes the appended special tokens as the dense embedding, incurring no extra forward computation beyond what is already performed in the causal pass.
What training objective does UEmbed use?
UEmbed is trained with a unified objective combining dense and sparse InfoNCE losses alongside FLOPS regularizers (with weights α_q = α_d = 1×10⁻⁴) to enforce sparsity; the sparse loss factor is λ=1.0, the dense retrieval temperature is τ=0.03, and the sparse retrieval temperature is τ_s=32.
Does joint training of dense and sparse modes cause negative transfer?
No; ablation studies show that the jointly trained UEmbed model closely matches the performance of single-mode specialists, indicating that the dual-mode capability incurs a negligible performance penalty, and Table 10 confirms that joint dense-sparse training preserves dense performance while delivering competitive sparse results.
What datasets and benchmarks are used to train and evaluate UEmbed?
Training uses 3.94 million query-document pairs merged from Echo-Embedding (broad domain coverage), MLDR (long-document retrieval), and MMEB (multimodal examples), covering text, image, video, and visually-rich document tasks. Evaluation uses the MMEB-v2 multimodal benchmark (with task-specific metrics) and the BEIR text benchmark (nDCG@10).
How does UEmbed handle hard negatives for multimodal training data?
Because multimodal datasets lack hard negatives, UEmbed uses the teacher model Qwen3-VL-Embedding-8B to retrieve the top-k non-relevant documents for each query, treating those retrieved documents as hard negatives during training.
What are the key quantitative results reported for UEmbed?
With N=16 special tokens, UEmbed-2B achieves the highest average MMEB-v2 score of 81.2 among the token-count ablations; moderate sparse temperatures of 0.8–0.9 yield the best average results of approximately 81.8. The paper states UEmbed matches or exceeds specialized models across image, video, and document retrieval on MMEB-v2, and achieves state-of-the-art performance in sparse multimodal retrieval while maintaining competitive dense retrieval scores.
How does UEmbed's hybrid scoring work?
Hybrid scoring combines dense cosine similarity and sparse inner product as s_hybrid(q,d) = α·s_dense(q,d) + β·s_sparse(q,d), with α fixed at 1.0 and modality-specific β values: 5×10⁻⁸ for Text, 5×10⁻⁴ for Image, 1×10⁻⁴ for Video, and 7×10⁻⁴ for VisDoc.
How does UEmbed preprocess the vocabulary before partitioning?
Before k-means partitioning, UEmbed compresses the vocabulary by stripping accents, lowercasing, and collapsing whitespace, merging variants such as 'Hello', 'HELLO', and 'héllò' into a single entry; this reduces the token count from 248,320 to 184,016.
What are the limitations and failure cases of UEmbed?
The paper identifies several modality-specific gaps: for images, UEmbed struggles with fine-grained categories and misleading lexical cues; for video, the lack of dedicated temporal supervision hurts action-recognition tasks; for document vision, a baseline using a rank-KL objective has an edge on multilingual and table-reasoning data. Additionally, N=32 special tokens causes a noticeable performance drop on VisDoc VDRv2 and overall VisDoc averages.
How does UEmbed compare to prior sparse retrieval methods?
Prior learned sparse retrieval methods rely on bidirectional encoder architectures and cannot handle multimodal inputs without auxiliary modules, whereas UEmbed uses a decoder-only backbone with learnable special tokens that natively supports text, image, video, and document modalities in a single model. The paper reports UEmbed achieves state-of-the-art performance in sparse multimodal retrieval while remaining competitive on dense retrieval benchmarks.
How does sparse retrieval scale compared to dense retrieval in UEmbed's evaluation?
A benchmark comparing a Faiss dense index to a Lucene sparse index shows that dense search attains slightly higher NDCG@5 on the current corpus, but sparse retrieval scales more favorably as corpus size grows, making it advantageous for large-scale deployments.
What are the implementation details and hardware requirements for training UEmbed?
UEmbed initializes from a pretrained multimodal backbone with N=16 learnable special tokens prepended; LoRA adapters are applied to attention and MLP projections while the visual encoder is frozen. Training uses bf16 mixed precision with DeepSpeed ZeRO and gradient checkpointing; UEmbed-2B and UEmbed-4B train on 16× A100 GPUs (per-device batch size 16), while UEmbed-9B uses 32× A100 GPUs (per-device batch size 8), all maintaining a total batch size of 256.
What role does text data play in stabilizing UEmbed's sparse training?
Because sparse scores are raw inner products with a dynamic range that can cause training instability, incorporating text data markedly stabilizes training, as evidenced by the rapid reduction of negative gaps shown in Figure 5 of the paper.
Who are the authors, and where was UEmbed published?
The paper does not explicitly state the author names or publication venue in the provided text; it is available on arXiv at https://arxiv.org/abs/2608.02583.
Key terms
- Learned Sparse Retrieval (LSR)
- A retrieval approach where a neural model learns to assign importance weights to vocabulary tokens for each document or query, enabling fast lexical matching via inverted indices while incorporating semantic understanding.
- decoder-only model
- A transformer architecture that processes tokens left-to-right using causal (unidirectional) attention, as opposed to bidirectional encoders that can attend to all tokens simultaneously.
- learnable special tokens
- Trainable embedding vectors appended to the input sequence that are optimized during training to summarize the input and produce task-specific representations, here used to generate sparse vocabulary weights.
- vocabulary partitioning
- The process of dividing the full vocabulary into N disjoint subsets so that each learnable special token is responsible for predicting sparse weights over only its assigned subset, reducing the projection burden per token.
- k-means clustering
- An unsupervised algorithm that groups items into k clusters by minimizing within-cluster variance, used here to partition vocabulary tokens into semantically coherent subsets assigned to each special token.
- InfoNCE loss
- A contrastive training objective that treats all other samples in a batch as negatives and maximizes the similarity between a query and its positive document relative to all in-batch negatives via a softmax formulation.
- FLOPS regularizer
- A penalty term added to the training loss that discourages activating many vocabulary dimensions simultaneously, enforcing sparsity in the learned sparse representations.
- EOS token
- The end-of-sequence token in a language model; in UEmbed, the hidden state of the EOS token (positioned just before the appended special tokens) is used as the dense embedding for semantic retrieval.
- inverted index
- A data structure that maps each vocabulary term to the list of documents containing it, enabling fast sparse retrieval by looking up only the terms present in a query.
- vLLM
- A high-throughput serving framework designed for large language models that leverages efficient memory management and batching, compatible with decoder-only architectures.
- LoRA (Low-Rank Adaptation)
- A parameter-efficient fine-tuning technique that inserts small trainable low-rank matrices into model layers, allowing adaptation of large pretrained models without updating all original weights.
- nDCG@10
- Normalized Discounted Cumulative Gain at rank 10, a standard information retrieval metric that measures ranking quality by giving higher credit to relevant documents appearing earlier in the top-10 results.
- BEIR
- A heterogeneous text retrieval benchmark comprising multiple datasets across diverse domains, used to evaluate the generalization of retrieval models.
- MMEB-v2
- A multimodal embedding benchmark that evaluates retrieval performance across image, video, and visually-rich document tasks using task-specific metrics.
- hybrid retrieval
- A retrieval strategy that linearly combines dense semantic similarity scores and sparse lexical scores to leverage the complementary strengths of both representation types.
- hard negatives
- Training examples that are semantically similar to a query but not relevant, retrieved by a strong model to make contrastive training more challenging and improve model discrimination.
- BM25
- A classical sparse retrieval algorithm that ranks documents based on term frequency and inverse document frequency, widely used as a strong lexical baseline in information retrieval.
- DeepSpeed ZeRO
- A distributed training optimization framework that partitions model states across GPUs to reduce memory usage, enabling training of large models across many accelerators.
- MLLM (Multimodal Large Language Model)
- A large language model extended with the ability to process and reason over multiple input modalities such as text, images, and video.
- information bottleneck
- The constraint that a single fixed-size vector representation must compress all relevant information from an input, which becomes problematic when projecting into a very large output space such as a full vocabulary.