PaperlayerView exact revision on arXiv

Source-linked AI summary

Uni-cot: Towards Unified Chain-of-Thought Reasoning Across Text and Vision

Luozheng Qin, Jia Gong, Yuqing Sun, Tianjiao Li, Mengping Yang, Xiaomeng Yang, Chao Qu, Zhiyu Tan, Hao Li

arXiv:2508.05606v3cs.CVcs.CL

TL;DR

Coherent multimodal reasoning remains difficult because visual-state transitions are hard to model and loosely integrated architectures can yield fragmented trajectories. Uni-CoT combines visual transitions with text reasoning in a unified model, using hierarchical reasoning and structured training; the paper reports strong performance across generation and understanding tasks. Fine-grained visual consistency and perfect text–image alignment remain open challenges.

  • Problem

    Multimodal reasoning methods struggle to represent visual-state transitions coherently, limiting support for visually grounded tasks.

    Source: §1 · p. 1

  • Method

    Uni-CoT combines structural visual transitions and text reasoning in one model, with macro-level planning and interleaved supervision plus micro-level MDP-style execution and auxiliary tasks.

    Source: §1 · pp. 1–2

  • Results

    Uni-CoT reports state-of-the-art image-generation results and superiority across reasoning-driven image-generation and editing benchmarks, alongside strong understanding results.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The authors present Uni-CoT as a scalable foundation for future multimodal reasoning systems.

    Source: §6 · p. 10

    Broader real-world applications requiring fine-grained visual consistency or perfectly aligned text and images remain an open challenge.

    Source: §6 · p. 10

Abstract

from arXiv · show

Chain-of-Thought (CoT) reasoning has been widely adopted to enhance Large Language Models (LLMs) by decomposing complex tasks into simpler, sequential subtasks. However, extending CoT to vision-language reasoning tasks remains challenging, as it often requires interpreting transitions of visual states to support reasoning. Existing methods often struggle with this due to limited capacity of modeling visual state transitions or incoherent visual trajectories caused by fragmented architectures. To overcome these limitations, we propose Uni-CoT, a Unified Chain-of-Thought framework that enables coherent and grounded multimodal reasoning within a single unified model. The key idea is to leverage a model capable of both image understanding and generation to reason over visual content and model evolving visual states. However, empowering a unified model to achieve that is non-trivial, given the high computational cost and the burden of training. To address this, Uni-CoT introduces a novel two-level reasoning paradigm: A Macro-Level CoT for high-level task planning and A Micro-Level CoT for subtask execution. This design significantly reduces the computational overhead. Furthermore, we introduce a structured training paradigm that combines interleaved image-text supervision for macro-level CoT with multi-task objectives for micro-level CoT. Together, these innovations allow Uni-CoT to perform scalable and coherent multi-modal reasoning. Furthermore, thanks to our design, all experiments can be efficiently completed using only 8 A100 GPUs with 80GB VRAM each. Experimental results on reasoning-driven image generation benchmark (WISE) and editing benchmarks (RISE and KRIS) indicates that Uni-CoT demonstrates SOTA performance and strong generalization, establishing Uni-CoT as a promising solution for multi-modal reasoning. Project Page and Code: https://sais-fuxi.github.io/projects/uni-cot/

1 INTRODUCTION

Uni-CoT targets coherent multimodal reasoning by combining visual-state transitions with text reasoning in one model. Its macro–micro design and split training approach address the computational and training burdens, and the paper reports state-of-the-art results on image-generation and understanding benchmarks.

1 INTRODUCTION

  • Multimodal reasoning remains challenging because text-focused methods struggle on visually grounded tasks, while loosely integrated visual generators can produce fragmented reasoning and incoherent transitions.Programmatic operations such as cropping approximate local transitions but may miss global structural changes.Source: §1 · p. 1
  • Uni-CoT integrates structural visual transitions with coherent text reasoning in a unified model capable of image understanding and generation.The unified model is intended to reduce discrepancies between reasoning and visual dynamics and support end-to-end training.Source: §1 · pp. 1–2
  • Macro-level planning decomposes tasks into subtasks, while micro-level execution solves subtasks locally and models them as an MDP to reduce redundant attention interactions.The paper says this restructuring lowers multimodal reasoning complexity from quadratic to nearly linear.Source: §1 · p. 2
  • Training separates macro-level interleaved text–image supervision from micro-level MDP-style execution with four auxiliary tasks.The auxiliary tasks include action generation and reward estimation; the decoupled paradigm supervises both global planning and local execution.Source: §1 · p. 2
  • Uni-CoT achieves state-of-the-art results on image-generation and understanding benchmarks while training on 8×A100 GPUs.The paper presents this as evidence of the framework’s potential for multimodal reasoning.Source: §1 · p. 2

