PaperlayerView exact revision on arXiv

Source-linked AI summary

TTE-Flash: Accelerating Reasoning-based Multimodal Representations via Think-Then-Embed Tokens

Jianpeng Cheng, Xian Wu, Jiangfan Zhang, Wentao Bao, Chaitanya Ahuja, Shlok Kumar Mishra, Hanchao Yu, Yang Gao, Fan Xia, Qi Guo, Shaodan Zhai, Xiangjun Fan, Jun Xiao

arXiv:2605.16638v1cs.AI

TL;DR

Explicit Chain-of-Thought improves multimodal embeddings but is computationally costly. TTE-Flash replaces it with latent think tokens and achieves stronger MMEB-v2 performance than explicit-CoT models while being 70x more efficient.

  • Problem

    Explicit Chain-of-Thought improves multimodal representations but its real-time computational cost is prohibitive, motivating efficient latent reasoning.

  • Method

    TTE-Flash uses decoupled latent think and embedding tokens in a register-based architecture trained for reasoning and representation learning in one pre-filling pass.

  • Results

    TTE-Flash-2B outperforms explicit-CoT models on MMEB-v2 while being 70x more efficient, with positive think-token scaling across 15 zero-shot video subsets.

  • Takeaways & Limitations

    Latent think tokens provide an efficient reasoning-aware alternative to explicit-CoT multimodal embeddings, with task-dependent think-token scaling supporting adaptive budget allocation.

Abstract

from arXiv · show

Recent research has demonstrated that Universal Multimodal Embedding (UME) benefits significantly from Chain-of-Thought (CoT) reasoning. In this paradigm, a generative model produces explicit reasoning traces for a multimodal query, with the final representation extracted from an <eos> embedding token attending to both the query and the reasoning. Despite its effectiveness, the computational overhead of generating explicit CoT traces is often prohibitive. In this work, we propose replacing explicit CoT with latent think tokens, which are interpreted as latent variables that can produce explicit CoT traces as observed variables. By optimizing think tokens using CoT generation loss and subsequent embedding tokens using contrastive loss, we produce high-performance, reasoning-aware representations at a constant inference cost. Our study investigates two key architectural designs: 1) how think and embeddings tokens should be extracted from the same LLM backbone. 2) how the tokens should be trained as two dependent tasks. We introduce TTE-Flash-2B, a reasoning-aware multimodal representation model that outperforms its explicit-CoT counterpart on the MMEB-v2 benchmark, while producing latent think tokens that are interpretable both textually and visually. Furthermore, zero-shot evaluation across 15 video datasets reveals scaling behavior as the number of think tokens increases, and motivating a pilot study of adaptive think budget allocation based on task requirements.

1 Introduction

TTE-Flash replaces costly explicit Chain-of-Thought traces in reasoning-enhanced Universal Multimodal Embedding with latent think tokens followed by embed tokens for constant-time inference. TTE-Flash-2B outperforms explicit-CoT reasoning-based UME baselines on MMEB-v2 while being 70x more efficient.

  • Architecture: The method uses a unified LLM backbone to generate think tokens before embed tokens, replacing explicit CoT generation with latent intermediate representations.This design targets the prohibitive real-time computational cost of generating explicit CoT traces.
  • Architecture: The architecture study compares looped latent reasoning, which is autoregressive and memory-bound, with register-based tokens processed in one pre-filling pass.The comparison addresses how to extract think and embed representations while balancing concurrency and representational fidelity.
  • Training: Think tokens are trained as latent variables that generate explicit CoT traces, while subsequent embed tokens learn multimodal representations through contrastive retrieval training.This information-bottleneck formulation produces think tokens aligned with multimodal instructions and ground-truth CoT traces.
  • Training: Overlapping think and embed tokens degrades both reasoning and representation performance, motivating separate but dependent objectives within the shared multimodal backbone.The preliminary finding supports treating reasoning and representation learning as related but distinct tasks.
  • Results: TTE-Flash-2B outperforms explicit-CoT reasoning-based UME baselines on MMEB-v2 while being 70x more efficient.The model is designed for constant-time inference using Think-Then-Embed tokens.
  • Interpretability: The latent think tokens support textual interpretation through decoded CoT traces and visual interpretation through an attached image-generation head.These interpretations are reported alongside zero-shot evaluation across 15 video datasets.

