PaperlayerView exact revision on arXiv

Source-linked AI summary

EndoCoT: Scaling Endogenous Chain-of-Thought Reasoning in Diffusion Models

Xuanlang Dai, Yujie Zhou, Long Xing, Jiazi Bu, Xilin Wei, Yuhong Liu, Beichen Zhang, Kai Chen, Yuhang Zang

arXiv:2603.12252v4cs.CVcs.CL

TL;DR

Diffusion models using MLLMs as text encoders still struggle with step-by-step reasoning tasks, while static one-pass conditioning leaves MLLM reasoning potential underused. EndoCoT iteratively refines latent thought states, grounds the final state in text, and reports 25-40% gains over prior work on complex visual reasoning benchmarks.

  • Problem

    Diffusion models with MLLM text encoders struggle with tasks requiring step-by-step logical reasoning, and static conditioning does not exploit MLLMs’ dynamic reasoning potential.

    Source: §1 · p. 1

  • Method

    EndoCoT iteratively refines MLLM latent thought states, aligns them with DiT denoising, and grounds the final reasoning state in textual supervision.

    Source: §1 · p. 2

  • Results

    25-40% gains over prior work on complex visual reasoning benchmarks; EndoCoT establishes a new state of the art across all evaluated visual reasoning benchmarks under task-specific training.

    Source: §1 · p. 3, §5.1 · p. 9

  • Takeaways & Limitations

    Experiments support that the reasoning capability arises from coupling the multimodal text encoder and DiT backbone, rather than from standard denoising alone.

    Source: §6 · p. 11

    EndoCoT requires manual tuning of the reasoning-step count and high-quality datasets with explicit intermediate supervision.

    Source: §6 · p. 11

Abstract

from arXiv · show

Recently, Multimodal Large Language Models (MLLMs) have been widely integrated into diffusion frameworks primarily as text encoders to tackle complex tasks such as spatial reasoning. However, this paradigm suffers from two critical limitations: (i) MLLMs text encoder exhibits insufficient reasoning depth. Single-step encoding fails to activate the Chain-of-Thought process, which is essential for MLLMs to provide accurate guidance for complex tasks. (ii) The guidance remains invariant during the decoding process. Invariant guidance during decoding prevents DiT from progressively decomposing complex instructions into actionable denoising steps, even with correct MLLM encodings. To this end, we propose Endogenous Chain-of-Thought (EndoCoT), a novel framework that first activates MLLMs' reasoning potential by iteratively refining latent thought states through an iterative thought guidance module, and then bridges these states to the DiT's denoising process. Second, a terminal thought grounding module is applied to ensure the reasoning trajectory remains grounded in textual supervision by aligning the final state with ground-truth answers. With these two components, the MLLM text encoder delivers meticulously reasoned guidance, enabling the DiT to execute it progressively and ultimately solve complex tasks in a step-by-step manner. Extensive evaluations across diverse benchmarks (e.g., Maze, TSP, VSP, and Sudoku) achieve an average accuracy of 92.1%, outperforming the strongest baseline by 8.3 percentage points. The code and dataset are publicly available at https://internlm.github.io/EndoCoT/.

1. Introduction

Diffusion models with MLLM text encoders struggle with step-by-step reasoning and generalization because their guidance is static. EndoCoT iteratively refines latent thoughts and grounds the final state, outperforming baselines across four reasoning tasks.

1. Introduction

  • Prior methods commit to solutions in early denoising steps, and DiffThinker [11] produces fragile solutions that fail on novel domains.The paper attributes the fragility to forced alignment that yields pattern matching rather than robust iterative problem-solving.Source: §1 · p. 1, Figure 1 · p. 2
  • MLLMs cannot encode all logical constraints in one pass, while DiTs cannot maintain alignment with complex constraints using static embeddings.The introduction identifies these as limited single-step reasoning and static-guidance failure, respectively.Source: §1 · pp. 1–2
  • 90% on Maze-32 and 95% on Sudoku-35, outperforming the strongest baseline by 25% and 40%, respectively; Figure 1(a) shows gains across all benchmarks.The four tasks are Maze, TSP, Visual Spatial Planning, and Sudoku.Source: Figure 1 · p. 2, §1 · p. 2
  • EndoCoT iteratively updates MLLM latent states to create CoT-like reasoning linked to DiT denoising, then grounds the final state in ground-truth answers.Terminal grounding is intended to prevent cumulative drift in the terminal output.Source: §1 · p. 2
  • The paper presents layer-wise sensitivity and attention-entropy analyses to localize reasoning and identify bottlenecks, and reports inference-time scaling and clearer, more controllable image-editing trajectories.It describes EndoCoT as the first diffusion framework to enable genuine chain-of-thought reasoning through iterative latent-state refinement rather than single-pass solution precomputation.Source: §1 · pp. 2–3

