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

• A unified backbone generates a predefined number of latent think tokens before embed tokens, replacing costly explicit reasoning generation.Earlier Think-Then-Embed approaches improve representation quality through CoT reasoning but incur prohibitive generation costs (Cui et al., 2025; Lan et al., 2025).Source: §1 · p. 1
• 8-token registers double looped-method throughput on MMEB but incur 3% relative retrieval-accuracy degradation; introducing registers at every transformer layer narrows the gap.Looped architectures recursively feed outputs back as inputs, making them autoregressive and memory-bound. Registers use a single pre-filling pass but lack context-dependence; per-layer registers retain forward efficiency.Source: §1 · p. 2
• Think tokens are trained as latent variables and an information bottleneck for generating explicit CoT, addressing the challenge of supervising intermediate hidden states.Out-of-domain MMEB evaluations show that decoded CoT traces have high fidelity and align with multimodal instructions and ground truth.Source: §1 · p. 2
• Overlapping think and embed tokens degrades both tasks, motivating specialized parameter sharing and decoupling that respects the sequential Think-Then-Embed dependency.Embed tokens learn multimodal representations through multi-vector retrieval with contrastive learning (Santhanam et al., 2022; Faysse et al., 2024).Source: §1 · p. 2
• 70x more efficient than explicit-CoT reasoning-based UME baselines, TTE-Flash-2B also outperforms them on MMEB-v2.The model uses Think-Then-Embed tokens for constant-time inference.Source: §1 · p. 2
• 15 video datasets reveal task-specific zero-shot scaling with think-token count, motivating a pilot study of adaptive budgets based on task difficulty.Think tokens also support textual interpretation through CoT decoding and visual interpretation through an attached image generation head.Source: §1 · p. 2

Universal Multimodal Embedding

• UME maps multimodal inputs into a shared space through a unified MLLM backbone (Lin et al., 2024; Gu et al., 2026, 2025).Unlike CLIP-style dual towers with disjoint encoders and late alignment (Radford et al., 2021; Zhai et al., 2023), LLM-centric fusion supports reasoning-aware representations (Jiang et al., 2024a; Zhang et al., 2024; Lin et al., 2024).Source: §2 · p. 3
• TTE (Cui et al., 2025), UME-R1 (Lan et al., 2025), and MMEmb-R1 (Wang et al., 2026) integrate explicit CoT before UME representation learning.Concurrent work PLUME (He et al., 2026) introduces latent reasoning in UME for an improved accuracy-efficiency trade-off on MMEB-v2 (Meng et al., 2025).Source: §2 · p. 3

Latent Reasoning

• Latent reasoning replaces explicit reasoning tokens with intermediate hidden states; TTE-Flash treats thought vectors as a bottleneck generating explicit CoT within the same pretrained LLM.Existing training strategies include curriculum learning (Hao et al., 2024; Deng et al., 2024), explicit-CoT distillation (Deng et al., 2023; Shen et al., 2025), end-to-end supervision, and reinforcement learning (Yue et al., 2025).Source: §2 · p. 3

Joint Contrastive-Generative Models

• TTE-Flash combines contrastive and generative objectives within a unified LLM, distinguishing its Think-Then-Embed task through dedicated token heads for reasoning and representation learning.CoCA (Yu et al., 2022) combines contrastive and captioning losses, while BLIP (Li et al., 2022) additionally uses image-text matching. InternVL (Chen et al., 2024) performs vision-language contrastive and generative training within a shared LLM backbone.Source: §2 · p. 3

3.1 Backbone

• TTE-Flash appends learnable think and embed registers to a unified causal LLM, extracting all representations in one pre-filling stage.Registers are shared across inputs. Think tokens attend to multimodal input, while embed tokens attend to both input and think tokens.Source: §3.1 · p. 3
• N decoding steps are theoretically avoided by registers, selected for efficiency despite weaker performance than loops in four-token and eight-token ablations.N counts think plus embed tokens: memory-bound loops reload the prompt KV cache N times, whereas compute-bound registers compute and load all KV once. Adding registers at every transformer layer closes the performance gap by increasing expert parameters for both tasks.Source: §3.1 · pp. 3–4

