T$^\star$: Progressive Block Scaling for MDM Through Trajectory Aware RL

TL;DR

The paper tackles how to scale block size in masked diffusion language models without sacrificing math-focused reasoning. It introduces T$^\star$, a progressive block-scaling curriculum built on trajectory-aware reinforcement learning (TraceRL) that expands block size from an initial $B_0$ to a target $\hat{B}$ through a three-step per-stage process (RL on current blocks, boundary shifting, and block expansion). A PPO-style objective with a clipped surrogate and trajectory-level credit assignment is used to train across denoising steps, with KL regularization for stability. Empirical results on SDAR-1.7B-Chat and SDAR-4B-Chat across MATH500, GSM8K, and AIME24 show that T$^\star$ maintains or improves performance while enabling larger blocks, while direct TraceRL can collapse at large block sizes; the work also reveals that trajectory-aware RL can yield non-canonical but effective denoising schedules. Overall, the approach offers a practical pathway to higher-parallel decoding in diffusion LMs and suggests trajectory-aware policy learning as a complement to external search scaffolds for reasoning tasks.

Abstract

We present T$^\star$, a simple \textsc{TraceRL}-based training curriculum for progressive block-size scaling in masked diffusion language models (MDMs). Starting from an AR-initialized small-block MDM, T$^\star$~transitions smoothly to larger blocks, enabling higher-parallelism decoding with minimal performance degradation on math reasoning benchmarks. Moreover, further analysis suggests that T$^\star$~can converge to an alternative decoding schedule $\hat{\rm S}$ that achieves comparable performance.

Figures (4)

• Figure 1: Validation accuracy during block scaling (1.7B). MATH500 validation accuracy over training epochs for T$^\star$ and a direct TraceRL baseline (dashed). Vertical dotted lines indicate stage transitions ($B{=}4 \rightarrow 8 \rightarrow 16$). Horizontal dashed lines show the accuracies of the original SDAR checkpoints trained at each block size.
• Figure 2: Performance vs. block size across model scales. Performance on MATH500, GSM8K, and AIME24 as a function of block size $B$ for SDAR models at 1.7B (left) and 4B (right), comparing the Base model, TraceRL trained at the same $B$, and our progressive strategy T$^\star$.
• Figure 3: Decoding schedule under TraceRL vs. T$^\star$. More results can be found in Appendix \ref{['app:case-study']}
• Figure 4: Case study: decoding schedule under TraceRL vs. T$^\star$. We visualize the token-level first-unmask step index (heatmaps; darker means decoded later) and the corresponding model solutions for a representative algebra problem, evaluated with block sizes $B\in\{8,16,32\}$. The top row shows a model trained with direct TraceRL at the same block size, and the bottom row shows the model obtained by T$^\star$.