2. Related Work

Prior work brings reasoning into multimodal and diffusion systems, but the paper identifies genuine iterative chain-of-thought within diffusion as unresolved. It builds on latent-space reasoning to pursue iterative reasoning inside the diffusion process.

Reasoning in Multimodal Large Language Models

  • CoT and test-time scaling [4] [36] [39] [43] extend from autoregressive LLMs to multimodal settings [1] [34], including image and video reasoning [6] [20] [23] [28] [30] [37] and interleaved text-visual latent generation [38].These approaches use visual content as evidence for multistep reasoning.Source: §2 · p. 3
  • Fixed-length text encoders [21] [22] limit multistep reasoning in diffusion, while VLM encoders [32] provide richer semantic representations for reasoning-conditioned generation.Source: §2 · p. 3

Reasoning in Diffusion Model

  • Existing diffusion frameworks still lack genuine iterative chain-of-thought within the diffusion process; this work addresses the gap through iterative latent-state refinement.The paper says this enables robust logical planning and spatial grounding.Source: §2 · p. 3
  • Diffusion approaches inject textual reasoning traces [15] [16], use video priors for state transitions [33] [40], or explore MMDiT reasoning [11], but remain short of iterative CoT within diffusion.Video-prior methods rely on temporal smoothness that may not transfer to discrete logical tasks; visual generation has also used next-token autoregression [7] [9].Source: §2 · p. 3

Latent Reasoning

  • Latent reasoning supports multistep inference in continuous spaces [10] [27] [41], compresses reasoning chains, and enables tree-structured exploration [24]; this work trains diffusion models end-to-end with latent tokens for test-time scaling.Source: §2 · p. 3

3. Analysis of Reasoning in Diffusion Models

Analysis of Qwen-Image-Edit-2511 [32] identifies two limits: the MLLM cannot encode all complex constraints in one pass, and static text conditioning fails to maintain DiT grounding in dense scenes. These findings motivate jointly fine-tuning the MLLM and DiT and using iterative latent reasoning with evolving conditioning.

Layer-Wise Sensitivity Analysis

  • Sensitivity peaks in the Vision Encoder and at the LLM–DiT boundary, indicating that MLLM reasoning is central and motivating joint fine-tuning of the MLLM and DiT.The analysis measures activation-response magnitude across the Vision Encoder, Language Model, and Diffusion Transformer blocks.Source: §3 · p. 3, Figure 2 · p. 4

Limited Single-Step Reasoning

  • A single MLLM pass leaves complex constraints unenforced: in 8×8 mazes, the DiT grounds trajectories but generated paths can pass through walls.The finding motivates iterative refinement of the MLLM’s understanding in latent space.Source: §3 · p. 4, Figure 2 · p. 4

Static-Guidance Failure in Complex Dynamic DiT Decoding

  • In 32×32 mazes, cross-attention entropy becomes globally high and diffuse, showing that static text embeddings cannot keep DiT spatial features anchored to logical tokens.The conditioning signal therefore needs dynamic updates throughout the reasoning process.Source: §3 · p. 4, Figure 2 · p. 4

4. Methods

EndoCoT enables diffusion models to reason step by step during generation by iteratively refining MLLM thought states, grounding the final state in textual supervision, and training progressively. Its inference process recursively updates latent thoughts to produce the final result.

4.2. EndoCoT

  • EndoCoT combines iterative latent-thought refinement, explicit textual grounding, and progressive training to enable diffusion models to perform autonomous chain-of-thought reasoning.The framework aims to build reasoning capability first and then refine output quality.Source: §4 · p. 5, §4.2 · p. 6

4.1. Preliminary

  • Flow matching constructs a linear probability path between clean data and Gaussian noise, with a neural network approximating the path’s ground-truth vector field.The conditioning variable may be included in the neural network input.Source: §4.1 · p. 5

4.2.1. Iterative Thought Guidance.

  • At each reasoning step, EndoCoT recursively updates a latent thought state and conditions a complete diffusion denoising trajectory on it to generate an intermediate visual output.Intermediate targets can be obtained by sequential ground-truth decomposition, and the MLLM and DiT are jointly adapted with LoRA.Source: §4.2.1 · pp. 6–7

4.2.2. Terminal Thought Grounding.

  • To address the modality gap and underuse of textual ground truth, terminal thought grounding aligns the final latent reasoning state with embedded ground-truth reasoning steps using an L2 semantic loss.The alignment loss is active only at the final reasoning step; the paper sets λ_align = 1 and reports that the term helps prevent latent-state drift.Source: §4.2.2 · p. 7

4.2.3. Progressive Training Strategy

  • Progressive training first supervises every reasoning step, then preserves intermediate steps in the forward pass while computing gradients only for the final output.The second stage uses short-cycle fine-tuning with limited iterations to help preserve the learned reasoning chains.Source: §4.2.3 · pp. 7–8