2 Related Work

UME uses a unified MLLM backbone and LLM-centric fusion to produce reasoning-aware shared representations, unlike CLIP-style dual towers. Related methods add explicit or latent reasoning, while joint contrastive-generative objectives unify understanding and generation.

  • Universal Multimodal Embedding: UME maps multimodal inputs into a shared space through a unified MLLM backbone and LLM-centric fusion, unlike CLIP-style dual-tower architectures [Radford et al., 2021; Zhai et al., 2023].LLM-centric fusion supports rich, reasoning-aware representations.
  • Universal Multimodal Embedding: Reasoning-enhanced UME frameworks TTE (Cui et al., 2025), UME-R1 (Lan et al., 2025), and MMEmb-R1 (Wang et al., 2026) integrate explicit CoT before representation learning.
  • Universal Multimodal Embedding: PLUME (He et al., 2026) introduces latent reasoning in UME to achieve a superior accuracy-efficiency trade-off on MMEB-v2 (Meng et al., 2025).
  • Latent Reasoning: Latent reasoning treats CoT as hidden-state computation by generating intermediate thought vectors instead of decoding explicit reasoning tokens.Coconut (Hao et al., 2024) recursively feeds back the last hidden state, CoLaR (Tan et al., 2025) compresses multiple CoT tokens, and LaDiR (Kang et al., 2025) applies diffusion.
  • Joint Contrastive-Generative Models: Foundation models commonly combine contrastive and generative objectives for unified understanding and generation, as illustrated by CoCA (Yu et al., 2022) and BLIP (Li et al., 2022).CoCA combines contrastive and captioning losses, while BLIP adds image-text matching to contrastive and captioning objectives.

3 TTE-Flash

TTE-Flash uses a unified causal LLM to encode multimodal inputs, generate latent think tokens, and extract embeddings, with register tokens enabling efficient parallel computation. Think tokens are trained through compressed CoT generation, while embeddings use contrastive learning with pairwise similarity scoring.

  • Architecture: A unified causal LLM lets think tokens attend to multimodal inputs and embedding tokens attend to both inputs and thoughts, mirroring explicit TTE.The architecture treats multimodal content, thinking, and representation learning as token sequences.
  • Architecture: Register tokens extract all think latents and embeddings in one pre-filling stage, avoiding the N sequential decoding steps required by looped architectures.Looped designs are memory-bounded, whereas registers compute and load the KV cache once and are therefore compute-bounded.
  • Think Loss: Think registers are optimized solely with CoT generation loss, forcing N latent vectors to compress multimodal information sufficient to generate an L-token reasoning trace.The bottleneck uses N = 1–32 vectors to represent information otherwise requiring approximately L ≈300 discrete tokens.
  • Embed Loss: Embed registers are trained with standard contrastive loss on query and target embeddings conditioned on the original input and think latents.The scoring function sums pairwise vector similarities; with N = 1, it reduces to standard single-vector retrieval.
  • Embed Loss: Pairwise similarity preserves positional correspondence between dependent causal register pairs, unlike “sum of maximum” scoring used in prior work (Santhanam et al., 2022; Faysse et al., 2024).The fine-tuned backbone performs one pre-filling pass, while the pretrained backbone decodes think latents for visualization or interpretation.

4 Experiments

