UniT: Unified Multimodal Chain-of-Thought Test-time Scaling

TL;DR

UniT extends unified multimodal models with test-time scaling by introducing an agentic data synthesis pipeline, unified model training, and budgeted inference to support multi-round multimodal reasoning and refinement. The approach yields emergent cognitive behaviors such as verification, subgoal decomposition, and content memory, and demonstrates that sequential chain-of-thought scaling outperforms best-of-N parallel sampling across compositional generation, multi-turn editing, and visual reasoning tasks. Models trained on short reasoning trajectories generalize to longer inference chains at test time, delivering substantial gains with more compute-efficient refinement. This work establishes multimodal chain-of-thought test-time scaling as a practical paradigm for advancing both generation and understanding in unified models, with robust improvements on OneIG-Bench, CompBench, ImgEdit, and MIRA benchmarks. The findings support broader adoption of iterative, memory-enabled reasoning in unified multimodal systems for complex tasks requiring both perception and generation.

Abstract

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.

Figures (9)

• Figure 1: Multimodal chain-of-thought enables test-time scaling through emergent cognitive behaviors. We propose the UniT framework for unified multimodal models, which induces subgoal decomposition for compositional tasks and unlocks content understanding and memory for multi-turn editing. Controlling the number of test-time images, chain-of-thought sequential scaling outperforms best-of-N parallel scaling across generation and reasoning benchmarks. User input Model output
• Figure 2: Agentic framework for synthesizing chain-of-thought training data. Starting from a user prompt, an image generation model generates an initial image. A vision-language model then performs verification - evaluating whether the output satisfies the prompt. When unsatisfactory, the VLM engages in explicit subgoal decomposition through thinking tokens, planning concrete improvements, and rewriting editing instructions. This iterative loop continues until verification succeeds, generating multi-turn reasoning trajectories that teach unified models to refine outputs through test-time computation. The explicit reasoning traces of the three models capture how cognitive behaviors emerge from the interplay between generation, verification, and planning.
• Figure 3: UniT enables iterative refinement for compositional instructions through multimodal chain-of-thought reasoning. UniT exhibits: (i) error verification and correction—identifying and fixing constraint violations that Bagel misses (top: correcting leash placement and dog action); (ii) subgoal decomposition with subject consistency—sequentially addressing instructions while maintaining subject identity across rounds (middle: preserving bear features through style transformation, bottom: skateboard consistency); (iii) quality preservation—maintaining visual fidelity through iterative refinement rather than degradation (top: reduced artifacts and haloing).
• Figure 4: Qualitative examples of chain-of-thought test-time scaling. Representative trajectories showing progressive refinement across different tasks and computational budgets. Examples demonstrate how explicit chain-of-thought reasoning enables the model to iteratively improve compositional generation.
• Figure 5: Training vs. inference round distribution demonstrates beyond-training generalization. The model is trained on trajectories averaging 3.6 refinement rounds, but effectively generalizes to longer inference chains averaging 4.7 rounds at test time. This distribution shift reveals the model's emergent ability to extend inference beyond its training distribution, a key property of effective test-time scaling.