4.2.4. Inference Process

  • At inference, specifying the number of reasoning steps lets EndoCoT recursively update latent thought states to generate the final result without decoding intermediate visual states.Reasoning steps are distinct from the diffusion model’s internal denoising timesteps.Source: §4.2.1 · p. 6, §4.2.4 · p. 8

5. Experiments

EndoCoT is evaluated on visual reasoning and image editing tasks, with experiments examining benchmark performance, training settings, design choices, inference scaling, and progressive editing.

5.1. Results on Visual Reasoning Tasks

  • The evaluation covers five multi-step reasoning benchmarks under zero-shot, task-specific, and unified settings, using Qwen-Image-Edit-2511 with LoRA fine-tuning.The setup uses LoRA rank 32, a 1 × 10^-4 learning rate, and five training epochs; generated task data are rendered at 512 × 512 resolution.Source: §5.1 · pp. 8–9, §5 · p. 8
  • 95% Sudoku-35 and 90% Maze-32 accuracy exceed DiffThinker’s 55% and 65%; VSP-Super-32 reaches 85% while generative baselines fail.Under task-specific training, the model establishes a new state of the art across evaluated visual reasoning benchmarks and maintains near-perfect scores on standard VSP levels.Source: Table 1 · p. 9, §5.1 · p. 9
  • Unified training remains competitive with DiffThinker while performing slightly below task-specific training.One model is trained on Maze, TSP, Sudoku, and VSP together and evaluated across all tasks simultaneously.Source: §5.1 · p. 9, Table 1 · p. 9

5.2. Ablation Study and Analysis

  • On Maze-32, removing semantic loss drops accuracy from 90% to 14% and path repetition rate from 98.13% to 67.24%.The auxiliary loss aligns continuous latent tokens with ground-truth reasoning steps, acting as a regularizer against implicit token drift.Source: §5.2 · pp. 9–10, Table 2 · p. 9
  • Explicit autoregressive reasoning reaches 34% accuracy on Maze-8 but 0% on Maze-32, where discrete-token error accumulation and mode collapse produce repetitive loops.The variant replaces continuous implicit tokens with generated discrete intermediate reasoning steps.Source: Table 2 · p. 9, §5.2 · p. 10
  • Increasing implicit reasoning budgets improves accuracy and path repetition, while higher resolution reduces relative computational cost without repeating expensive DiT denoising.Inference time trades off against accuracy as token budgets increase; latency remains stable and predictable across resolutions.Source: §5.2 · p. 10, Table 4 · p. 10, Figure 6 · p. 10
  • On Maze-32, DiT-only accuracy is 18%, while MLLM-only fails, showing that reliable spatial reasoning requires both cognitive planning and visual grounding.The MLLM-only baseline produces erratic, wandering paths that can become trapped in dead ends.Source: §5.2 · p. 11, Figure 8 · p. 11

Results on Image Editing Tasks

  • For image editing, the model progressively adds a stone lantern and deer, transforms the deer into a sheep, and uses reasoning step τ to control editing operations.Given step-by-step instructions, it iteratively plans and executes each modification during generation.Source: Results on Image Editing Tasks · p. 11, Figure 9 · p. 11

6. Conclusion · Appendix · A. Overview

EndoCoT combines iterative latent-thought refinement with textual grounding to support logical planning and visual generation. The conclusion attributes its reasoning capability to coupling the multimodal text encoder with the DiT backbone, while noting supervision and tuning requirements.

6. Conclusion

  • EndoCoT enables diffusion models to perform endogenous Chain-of-Thought reasoning by iteratively refining latent thought states and grounding the final output in textual supervision.The paper presents this as bridging high-level logical planning and precise visual generation.Source: §6 · p. 11
  • Experiments across diverse reasoning tasks indicate that this capability arises from synergistic coupling of the multimodal text encoder and DiT backbone, rather than standard denoising alone.Source: §6 · p. 11
  • The approach currently requires manual tuning to select the optimal number of reasoning steps.Source: §6 · p. 11
  • EndoCoT also relies on high-quality datasets with explicit intermediate supervision.Source: §6 · p. 11
  • Future work will explore adaptive mechanisms for automatic reasoning-depth control.Source: §6 · p. 11
  • Future work will also extend the framework to broader general-purpose tasks.Source: §6 · p. 11

Appendix · A. Overview

  • The supplementary material covers dataset statistics, task prompts and additional results, additional ablations, EndoCoT insights and observations, and editing datasets.Source: A. Overview · p. 15

B. Data Statistics … D. More Ablations

The appendices describe a 182.4K-instance dataset spanning Maze, TSP, Sudoku, and VSP, detail how tasks and intermediate reasoning steps are generated, and provide evaluation prompts and qualitative comparisons.