Experiments show that TTE-Flash benefits from loop-based and decoupled architectures, pairwise similarity, and more think tokens, culminating in MMEB-V2 gains over explicit-CoT and larger baselines. Additional studies examine adaptive budgets and whether think tokens encode textual and visual semantics.

  • Architecture comparison: Loop-based approaches consistently outperform register-based approaches on MMEB-V1, while per-layer registers reduce the performance gap under restricted computational constraints.The comparison averages 20 checkpoints with a total token budget of 8; latency and throughput are analyzed separately in Table 1.
  • Similarity function: Sum of pairwise similarities outperforms sum of maximum similarities because position-aware embeddings are causally dependent.The study compares maximum query-target matching with same-position pairwise similarity aggregation.
  • Think/embed decoupling: Decoupling think and embedding tokens outperforms sharing them on MMEB-V1 retrieval and CoT-generation similarity.The ablation uses 8 tokens and evaluates retrieval alongside cosine similarity between decoded CoT and ground truth.
  • Think-token scaling: Performance generally increases with the number of think tokens, with positive scaling also observed on 15 zero-shot video datasets.The think-token ablation fixes the number of embedding tokens to 1, and the video evaluation applies an image-trained checkpoint to MMEB-V2 tasks.
  • Final benchmark results: TTE-Flash 2B with 32 think tokens and 1 embedding token outperforms explicit-CoT TTE-V1 2B variants and the 7B VLM2Vec baseline on MMEB-V2.The final configuration uses per-layer registers, shared LoRA adapters, and task-decoupled registers; Table 2 summarizes the benchmark comparison.
  • Adaptive thinking budget: Adaptive think with a maximum of 32 tokens scores 66, 49.4, and 65.5 on image, video, and visdoc, below fixed-32 scores of 68.3, 50.6, and 68.1.The adaptive mechanism predicts a first-K token mask using input-conditioned budget logits and Gumbel-Softmax; it allocates more tokens to complex VQA tasks, including GQA at 13.7 and OK-VQA at 13.0 tokens.

5 Conclusion … A.2 Budget Regularization

TTE-Flash uses latent think tokens in a single causal pre-filling pass to produce efficient reasoning-aware embeddings, outperforming explicit-CoT TTE-2B variants and VLM2vec 7B on MMEB-V2. An adaptive mechanism allocates shorter or longer reasoning prefixes per input difficulty through budget prediction, differentiable masking, and regularization.

  • 5 Conclusion: TTE-Flash-2B outperforms explicit-CoT TTE-2B variants and the larger VLM2vec 7B baseline on MMEB-V2 by scaling the number of latent think tokens.The model uses register tokens to generate think and embedding representations within one causal pre-filling pass.
  • 5 Conclusion: Decoded CoT traces are compared across think-token budgets of 2, 4, 8, and 16.The comparison examines how latent reasoning outputs vary as the think budget increases.
  • A Adaptive Think Budget Allocation: The adaptive think-budget mechanism routes easy inputs to shorter reasoning prefixes and harder inputs to more think tokens.Short classification queries are examples of easy inputs, while compositional VQA is an example of a harder input.
  • A.1 Budget Predictor: A lightweight predictor maps pooled input representations to budget logits over N candidate positions and samples a one-hot budget with straight-through Gumbel-Softmax.The estimator uses temperature τ and enables differentiable budget selection.
  • A.1 Budget Predictor: The one-hot budget is converted into a contiguous prefix mask whose active think tokens precede the embedding position, with gradients flowing through the mask.The mask is applied multiplicatively to the input embeddings, enabling end-to-end training.
  • A.2 Budget Regularization: An ℓ1 penalty on the average number of active think tokens discourages trivial all-on budget solutions, using λb = 0.01.The total objective combines contrastive, generation, and budget-regularization losses.
  • A.2 Budget Regularization: Contrastive and generation losses activate think tokens when useful, while regularization turns them off when unnecessary, adapting budgets to per-sample difficulty.This produces an equilibrium between reasoning utility and token economy.

B Experiments Details · B.1 Datasets Details

