💬 NLP

T $^\star$ : Progressive Block Scaling for Masked Diffusion Language Models Through Trajectory Aware Reinforcement Learning

The paper introduces T $^\star$ , a TraceRL-based training curriculum that enables masked diffusion language models to progressively scale from small to large decoding blocks, achieving high-parallelism decoding with minimal performance loss on math reasoning tasks.

Original authors: Hanchen Xia, Baoyou Chen, Yutang Ge, Guojiang Zhao, Siyu Zhu

Published 2026-03-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Hanchen Xia, Baoyou Chen, Yutang Ge, Guojiang Zhao, Siyu Zhu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

🧱 The Big Idea: Building a Better "Guessing Game"

Imagine you are trying to solve a complex math problem, but instead of writing the answer one word at a time (like a standard AI), you have a magical whiteboard where you can erase and rewrite chunks of the answer simultaneously.

This is how Masked Diffusion Models (MDMs) work. They don't just write left-to-right; they look at the whole sentence, guess which parts are wrong, erase them, and fill them in all at once. This is super fast (parallel) but can be messy if the chunks are too big.

The paper introduces T⋆, a clever training method that teaches these models how to handle bigger chunks without getting confused.

🚗 The Problem: The "Speed vs. Control" Dilemma

Think of an AI model like a car driver.

Small Blocks (Block Size = 4): The driver is very careful. They check the road, make a tiny adjustment, check again, and move forward. They are slow but very accurate.
Large Blocks (Block Size = 32): The driver tries to make huge, sweeping turns to get to the destination faster. They can cover more ground at once (high parallelism), but if they turn too sharply, they might crash or spin out of control.

The Issue: When researchers tried to train these models to use "Large Blocks" immediately, the models crashed. Their math skills dropped because they were trying to guess too many words at once without enough practice. It's like asking a beginner driver to race on a track with 32 lanes open at once—they panic and fail.

🪜 The Solution: T⋆ (The "Staircase" Training)

The authors created a training curriculum called T⋆. Instead of throwing the model into the deep end, they built a staircase.

Start Small: They begin by training the model on tiny blocks (like 4 words at a time). The model learns the logic of math problems perfectly here.
The "TraceRL" Coach: They use a special reinforcement learning technique (called TraceRL) that acts like a coach. The coach doesn't just say "Good job" at the end; it watches every single step the model takes and gives feedback on how it unmasked the words.
Climb the Stairs: Once the model is a pro at small blocks, the training slowly increases the block size (4 → 8 → 16 → 32).
- Crucial Step: Before moving to the next size, the model practices on the current size using the coach's feedback. This ensures the model doesn't forget how to be accurate before it tries to be faster.

The Analogy: Imagine learning to juggle.

Old Way: Try to juggle 10 balls immediately. You drop them all.
T⋆ Way: Start with 1 ball. Master it. Then 2. Then 3. By the time you reach 10, your brain has built the muscle memory to handle the speed without dropping the balls.

🧠 The Secret Sauce: "Trajectory Awareness"

Why does this work so well?

In standard training, the AI only gets a "Yes/No" at the very end of the math problem.

AI: "Here is the answer."
Teacher: "Wrong."
AI: "Okay, I guess I'll try something else." (This is inefficient).

T⋆ (using TraceRL) is different. It looks at the journey (the trajectory).

Teacher: "You got the first step right, but you erased the wrong word in step 3. Let's fix that specific move."

By rewarding the model for making good intermediate moves, the model learns a better "unmasking schedule." It learns when to finalize a word and when to keep guessing, even when dealing with huge blocks of text.

🏆 The Results: Fast AND Smart

The paper tested this on hard math benchmarks (like MATH500 and GSM8K).

