EndoCoT: Teaching Diffusion Models to Think Before They Draw

Paper: 2603.12252 Authors: Xuanlang Dai, Yujie Zhou, Long Xing, Jiazi Bu, Xilin Wei, Yuhong Liu, Beichen Zhang, Kai Chen, Yuhang Zang Categories: cs.CV, cs.CL

The Gap

The standard playbook for hard visual reasoning tasks is: plug an MLLM into a diffusion transformer (DiT) as a text encoder, let it encode the prompt, and hand the embedding off to the denoising process. Works fine for “a red apple on a table.” Falls apart for “solve this 9x9 Sudoku grid.”

Two things go wrong. First, single-pass encoding never triggers chain-of-thought. MLLMs are known to reason better when they think step-by-step, but a one-shot encode call skips all of that — you get the model’s first instinct, not its considered answer. Second, even if the encoding were perfect, the guidance is frozen. The DiT denoises over dozens of steps, but it’s steering by the same fixed vector the whole time. Complex tasks need progressive decomposition — early steps should lay out structure, later steps should fill in details — but a static embedding can’t do that.

Prior work (standard MLLM-DiT pipelines) treated the MLLM as a smarter tokenizer. This paper treats it as a reasoner that needs to stay in the loop.

Problem: MLLM used as static encoder
    |
    v
Assumption 1: CoT requires iterative state updates, not one-shot encoding
Assumption 2: DiT denoising benefits from evolving guidance, not fixed vectors
    |
    v
Method:
  [Iterative Thought Guidance] --> refine latent thought states over steps
  [Terminal Thought Grounding]  --> anchor final state to ground-truth text
    |
    v
Evidence: Maze / TSP / VSP / Sudoku benchmarks, 92.1% avg, +8.3pp over best baseline
    |
    v
Conclusion: Endogenous CoT in the encoder + dynamic guidance in the decoder = better structured reasoning

The Increment

One sentence: Before this paper, MLLMs in diffusion pipelines were mute oracles consulted once; after it, they’re active reasoners whose thinking evolves in sync with the denoising process.

Core Mechanism

The framework has two interlocking pieces. The first is the Iterative Thought Guidance (ITG) module. Instead of encoding the prompt once, ITG maintains a latent “thought state” — a vector that gets updated across multiple rounds. At each round, the MLLM takes the current thought state plus the original prompt and produces a refined state. This is essentially unrolling the chain-of-thought process into the embedding space: each iteration is one reasoning step, and the state carries the accumulated reasoning forward.

The second piece is Terminal Thought Grounding (TTG). After the iterative refinement converges, you have a final thought state. TTG aligns this final state with the ground-truth answer text using a contrastive or alignment loss. This prevents the iterative process from drifting into some arbitrary latent region — it keeps the reasoning trajectory tethered to what the answer actually looks like in text space.

These two thought states — the evolving intermediate ones and the grounded final one — are then injected into the DiT’s cross-attention layers at different denoising timesteps. Early timesteps get earlier thought states (coarse reasoning), later timesteps get later ones (refined reasoning). The guidance is no longer a static vector; it’s a schedule.

Prompt
  |
  v
[MLLM Encoder]
  |
  v
Thought State t=0
  |
  +---> [ITG: MLLM(state_t, prompt) --> state_t+1] --+
  |                                                   |
  +<--------------------------------------------------+
  |   (repeat N times)
  v
Final Thought State
  |
  +---> [TTG: align(final_state, ground_truth_text)]
  |
  v
Thought State Schedule [s_0, s_1, ..., s_N]
  |
  v
DiT Denoising:
  step 1 <-- s_0 (coarse guidance)
  step 2 <-- s_1
  ...
  step T <-- s_N (fine-grained guidance)
  |
  v
Output Image / Solution

Think of it like a detective briefing a sketch artist. In the old approach, the detective describes the suspect once, the artist draws, done. In EndoCoT, the detective keeps refining their description as the sketch develops — “actually, the jaw was sharper, the eyes were closer together” — and the artist updates the drawing accordingly. The TTG module is like the detective checking their notes against a known photo at the end to make sure they haven’t drifted into fiction. The ITG module is the back-and-forth conversation. The thought state schedule is the sequence of revised briefs handed to the artist at each stage of the drawing.

Key Concepts

