Visual-Aware CoT: Achieving High-Fidelity Visual Consistency in Unified Models

TL;DR

VACoT addresses the gap in multimodal generation by enforcing visual context consistency alongside text alignment within a unified model. It introduces Adaptive Visual Planning to produce structured visual-checklists and Iterative Visual Correction for self-refinement, guided by a visual-consistency reward implemented through flow-GRPO. A two-stage training pipeline—visual planning/correction via supervised fine-tuning and optimization with a visual-consistency reward—enables robust, high-fidelity multi-reference generation. Empirical results on OmniContext and GenEval demonstrate substantial gains in visual fidelity and identity preservation without sacrificing text-grounded generation quality, signaling strong practical potential for visually aware unified models.

Abstract

Recently, the introduction of Chain-of-Thought (CoT) has largely improved the generation ability of unified models. However, it is observed that the current thinking process during generation mainly focuses on the text consistency with the text prompt, ignoring the \textbf{visual context consistency} with the visual reference images during the multi-modal generation, e.g., multi-reference generation. The lack of such consistency results in the failure in maintaining key visual features (like human ID, object attribute, style). To this end, we integrate the visual context consistency into the reasoning of unified models, explicitly motivating the model to sustain such consistency by 1) Adaptive Visual Planning: generating structured visual check list to figure out the visual element of needed consistency keeping, and 2) Iterative Visual Correction: performing self-reflection with the guidance of check lists and refining the generated result in an iterative manner. To achieve this, we use supervised finetuning to teach the model how to plan the visual checking, conduct self-reflection and self-refinement, and use flow-GRPO to further enhance the visual consistency through a customized visual checking reward. The experiments show that our method outperforms both zero-shot unified models and those with text CoTs in multi-modal generation, demonstrating higher visual context consistency.

Figures (16)

• Figure 1: Comparison of Chain-of-Thought approaches for generation. (a) Text CoT uses text-only understanding modules for analysis and planning. (b) Text-align Multi-modal CoT iteratively refines generation through text-alignment evaluation. (c) Visual-Aware Multi-modal CoT (Ours) incorporates visual planning and visual-aware evaluation for improved visual consistency.
• Figure 2: Adaptive Visual Planning and Iterative Visual Correction Process of Our Method.
• Figure 3: Dataset construction for the planning and correction process.
• Figure 4: Our training pipeline of Visual-Aware CoT.
• Figure 5: Qualitative Comparison on Multi-Reference Generation. Our VACoT demonstrates superior identity consistency and visual coherence compared to baseline methods. BAGEL often fails to maintain character consistency across references, while UiG and UniCoT focus on text alignment rather than visual consistency. Our approach effectively preserves identity features and generates more accurate multi-reference images.