LaME: Learning to Think in Latent Space for Multimodal Embedding via Information Bottleneck

Peixi Wu; Biao Yang; Feipeng Ma; Bosong Chai; Bo Lin; Wei Yuan; Fan Yang; Tingting Gao; Hebei Li; Xiaoyan Sun

Source-linked AI summary

LaME: Learning to Think in Latent Space for Multimodal Embedding via Information Bottleneck

Peixi Wu, Biao Yang, Feipeng Ma, Bosong Chai, Bo Lin, Wei Yuan, Fan Yang, Tingting Gao, Hebei Li, Xiaoyan Sun

arXiv:2606.13061v3Jun 11, 2026cs.CV

TL;DR

CoT-based multimodal embedding is costly at inference and dependent on annotation quality, motivating reasoning directly in latent space. LaME uses a weakly supervised information bottleneck with K learnable reason tokens, dual-head supervision, and two-stage training. It achieves competitive benchmark performance while providing 60× faster inference than explicit CoT methods and 2× faster inference than iterative latent baselines.

Problem
CoT-based embedding incurs high inference cost and depends strongly on the quality of textual CoT annotations, limiting practical low-latency and large-scale training.
Method
LaME performs latent reasoning through K learnable reason tokens as a fixed-capacity information bottleneck, using weak dual-head supervision and two-stage training.
Results
LaME achieves competitive performance on MMEB-v2 and MRMR, with 60× faster inference than explicit CoT methods and 2× faster inference than iterative latent baselines.
Takeaways & Limitations
Embedding-oriented reasoning can be performed in latent space within a single forward pass without relying on explicit CoT annotations.
Takeaways & Limitations
LaME reduces interpretability relative to CoT-based methods and its two-stage training adds engineering complexity; simplifying training remains open.

Abstract

from arXiv · show

Reasoning-driven universal multimodal embedding has advanced rapidly by introducing Chain-of-Thought (CoT) reasoning into the embedding pipeline. Despite the strong performance across both general and complex tasks, this paradigm suffers from two core limitations: (i) autoregressive CoT reasoning incurs high computational cost, making it impractical for low-latency retrieval; and (ii) embedding performance is heavily coupled with CoT annotation quality, making large-scale training unreliable. These raise fundamental questions: Is textual CoT the optimal form of reasoning for embedding, and can effective embedding reasoning be accomplished in latent space? To this end, we propose LaME (Latent Reasoning Multimodal Embedding), which formulates embedding-oriented latent reasoning as a weakly supervised information bottleneck. LaME employs K learnable reason tokens as a fixed-capacity bottleneck, completing all reasoning within a single forward pass. The two weak supervision signals structurally decouple contrastive from autoregressive objectives and eliminate dependence on CoT annotations, while a two-stage training pipeline ensures stable convergence. Experiments on MMEB-v2 and MRMR show that LaME achieves competitive performance, surpassing some explicit CoT-based models, while delivering 60x faster inference than explicit CoT methods and 2x faster than latent baselines with throughput comparable to discriminative embedding models. Code is available at https://github.com/PeppaWu/LaME.

1 Introduction

LaME addresses the cost and annotation dependence of CoT-based multimodal embedding by performing weakly supervised latent reasoning through a fixed-capacity bottleneck in one forward pass. Its dual-head supervision and two-stage training support competitive retrieval performance with substantially faster inference.

1 Introduction: CoT-based embedding is limited by high inference cost and strong dependence on potentially defective textual reasoning annotations.These limitations motivate investigating whether textual CoT is necessary for embedding-oriented reasoning.
1 Introduction: A two-stage pipeline first stabilizes bottleneck initialization and then jointly optimizes the model for embedding-oriented latent reasoning.Warm-up trains reasoning tokens and supervision heads before broader parameter updates, reducing instability during convergence.
1 Introduction: LaME formulates embedding-oriented latent reasoning as a weakly supervised information bottleneck using K learnable reason tokens.The fixed-capacity bottleneck confines latent reasoning without requiring explicit CoT formats.
1 Introduction: Dual-head bottleneck supervision separates contrastive and generative objectives, while avoiding dependence on intermediate CoT annotations.The supervision signals probe latent reasoning and support embedding learning through distinct objectives.
1 Introduction: 60× faster inference than explicit CoT methods and 2× faster than iterative latent rollout baselines while achieving competitive performance on MMEB-v2 and MRMR.The method confines reasoning to prefill tokens within a single forward pass.

2 Related Works

Multimodal embedding research spans contrastive MLLM fine-tuning, explicit CoT generation, and latent hidden-state reasoning. These approaches trade representational capability against inference cost, annotation dependence, or limited reasoning flexibility.