• Chain-of-Thought in embedding space: Normally, CoT means generating intermediate text tokens (“first, I notice… then, I conclude…”). Here, CoT happens in the continuous latent space — no tokens are generated, but the thought state vector is iteratively updated in a way that mimics the reasoning trajectory. The intuition: if you run the MLLM’s attention layers repeatedly over the same input while carrying forward a state, you’re approximating what would happen if it wrote out its reasoning step by step.

• Dynamic guidance scheduling: Standard diffusion conditioning gives the DiT one context vector for all timesteps. Dynamic scheduling means different timesteps receive different conditioning vectors. This matters because diffusion models are known to handle different aspects of generation at different noise levels — coarse structure early, fine details late. Matching the reasoning depth of the guidance to the denoising stage is the key structural insight here.

• Terminal grounding: Without this, iterative refinement in latent space can wander. TTG is a regularizer that says “your final thought state must be close to the embedding of the correct answer.” It’s the difference between a student who thinks hard but goes off on a tangent versus one who keeps checking their work against the answer key.

Framework Shift

Before (mainstream approach):        After (this paper):

Prompt                               Prompt
  |                                    |
  v                                    v
[MLLM: encode once]                  [MLLM: iterate N times]
  |                                    |   state_0 -> state_1 -> ... -> state_N
  v                                    |                                   |
Fixed embedding e                      |                              [TTG: ground]
  |                                    v
  +---> DiT step 1                   Thought schedule [s_0..s_N]
  |                                    |
  +---> DiT step 2                     +---> DiT step 1 (guided by s_0)
  |                                    |
  +---> DiT step T                     +---> DiT step T (guided by s_N)
  |                                    |
  v                                    v
Output                               Output

Same e at every step                 Different s_i at each step

From static oracle to active co-reasoner, the core shift is: the MLLM’s reasoning and the DiT’s denoising are now temporally coupled, not sequentially decoupled.

Expert Assessment

Problem choice: This is a real gap. The mismatch between single-pass encoding and multi-step denoising has been quietly ignored because most diffusion tasks don’t need deep reasoning. As the field pushes into structured problem-solving (routing, constraint satisfaction, spatial planning), this mismatch becomes load-bearing. The paper is well-positioned on the field’s trajectory.

Method maturity: The core idea is clean and the two-module design is well-motivated. That said, the iterative thought guidance is essentially repeated MLLM forward passes, which is expensive. The paper doesn’t deeply engage with whether a lighter mechanism — say, a small recurrent module trained on top of frozen MLLM embeddings — could achieve similar results at a fraction of the cost. The “endogenous” framing is a bit generous; you’re still running the MLLM multiple times, which is more “iterative exogenous” than truly endogenous.

Experimental integrity: The benchmark suite (Maze, TSP, VSP, Sudoku) is reasonable for structured reasoning, but these are all combinatorial/algorithmic tasks. It’s not clear the gains transfer to more naturalistic spatial reasoning (e.g., compositional scene generation). The +8.3pp improvement is solid, but the paper should show ablations that isolate ITG vs TTG contributions more cleanly — if TTG alone gets you 6pp, the iterative part is less impressive than advertised.

Writing quality: The motivation section is strong. The weakest part is the analysis of why the method works — there’s no visualization of how thought states evolve across iterations, which would be the most convincing evidence that CoT is actually happening in latent space rather than just “more compute = better results.” Adding that would elevate the paper from “it works” to “here’s why it works.”

Verdict: weak accept — solid engineering contribution with a clean story, but needs stronger mechanistic evidence that the iterative reasoning is doing what the authors claim.

Takeaways

A few concrete things worth stealing:

• The thought state schedule idea — mapping reasoning depth to denoising timestep — is transferable to any conditional generation task where the condition has internal structure. If you’re conditioning on a program, a proof, or a multi-step plan, you could similarly schedule which “stage” of the condition guides which denoising step.

• TTG as a grounding regularizer is a lightweight trick: align your final latent state to a known-good text embedding before using it downstream. This is applicable anywhere you have iterative latent refinement and want to prevent drift.

• The framing of “MLLM as reasoner, not tokenizer” is worth internalizing. If you’re building any pipeline that uses an LLM/MLLM as an encoder, ask whether single-pass encoding is actually sufficient for your task’s reasoning demands.