2 PRELIMINARY: THE UNIFED MODEL - BAGEL

BAGEL is the unified model underlying Uni-CoT, combining image-understanding and image-generation experts through shared multimodal self-attention. Its visual token load makes step-wise multimodal reasoning computationally expensive and challenging to optimize.

2 PRELIMINARY: THE UNIFED MODEL - BAGEL

  • BAGEL uses separate image-understanding and image-generation experts linked by unified self-attention over shared multimodal token sequences.The experts share an architecture but have different parameters.Source: §2 · p. 3
  • The understanding pathway uses ViT to encode images into 4,900 tokens, while the generation pathway uses VAE to produce 4,096 tokens.The ViT maps images to a 70×70 latent grid; the VAE uses a 64 × 64 latent grid.Source: §2 · p. 3
  • Expert routing activates ViT for grounded image understanding and VAE-based generation for image creation or editing.The understanding expert fuses image tokens with text for next-token prediction; the generation expert predicts VAE latent features through Rectified Flow.Source: §2 · p. 3
  • BAGEL’s dual-pathway design supports fine-grained visual generation and semantic grounding within one model.The framework is described as handling diverse tasks under a single model.Source: §2 · p. 3
  • Each multimodal reasoning step requires nearly 9,000 visual tokens in addition to about 300 text tokens, making training and inference prohibitively expensive.The passage also identifies optimization challenges that can hinder convergence and limit generalization, especially for long compositional reasoning chains.Source: §2 · p. 3

3 METHOD

Uni-CoT replaces monolithic multimodal CoT with a hierarchical macro–micro trajectory: planning and synthesis surround localized subtask execution. Its attention restrictions and MDP-style refinement target the high cost of long multimodal histories.

3.1 OVERVIEW OF PROBLEM AND OUR SOLUTION

  • The framework decomposes multimodal reasoning into macro planning, localized micro execution, and macro summarization, reducing full-history quadratic cost to O(T) with Markovian transitions.A multimodal state can contain approximately 10,000 tokens; partitioning the trajectory and conditioning each micro transition on the previous state and instruction reduces computation.Source: §3.1 · p. 4

3.2 MACRO-LEVEL REASONING: PLANNING AND SUMMARIZATION

  • Macro reasoning plans sequential or concurrent subgoals and synthesizes their results, while its attention mask exposes only the input, subgoals, intermediate results, and final answer.In the sequential form, each subgoal follows completion of the previous one; in the parallel form, subgoals execute concurrently.Source: §3.2 · pp. 4–5

3.3 MICRO-LEVEL REASONING: MDP-GUIDED SELF-REFLECTION

  • Micro CoT begins with a subtask attempt, evaluates it, and revises inconsistencies or cross-modal mismatches in a closed loop, conditioning each new state only on the previous state and fixed subgoal.This MDP formulation reduces the trajectory cost toward linear complexity and supports scaling self-reflective reasoning to long multimodal trajectories.Source: §3.3 · p. 6

3.4 TRAINING PARADIGM

  • Training supervises macro planning and synthesis with joint text-image loss, while micro learning adds text and image action generation, next-state prediction, and reward estimation.The macro joint loss applies cross-entropy to text generation and mean squared error to image generation.Source: §3.4 · p. 6

4 EXPERIMENTS

The experiments evaluate Uni-CoT on image-generation and image-understanding benchmarks and compare its hierarchical variant with Uni-CoT Raw for inference complexity and training dynamics. Reported results favor Uni-CoT on reasoning-driven generation and structured understanding, with lower inference cost and faster convergence than Uni-CoT Raw.

4.1 IMPLEMENTATION DETAILS

  • The dataset contains approximately 30K reasoning samples, including 11K macro-planning and 20K micro-refinement pairs; fine-tuning used 31K samples and 8 A100 GPUs for one week.Source: §4.1 · p. 7

