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 systems using MLLMs as static text encoders struggle with step-by-step visual reasoning and generalization to novel configurations. EndoCoT iteratively refines latent thought states, connects them to DiT denoising, and grounds the terminal state in textual supervision, achieving 25-40% gains over prior work across complex visual reasoning benchmarks.

  • Problem

    MLLMs used as static conditional encoders do not perform genuine reasoning during generation, limiting step-by-step solutions and robustness on novel domains.

  • Method

    EndoCoT iteratively updates MLLM latent states, establishes correspondence with DiT denoising, and grounds the final reasoning state in ground-truth textual supervision.

  • Results

    25-40% gains over prior work are reported across complex visual reasoning benchmarks, with 95% Sudoku-35 and 90% Maze-32 accuracy under task-specific training.

  • Takeaways & Limitations

    EndoCoT bridges high-level logical planning and precise visual generation through synergistic coupling of the multimodal text encoder and DiT backbone.

  • Takeaways & Limitations

    The approach requires manual tuning of reasoning steps and high-quality datasets with explicit intermediate supervision.

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 systems that use MLLMs as static text encoders struggle with sequential constraints and novel configurations because single-pass reasoning is insufficient and guidance does not evolve during denoising. EndoCoT addresses these bottlenecks by iteratively refining latent thought states, grounding the terminal state in textual supervision, and linking reasoning to DiT denoising.

  • Diffusion models with MLLM text encoders struggle on mazes, TSP routes, and Sudoku requiring sequential constraints and generalization to novel configurations.
  • Single-step MLLM encoding cannot represent all necessary logical constraints, while static embeddings prevent DiTs from maintaining alignment during complex denoising.
  • EndoCoT iteratively updates MLLM latent states, couples them to DiT denoising, and grounds the final reasoning state with ground-truth answers.
  • EndoCoT evaluates Maze, TSP, VSP, and Sudoku tasks and outperforms baselines across all benchmarks while maintaining accuracy as task complexity increases.
  • 25-40% gains over prior work are reported on complex visual reasoning benchmarks, alongside controllable inference-time scaling and clearer transformation trajectories in image editing.

2. Related Work

Prior work extends chain-of-thought and richer multimodal representations to diffusion, but diffusion systems commonly retain fixed or decoupled conditioning that limits multi-step reasoning. Latent-space reasoning motivates training diffusion models with latent tokens for test-time scaling within the diffusion process.

  • Chain-of-thought and test-time scaling strategies have expanded from autoregressive LLMs to multimodal settings, while diffusion models remain constrained by fixed-length text encoders.
  • Reasoning-aware diffusion methods inject textual traces into conditioning, but often treat MMDiT as a conditional decoder and MLLMs mainly as prompt enhancers.
  • Latent-space reasoning compresses long chains and supports tree-structured exploration, motivating end-to-end diffusion training with latent tokens for test-time scaling.

3. Analysis of Reasoning in Diffusion Models

Analysis identifies that reasoning is concentrated in the MLLM and its junction with DiT, but single-pass encoding and static conditioning fail on complex visual reasoning. These findings motivate iterative latent reasoning and progressive training that updates latent representations while preserving learned dynamics.

  • Layer-Wise Sensitivity Analysis: Layer-wise sensitivity peaks in the Vision Encoder and at the junction between terminal LLM layers and initial DiT layers, motivating joint MLLM–DiT fine-tuning.
  • Limited Single-Step Reasoning: On lower-complexity 8×8 mazes, DiT grounds trajectories spatially but paths violate physical constraints, showing that one MLLM forward pass is insufficient.
  • Static-Guidance Failure in Complex Dynamic DiT Decoding: In high-complexity 32×32 mazes, cross-attention entropy becomes diffuse, weakening alignment between spatial patches and logical text tokens during denoising.
  • Summary: The analysis concludes that MLLMs need iterative refinement and DiTs need dynamically evolving conditioning to preserve alignment with complex logical constraints.
  • EndoCoT Overview: EndoCoT's progressive training fits intermediate and final states before optimizing the terminal state, while inference iteratively updates latent representations.

4. Methods