B. Data Statistics

  • 182.4K instances across Maze, TSP, Sudoku, and VSP comprise the constructed dataset, including intermediate reasoning steps.Table 5 summarizes the sample size for each task.Source: Table 5 · p. 15, B. Data Statistics · p. 15

B.1. Algorithmic Formulation

  • Maze paths are generated with DFS-carved connected, loop-free grids and BFS shortest paths, while TSP tours use Held-Karp dynamic programming on randomly placed cities.The resulting paths or tours are recorded sequentially as intermediate steps.Source: §B.1 · p. 15
  • Sudoku boards are constructed by backtracking, clues are removed while preserving a unique solution, and reasoning steps reverse that process by filling empty cells.The boards satisfy standard row, column, and block constraints.Source: §B.1 · p. 15
  • VSP uses a FrozenLake grid and Dijkstra’s algorithm to compute a safe path, recording state transitions to capture perception and hazard avoidance.The grid contains safe tiles, hazard holes, a start, and a goal.Source: §B.1 · p. 15

B.2. Pseudocode for Dataset Generation

  • The dataset-generation pseudocode summarizes DFS/BFS for mazes, Held-Karp for TSP, Dijkstra for VSP, and clue removal from a valid Latin square for Sudoku.The procedures are shown in Figure 14.Source: §B.2 · p. 15

C.1. Evaluation Prompts

  • The evaluation prompts ask models to fill Sudoku cells from top left to bottom right and to draw a stepwise path from start to goal while avoiding holes.The path prompt specifies D, U, R, and L as direction labels and a green dot for the current endpoint.Source: §C.1 · p. 16

C.2. Qualitative Examples

  • Additional qualitative examples compare EndoCoT with ThinkGen [15] and ChronoEdit [33] across reasoning tasks.The examples appear in Figures 15–18.Source: §C.2 · p. 16

D.1. Effect of Two-Stage Training

On Maze, the two-stage training analysis indicates that Terminal Consolidation enhances the model’s awareness of the final stage.

D.1. Effect of Two-Stage Training

  • On Maze, Table 6 analyzes single-stage versus multi-stage training, and the analysis indicates Terminal Consolidation enhances awareness of the final stage.The supplied passage does not report specific values or a quantitative comparison.Source: Table 6 · p. 16, §D.1 · p. 16

D.2. Impact of Terminal Consolidation Training Duration … E.1. Training Settings

EndoCoT’s Sudoku predictions generalize across resolutions and font styles, while excessive Terminal Consolidation training makes intermediate reasoning steps sparse and shifts inference toward the target state. Training hyperparameters and LoRA configurations are summarized in Table 7.

E.1. Training Settings

D.2. Impact of Terminal Consolidation Training Duration

  • With a constant number of reasoning steps τ, checkpoints with more Stage 2 training steps converge more toward the target state at inference.Source: §D.2 · p. 16, Figure 10 · p. 17

D.3. Generalization under OOD Settings

  • EndoCoT generalizes across Sudoku resolutions and font styles, outperforming baseline methods in the settings described.The passage reports stronger generalization than baselines but provides no numerical comparison.Source: §D.3 · p. 16
  • For high-resolution Sudoku inputs, EndoCoT consistently predicts the correct 9x9 grid at the center.Source: §D.3 · p. 16
  • For low-resolution inputs, EndoCoT recovers partially cropped correct answers, whereas DiffThinker fails to generate valid solutions.Source: §D.3 · p. 16

E.2. Further Analysis

Gradient comparisons between reasoning and non-reasoning datasets show that reasoning-related signals are concentrated primarily in MLLM layers.

E.2. Further Analysis

  • Gradient magnitudes differ most consistently in MLLM layers when comparing reasoning and non-reasoning datasets across the MLLM and DiT.The analysis supports the conclusion that reasoning capabilities are primarily mediated by the MLLM.Source: §E.2 · p. 17, Figure 12 · p. 18

E.3. Experimental Observations · F. Image Editing Details

The experiments find that fine-grained step supervision, sparse semantic supervision, and latent-token counts matched to task complexity matter for reasoning performance. The image-editing details describe a multi-step dataset and procedures for generating puzzle trajectories and images.

E.3. Experimental Observations

  • Fine-grained intermediate-step supervision is necessary because inference uses fewer reasoning steps than training supervision.The paper says this finer supervision helps the model learn latent reasoning processes.Source: §E.3 · p. 18
  • Applying semantic loss at intermediate steps disrupts DiT–MLLM synergy and performs worse than EndoCoT’s sparse semantic supervision strategy.Source: §E.3 · p. 18
  • The number of latent tokens required for a puzzle depends on both its algorithmic and visual complexity.Source: §E.3 · p. 18

F. Image Editing Details

Related papers

Loading 2603.12252v4…