MMEBv2 is a broad multimodal embedding benchmark spanning 9 meta-tasks and 78 tasks across text, image, video, and document modalities. Its unified retrieval-based protocol evaluates models consistently across diverse multimodal understanding capabilities using the benchmark’s prescribed training and evaluation split.

  • B.1 Datasets Details: MMEBv2 comprises 9 meta-tasks and 78 individual tasks spanning text, image, video, and document modalities.
  • B.1 Datasets Details: Image evaluation covers 10 classification tasks and 10 visual question answering tasks, including ImageNet-1K, HatefulMemes, OK-VQA, DocVQA, and GQA.
  • B.1 Datasets Details: Image-level retrieval spans 12 tasks, while visual grounding spans 4 tasks including MSCOCO, RefCOCO, and Visual7W-Pointing.
  • B.1 Datasets Details: Video evaluation includes 5 video-classification tasks and 5 video-question-answering tasks, covering datasets such as Kinetics-700, Video-MME, MVBench, and EgoSchema.
  • B.1 Datasets Details: Video-level retrieval comprises 5 tasks, and moment retrieval comprises 3 tasks: QVHighlights, Charades-STA, and MomentSeeker.
  • B.1 Datasets Details: All tasks use a retrieval-based framework in which multimodal queries retrieve correct targets from candidate pools, following the MMEBv2 (Meng et al., 2025) training/evaluation split.

B.2 Training and Evaluation Details

The model is trained from Qwen3-VL-2B-Instruct on MMEB-V2, with task performance measured by Hit@1 and the overall score computed as a hierarchical weighted average.

  • Training Setup: Training uses Qwen3-VL-2B-Instruct on the MMEB-V2 training set, with hyperparameters specified in Table 3.Table 3 presents the training configuration.
  • Evaluation: Individual tasks are evaluated using Hit@1, while the overall benchmark score is a hierarchical weighted average.The overall aggregation follows Meng et al. (2025).

C Visual Decoding: Methodology … C.3 DiTDH-XL: Diffusion Transformer

The visual-decoding method uses a frozen TTE-Flash backbone to condition latent diffusion, translating thinking-token representations into images through a semantic autoencoder and DiTDH-XL transformer. The pipeline preserves spatial conditioning through Perceiver resampling and per-token denoising operations.

  • C Visual Decoding: Methodology: The decoder trains a latent diffusion model conditioned on frozen TTE-Flash representations to generate images from free-form vision-language inputs.This methodology follows the premise that semantically rich representation encoders co-occur with strong generative capabilities.
  • C.1 Overview: The two-stage pipeline encodes target images into DINOv2 semantic latents with a frozen encoder, trains a ViT-XL decoder, and decodes generated latents into images.The RAE uses latents in R768×16×16, while only the decoder is trained; DiTDH-XL generates these latents from thinking-token representations.
  • C.2 Conditioning via TTE-Flash Representations: For visualization, frozen TTE-Flash produces 8 thinking tokens with 2048-dimensional embeddings, which condition the diffusion model.A Perceiver resampler uses 256 learned 768-dimensional queries and 8-head cross-attention to project these features into a fixed spatial conditioning signal.
  • C.2 Conditioning via TTE-Flash Representations: The Perceiver output is linearly projected from 2048 to 768 dimensions and reshaped into c ∈R768×16×16, matching DiT spatial latent dimensions.The conditioning is computed as LayerNorm(Q + CrossAttn(Q, Linear(Hthink), Linear(Hthink))).
  • C.3 DiTDH-XL: Diffusion Transformer: DiTDH-XL combines a 28-layer, 1152-hidden-size encoder with a 2-layer, 2048-hidden-size DDT decoder head for spatially-aware denoising.The decoder receives encoder outputs as per-token conditioning and applies thinking-token-conditioned scale, shift, and gate operations to re-embedded noisy inputs.
  • C.3 DiTDH-XL: Diffusion Transformer: Training uses flow matching with linear interpolation between Gaussian noise and DINOv2 latents, predicting velocity with MSE under logit-normal timestep sampling.The path is xt = (1−t) x1 +t x0, with target velocity ut = x0 −x1.
  • C.3 DiTDH-XL: Diffusion Transformer: Images are sampled with 50-step Euler ODE integration and classifier-free guidance at scale 4.0, enabled by 10% caption dropout during training.The diffusion transformer is trained for 200 epochs on 243K MMEB samples with effective batch size 1024 and AdamW optimization.

C.4 Qualitative Results · D More CoT Examples Generated from Think Tokens

