Source-linked AI summary
Machine Mental Imagery: Empower Multimodal Reasoning with Latent Visual Tokens
Zeyuan Yang, Xueyang Yu, Delin Chen, Maohao Shen, Chuang Gan
TL;DR
VLMs remain limited on multimodal reasoning that requires visual imagination, and image-generation training can hinder reasoning quality. Mirage interleaves latent visual tokens with text, trains them through visual grounding and text-only supervision, and refines the trajectory with reinforcement learning. Across four spatial-reasoning benchmarks, it consistently outperforms text-only baselines without generating pixel-level images.
Problem
VLMs struggle with multimodal reasoning that requires coherent manipulation of visual elements, while combining reasoning with pixel-level image synthesis can degrade reasoning quality.
Method
Mirage interleaves hidden-state-derived latent visual tokens with text and trains them through joint visual grounding, text-only latent relaxation, and reinforcement-learning refinement.
Results
Across four spatial-reasoning benchmarks, Mirage consistently outperforms text-only baselines without generating pixel-level images.
Takeaways & Limitations
Latent visual tokens can support stronger multimodal reasoning while embodying meaningful visual cues, without requiring explicit image generation.
Takeaways & Limitations
Performance depends on the quality of generated multimodal trajectories, and evaluation is currently limited beyond spatial reasoning tasks.
Abstract
from arXiv · showhide
Vision-language models (VLMs) excel at multimodal understanding, yet their text-only decoding forces them to verbalize visual reasoning, limiting performance on tasks that demand visual imagination. Recent attempts train VLMs to render explicit images, but the heavy image-generation pre-training often hinders the reasoning ability. Inspired by the way humans reason with mental imagery-the internal construction and manipulation of visual cues-we investigate whether VLMs can reason through interleaved multimodal trajectories without producing explicit images. To this end, we present a Machine Mental Imagery framework, dubbed as Mirage, which augments VLM decoding with latent visual tokens alongside ordinary text. Concretely, whenever the model chooses to ``think visually'', it recasts its hidden states as next tokens, thereby continuing a multimodal trajectory without generating pixel-level images. Begin by supervising the latent tokens through distillation from ground-truth image embeddings, we then switch to text-only supervision to make the latent trajectory align tightly with the task objective. A subsequent reinforcement learning stage further enhances the multimodal reasoning capability. Experiments on diverse benchmarks demonstrate that Mirage unlocks stronger multimodal reasoning without explicit image generation.
1 Introduction
VLMs struggle with multimodal reasoning that requires manipulating visual elements, while explicit image generation can degrade reasoning. Mirage addresses this by interleaving latent visual tokens with text and training them through staged supervision and reinforcement learning.
- VLMs still stumble on spatial reasoning tasks that require coherent understanding and manipulation of visual elements.
- Human-inspired mental imagery motivates reasoning with compact, task-relevant visual embeddings instead of photorealistic images.
- Mirage reuses a model’s hidden state as a compact visual embedding when a special token triggers visual reasoning, then appends it to the context.
- Its second stage removes direct latent-vector supervision and optimizes text tokens, allowing generated latent embeddings to act as flexible priors for word generation.
- Mirage enables interleaved reasoning trajectories that mix latent visual tokens with ordinary text without external visual decoders.
- The two-stage training paradigm produces stable yet flexible interleaved reasoning, while reinforcement learning further boosts performance.
- Mirage delivers consistent gains across diverse multimodal reasoning benchmarks, with latent tokens embodying meaningful visual cues.
2 Related Work
Related work extends chain-of-thought reasoning with visual evidence, unified image-text generation, latent reasoning representations, and specialized control tokens. Mirage is positioned among these approaches as a latent-token alternative focused on interleaved multimodal reasoning.
- Multimodal chain-of-thought methods interleave visual evidence, such as attention-selected image crops, with text to improve visual question answering.
- Unified token-based models can emit arbitrary text and image sequences, but require large-scale pixel-level supervision and heavier decoding.
- Specialized tokens such as <plan> and <pause> have been introduced to control reasoning stages and stabilize multi-step reasoning.
- Latent-reasoning approaches distill chain-of-thought behavior into hidden states, sometimes removing explicit intermediate steps to improve inference efficiency.
3 Multimodal Reasoning with Latent Visual Tokens
Mirage trains VLMs to reason through interleaved text and latent visual tokens, using helper-image supervision before relaxing latent constraints. The framework then uses text-only optimization and reinforcement learning to support adaptive multimodal reasoning without explicit image generation.
- Framework: Mirage interleaves compact latent visual embeddings with ordinary text, allowing VLMs to continue multimodal reasoning without an external image decoder.The model reuses hidden states as visual embeddings when it emits a special visual-reasoning token.
- Data Generation: Synthetic training data pairs each task input with a task-specific helper image and reasoning chain that incorporates the image.Helper images encode precisely the visual cues needed for supervision, such as annotated paths or composite jigsaw fragments.
- Joint Supervision: Stage 1 compresses helper-image patch features into salient vectors and jointly trains latent-token visual alignment with surrounding text prediction.Latent tokens use cosine similarity to target embeddings, while text uses cross-entropy; the combined objective is L1 = Lvisual + γLtext.
- Latent Relaxation: Stage 2 removes the latent cosine loss and optimizes text tokens, allowing self-generated differentiable embeddings to act as flexible priors for subsequent generation.Gradients from later text prediction flow through the latent tokens, adapting them toward task-focused reasoning.
- Results: Experiments report consistent gains across diverse multimodal reasoning benchmarks over text-only decoding, with latent tokens embodying meaningful visual cues.The framework’s two stages are empirically validated as more effective than naive text-only decoding.
- Reinforcement Learning: After supervised training, Mirage applies group relative policy optimization to further improve reasoning performance.The reinforcement-learning stage samples multiple responses for each training query.
4 Experiments
Experiments evaluate Mirage across spatial reasoning benchmarks, model scales, training stages, hyperparameters, and helper-image settings. Results generally favor interleaved latent visual reasoning, while revealing sensitivity to synthesized reasoning quality and excessively long latent sequences.
- Benchmark Results: Mirage improves over direct fine-tuning by 3% on VSP spatial reasoning and 11% on VSP spatial planning.
- Benchmark Results: Mirage also improves over CoT SFT + GRPO by 2% on spatial reasoning and 7% on spatial planning.
- Benchmark Results: Across COMT, Jigsaw, and SAT, interleaving compact visual cues produces consistent performance gains on the evaluated tasks.
- Limitations: On VSP spatial planning, synthesized reasoning thoughts perform worse than direct answer-label training, partly because imperfect generated thoughts can propagate errors.
- Training Enhancements: GRPO adds a further +2% accuracy gain on VSP tasks when latent visual tokens are interleaved within text trajectories.
- Ablation Study: The two training stages play complementary roles: the first grounds latent tokens, while the second lets them adapt to the task objective.
- Ablation Study: Varying latent size from 2 to 6 maintains strong performance, whereas increasing it to 8 causes an approximately 13% drop.
5 Analysis
Analysis shows that helper images can strongly support spatial reasoning, while Mirage’s latent tokens retain visual structure after direct supervision is removed. The framework also generalizes to smaller models, with larger gains over text-only baselines on two tasks.
- Generalization to Smaller Models: Mirage achieves larger gains than text-only baselines on Qwen2.5-VL 3B, improving Jigsaw by 5% and SAT Real by 10%.The gains are reported across both evaluated tasks on the smaller model.
- Synthesized Data Quality: Helper images yield nearly 100% accuracy on both VSP tasks when supplied as prior input in zero-shot and fine-tuned settings.Zero-shot gains are substantial for spatial reasoning but limited to simpler map layouts for spatial planning.
- Synthesized Data Quality: The helper-image results suggest that synthesized visual information is informative for VLM reasoning and may define an upper bound if fully internalized by latent thoughts.The passage explicitly frames complete internalization of helper-image information as a strong performance upper bound for Mirage.
- Latent Behavior Analysis: The latent-behavior analysis investigates whether tokens still encode visual representations after the second stage removes their direct supervision.The first stage anchors latent tokens in the visual subspace, while the second stage lets them adapt without predetermined embeddings.
- Latent Behavior Analysis: After text-only second-stage training, Mirage’s latent tokens remain clustered near but just outside the visual representation subspace.The t-SNE visualization is consistent with the two-stage training design despite the removal of direct latent-vector supervision.
6 Conclusion
Mirage enables VLMs to reason with compact latent visual tokens interleaved with text, avoiding pixel-level image generation. Across four spatial-reasoning benchmarks, it consistently outperforms text-only baselines, while broader task generalization and dependence on synthetic trajectory quality remain open limitations.
- Conclusion: Mirage interleaves compact latent visual tokens with text so VLMs can reason multimodally without generating pixel-level images.Its training combines visual-embedding supervision, text-only adaptation, and reinforcement-learning refinement.
- Conclusion: Across four spatial-reasoning benchmarks, Mirage consistently outperforms text-only baselines.The conclusion presents this as evidence for the effectiveness and potential of latent visual reasoning.
- Limitations and Future Works: The framework’s interleaved reasoning depends on the quality of generated multimodal trajectories, making high-quality dataset curation an important next step.The paper also leaves extension beyond spatial-reasoning benchmarks as an open direction.
A Datasets
The dataset pipeline creates task-specific helper images and uses them with generated reasoning traces to supervise fine-tuning. The images encode maps, paths, candidate-piece composites, or target views for spatial and visual reasoning tasks.
- Help Image Generation: Task-specific tools generate helper images for the VSP, BLINK Jigsaw, and SAT datasets.The section describes separate generation pipelines for each task.
- VSP Spatial Reasoning: VSP Spatial Reasoning uses a rendered map state showing the agent’s position after part of an action trajectory.The state is rendered with OpenAI Gym from the initial map and action sequence.
- VSP Spatial Planning: VSP Spatial Planning supplies a map annotated with a valid ground-truth path as a red arrow from the agent’s start position to the goal.The task is converted into reading the highlighted trajectory.
- BLINK Jigsaw: BLINK Jigsaw helper images insert a randomly chosen candidate patch into the masked reference region so the model can judge whether the composite is seamless.A smooth blend indicates the correct candidate; otherwise the other candidate should be selected.
- SAT: SAT helper images provide target views for GoalAim and ObjM camera-pose movement, generated with CogVideoX-5B under three action primitives.The supplied passages identify forward movement, left turns, and a third truncated primitive condition.
- Textual Thoughts Generation: For each task, Qwen2.5-VL 32B generates textual thoughts conditioned on the helper image and ground-truth answer, and these thoughts supervise fine-tuning.Task-specific prompts are used, with one illustrative example per task provided in the referenced tables.
A.3 Data Configuration
The data configuration partitions benchmark examples among evaluation, fine-tuning, and reinforcement learning, while synthesizing training data for VSP because it has no training set.
- Data Configuration: COMT Mathematical Geometry uses 200 evaluation examples and 820 fine-tuning examples.For other benchmarks, the configuration uses 1,000 fine-tuning samples and 2,000 non-overlapping reinforcement-learning samples.
- Data Configuration: VSP has no training set, so the authors follow its data-generation recipe to synthesize their own data.The passage also states that BLINK-Jigsaw and SAT samples are randomly drawn from official training splits.
- Data Configuration: Table 8 reports the dataset statistics for the evaluated data configuration.
B.1 Implementation Details
Mirage uses Qwen2.5-VL-7B-Instruct with a trainable language/multimodal model except for the vision encoder, combining token prediction and latent-token alignment during fine-tuning. Reinforcement learning uses GRPO with format and correctness rewards, KL regularization, and no latent-token KL divergence.
- Supervised Fine-tuning: Qwen2.5-VL-7B-Instruct serves as Mirage’s base VLM, with all components trainable except the vision encoder.
- Supervised Fine-tuning: Fine-tuning combines cross-entropy next-token prediction with cosine similarity alignment for latent visual tokens, using visual-loss weight γ = 0.1.
- Reinforcement Learning: Reinforcement learning uses VERL with Group Relative Policy Optimization (GRPO).
- Reinforcement Learning: The reward function combines format and correctness rewards weighted by σf and σc, respectively, while KL regularization is applied.
- Reinforcement Learning: Entropy regularization is disabled by setting λen = 0, and KL divergence on latent visual tokens is omitted during reinforcement learning.
B.2 Efficiency Analysis
Mirage’s two training stages require several hours on a single NVIDIA H100 GPU, with reinforcement learning taking longer than supervised fine-tuning for the example task.
- 3.5 hours for Stage 1 and 7.2 hours for Stage 2 are required on one NVIDIA H100 GPU for the VSP spatial reasoning task.
- Text-only CoT supervised fine-tuning takes about 5.5 hours on the same hardware, between Mirage’s Stage 1 and Stage 2 durations.