2 Related Works: MLLM-based embedding methods extend contrastive dual-encoder learning by using vision-language fusion and instruction-following backbones.Representative systems include VLM2Vec, E5-V, MM-Embed, GME, UniME, UniME-V2, and RzenEmbed.
2 Related Works: Explicit CoT methods generate textual reasoning before embedding, improving reasoning capacity but coupling retrieval learning to autoregressive decoding and annotations.This coupling creates substantial inference overhead and dependence on retrieval-friendly CoT supervision.
2 Related Works: Latent reasoning methods replace some textual thinking with hidden-state iteration or additional computation slots, but iterative rollouts retain inference overhead.Prior work includes Coconut, Pause tokens, Implicit CoT, CODI, CoLaR, and PLUME.
2 Related Works: Information bottleneck methods compress input information while preserving task-relevant information, motivating finite-capacity latent representations for multimodal embedding.The IB objective is expressed as minimizing I(Z, X) −β I(Z, Y ), with related bottleneck ideas used in representation learning.

3 Our Framework

LaME confines multimodal latent reasoning to K learnable reason tokens and supervises the bottleneck with retrieval-oriented and target-reconstruction signals. A two-stage pipeline stabilizes training while separating generative probing from final embedding optimization.

3.2 Reason Tokens as Information Bottleneck: LaME confines latent reasoning to K learnable reason tokens within a single forward pass, using their finite capacity as an information bottleneck.The bottleneck is intended to retain task-relevant information while discarding noise.
3.3 Dual-Head Supervision: The Decoder Head reconstructs retrieval targets, answers, or keywords from the first K_r reason tokens without supervising intermediate CoT steps.This target-only weak signal decouples training from textual CoT annotation quality.
3.3 Dual-Head Supervision: The Embedding Head mean-pools the remaining K_e reason-token states and trains the resulting reason embedding with InfoNCE.Its non-autoregressive path structurally separates contrastive optimization from generative decoding.
3.3 Dual-Head Supervision: The joint objective combines contrastive, decoder, embedding, and diversity losses, with the diversity regularizer penalizing pairwise similarity among reason tokens.A separate learnable embed token is projected into the final retrieval embedding space.
3.4 Two-Stage Training: Two-stage training first warms up reason tokens and supervision heads with the backbone frozen, then jointly optimizes the model while adding contrastive supervision through an embed token.Warm-up gives randomly initialized reason tokens semantic structure before backbone updates, preventing degenerate solutions and accelerating convergence.
3.4 Two-Stage Training: LaME is evaluated on MMEB-V1 and the full MMEB-V2 benchmark, whose tables average performance across 36 and 78 subtasks respectively.MMEB-V2 covers image, video, and visual-document modalities and uses multiple retrieval metrics by scenario.

4 Experiments

LaME is evaluated on multimodal retrieval benchmarks spanning image, video, visual-document, and reasoning-intensive tasks, using compact latent bottleneck configurations. The experiments compare its performance with discriminative and reasoning-based baselines and examine core design choices.

Evaluation Setup: LaME is evaluated on MMEB-V2 image, video, visual-document, and MRMR reasoning-intensive retrieval benchmarks.MMEB-V2 uses Hit@1 for image and video tasks and NDCG@5 for visual-document tasks, while MRMR uses nDCG@10.
Ablation Setup: Table 4 compares LaME’s core bottleneck components against a discriminative baseline on MMEB-V2.The supplied passage identifies the ablation scope but does not provide the table’s cell values.
Ablation Setup: The experiments also vary the number of latent reason tokens and measure throughput on a single GPU.This tests how bottleneck capacity affects MMEB-V2, MRMR, and efficiency.

K Image Video VisDoc MRMR Throughput(samples/s)

LaME achieves strong retrieval results across multimodal and reasoning-intensive benchmarks, while its ablations show that two-stage training, dual-head supervision, and moderate bottleneck capacity are important. Qualitative examples illustrate single-pass latent reasoning for world knowledge and mathematical queries.

Main Results: 64.4 is LaME’s best overall score among 2B MMEB-V2 models, leading Ops-MM-Embed by 1.4 and PLUME by 2.8 points; its 7B score is 68.8, just 0.1 below Ops-MM-Embed.At 7B, LaME achieves the best image and visual-document averages, supporting generalization across modalities and model scales.
Main Results: 49.8 is LaME’s MRMR overall score at 7B, competitive with RIME at 50.2 and above Ops-MM-Embed at 48.1.LaME ranks first on Science 73.8, Math 29.5, Physics 44.4, and Engineering 36.4.
Ablation Studies: 64.4 is the full core-component configuration’s MMEB-V2 score, versus 63.8 for the discriminative baseline and 63.3 for naive joint attachment without two-stage training.Using only the decoder or embedding head with two-stage training reaches 64.0 and 64.1, respectively.
Ablation Studies: Increasing K from 0 to 8 improves Image from 68.5 to 69.3, Video from 43.9 to 44.5, VisDoc from 71.3 to 72.1, and MRMR from 48.0 to 49.8.Increasing K to 16 provides only marginal MRMR gains while degrading standard retrieval, indicating that excessive capacity weakens the bottleneck constraint.
Ablation Studies: Mean pooling reaches 64.4 versus 63.8 for attention pooling, while the Qwen3-0.6B decoder reaches 64.4 and the larger Qwen3-1.7B decoder drops to 64.1.These results support lightweight decoder supervision and aggregated embedding-head compression.
Qualitative Analysis: Qualitative examples show latent tokens supporting single-pass world-knowledge reconstruction and mathematical reasoning, including a pyramid slant-height answer of 4.3.The examples associate the latent process with racing-related knowledge and internalized Pythagorean reasoning.