4.2 EXPERIMENTAL SETUP

  • Evaluation spans GenEval and WISE for image generation, plus MME, MMMU, MathVista, MMBench, and Jigsaw-R1 for image understanding.Source: §4.2 · p. 7

4.3 RESULTS FOR IMAGE GENERATION

  • Uni-CoT surpasses Bagel on GenEval, mainly through macro decomposition, and achieves state-of-the-art results across WISE domains, where self-reflection corrects initial errors.Figure 3 also shows coherent planning that turns an unnatural prompt into natural intermediate prompts.Source: §4.3 · p. 7, Table 1 · p. 8, Table 2 · p. 8

4.4 RESULTS FOR IMAGE UNDERSTANDING

  • Uni-CoT scores slightly above Bagel on MME, MathVista, and MMBench, and substantially outperforms compared open-source models on Jigsaw-R1, particularly in perception-heavy tasks.Figure 3 illustrates parallel decomposition for a Jigsaw-R1 case: candidate options are examined concurrently and then integrated.Source: §4.4 · pp. 8–9, Table 3 · p. 9

4.5 COMPLEXITY ANALYSIS

  • 2.24×, 2.95×, and 11.26× fewer average token interactions at 2, 3, and 10 reasoning steps, respectively, accompany near-linear scaling versus Uni-CoT Raw’s almost quadratic growth.Uni-CoT reaches comparable loss in 6,000 steps versus Raw’s 12,000; after 12,000 steps on approximately 10,000 samples, WISE scores are 0.73 vs. 0.70.Source: Figure 4 · p. 9, §4.5 · p. 9, Figure 5 · p. 9

5 RELATED WORK

Prior multimodal CoT methods struggle to model visually grounded state transitions coherently. Uni-CoT’s related-work positioning emphasizes reducing multimodal reasoning complexity while preserving reasoning, rather than only changing attention computation.

Multi-modal CoT with Visual Reasoning

  • RL-based approaches strengthen textual reasoning but struggle in visually intensive domains, motivating methods that align reasoning with visual content.Prior work also extends CoT to multimodal language models for vision-language tasks.Source: §5 · p. 9
  • Programmatic visual operations capture local changes but miss global structural transformations, while loosely coupled generators or reward models can fragment reasoning flows.The cited tasks include navigation and rearrangement.Source: §5 · p. 10

Hierarchical Reasoning

  • Unlike prior hierarchical methods focused on reasoning quality, controllability, or planning efficiency, this work uses hierarchy to reduce multimodal CoT training and inference complexity while preserving reasoning ability.Source: §5 · p. 10

Efficient Attention

  • Unlike efficient-attention methods that reduce transformer attention cost, this formulation targets reasoning complexity while preserving CoT integrity.Source: §5 · p. 10

6 CONCLUSION

Uni-CoT presents a unified multimodal CoT framework with hierarchical reasoning, reporting strong performance across generation and understanding tasks. The authors identify broader real-world applications requiring fine-grained visual consistency or perfect text–image alignment as an open challenge.

6 CONCLUSION

  • Uni-CoT reports superiority in performance and interpretability across reasoning-driven image generation and editing benchmarks.Source: §6 · p. 10
  • The framework combines vision and language within one model and uses two-level hierarchical reasoning to reduce computational complexity and improve reasoning efficiency.Source: §6 · p. 10
  • Broader real-world applications requiring fine-grained visual consistency or perfectly aligned text and images remain an open challenge.The authors identify improved trajectory modeling and stronger visual–textual alignment as future directions.Source: §6 · p. 10

8 REPRODUCIBILITY STATEMENT

The paper describes reproducibility materials, discloses limited LLM assistance in writing, and specifies distinct attention access for macro and micro reasoning.

8 REPRODUCIBILITY STATEMENT

  • Architecture, objectives, datasets, preprocessing, evaluation protocols, most hyperparameters, and inference scripts are documented or released to support independent verification.Source: §8 · p. 10

A THE USE OF LARGE LANGUAGE MODELS (LLMS)

  • An LLM was used only to polish language, improve clarity, and organize sections, not for core research activities.The stated exclusions include problem formulation, experimental design, data analysis, and result interpretation.Source: A THE USE OF LARGE LANGUAGE MODELS… · p. 17