The "Direct" Approach: If you just tell a model to use big blocks, its accuracy crashes (like a car spinning out).
The T⋆ Approach: The model keeps its high accuracy (it's still a math genius) but gains the speed of big blocks.

The "Magic" Finding:
Usually, when you speed up, you lose precision. T⋆ found a "sweet spot" where the model learned a new way of thinking. It didn't just go back to writing slowly (left-to-right); it found a unique, non-linear way of solving problems that is both fast and accurate.

🌟 Summary in One Sentence

T⋆ is a "training ladder" that teaches AI models to solve complex math problems by starting with small, careful steps and gradually teaching them to handle massive chunks of text simultaneously, without losing their accuracy.

1. Problem Statement

Masked Diffusion Language Models (MDMs) offer a promising alternative to autoregressive (AR) models by enabling parallel decoding, which can significantly reduce inference latency. However, a critical challenge arises when scaling the block size (the number of tokens denoised in parallel):

Performance Degradation: As block size increases, reasoning capabilities (particularly in mathematics) degrade significantly. Larger blocks require the model to resolve more masked tokens simultaneously, increasing the complexity of the denoising task.
Training Instability: Directly applying Reinforcement Learning (RL) to large-block MDMs often leads to training collapse. The authors attribute this to "Lazy Likelihood Displacement," where the model starts from a low-confidence regime in large blocks, leading to noisy advantage estimates and likelihood drift during policy updates.
Lack of Optimal Schedules: Standard supervised fine-tuning (SFT) only provides the final target sequence, not the optimal intermediate "denoising schedule" (i.e., which tokens should be finalized at each step).

2. Methodology: T⋆

The authors propose T⋆, a progressive training curriculum designed to scale block sizes while maintaining reasoning performance. The method combines Trajectory-Aware Reinforcement Learning (TraceRL) with a staged block-expansion strategy.

A. Trajectory-Aware RL (TraceRL) Backbone

Instead of treating the final output as the only reward signal, T⋆ treats the denoising process as a multi-step trajectory $\tau$ .

Action Space: Each finalized token is treated as an action.
Reward Shaping: The method aggregates token-level rewards (e.g., correctness of the final answer) into step-level advantages using TD/GAE (Generalized Advantage Estimation). This allows credit assignment to propagate back through the entire denoising trajectory, not just the final output.
Objective: A PPO-like objective optimizes the policy to maximize the expected reward over the trajectory while maintaining a KL-divergence penalty to prevent drifting too far from the base model.

B. Progressive Block Scaling (The T⋆ Curriculum)

To overcome the instability of training large blocks from scratch, T⋆ employs a stage-wise expansion strategy:

Initialization: Start with a small block size (e.g., $B=4$ ) initialized from an AR-pretrained model.
Stabilization: Train the model at the current block size $B$ $B$ using TraceRL. Crucially, the training involves two partition updates per step:
- Standard block partition.
- Shifted partition: An offset of $\Delta = B/2$ is applied to the block boundaries to ensure robustness against boundary artifacts.
Expansion: Once the policy is stable at size $B$ , the block size is doubled ( $B \leftarrow 2B$ ). The model continues training with the learned policy from the smaller block as a warm start for the larger block.
Iteration: This process repeats until the target block size is reached.

3. Key Contributions

T⋆ Algorithm: A novel, easy-to-implement curriculum that enables MDMs to scale block sizes (from 4 to 32+) with minimal performance loss, effectively bridging the gap between AR-like causality and high-parallelism masked updates.
Stability Analysis: The paper identifies that direct TraceRL fails on large blocks due to initialization issues (low confidence regimes). T⋆ mitigates this by "warm-starting" the larger block policy with the refined policy from smaller blocks.
Schedule Discovery: Analysis reveals that T⋆ does not simply revert to a standard left-to-right (AR) schedule. Instead, it learns a non-canonical denoising schedule that optimizes token finalization order for the specific block size, improving reasoning without external search scaffolds.
Efficiency Gains: The method preserves the parallelism benefits of larger blocks, significantly increasing Tokens Per Forward (TPF) compared to AR baselines.

4. Experimental Results

Experiments were conducted on SDAR-1.7B and SDAR-4B models using mathematical reasoning benchmarks (MATH500, GSM8K, AIME24).

Performance Stability:
- Direct TraceRL: Exhibited severe performance collapses when applied directly to larger blocks (e.g., MATH500 accuracy dropped from ~56% to ~30% for the 1.7B model at $B=16$ ).
- T⋆: Maintained stable performance across all block sizes. For the 1.7B model, T⋆ achieved 63.4% on MATH500 at $B=8$ and 59.0% at $B=32$ , significantly outperforming both the base model and direct TraceRL at larger scales.
Scaling Trends:
- Base models showed a downward trend in accuracy as block size increased.
- T⋆ consistently matched or exceeded the performance of base models at smaller blocks and achieved the strongest results at larger blocks.
Decoding Efficiency:
- As block size increased from 8 to 32, TPF increased by ~29% (from 2.91 to 3.74), indicating a tangible reduction in forward passes required for generation.
- T⋆ maintained these efficiency gains without sacrificing reasoning accuracy.
Schedule Analysis:
- The LocalStrict metric (measuring adherence to left-to-right order) showed that while T⋆ remains largely monotone, it learns flexible schedules distinct from the canonical AR order, adapting to the block constraints.

5. Significance

Practical Deployment: T⋆ provides a practical recipe for deploying MDMs in real-world scenarios where inference speed (parallelism) is critical, without sacrificing the high-level reasoning capabilities typically reserved for AR models.
Internal Policy Learning: The work demonstrates that complex reasoning improvements can be achieved by optimizing the model's internal denoising schedule via RL, rather than relying on external, computationally expensive search procedures (like Tree-of-Thoughts).
Mitigation of RL Collapse: The paper offers a concrete solution to the "death spiral" of likelihood displacement in RL for diffusion models, suggesting that progressive curriculum learning is essential for scaling RL in generative models.

In conclusion, T⋆ successfully decouples the trade-off between parallelism and reasoning quality in masked diffusion models, enabling them to scale to larger block sizes while preserving the strong reasoning capabilities inherited from autoregressive initialization.