5 Conclusion

LaME frames embedding-oriented latent reasoning as an information bottleneck implemented with learnable reason tokens and guided by decode and contrastive losses. The conclusion reports strong MMEB-v2 performance and inference speedups while identifying adaptive capacity and broader retrieval tasks as future directions.

Conclusion: LaME performs single-pass latent reasoning with learnable reason tokens and requires no supervision for the structure or content of its latent thinking process.The bottleneck is guided only by decode and contrastive losses.
Conclusion: LaME achieves state-of-the-art MMEB-v2 performance with significant inference speedups over explicit CoT and iterative methods.The conclusion summarizes the method’s reported benchmark and efficiency advantages without supplying additional numerical values.
Conclusion: Future work will explore adaptive bottleneck capacity and extension to broader and more complex retrieval tasks.

Limitations

The paper identifies reduced interpretability without explicit CoT traces and extra engineering complexity from two-stage training as limitations of LaME.

Limitations: Without explicit CoT traces, LaME’s latent reasoning is less interpretable than CoT-based methods.The paper frames this as a reduction in interpretability rather than a failure of retrieval performance.
Limitations: The two-stage training pipeline adds engineering complexity compared with single-stage approaches, motivating work to simplify training without sacrificing performance.

Ethics Statement

The work concerns multimodal retrieval and representation learning without human subjects, clinical data, or direct user interaction. It cautions that retrieval systems may inherit dataset biases and require curation and safety filtering.

The study involves multimodal retrieval and representation learning, not human subject recruitment, clinical data, or direct user interaction.
Source-dataset biases, noise, and content imbalance may carry into retrieval outputs, motivating dataset curation and application-specific safety filtering.The rewrite-based design is intended to reduce redundant reasoning and semantic distortion but can still produce imperfect or biased outputs.
Downstream deployment should account for potentially imperfect or biased outputs from the retrieval system.

A.2 Training Data Composition

Training uses approximately 1.55 million pairs from the VLM2VEC-v2 corpus across image, video, and visual-document modalities. The composition spans matching, question answering, and retrieval scenarios.

Approximately 1.55 million training pairs are drawn from the VLM2VEC-v2 training corpus across images, videos, and visual documents.
Image data includes MSCOCO, ImageNet-1K, ChartQA, DocVQA, CIRR, and other datasets after filtering.The passage reports approximately 677K image samples.
The training composition covers image-text matching, visual question answering, and retrieval scenarios.

B.1 Ablation on Decoding Strategy

Decoding ablations show that concise answer supervision is preferable to full CoT decoding, while compressed reason embeddings remain close to final retrieval embeddings. Diversity regularization also prevents latent-token collapse.

B.1 Ablation on Decoding Strategy: 69.3 / 44.5 / 72.1: Only-Answer achieves the best results, surpassing CoT decoding and No Decoder at 63.8 average.
B.1 Ablation on Decoding Strategy: Long CoT supervision lets easy-to-predict tokens dominate gradients, whereas concise answer tokens preserve discriminative bottleneck information.No Decoder’s degradation indicates that the decoder still provides essential training-time reasoning supervision.
B.2 Embedding Comparison: 64.1 overall: the reason embedding trails the final retrieval embedding by only 0.4 points, indicating that reason tokens encode retrieval-relevant information.
B.2 Embedding Comparison: The reason-versus-retrieval embedding gap is 0.9 on VisDoc, compared with 0.2 on Image and 0.4 on Video.The passage attributes this pattern to fine-grained document understanding benefiting from the complete hidden-state sequence.
B.3 Ablation on the Diversity Regularizer: 64.4: increasing λDiv from 0 to 0.05 improves performance from the lowest 63.4 result, preventing degenerate reason-token representations.

C Theoretical Analysis of Information Bottleneck

LaME treats latent reasoning as an information bottleneck that compresses multimodal inputs while preserving decoder and embedding supervision. Its fixed token capacity and two-stage optimization jointly constrain and train the latent representation.

LaME defines X as interleaved text and visual tokens, Z as reason-token hidden states, and Y as supervision from decoder and embedding heads.
The structural bottleneck supports Z on at most K tokens of dimension d, imposing I(Z; X) ≤ K·d·log(1+SNR).Unlike variational IB, token count K itself serves as the bottleneck radius without an auxiliary prior or Monte Carlo estimation.
Decoder and embedding heads provide complementary predictive channels over separate reason-token subsets, jointly trading reconstructive against discriminative information.
Two-stage training first stabilizes reason-token representations with a frozen backbone, then enables backbone co-adaptation while preserving the [EMBED] pathway.