B DETAILS OF ATTENTION MASK

  • Macro attention exposes system prompts, high-level plans, and final subtask outcomes, whereas micro attention exposes only the immediately preceding state and current subtask instruction.Macro attention masks intermediate execution traces; the attention schemes are shown in Figure S1.Source: B DETAILS OF ATTENTION MASK · pp. 17–18, Figure S1 · p. 17

C.1 TRAINING PARADIGM

Uni-CoT separates global planning and synthesis from local subtask execution and reflection, with training objectives tailored to each level. Its training setup combines multimodal supervision, a curated reasoning dataset, and efficient implementation practices.

C.1 TRAINING PARADIGM

  • Training separates Macro-Level CoT for global planning and final synthesis from Micro-Level CoT for subtask execution and self-reflection.Source: §C.1 · p. 18
  • The macro planner treats task decomposition as multimodal understanding and learns a structured textual plan with cross-entropy loss.Source: §C.1 · p. 18
  • The macro summarizer integrates the global plan and intermediate multimodal outputs into a final solution, producing text for understanding or text and images for generation.Table S1 summarizes the macro-level training tasks.Source: §C.1 · p. 18, Table S1 · p. 19
  • For image generation, Rectified Flow linearly interpolates clean latent x0 and Gaussian noise x1 into xt, and the model learns to predict the velocity field with mean-squared error.The model’s predicted velocity is conditioned on noisy latent xt and context c.Source: §C.1 · p. 18
  • Subtask completion uses a joint loss combining cross-entropy for text and mean-squared error for image predictions; the text loss compares predicted probabilities with ground-truth tokens.The joint loss uses λ to balance its two losses.Source: §C.1 · p. 18
  • Micro-level self-reflection uses an MDP framework with four auxiliary objectives: text-action generation, image-action generation, next-state prediction, and reward estimation.The objectives model editing instructions, visual modifications, modified-image status, and feedback on intermediate-step quality.Source: §C.1 · pp. 18–19

C.2 TRAINING HYPER PARAMETERS

  • Training uses Adam with a constant learning rate of 2e-5, a 200-step linear warmup, and packed sequences targeting 32,768 tokens; experiments run on 8 NVIDIA A100 GPUs.The implementation also uses FlashAttention, FSDP, and mixed precision.Source: §C.2 · p. 19
  • The data-curation pipeline creates macro planning traces by enhancing prompts and decomposing them into subtasks, then creates micro trajectories through image generation and iterative evaluation and editing.Source: Figure S2 · p. 20

D DETAILS OF DATASET.

The appendix describes scalable data pipelines and interleaved supervision for macro- and micro-level reasoning. It also identifies limited scale and synthetic composition as constraints on real-world generalization, motivating a hybrid data-scaling strategy.

Macro-level reasoning data

  • Macro-level data pairs enhanced prompts, 2–3 decomposed subtasks, execution and refinement with interleaved text and intermediate images.The process uses models such as GPT-4o or Qwen-plus for decomposition and GPT-4o for evaluation and refinement.Source: D DETAILS OF DATASET · p. 20

Data curation process

  • The pipelines draw on robotics/VLA datasets, scientific videos, and movies to connect visual transitions with textual guidance.These sources cover spatial and physical interactions, scientific or procedural phenomena, and daily-life causal or behavioral transitions.Source: Figure S3 · p. 21, D DETAILS OF DATASET · p. 22

Micro-level reasoning data

  • Micro-level data records repeated image evaluation and editing cycles as interleaved text-image supervision for self-reflection.VLMs assess image quality and instruction-following, then image-editing models apply refinement instructions.Source: D DETAILS OF DATASET · pp. 20–21

Data information

  • Approximately 11K long-form trajectories yield 11K macro-level and 20K micro-level samples, alongside additional generation, editing, and understanding data.The additional sources include 114K text-to-image samples, 68K Echo4o samples, 46K ShareGPT4o-image samples, and approximately 100K LLaVA-OV samples.Source: D DETAILS OF DATASET · p. 21

Future Works for Data Collection

  • The current dataset’s relatively small scale and predominantly synthetic composition limit generalization to complex real-world scenarios.The authors propose combining high-quality real-world data with structured synthetic data to expand diversity and regulate reasoning complexity cost-effectively.Source: D DETAILS OF DATASET · pp. 21–22

Data curation process

  • Most images require only lightweight post-processing, so GPU use and API budgets mainly go toward formulating textual guidance between visual states.The authors estimate that the pipelines could scale to 100K–500K samples.Source: D DETAILS OF DATASET · p. 22

