Source-linked AI summary
UniT: Unified Multimodal Chain-of-Thought Test-time Scaling
Leon Liangyu Chen, Haoyu Ma, Zhipeng Fan, Ziqi Huang, Animesh Sinha, Xiaoliang Dai, Jialiang Wang, Zecheng He, Jianwei Yang, Chunyuan Li, Junzhe Sun, Chu Wang, Serena Yeung-Levy, Felix Juefei-Xu
TL;DR
Unified multimodal models largely operate in a single pass, leaving open how to support scalable iterative generation, reflection, and refinement. UniT combines agentic data synthesis, unified training, and test-time inference, showing longer-chain generalization, efficient sequential scaling, and broad multimodal improvements.
Problem
Unified multimodal models typically lack explicit mechanisms for evaluating, reflecting on, and refining outputs during complex multimodal tasks.
Method
UniT combines agentic multi-round data synthesis, unified model training, and test-time iterative refinement to elicit multimodal reasoning behaviors.
Results
UniT improves multimodal generation, editing, instruction following, and visual reasoning, while sequential scaling achieves comparable performance to parallel sampling with 2.5× less computational cost.
Takeaways & Limitations
Models trained on shorter trajectories generalize to longer inference chains, supporting multimodal chain-of-thought test-time scaling across generation and understanding tasks.
Takeaways & Limitations
Untrained Bagel cannot practically perform multimodal chain-of-thought reasoning because image quality degrades with additional context images and verification is unreliable.
Abstract
from arXiv · showhide
Unified models can handle both multimodal understanding and generation within a single architecture, yet they typically operate in a single pass without iteratively refining their outputs. Many multimodal tasks, especially those involving complex spatial compositions, multiple interacting objects, or evolving instructions, require decomposing instructions, verifying intermediate results, and making iterative corrections. While test-time scaling (TTS) has demonstrated that allocating additional inference compute for iterative reasoning substantially improves language model performance, extending this paradigm to unified multimodal models remains an open challenge. We introduce UniT, a framework for multimodal chain-of-thought test-time scaling that enables a single unified model to reason, verify, and refine across multiple rounds. UniT combines agentic data synthesis, unified model training, and flexible test-time inference to elicit cognitive behaviors including verification, subgoal decomposition, and content memory. Our key findings are: (1) unified models trained on short reasoning trajectories generalize to longer inference chains at test time; (2) sequential chain-of-thought reasoning provides a more scalable and compute-efficient TTS strategy than parallel sampling; (3) training on generation and editing trajectories improves out-of-distribution visual reasoning. These results establish multimodal test-time scaling as an effective paradigm for advancing both generation and understanding in unified models.
1 Introduction
UniT addresses the gap between single-pass unified multimodal models and tasks requiring iterative generation, verification, and refinement. It combines agentic data synthesis, unified training, and test-time scaling, yielding longer-chain generalization and broad multimodal gains.
- Unified multimodal models generally produce outputs in a single pass, limiting explicit evaluation, reflection, and refinement for complex multimodal tasks.
- The open question is how to enable unified models to iteratively generate, reflect, and refine with scalable multimodal inference.
- Sequential chain-of-thought scaling substantially outperforms best-of-N parallel sampling while achieving comparable performance with 2.5× less computational cost.
- UniT improves compositional editing, multi-turn editing, instruction following, and out-of-distribution visual reasoning across multimodal benchmarks.Reported improvements include 5.56% on CompBench, 2.95 human preference scores on ImgEdit, 10.34% on OneIG, and 53.33% on MIRA.
- UniT integrates agentic data synthesis, unified model training, and test-time scaling to induce verification, subgoal decomposition, and content memory.
- Models trained on shorter trajectories averaging 3.6 rounds generalize to longer inference chains averaging 4.7 rounds at test time.
2 Related Works
Related work covers test-time scaling, unified multimodal models, and multimodal chain-of-thought. UniT extends these directions with an agentic framework that synthesizes iterative reasoning trajectories for unified models.
- Test-time scaling allocates additional inference computation through parallel candidate sampling or sequential critique and refinement.
- Unified multimodal models combine understanding and generation within architectures spanning text and image representations.
- Multimodal chain-of-thought methods extend reasoning across visual representations and interleaved text-image modalities.
- UniT’s agentic framework synthesizes training trajectories by coupling generation, verification, subgoal decomposition, and iterative editing.
3 Method
UniT trains a unified model on agentically synthesized reflection-editing trajectories and controls inference compute through iterative image-generation rounds. The resulting system supports multimodal cognitive behaviors and generalizes beyond its training trajectory length.
- Multimodal Chain-of-Thought Data: The agentic pipeline starts from prompts, generates images, verifies them, plans corrections, applies edits, and repeats until the output satisfies the prompt.
- Multimodal Chain-of-Thought Data: The pipeline induces verification, subgoal decomposition, and content memory across multimodal refinement rounds.
- Multimodal Chain-of-Thought Data: Quality, relevance, visual-change, and benchmark-deduplication filters retain 12k high-quality training trajectories.
- Training and Inference: At inference, nested text and image classifier-free guidance separately control prompt adherence and visual consistency.Text guidance uses st=4.0 and image guidance uses si=2.0.
- Training and Inference: Sequential scaling controls compute by the number of image-generation rounds, with each round combining chain-of-thought reasoning and image generation or editing.
- Training and Inference: Training on trajectories averaging 3.6 rounds generalizes to inference chains averaging 4.7 rounds, while untrained Bagel degrades quickly as context images increase.
4 Experiments
UniT is evaluated across generation, editing, and visual reasoning benchmarks using iterative multimodal chain-of-thought inference. The experiments show gains with additional refinement rounds, including substantial improvements in multi-turn editing and out-of-distribution visual reasoning.
- UniT is evaluated on text-to-image generation, compositional editing, multi-turn editing, and visual reasoning benchmarks.
- 10.34% improvement over the base model is achieved on OneIG-Bench at C=10, while CompBench improves 5.56% from C=1 to C=10.
- 225.19% improvement is achieved on ImgEdit from C=1 to C=4, with maintained context and reasoning chains supporting multi-turn interactions.
- 53.33% improvement is achieved on MIRA from C=1 to C=10, although the remaining gap to frontier models reflects base-model scale and proprietary-data differences.
- Iterative visual reasoning progressively refines outputs through verification and subgoal decomposition, with generation errors corrected across rounds and cognitive behaviors transferring to MIRA tasks.
5 Discussion
The discussion finds sequential chain-of-thought scaling more effective than independent parallel sampling across multimodal tasks, while different cognitive behaviors contribute to different task types. It also identifies deployment trade-offs and data-quality sensitivities.
- Sequential versus parallel scaling: Sequential scaling shows steeper gains and continued improvement up to C=10, whereas parallel scaling plateaus after a few samples.
- Sequential versus parallel scaling: Sequential refinement accumulates successful edits and prior textual context, while parallel sampling generates independent outputs without inter-sample learning.
- Deployment trade-offs: Sequential scaling prioritizes performance and parallel scaling prioritizes latency, with speculative decoding, KV-cache reuse, and early stopping potentially narrowing the latency gap.
- Unified versus modular pipelines: The unified model offers faster inference and single-model deployability than the multi-model teacher pipeline, although the teacher scores slightly higher.
- Cognitive behavior ablation: Removing subgoal decomposition produces 3.8% and 2.5% drops on OneIG-Bench and CompBench, respectively, showing task-specific sensitivity to planning.
- Cognitive behavior ablation: Removing content memory causes a 42.5% relative drop on ImgEdit, while removing verification reduces MIRA performance by 1.9%.
- Data quality: Removing relevance filtering degrades OneIG-Bench and CompBench by 3.1% and 2.5%, while removing the minimal-visual-changes filter hurts ImgEdit by 1.16 points.
6 Conclusion
UniT extends test-time compute from text-only reasoning to unified models that perform visual understanding and generation. Its agentic data, budget forcing, and sequential reasoning produce gains across complex multimodal tasks.
- UniT extends inference-time compute to unified models handling both visual understanding and generation.
- The framework combines agentic cognitive-behavior induction, budget forcing for beyond-training generalization, and evidence that sequential reasoning outperforms parallel sampling.
- Iterative refinement through explicit reasoning improves compositional generation, multi-turn editing, and visual reasoning.
A Data Synthesis Pipeline
UniT’s data synthesis pipeline automatically creates multimodal chain-of-thought trajectories through an iterative loop of generation, verification, planning, and editing. The resulting examples explicitly induce verification, subgoal decomposition, and content memory.
- The pipeline automatically generates 12K multimodal chain-of-thought training trajectories.
- An image-generation model creates initial images, a vision-language model verifies and plans refinements, and an image-editing model applies those instructions in an iterative loop.
- Prompt generation covers compositional attributes, spatial relations, and multi-object tasks using 20K diverse prompts.
- During verification, Qwen3-VL identifies deficiencies, plans improvements, and produces explicit editing instructions when outputs fail to satisfy prompts.
- A bookshelf trajectory demonstrates verification, sequential subgoal decomposition, and content memory by tracking cumulative changes across Images #1, #2, and #3.
A.4 Training Data Statistics
The training set contains 12K quality-filtered multimodal reasoning trajectories generated through structured visual analysis and iterative refinement. These trajectories average 3.6 refinement rounds and induce verification, subgoal decomposition, and content memory.
- 12K quality-filtered trajectories form the training data, averaging 3.6 refinement rounds across a range of 1–8 rounds.Training on this data requires 700 H100 GPU hours.
- A bookshelf example shows verification of unwanted books, sequential removal and frame addition, and references to prior content across three refinement rounds.
- The VLM prompt guides image description, requirement comparison, and decisions among editing, backtracking, or completion.This structure is designed to produce verification, subgoal decomposition, and content memory.
B Additional Qualitative Results
Additional qualitative results show explicit multimodal reasoning in which reflection explains shortcomings and guides iterative improvement. The examples span different task types and computational budgets.
- Reflection produces detailed reasoning about why an output falls short and how to improve it, rather than only issuing new instructions.
- Figure 8 presents representative trajectories across different task types and computational budgets.The examples illustrate qualitative chain-of-thought test-time scaling results.
C Generalization Preservation
Fine-tuning on reasoning-heavy trajectories preserves and can improve base-model capabilities without test-time scaling, while evaluation is bounded by memory constraints and extended editing can degrade image quality. The framework also uses structured verification and planning to support these behaviors.
- Evaluation is capped at C=10 because of GPU memory constraints.
- Image quality collapses when consecutive editing rounds make minimal visual changes, defined here as LPIPS < 0.03.The degradation is attributed to accumulated autoregressive noise; such steps are described as infrequent.
- The VLM prompt uses image description, comparison analysis, and action decisions to induce verification, subgoal decomposition, and content memory.
- Beyond C=10, the authors expect test-time-scaling performance to saturate or degrade once quality collapse dominates.The exact inflection point depends on the base generation and editing models.
E Failure Analysis
Failure analysis identifies weaknesses in precise physical and spatial reasoning, verification, highly constrained compositions, and the computational cost of inference-time scaling. The authors propose broader modality coverage and stronger reflection mechanisms as future directions.
- Failure Cases: Precise physical or fine-grained spatial tasks can fail because refinement may not correct physics violations or attribute-binding errors inherited from base models.Examples include incorrect leash-dog assignments and helmet placement or sizing.
- Failure Cases: Verification hallucinations can trigger unnecessary edits that reduce quality when the VLM misidentifies nonexistent visual issues.
- Failure Cases: Highly complex prompts can create subgoal conflicts in which satisfying one compositional constraint violates another.
- Limitations: Test-time scaling requires additional inference computation, motivating more efficient reflection and adaptive budget allocation.
- Future Directions: Future work includes adding audio and video, explicit physical reasoning, reinforcement learning from human feedback, and combinations with other inference-time techniques.
- Failure Cases: Figure 9 illustrates failures involving exact object counts, geometric arrangements, and layout changes across intermediate images.