Papers
Topics
Authors
Recent
Search
2000 character limit reached

Block-R1: Rethinking the Role of Block Size in Multi-domain Reinforcement Learning for Diffusion Large Language Models

Published 12 May 2026 in cs.LG | (2605.11726v1)

Abstract: Recently, reinforcement learning (RL) has been widely applied during post-training for diffusion LLMs (dLLMs) to enhance reasoning with block-wise semi-autoregressive generation. Block size has therefore become a vital factor in dLLMs, since it determines the parallel decoding granularity and affects the rollout trajectories during RL optimisation, e.g., GRPO. Instead of investigating the effect of block size during inference on individual domains, this paper studies block size from a domain conflict perspective for dLLM RL post-training in multi-domain scenarios. The main contributions are: (1) a formulation of domain block size conflict in multi-domain RL for dLLMs, which will largely affect the post-training effectiveness for rollout-based RL methods; (2) a novel dataset, Block-R1-41K is constructed with a best-improved training block size for each sample, which also induces a Block Size Conflict Score to quantitatively measure the domain conflict; (3) a new benchmark, Block-R1, for flexible RL post-training for dLLMs in both single and cross domain; and (4) a simple yet powerful cross-domain post-training method with sample-level best-improved training block sizes. Extensive experiments on 13 distinct datasets, 7 latest RL algorithms, and various different dLLM backbones are covered in Block-R1. The benchmark is open-sourced at https://github.com/YanJiangJerry/Block-R1, with the dataset released at https://huggingface.co/datasets/dLLM-R1/Block-R1-41K.

Authors (3)

Summary

  • 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

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 cc 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\mathcal{D}_i and Dj\mathcal{D}_j, block size preferences are disjoint if

argโกmaxโกcโˆˆBJi(ฮธ,c)โˆฉargโกmaxโกcโˆˆBJj(ฮธ,c)=โˆ…\arg\max_{c \in B} J_i(\theta, c) \cap \arg\max_{c \in B} J_j(\theta, c) = \emptyset

where Jk(ฮธ,c)J_k(\theta, c) is the expected reward on domain Dk\mathcal{D}_k under block size cc. 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:

  1. Cross-Domain Source Selection: High-quality reasoning datasets spanning diverse domains are curated.
  2. 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.
  3. 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.
  4. 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.
  5. Balanced Multi-Domain Assembly: The final corpus balances sample counts from each domain to prevent data skew. Figure 2

    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 3

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

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

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.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Collections

Sign up for free to add this paper to one or more collections.

Tweets

Sign up for free to view the 1 tweet with 2 likes about this paper.