3.2 Training

• 1 to 32 think vectors encode information otherwise requiring approximately 300 discrete tokens, with the fixed budget imposing an information bottleneck.The LoRA-equipped backbone and think registers produce deterministic latents shaped solely by CoT generation loss. Decoding uses the same backbone, frozen and without LoRA, preserving its pretrained generation capability.Source: §3.2 · pp. 4–5, Figure 2 · p. 4
• Embedding registers learn through contrastive loss, with multi-vector retrieval scored by summing corresponding pairwise similarities rather than taking maxima.Training uses in-batch negatives and temperature τ; N = 1 recovers single-vector retrieval. Pairwise scoring preserves positional correspondence because causal attention produces dependent vectors.Source: §3.2 · p. 5

3.3 Inference

• Inference extracts embeddings in a single pre-filling pass; interpreting think tokens instead passes their latents through the pretrained backbone for decoding.Embedding extraction uses the fine-tuned backbone with both think and embed registers.Source: §3.3 · p. 5

4.1 Ablations

• Loops outperform registers on MMEB-V1, but registers provide lower latency and higher throughput; adding per-layer registers narrows the performance gap.Increasing register count has negligible latency impact. Position-specific registers added to each transformer layer supply task-specific parameters, adding 55K parameters for one token and up to 1.7M for 32 tokens.Source: §4.1 · p. 5, Figure 4 · p. 6, Table 1 · p. 6
• Position-aligned pairwise similarity yields a higher mean and lower variance than sum-of-max, while decoupled think and embed tokens improve retrieval and CoT generation.Similarity functions achieve comparable overall performance across 20 runs with an eight-token budget; the authors hypothesize that positional alignment explains the pairwise advantage. The eight-token decoupling ablation evaluates retrieval and cosine similarity between decoded and ground-truth CoT.Source: §4.1 · p. 6, Figure 5 · p. 6, Figure 6 · p. 6
• 8 think tokens largely eliminate the incremental retrieval gains from additional embedding tokens, whereas increasing embedding tokens helps when no think tokens are present.With eight think tokens, multi-vector retrieval performs close to single-vector retrieval on MMEB-V1. The authors also confirm diminishing gains on longer-input ViDoRe retrieval tasks (Faysse et al., 2024).Source: §4.1 · p. 7, Figure 9 · p. 7, Figure 8 · p. 7

4.2 Full Results

• 70x more efficient than TTE-V1, TTE-flash 2B also outperforms explicit-CoT TTE-V1 2B variants, VLM2vec 7B, and latent-reasoning PLUME 2B on MMEB-V2.The final configuration uses 32 think tokens, one embedding token, per-layer registers, shared LoRA adapters, and task-decoupled registers. MMEB-V2 evaluation covers image, video, and visual-document tasks.Source: §4 · p. 5, Table 1 · p. 6, Table 2 · p. 7

4.1 Ablations

• 15 zero-shot video datasets show positive think-token scaling, with stronger upward trends on complex reasoning tasks and task-dependent optimal budgets.The video evaluation uses image-trained checkpoints. MMEB-V1 performance also generally improves as think-token count increases while the embedding-token count remains fixed at one.Source: §4.1 · pp. 6–7, Figure 7 · p. 7, Figure 10 · p. 8

4.3 Adaptive Think

• 50% overall token activation accompanies lower scores than fixed-budget inference, while adaptive budgets reflect reasoning complexity, retrieval direction, and temporal demands.With a maximum of 32 think tokens, adaptive image/video/visdoc scores are 66/49.4/65.5, versus 68.3/50.6/68.1 with fixed 32-token inference. An input-conditioned two-layer MLP selects a first-K token mask using hard Gumbel-Softmax sampling (Jang et al., 2016), with an additional loss penalizing K.Source: §4.3 · pp. 8–9

