- The paper demonstrates that dynamic, sample-level block size adaptation in RL significantly outperforms fixed block size methods in diverse reasoning domains.
- It introduces Block-R1-41Kโa curated, multi-domain dataset with empirical annotations for optimal block sizesโto fine-tune diffusion LLMs.
- Empirical evaluations reveal that block-conditioned RL improves benchmarks such as Countdown and GSM8K by 20-30 percentage points compared to static approaches.
Block-R1: Sample-Level Block Size Adaptation for Multi-Domain RL in Diffusion LLMs
Introduction
Block-R1 introduces a principled framework for addressing block size conflict in multi-domain reinforcement learning (RL) post-training of diffusion LLMs (dLLMs). While prior RL approaches for dLLMs typically fix the block size throughout training, this work establishes that such rigidity significantly impairs cross-domain generalization because different reasoning domainsโincluding code generation, mathematical reasoning, and logical puzzlesโexhibit variable optimal block granularity. The core contribution is a dataset and benchmark where each training sample is annotated with its empirically optimal block size, enabling block-conditioned RL policies that adaptively align rollout generation to the local sample structure.
Figure 1: RL post-training on heterogeneous domains requires different block sizes for optimal performance; uniform block size across domains causes performance degradation compared to single-domain RL or the base model.
Block Size Conflict in Multi-Domain dLLM RL
The central argument is that block size should not be treated as a static decoding hyperparameter in RL post-training for dLLMs. The model decodes sequences via blockwise, semi-autoregressive denoising, with block size c directly controlling the parallel granularity of token generation. Empirical analysis (Figure 1 and related quantitative results) demonstrates that optimal block sizes differ markedly across domainsโfor example, smaller blocks benefit mathematical reasoning due to fine-grained verification, while larger blocks are favored in tasks like Sudoku for structural consistency.
The work formalizes the domain block size conflict: for two domains Diโ and Djโ, block size preferences are disjoint if
argcโBmaxโJiโ(ฮธ,c)โฉargcโBmaxโJjโ(ฮธ,c)=โ
where Jkโ(ฮธ,c) is the expected reward on domain Dkโ under block size c. Theorem 1, under objective alignment and preference divergence assumptions, strictly bounds the performance of any fixed block size versus sample- or domain-adaptive selection, confirming the inherent sub-optimality of the former in the multi-domain regime.
Block-R1-41K: Dataset Construction and Sample-Level Block Size Optimization
Block-R1-41K is constructed through a stringent pipeline:
- Cross-Domain Source Selection: High-quality reasoning datasets spanning diverse domains are curated.
- Unified Reward Design: A composite reward function is applied, integrating format-, accuracy-, and constraint-based terms to cover heterogeneous tasks while allowing meaningful reward comparison across domains.
- Teacher-Student Filtering: Samples are filtered based on performance gaps between a strong (teacher) and weak (student) model across all candidate block sizes, retaining only those with significant, attainable improvement space.
- Sample-Level Best-Improved Block Size Selection: For each sample, the block size maximizing the teacher-student improvement gap is selected, providing a direct annotation of its optimal training granularity.
- Balanced Multi-Domain Assembly: The final corpus balances sample counts from each domain to prevent data skew.
Figure 3: Five-stage construction pipeline for Block-R1-41K, including reward unification, sample filtering, and sample-level block size annotation.
This sample-level conditioning induces an empirical block size preference distribution per domain, used to compute domain-wise Block Size Conflict Score (BCS) via the Wasserstein distance. A higher BCS between two domains predicts greater performance degradation when post-training with a shared fixed block size.
Figure 4: (a) Correlation between BCS and performance drop in multi-domain RL; (b) Heatmap of pairwise domain BCS; (c) Distribution of best-improved block sizes per domain, highlighting domain-dependent block size preference.
Empirical Evaluation and Benchmarking
Block-R1 serves as a benchmark for RL post-training across 13 reasoning benchmarks, 7 RL algorithms, and 10 dLLM backbones. Comparative analyses strongly underline two observations:
- Single-domain RL achieves domain-specific gains but fails to generalize. Training on one domain yields significant improvement locally (e.g., up to ~81% on GSM8K), but frequently degrades or stagnates on other domains due to overfitting to domain-specific parallelism structures.
- Vanilla multi-domain RL with a global block size performs poorly. Directly mixing data with a fixed block size leads to substantial degradation (e.g., Countdown drops from ~59% to ~30%), exposing the adverse effects of structural incompatibility across domains.
Block-R1โs sample-level block-conditioned RL robustly outperforms both baselines. Multi-domain post-training with Block-R1 achieves simultaneous improvements on all domains (e.g., Countdown: 30.08 โ 62.11; KK: 30.14 โ 50.14; GSM8K: 57.24 โ 80.74).
Figure 5: Average reward improvement as a function of training block size, demonstrating highly domain-dependent variation and the non-neutral role of block size in reward maximization.
Block-R1 also remains compatible with dynamic inference-time block size selection (e.g., via b1), which yields further incremental gains when combined.
Theoretical and Practical Implications
The results affirm that block size is a structural factor in RL for dLLMsโnot merely an implementation convenience. Fixed block sizes undermine the ability to exploit the expressive flexibility of blockwise dLLMs, especially as models scale to multi-domain contexts with compositional and architectural heterogeneity. Adopting sample-level block conditioning ensures that optimal granularity is aligned with each sampleโs latent structure, both maximizing reachable reward improvement during training and enabling effective generalization.
Practically, the Block-R1 dataset and methodology set a new standard for evaluating and improving dLLM reasoning across diverse domains. The datasetโs annotation process provides a reusable diagnostic for block size conflict and establishes an empirical foundation for future sample-wise block size adaptation at both training and inference.
Future Directions
Block-R1 outlines unexplored terrain for adaptive sequence modeling in diffusion frameworks. Potential research avenues include:
- End-to-end differentiable block size selection mechanisms, possibly through reinforcement learning or meta-learning.
- Extension of block-aware adaptation into non-block-based or fully non-autoregressive architectures.
- Coupling block size adaptation with dynamic reward shaping or preference modeling to further individualize learning signals.
- Integration with inference-time adaptation methods to bridge training-inference divides holistically.
Conclusion
Block-R1 shows that brute-force uniformity in RL post-training for diffusion LLMs is fundamentally mismatched to cross-domain generalization. Explicitly modeling and exploiting sample-level block size heterogeneity is criticalโyielding substantial improvements in reasoning benchmarks and establishing empirically justified best practices for dLLM design and post-training. This framework guides both theoretical development and experimental evaluation of block-conditioned learning in sequence modeling.
Figure 6: Detailed domain-domain BCS pair visualization, quantitatively linking cross-domain block size divergence to observed performance degradation in fixed-block multi-domain RL.