EndoCoT enables diffusion models to reason through complex visual tasks by iteratively refining MLLM latent thought states, conditioning DiT denoising on each state, and grounding the final state with textual supervision.

  • Core framework: EndoCoT iteratively refines latent thought states in the MLLM and connects them to the DiT denoising process.The framework includes iterative thought guidance, terminal thought grounding, and progressive training.
  • Iterative Thought Guidance: Each reasoning step updates the conditioning signal and visual output, rather than relying on static text embeddings throughout generation.The model performs T reasoning steps, with each step refining both the thought state and intermediate visual result.
  • Iterative Thought Guidance: The MLLM recursively updates h_τ from fixed prompt and image prefix embeddings, using the previous thought state as input to the next step.The previous latent state bypasses the discrete embedding lookup and directly enters the MLLM input.
  • Conditional Flow Generation: At every reasoning step, the current thought state conditions a complete flow-denoising trajectory that produces an intermediate visual output.Intermediate targets can represent decomposed ground-truth steps such as partial maze path segments.
  • Terminal Thought Grounding: Terminal thought grounding aligns the final reasoning state with embedded ground-truth reasoning steps through an auxiliary alignment objective.The alignment term is activated at the final reasoning step and is intended to prevent latent reasoning drift.
  • Progressive Training Strategy: Progressive training first supervises all reasoning steps, then concentrates gradients on the final output while preserving intermediate reasoning in the forward pass.The second stage uses limited iterations to reduce degradation of learned reasoning chains.
  • Progressive Training Strategy: Figure 5 illustrates incremental refinement from initial inputs through intermediate steps to final solutions for Maze, TSP, Sudoku, and VSP.The figure presents the initial input at left and the final optimal solution at right for each task.

5. Experiments

Experiments evaluate EndoCoT across visual reasoning tasks, training settings, ablations, inference-time scaling, efficiency, and image editing. The results show strong task-specific performance, transferable unified reasoning, and benefits from semantic grounding and iterative implicit tokens.

  • Task-Specific Training: 95% Sudoku-35 and 90% Maze-32 accuracy substantially outperform DiffThinker’s 55% and 65%, respectively.EndoCoT also reaches 85% on VSP-Super, while ThinkGen and ChronoEdit fail completely there.
  • Unified Training: Unified training preserves competitive performance across all tasks, although it performs slightly below task-specific training.A single model is trained on Maze, TSP, Sudoku, and VSP simultaneously.
  • Ablation Study: Removing semantic loss reduces Maze-32 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 and constrains their trajectories.
  • Ablation Study: Explicit tokens achieve 34% accuracy on Maze-8 but 0% on Maze-32, with long-horizon reasoning degrading through error accumulation and repetitive token loops.The experiment replaces implicit tokens with discrete autoregressive intermediate reasoning steps.
  • Inference-Time Scaling: Increasing the implicit reasoning budget τ consistently improves accuracy and path repetition rate, especially on Maze-32.Figure 6 examines the associated trade-off between accuracy and overall inference time.
  • Efficiency: As resolution increases, EndoCoT’s relative computational cost decreases because it avoids repeating expensive DiT denoising steps and maintains stable inference latency.The experiments compare execution time across resolutions and token budgets.
  • Component Analysis: Using only DiT yields 18% accuracy on Maze-32, while using only the MLLM produces erratic wandering paths that become trapped in dead ends.The results indicate that logical reasoning and visual grounding are both needed for spatial reasoning.
  • Image Editing: EndoCoT supports progressive image editing by iteratively planning and executing object additions and transformations, with τ controlling the number of editing operations.The qualitative comparison includes sequentially adding a stone lantern and deer, then transforming the deer into a sheep.

6. Conclusion

EndoCoT enables endogenous chain-of-thought reasoning by refining latent thought states and grounding the final output in textual supervision. Experiments attribute its capability to coupling the MLLM text encoder with the DiT backbone, while requiring manual reasoning-step tuning and high-quality datasets.

  • Conclusion: EndoCoT performs endogenous chain-of-thought reasoning by iteratively refining latent thought states and grounding the final output in textual supervision.The framework bridges high-level logical planning and precise visual generation.
  • Conclusion: Its reasoning capability arises from synergistic coupling between the multimodal text encoder and DiT backbone rather than standard denoising alone.This conclusion connects the model’s planning and visual generation components.
  • Conclusion: The approach currently requires manual tuning of the optimal reasoning-step count and high-quality datasets.

B. Data Statistics