4.4 Visualization of Thinking Output

• Decoded reasoning improves with more think tokens, while a diffusion decoder reveals visual semantics through caption-grounded generation and image-plus-text composition.Human evaluation favors larger think budgets, and Claude 4.7-as-judge identifies CoT-16 as superior in coherence and relevance. A DiTDH-XL decoder (Peebles and Xie, 2023), conditioned on frozen TTE-Flash think embeddings, learns to reconstruct DINOv2 (Oquab et al., 2024) image latents.Source: §4.4 · p. 9

5 Conclusion

• TTE-flash-2B outperforms both explicit-CoT TTE-2B variants and the larger VLM2vec 7B baseline on MMEB-V2 as the think-token budget scales.The supplied conclusion reports this comparison without numerical benchmark scores.Source: §5 · p. 9
• 15 zero-shot video subsets confirm positive scaling patterns as the number of think tokens increases.Task-dependent thinking-depth requirements motivate the adaptive-budget pilot.Source: §5 · p. 9
• Latent think tokens map to explicit CoT traces, while register tokens produce all think and embedding representations in one causal pre-filling pass.The authors favor this register-token-based architecture.Source: §5 · p. 9

A Adaptive Think Budget Allocation

• Adaptive think-budget allocation learns input-specific token counts, aiming to assign shorter reasoning prefixes to easy inputs and more think tokens to harder inputs.Short classification queries and compositional VQA illustrate the intended easy-versus-hard routing.Source: A Adaptive Think Budget Allocation · p. 13

A.1 Budget Predictor

• A lightweight predictor maps pooled input embeddings to budget logits over N positions, then samples a one-hot budget using straight-through Gumbel-Softmax with temperature τ.The selected position encodes the chosen budget.Source: §A.1 · p. 13
• A reverse cumulative sum creates a contiguous prefix mask whose final active position supplies the embedding and preceding active positions supply think tokens.The mask multiplies input embeddings; Gumbel reparameterization allows gradients to flow through it for end-to-end training.Source: §A.1 · p. 13

A.2 Budget Regularization

• 0.01-weighted ℓ1 budget regularization discourages all-on allocations, balancing contrastive and generation losses to adapt token usage to sample difficulty.The penalty targets the average active think-token count; task losses reward useful tokens, while regularization pushes unnecessary tokens off.Source: §A.2 · p. 13

B.1 Datasets Details

• 78 individual tasks across 9 meta-tasks comprise MMEBv2 (Meng et al., 2025), spanning text, image, video, and document modalities.Evaluation embeds each query and retrieves the correct target from a candidate pool, enabling consistent comparison across capabilities. Experiments follow the benchmark’s training/evaluation split.Source: §B.1 · p. 14
• 10 Image Classification tasks include ImageNet-1K, ObjectNet, HatefulMemes, ImageNet-A, and ImageNet-R.The remaining datasets are VOC2007, N24News, SUN397, Country211, and Place365.Source: §B.1 · p. 14
• 10 Visual Question Answering tasks include OK-VQA, A-OKVQA, DocVQA, InfoVQA, ChartQA, Visual7W, ScienceQA, GQA, TextVQA, and VizWiz.Source: §B.1 · p. 14
• 12 Image-level Retrieval tasks include bidirectional MSCOCO retrieval, VisDial, CIRR, NIGHTS, WebQA, EDIS, OVEN, Wiki-SS-NQ, FashionIQ, and bidirectional VisualNews retrieval.Source: §B.1 · p. 14
• 4 Visual Grounding tasks comprise MSCOCO, RefCOCO, RefCOCO-Matching, and Visual7W-Pointing.Source: §B.1 · p. 14
• 5 tasks each cover Video Classification, Video Question Answering, and Video-level Retrieval, while 3 tasks cover Moment Retrieval.Examples include UCF101 and Kinetics-700 for classification, Video-MME and EgoSchema for question answering, and MSR-VTT and YouCook2 for video-level retrieval. Moment Retrieval comprises QVHighlights, Charades-STA, and MomentSeeker.Source: §B.1 · p. 14

