Summary (Overview)

• UEmbed is a decoder-only multimodal embedding model that produces both sparse lexical and dense representations in a single causal forward pass, eliminating the need for bidirectional encoders or auxiliary cross-modal modules.
• It appends $N=16$ learnable special tokens to the input, partitions the vocabulary into $N$ disjoint subsets via $k$-means clustering, and each token predicts sparse weights over its assigned subset.
• On MMEB-v2, UEmbed-9B achieves 71.8 (dense) and 71.0 (sparse), leading models trained on public data in dense retrieval and setting a new state-of-the-art for sparse multimodal retrieval.
• On BEIR (9 datasets, nDCG@10), UEmbed remains competitive with strong baselines like SPLADE-v3 and Echo-Mistral-SPLADE.
• Practical advantages include hybrid scoring (improving text and visual-document retrieval), compatibility with high-throughput serving (vLLM) and inverted indices, and cost savings in agentic search (BrowseComp-Plus).

Introduction and Theoretical Foundation

Information retrieval (IR) underpins applications like web search and question answering. Sparse lexical methods (e.g., BM25) are efficient and interpretable but limited to exact term matching. Learned Sparse Retrieval (LSR) methods like SPLADE improve semantic matching via neural token weighting and vocabulary expansion. However, LSR has three key limitations:

1. Architectural constraint: LSR relies on encoder-style bidirectional architectures (e.g., BERT). Decoder-only models cannot directly use max-pooling over all hidden states due to causal masking.
2. Limited modality: Most LSR methods are text-only; multimodal extensions require auxiliary cross-modal modules.
3. Unexplored practical utility: Sparse retrieval is rarely evaluated beyond accuracy on traditional benchmarks.

The primary challenge in deriving sparse representations from causal models is the information bottleneck: relying on a single token (e.g., EOS) to project into a massive $|V|$-dimensional vocabulary space limits capacity. The paper proposes UEmbed to overcome this by appending $N$ learnable special tokens, each responsible for a disjoint vocabulary subset, effectively distributing the projection.

Methodology

Preliminaries

InfoNCE Loss encourages high similarity for positive pairs:

$$\mathcal{L}_{\text{InfoNCE}} = -\log \frac{\exp(\text{sim}(q, d^+)/\tau)}{\sum_{d \in \{d^+\} \cup \mathcal{N}} \exp(\text{sim}(q, d)/\tau)}$$

where $\text{sim}(\cdot,\cdot)$ is cosine similarity (dense) or inner product (sparse), and $\tau$ is temperature.

Learned Sparse Retrieval (SPLADE) produces sparse weights by max-pooling over all token hidden states:

$$w_t = \max_{i=1}^L \log\left(1 + \text{ReLU}(\mathbf{h}_i^\top \mathbf{e}_t)\right)$$

where $\mathbf{h}_i$ is the $i$-th token hidden state, $\mathbf{e}_t$ is the embedding of term $t$, and $L$ is sequence length.

UEmbed Method

1. Vocabulary Compression: Reduce vocabulary from 248,320 to 184,016 via accent stripping, lowercasing, and whitespace collapsing.

2. Partitioned Sparse Heads: Append $N=16$ learnable special tokens $\langle s_1\rangle,\dots,\langle s_N\rangle$. Vocabulary $\mathcal{V}$ is partitioned into $N$ disjoint subsets $\mathcal{V}_1,\dots,\mathcal{V}_N$ via $k$-means clustering. Each special token $\langle s_k\rangle$ predicts sparse weights over $\mathcal{V}_k$:

   $$w_t^{(k)} = \log\left(1 + \text{ReLU}\left(\mathbf{W}_t^{(k)\top}\mathbf{h}_{s_k} + b_t^{(k)}\right)\right) \quad \forall t \in \mathcal{V}_k$$

   where $\mathbf{h}_{s_k}$ is the final hidden state of $\langle s_k\rangle$. The full sparse vector is the concatenation:

   $$\mathbf{w} = \bigoplus_{k=1}^N \left[ w_t^{(k)} \right]_{t \in \mathcal{V}_k}$$

3. Dense Representation: Use the hidden state of the EOS token preceding the special tokens as the dense embedding $\mathbf{d} \in \mathbb{R}^D$.

4. Unified Objective:

   $$\mathcal{L} = \mathcal{L}_{\text{InfoNCE}}^{\text{dense}} + \lambda \mathcal{L}_{\text{InfoNCE}}^{\text{sparse}} + \alpha_q \mathcal{L}_{\text{FLOPS}}^q + \alpha_d \mathcal{L}_{\text{FLOPS}}^d$$

   where $\mathcal{L}_{\text{FLOPS}}$ regularizers encourage sparsity.

Data Curation

• Sources: 3.94M query-document pairs from Echo-embedding, MLDR, and MMEB training sets.
• Hard Negative Mining: Use Qwen3-VL-Embedding-8B as teacher to retrieve top-$k$ non-relevant documents for each query.

Empirical Validation / Results

Multimodal Results (MMEB-v2)

Table 1 reports aggregate scores across image, video, and visual document tasks.

Key findings:

• UEmbed-9B (dense) leads models trained on public data, outperforming RzenEmbed-V2-7B (71.1) and Ops-MM-Embed-7B (67.1).
• Sparse performance is within 1.0 point of dense at all scales, with minimal gap on VisDoc tasks.

Text Results (BEIR)

UEmbed-9B matches Echo-Mistral-SPLADE (55.2) while also supporting dense retrieval and multimodal inputs.

Ablation Studies

• Joint vs. single-mode training: Joint training matches specialists (Figure 2a).
• Sparse temperature: $\tau_s = 32$ gives best sparse performance without harming dense (Figure 2b).
• Number of special tokens: $N=16$ is optimal; $N=32$ degrades due to tiny vocabulary subsets (Figure 2c).
• Vocabulary partitioning: Semantic clustering (k-means) outperforms random and max-distance strategies (Table 4).
• Comparison to bidirectional SPLADE: UEmbed-2B surpasses SPLADE baseline by +3.2 (dense) and +2.1 (sparse) on MMEB-v1 image subset (Table 3).

Practical Advantages

1. Hybrid scoring: Linear interpolation of dense and sparse similarities improves Text (+0.3) and VisDoc (+0.5) on MMEB-v2 (Table 5).
2. Efficiency: Fully autoregressive, compatible with vLLM and inverted indices.
3. Agentic search: On BrowseComp-Plus, sparse mode reduces average search rounds while maintaining recall (Table 6).

Theoretical and Practical Implications

• Theoretical: The paper demonstrates that sparse retrieval can be natively integrated into decoder-only MLLMs by overcoming the causal attention bottleneck via partitioned vocabulary heads. This unifies dense and sparse retrieval within a single causal forward pass, eliminating reliance on bidirectional encoders and auxiliary modules.
• Practical: UEmbed offers a drop-in replacement for both dense and sparse retrieval systems, with hybrid scoring, high-throughput serving, and cost savings in iterative reasoning (agentic search). Its interpretability via activated vocabulary terms aids bias auditing.

Conclusion

UEmbed is a decoder-only multimodal embedding model that unifies dense and sparse retrieval in one causal forward pass. It achieves state-of-the-art sparse multimodal retrieval on MMEB-v2 and competitive text retrieval on BEIR. The model is efficient, compatible with modern serving stacks, and demonstrates practical advantages in hybrid scoring and agentic search. Limitations include language bias (English/Chinese dominance), occasional anomalous token activations, and a slight performance gap in video tasks. Future work can address vocabulary stability, multilingual support, and modality-specific capacity.