E.1 RESULTS FOR IMAGE EDITING BENCHMARK

The image-editing evaluation covers general editing, reasoning-informed editing, and diagnostic task categories. Uni-CoT is reported to surpass Gemini 2.0 overall on KRIS and perform comparably to it on RISE.

Experiment Setup

  • The evaluation uses GEdit-Bench for general real-world editing, RISE for reasoning-informed dimensions, and KRIS for diagnostic knowledge categories.RISE covers temporal, causal, spatial, and logical dimensions; KRIS separates factual, conceptual, and procedural knowledge.Source: §E.1 · p. 22

Quantitative Results

  • On GEdit-Bench, Uni-CoT achieves competitive performance against other approaches, with higher scores preferred.The benchmark covers diverse real-world editing tasks.Source: §E.1 · p. 22
  • On KRIS, Uni-CoT surpasses all open-source baselines across perception, conceptual, and procedural categories, and exceeds Gemini 2.0 in overall score.On RISE, it performs comparably to Gemini 2.0 on overall results across four reasoning categories and on instruction reasoning, appearance consistency, and visual plausibility.Source: §E.1 · p. 22

E.2 NANO-BANANA LIKE GENERATION

The section turns abstract geospatial inputs into landscape images through explicit reasoning steps. In a geography test, the authors report strong generation with reasoning decomposition, while direct fine-tuning is unstable.

Reasoning-driven generation

  • Nano-Banana is described as generating photorealistic landscapes from abstract geospatial inputs, while the paper seeks to make such mappings explicit through interpretable subtasks.The cited examples include satellite imagery and isohypse line maps.Source: §E.2 · pp. 22–23
  • The map-to-landscape process uses three sequential steps: convert 2D contours to 3D terrain, select a 3D camera view, then render a photorealistic landscape.The decomposition is described as mirroring how geographers or cartographers reconstruct terrain and imagine a viewpoint.Source: §E.2 · p. 23

Geography reasoning evaluation

  • On 3K GeoPose3K geography samples, refinement for around 12,000 steps produced strong generation, while direct fine-tuning on raw map-image pairs yielded unstable training and incoherent or distorted outputs.The dataset is aligned to the three-step decomposition.Source: §E.2 · p. 23

Training dynamics

  • Explicit multimodal CoT converges smoothly within ∼2k iterations, whereas direct fine-tuning shows no stable convergence even after 12k iterations.The comparison uses mean squared error loss curves and concerns mapping 2D contour modalities to realistic 3D landscape imagery.Source: §E.2 · p. 23

E.3 ABLATION STUDY

The ablations report faster convergence and better WISE results than direct long-chain fine-tuning, assess macro- and micro-CoT across benchmarks, and test the framework on BLIP3o. Additional figures show editing, geography, and reasoning-trajectory examples.

E.3 ABLATION STUDY

  • Macro-CoT matters more on GenEval, while micro-CoT contributes more on WISE, a contrast the authors relate to differences in prompt style and complexity.GenEval prompts are programmatically generated and often unlike natural phrasing; WISE prompts are abstract but relatively simple.Source: §E.3 · p. 24, Table S6 · p. 24
  • The hierarchical training method reaches comparable loss in 6,000 steps versus 12,000 for Uni-CoT Raw and outperforms it across WISE metrics.The authors note that limited-compute results in Table S7 do not represent the full model’s final performance.Source: §E.3 · p. 24, Table S7 · p. 24
  • With BLIP3o, fine-tuning with Uni-CoT self-reflection improves performance by approximately 0.8 points across all metrics over fine-tuning without those signals.The setup uses BLIP3o-NEXT-edit-VAE fine-tuned for 100,000 steps on 8×A100 GPUs; direct dataset fine-tuning also improves over the original model.Source: §E.3 · p. 25, Table S8 · p. 25

F MORE VISUALIZATION

  • Supplementary figures show Nano-Banana-like generation, reliable editing, geography reasoning, and macro-, micro-, hybrid-CoT, and BLIP3o reasoning examples.Figure S6 specifically depicts textual editing instructions followed by modification of the current visual state.Source: F MORE VISUALIZATION · p. 25, Figure S4 · p. 26, Figure S6 · p. 27
Loading 2508.05606v3…