B.2 Training and Evaluation Details

• Training uses Qwen3-VL-2B-Instruct as the base model and the MMEB-V2 training set.Table 3 documents the complete training hyperparameters.Source: §B.2 · p. 14, Table 3 · p. 14
• Each individual task is evaluated using Hit@1, defined as accuracy at rank 1.Source: §B.2 · p. 14
• The overall benchmark score is computed as a hierarchical weighted average (Meng et al., 2025).Source: §B.2 · p. 14

C Visual Decoding: Methodology

• A latent diffusion model decodes TTE-Flash thinking-token representations into images while keeping the TTE-Flash backbone frozen.The model accepts free-form vision-language inputs.Source: C Visual Decoding: Methodology · p. 15

C.1 Overview

• DINOv2-B (Oquab et al., 2024) encodes reconstruction targets, while DiTDH-XL (Peebles and Xie, 2023) generates latents that ViT-XL decodes into images.The reconstruction autoencoder uses a 768×16×16 semantic latent space. Within that autoencoder, only the ViT-XL decoder is trained; the DINOv2-B encoder remains frozen.Source: §C.1 · p. 15

C.2 Conditioning via TTE-Flash Representations

• 8 thinking tokens, each with 2048 dimensions, are produced by the frozen TTE-Flash backbone to condition diffusion in the visualization experiments.The token count is fixed for these experiments.Source: §C.2 · p. 15
• 256 learned queries resample projected thinking tokens into a 768×16×16 spatial conditioning map through eight-head cross-attention.A linear projection reduces token features from 2048 to 768 dimensions. The resampler applies a residual query connection and LayerNorm before reshaping its output to match the diffusion latent dimensions.Source: §C.2 · p. 15

C.3 DiTDH-XL: Diffusion Transformer

• 28 encoder layers process patchified noisy latents using 2D rotary positional embeddings and Adaptive LayerNorm modulated by summed timestep and conditioning embeddings.The encoder has hidden size 1152 and 16 attention heads, with SwiGLU feed-forward networks and RMSNorm.Source: §C.3 · p. 15
• 2 DDT decoder layers receive per-token encoder conditioning rather than a single global vector, enabling spatially aware denoising.The decoder uses hidden size 2048 and 16 heads. It re-embeds the original noisy input and applies thinking-token conditioning through per-token scale, shift, and gate operations.Source: §C.3 · p. 15
• Flow matching predicts velocity x_0 − x_1 with MSE along a linear data-to-Gaussian-noise path, sampling logit(t) from N(0, 1).Training runs for 200 epochs on 243K MMEB samples spanning MSCOCO, VisualNews, CIRR, and NIGHTS, using AdamW and an effective batch size of 1024. Sampling uses 50-step Euler ODE integration and classifier-free guidance at scale 4.0, enabled by 10% caption dropout during training.Source: §C.3 · pp. 15–16

C.4 Qualitative Results

• For MSCOCO and VisualNews text-to-image tasks, the model receives only a caption and generates an image grounded in its described content.Source: §C.4 · p. 16
• Three MMEB tasks are represented by successful and failure cases in Appendix E tables.Source: §C.4 · p. 16
• Successful cases demonstrate that TTE-Flash thinking tokens encode rich visual semantics beyond what retrieval alone requires.The qualitative evidence comes from image generation grounded in captions or combined image-and-text inputs.Source: §C.4 · p. 16
• Failure cases reveal limitations in fine-grained spatial reasoning.These failures qualify the visual understanding demonstrated by successful generations.Source: §C.4 · p. 16
• Failure cases also expose limited faithfulness to complex compositional instructions.Composed retrieval requires generating an image that combines a reference image with a textual modification instruction.Source: §C.4 · p. 16
• For CIRR composed retrieval, the model receives a reference image and textual modification instruction, then generates an image combining both signals.Source: §C.4 · p. 16