Appendix E presents representative successful and failure cases across three MMEB tasks, showing that TTE-Flash thinking tokens encode rich visual semantics while exposing limitations in fine-grained spatial reasoning and complex compositional instruction faithfulness.

  • C.4 Qualitative Results: Appendix E covers text-to-image retrieval on MSCOCO and VisualNews and composed retrieval on CIRR, using captions or image–instruction pairs to generate grounded or compositional images.MSCOCO and VisualNews provide only text captions, whereas CIRR combines a reference image with a textual modification instruction.
  • C.4 Qualitative Results: Successful cases show rich visual semantics in TTE-Flash thinking tokens, while failures expose weaknesses in fine-grained spatial reasoning and faithfulness to complex compositional instructions.The qualitative examples indicate that the tokens represent more than retrieval alone requires.

Dataset: IMAGENET–1K

IMAGENET–1K examples pair image-classification representation prompts with retrieval targets and generated reasoning traces. The traces identify objects using visual attributes, labels, context, and functional features across varying CoT lengths.

  • Classification retrieval prompts: The dataset frames image representation as classification retrieval for targets including “African grey, African gray, Psittacus erithacus,” “beer bottle,” “drum, membranophone, tympan,” and “fire engine, fire truck.”These prompts explicitly pair each image with its intended retrieval label.
  • CoT-16 reasoning traces: CoT-16 traces classify images by combining visual appearance with contextual or functional evidence, such as plumage, labels, object form, and firefighting equipment.Examples identify an African grey parrot, beer bottle, drum, and fire engine from these cues.
  • Shorter reasoning traces: Shorter traces also use distinctive visual or textual cues, identifying beer through a “Bier” label and fire engines through red coloring, ladders, and “FIRE DEPT” text.The supplied examples include CoT-8 and CoT-4 generations for these classifications.

Dataset: MSCOCO … Dataset: VIDORE_SYNTHETICDOCQA_AI

The MSCOCO examples span object cropping, image-to-text caption retrieval, visual question answering, and multimodal matching prompts. Their generated CoT traces describe salient objects, actions, environments, and contextual cues at varying reasoning lengths.

  • Dataset: MSCOCO: MSCOCO includes object-isolation prompts for broccoli, pizza, and sandwich images.These prompts explicitly request cropping the image to isolate the labeled object.
  • Dataset: MSCOCO: MSCOCO also contains multimodal matching traces that focus on object appearance, ingredients, actions, and surrounding context.Pizza, broccoli, and sandwich examples emphasize features such as color, toppings, food preparation, and nearby people or surfaces.
  • Dataset: MSCOCO: Reasoning examples vary from CoT-2 through CoT-16 and progressively describe different levels of visual detail.Shorter traces identify central objects or activities, while longer traces provide richer environmental and attribute descriptions.
  • Dataset: MSCOCO: Caption-retrieval examples target scenes such as zebras in a field, people in a courtyard, hot dogs, girls eating, motorcycles, bathrooms, and skiing.The examples ask for captions describing everyday images and provide retrieval targets ranging from animals and people to places and activities.
  • Dataset: MSCOCO: CoT-16 traces identify visual subjects and context before producing captions or matching descriptions.Examples analyze zebras and grasslands, sandwiches on griddles, people near a school, hot dogs, girls eating, and motorcycles in streets.
  • Dataset: MSCOCO: Visual question-answering prompts ask for semantic attributes including sea name, building type, vehicle brand, season, weather, and ship type.Targets include pacific ocean, warehouse, lincoln, fall, sunny, and cruise.
  • Dataset: MSCOCO: Generated traces use scene cues to infer answers, including surfing conditions, warehouse structure, autumn foliage, vehicle logos, sunny weather, and vessel design.The reasoning connects visible objects and environments to the requested semantic labels.

Dataset: VISRAG_INFOVQA … Dataset: VisualNews_t2i – Failure Cases

