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

TL;DR

T$^ ext{*}$ introduces a TraceRL-based curriculum for progressive block-size scaling in masked diffusion language models, improving decoding efficiency with minimal performance loss.

Contribution

It proposes a novel training method that gradually scales block sizes in MDMs using reinforcement learning, enhancing decoding parallelism.

Findings

Enables higher-parallelism decoding with minimal performance loss.

Suggests convergence to an alternative decoding schedule with comparable performance.

Abstract

We present T$^\star$, a simple 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$ may actually converge to an alternative decoding schedule that achieves comparable performance.