Dataset: IMAGENET–1K

• CoT-16 identifies the African grey target, whereas CoT-8 says black macaw and CoT-2 says cat.CoT-4 mixes African Grey Parrot with Barbary Ape in its identification.Source: Dataset: IMAGENET–1K · p. 17
• For the beer bottle target, CoT-16 and CoT-8 output bottle, CoT-4 outputs beer, and CoT-2 describes a bottle labelled VODKA.CoT-8 explicitly interprets Bier as beer but still concludes with the generic label bottle.Source: Dataset: IMAGENET–1K · p. 17
• CoT-16 matches the drum target, while CoT-8 identifies a drumstick or mallet and CoT-4 and CoT-2 identify a bass drum.The reasoning also changes scene descriptions: CoT-16 describes an armored person striking a drum, whereas CoT-4 describes a complete drum kit.Source: Dataset: IMAGENET–1K · p. 18
• CoT-16, CoT-8, and CoT-4 identify the fire engine or fire truck target, while CoT-2 outputs the less specific label truck.CoT-16 uses water-tank, hose, and FIRE DEPT cues; CoT-2 instead emphasizes a boxy vehicle with an open cargo area.Source: Dataset: IMAGENET–1K · p. 18

Dataset: MSCOCO

• Object-cropping prompts request broccoli, pizza, and sandwich; generations include relevant pizza and sandwich descriptions but also an unrelated coffee cup.CoT-16 descriptions emphasize pizza features and sandwich preparation, while the supplied CoT-2 description focuses on a coffee cup.Source: Dataset: MSCOCO · pp. 19–20
• CoT-16 captions describe the target zebras and hot dogs, whereas CoT-2 generations instead describe elephants and a decorated cake.These examples demonstrate subject changes between generation settings, not aggregate retrieval-performance differences.Source: Dataset: MSCOCO · pp. 20–21
• CoT-16 is not uniformly target-consistent: the courtyard example becomes a school scene, although the motorcycle example retains parked motorcycles.At CoT-4, the courtyard example describes students taking pictures, while the motorcycle example describes a street with bicycles rather than motorcycles.Source: Dataset: MSCOCO · pp. 21–22
• Question-conditioned generations change from Pacific Ocean and warehouse at CoT-16 to surfboard and storage at CoT-4.The supplied retrieval targets are pacific ocean and warehouse; the shorter-setting answers change either the requested answer category or its specificity.Source: Dataset: MSCOCO · p. 24
• Generations range from Lincoln (CoT-16) and Suzuki (CoT-8) to a 1950s American car without a named make (CoT-2), while other examples identify a photograph (CoT-4) or birdcage (CoT-2).Source: Dataset: MSCOCO · pp. 25, 27

Dataset: WEBQA

• WEBQA examples show prompt-specific interpretations, generic descriptions, and conflicting identifications across CoT settings.For the RaiNY prompt, CoT-16 invokes Rai and NYC signs, while CoT-8 describes a clock tower. For the 1934 Cadillac prompt, CoT-16 describes a 1934 Cadillac Convertible, whereas CoT-2 names a 1957 Ford Thunderbird.Source: Dataset: WEBQA · p. 28

Dataset: MMLONGBENCH–DOC

• MMLONGBENCH–DOC generations describe academic-reference lists at CoT-16, CoT-8, and CoT-4, but switch to an El Nio event table at CoT-2.The longer traces discuss citation matching, publication locations, or title-and-year verification; CoT-2 instead describes spectral and temporal characteristics.Source: Dataset: MMLONGBENCH–DOC · p. 30

Dataset: VIDORE_SYNTHETICDOCQA_AI