The dataset contains 182.4K instances across four reasoning tasks: Maze, TSP, Sudoku, and VSP. Each task records or renders step-by-step solutions for spatial or combinatorial reasoning.

  • Dataset Composition: 182.4K instances comprise the constructed dataset across Maze, TSP, Sudoku, and VSP reasoning tasks.The dataset includes intermediate reasoning steps.
  • Maze: Maze instances use DFS for loop-free maze construction and BFS to compute shortest paths between randomly assigned starts and goals.Paths are recorded iteratively to simulate step-by-step spatial exploration.
  • TSP: TSP instances randomly place N cities on a 2D grid and use Held-Karp dynamic programming to derive optimal tours.The optimal node-to-node trajectory is rendered sequentially.
  • Sudoku: Sudoku instances use valid 9 × 9 boards, remove clues while preserving unique solutions, and record reverse filling steps from ground-truth solutions.The task prompt requires identifying and filling empty cells sequentially.
  • VSP: VSP instances generate FrozenLake grids with safe tiles, hazards, starts, and goals, then use Dijkstra’s algorithm to compute safe paths.State transitions are recorded step-by-step to capture environment perception and hazard avoidance.
  • Qualitative Examples: Qualitative examples compare EndoCoT with ThinkGen and ChronoEdit across the reasoning tasks.These examples are presented in Figs. 15–18.

D.1. Effect of Two-Stage Training

The Maze analysis compares single-stage and multi-stage training, with Table 6 highlighting Terminal Consolidation’s effect on awareness of the final stage.

  • Training Comparison: Table 6 presents a comparison between single-stage training and multi-stage training on Maze.
  • Training Comparison: Terminal Consolidation effectively enhances the model’s awareness of the final stage on the Maze task.Table 6 provides an illustrative analysis of this effect.

D.2. Impact of Terminal Consolidation Training Duration

Terminal Consolidation training duration presents a trade-off: stronger final-stage convergence can make intermediate reasoning steps sparse. EndoCoT also generalizes across Sudoku resolutions and font styles, including partially cropped low-resolution answers.

  • Training Duration: Excessive Stage 2 training causes intermediate reasoning steps to become sparse.This behavior is illustrated in Fig. 10.
  • Training Duration: With a constant number of reasoning steps τ, more Stage 2 training makes inference behavior converge more toward the target state.
  • Generalization: EndoCoT demonstrates stronger Sudoku generalization than baseline methods across varying resolutions and font styles.For low-resolution inputs, it can recover partially cropped correct answers, whereas DiffThinker fails to generate valid solutions.
  • Training Configuration: Training hyperparameters and LoRA configurations are summarized in Table 7.

E.2. Further Analysis

Further analysis finds that reasoning-related gradient differences and peak sensitivity concentrate in MLLM layers and near the MLLM–DiT junction. Sudoku tests additionally examine generalization across resolutions and font styles.

  • Gradient Analysis: The largest gradient differences between reasoning and non-reasoning datasets are consistently concentrated within MLLM layers.The analysis computes layer-wise gradient magnitude differences across the MLLM and DiT.
  • OOD Evaluation: The Sudoku out-of-distribution test compares performance across different resolutions and font styles.The figure caption states that EndoCoT shows stronger generalization than baseline methods.
  • Sensitivity Analysis: Peak sensitivity is concentrated within the Vision Encoder and at the junction between terminal LLM layers and initial DiT layers.This observation motivates jointly fine-tuning the MLLM and DiT components in EndoCoT.

E.3. Experimental Observations

Experiments identify supervision and representation choices that shape reasoning performance: fine-grained latent-step supervision is necessary, while excessive intermediate semantic supervision can hurt. The required latent-token count reflects both algorithmic and visual complexity.

  • Experimental Observations: Fine-grained intermediate-step supervision is necessary because inference typically uses fewer reasoning steps than training supervision.The authors report that finer-grained supervision helps the model learn latent reasoning processes.
  • Experimental Observations: Applying semantic loss to intermediate steps can reduce performance relative to EndoCoT’s sparse semantic supervision strategy.The reported explanation is that extra semantic supervision can disrupt synergy between DiT rendering and MLLM understanding.
  • Experimental Observations: Latent-token requirements depend on both the task’s algorithmic complexity and its visual complexity.The number of tokens needed to solve a puzzle is not determined by algorithmic complexity alone.
  • Dataset Generation: TSP instances sample cities, compute pairwise distances, solve for an optimal tour with Held-Karp dynamic programming, and render progressively longer trajectories.Each instance returns the city set, optimal tour, and rendered trajectory sequence.
  • Dataset Generation: Sudoku generation removes cells while preserving a unique solution, then simulates solving to produce a dataset instance.The procedure reverts a removal when multiple solutions appear and returns the puzzle, solution, and solving sequence.
Loading 2603.12252v4…