The supplementary examples illustrate multimodal reasoning traces, think-token-conditioned image generation, and successful or failed image retrieval across document, caption, and news-caption datasets. They span infographic interpretation, everyday-scene matching, news-image retrieval, and a documented failure case.

  • Dataset: VISRAG_INFOVQA: VISRAG_INFOVQA examples show CoT generations interpreting document titles, safety measures, infection-prevention actions, and infographic sections with causes of death.The examples use CoT-16, CoT-4, and CoT-2 generations to identify textual and visual information in document images.
  • Dataset: VISRAG_INFOVQA: A VISRAG_INFOVQA example identifies that face masks reduce infection risk by 70% from an infographic describing intervention impacts.The generation locates the “Infection” entry under “Face masks” and reports the stated percentage.
  • E Images Generated from Think Tokens: Visual decoding examples pair text or image-modification queries with ground-truth targets and DiTDH-XL images conditioned on TTE-Flash thinking tokens.These examples demonstrate visual outputs generated from the model’s latent thinking tokens.
  • Dataset: MSCOCO_t2i – Successful Cases: MSCOCO_t2i successful cases cover caption-to-image matching for motorcycles, birthday candles, kitchens, bathrooms, and bicyclists.The prompts include varied everyday scenes such as riders near trains, kitchen utensils, bathroom fixtures, and groups cycling in streets.
  • Dataset: MSCOCO_t2i – Successful Cases: Additional MSCOCO_t2i prompts express paraphrased descriptions of the same everyday scenes, including candles, kitchens, bathrooms, and bicycles.These cases vary wording while preserving the underlying scene-matching task.
  • Dataset: VisualNews_t2i – Successful Cases: VisualNews_t2i successful cases retrieve images for news captions about police, athletes, musicians, fashion, aviation, universities, politics, and conflict.The prompts include both short named-entity captions and longer event descriptions.
  • Dataset: VisualNews_t2i – Failure Cases: The VisualNews_t2i failure-case section contains a retrieval prompt about Huma Abedin photographed at the White House, labeled “GT Generated.”This example explicitly marks the generated result as a failure case.

Dataset: CIRR – Successful Cases

The successful CIRR cases demonstrate retrieval under varied image-editing instructions, including object counts, removals, additions, composition, lighting, color, and viewpoint changes. Examples also cover fine-grained subject focus and scene-content modifications.

  • Object and quantity changes: Successful cases retrieve images matching object-count and quantity changes, such as three soft-drink bottles, fewer paper towels, or an added deer.These examples show both categorical additions and count-based modifications.
  • Subjects and actions: Other cases specify subject relationships and actions, including same-breed dogs with a puppy, two same-breed dogs, barking dogs, and deer looking toward the camera.The examples require retrieving images that satisfy both identity-related and behavioral descriptions.
  • Composition and viewpoint: Successful retrievals also follow composition and viewpoint changes, including visible surfaces, white-background presentation, rear-facing deer views, and focused views of dogs’ heads.These prompts alter framing, viewpoint, background, or subject emphasis while retaining the scene’s main content.
  • Object and scene edits: The cases include object removal and insertion instructions, such as removing a lemon or seashells, adding a wooden log, and changing water color.The requested edits combine object-level changes with scene-property changes.
  • Environment and atmosphere: The successful examples further include environmental and atmospheric modifications such as brighter sunshine, added sunlight, green water, and indoor scenes with food.These prompts target illumination, color, setting, and scene context.

Dataset: CIRR – Failure Cases

The CIRR failure cases show generated retrieval prompts requiring controlled changes to animals, humans, environments, and composition while preserving similarity to the input image.

  • CIRR – Failure Cases: Generated cases request cross-species substitutions and human-presence constraints, including two different animal species without humans and a human with one different-species animal.These prompts test simultaneous identity and presence changes relative to the input or prior images.
  • CIRR – Failure Cases: Additional cases modify human appearance or animal type while retaining the intended scene relationship, including less gear and preserved environment.These prompts vary subject attributes and animal identity while specifying what should remain unchanged.
  • CIRR – Failure Cases: Other generated cases require preserving or altering specific visual attributes, such as showing a whole dog or animal standing on a dirt floor.The prompts isolate composition and environmental conditions involving the animal’s full body and ground surface.
Loading 2605.16638v1…