• One VIDORE_SYNTHETICDOCQA_AI example produces different document topics across settings: Zip sharing, Google assignment instructions, digital-media management, and a definition of data.These topics correspond respectively to CoT-16, CoT-8, CoT-4, and CoT-2; the CoT-4 trace names PDF and HTML files.Source: Dataset: VIDORE_SYNTHETICDOCQA_AI · p. 31
• The supplied CoT generations describe graph traversal at CoT-16, the Traveling Salesman Problem at CoT-8, Backward Search at CoT-4, and limitations of a linear-model assumption at CoT-2.Source: Dataset: VIDORE_SYNTHETICDOCQA_AI · p. 32

Dataset: VISRAG_INFOVQA

• A second VISRAG_INFOVQA example describes COVID-19 protective measures at CoT-16/8/4 but switches to digital transformation at CoT-2.The CoT-2 generation instead describes organizational change and components including data, technology, people, processes, and culture.Source: Dataset: VISRAG_INFOVQA · p. 34

E Images Generated from Think Tokens

• DiTDH-XL generates images conditioned on TTE-Flash think tokens, with examples showing the query, ground-truth target, and generated image.Queries are text prompts or reference images accompanied by modification instructions.Source: E Images Generated from Think Tokens · p. 35

Dataset: MSCOCO_t2i – Successful Cases

• MSCOCO_t2i successful-case prompts cover motorcycle riders, birthday candles, kitchens, bathrooms, and urban bicyclists.These examples request everyday images matching captions.Source: Dataset: MSCOCO_t2i – Successful Cases · pp. 35–37

Dataset: VisualNews_t2i – Successful Cases

• VisualNews_t2i successful-case prompts include a Maryland police cruiser, Ashton Eaton competing in the decathlon, and Obama meeting young businesspeople in Nairobi.Each prompt requests an image matching a news caption.Source: Dataset: VisualNews_t2i – Successful… · pp. 38–39

Dataset: VisualNews_t2i – Failure Cases

• The VisualNews_t2i failure example uses a caption identifying Huma Abedin photographed at the White House on August 10, 2011.The passage labels ground-truth and generated images but does not explain the visual mismatch.Source: Dataset: VisualNews_t2i – Failure Cases · p. 40

Dataset: CIRR – Successful Cases

• Quantity-change requests specify three soft-drink bottles or fewer paper towels per pack.These examples target changes in object count or package contents.Source: Dataset: CIRR – Successful Cases · p. 40
• Framing requests emphasize an animal’s head, reveal a surface through an angled shot, or show a deer’s back while it looks toward the camera.One head-focused request explicitly requires preserving the dog’s breed.Source: Dataset: CIRR – Successful Cases · pp. 41–43
• Dog-focused requests specify a same-breed dog running with its puppy, two same-breed dogs on the floor, or a barking dog with an unobstructed mouth.Source: Dataset: CIRR – Successful Cases · pp. 41–42
• Object and scene edits include removing a lemon, adding a wooden log, changing water color, and placing fries on a clean white plate.One request combines seashell removal with green water; another specifies blue water.Source: Dataset: CIRR – Successful Cases · pp. 41–42
• A compound request combines adding one more deer with adding sunlight.The requested modification changes both animal count and illumination.Source: Dataset: CIRR – Successful Cases · p. 43

Dataset: CIRR – Failure Cases

• Human-related edits request two animals of a different species without people, or a human facing the camera with less gear.The two-animal prompt specifies a species different from the original animal.Source: Dataset: CIRR – Failure Cases · p. 44
• Pose and framing requests specify either a whole dog standing on a dirt floor or an animal standing on dirt.The dog-specific prompt additionally requires the entire dog to be visible.Source: Dataset: CIRR – Failure Cases · pp. 44–45
• Species-change requests include a human with one animal unlike the previous two animals, or changing animal type while preserving the environment.The human-and-animal prompt explicitly requires a different species.Source: Dataset: CIRR – Failure Cases · p. 44