Unlike autoregressive language models, which terminate variable-length generation upon predicting an End-of-Sequence (EoS) token, Diffusion Language Models (DLMs) operate over a fixed maximum-length context window for a predetermined number of denoising steps. However, this process is independent of the required response length, resulting in computational waste for the majority of short responses common in reasoning and chat tasks. To address this problem, we conjecture that the latent prompt representation contains sufficient information to estimate the required output length. We provide empirical evidence for this phenomenon and propose a zero-shot mechanism to dynamically crop the context window before generation begins, leading to fewer diffusion steps and substantial computational savings. We evaluate our approach on four benchmarks with diverse tasks -- GSM8K (reasoning), HumanEval (code generation), IfEval (instruction following), and LongFormQA (question answering) -- revealing massive efficiency gains at minimal performance impact. We report significant reductions in FLOPs across all tasks, with no statistically significant performance degradation, and significant performance improvements in 2 out of 4 tasks.
Figures & tables
Figure 1: Predicted Length Distributions. Our SmartCrop ( τ=0.9 ) method successfully predicts task-specific output lengths across four benchmark datasets. The abrupt truncations observed in certain distributions correspond to context length constraints (refer to Section 4 for details).
Benchmark
Max new tokens ( Lnew )
Steps ( T )
Content
GSM8K
256
256
Math
HumanEval
512
512
Code
IfEval
1280
1280
Instruction
LongFormQA
512
64
QA
Table 1: Experimental Configurations. Summary of hyperparameters across benchmark datasets. Max new tokens ( Lnew ) denotes the number of masked tokens appended to the prompt in the diffusion canvas. Steps ( T ) indicates the total number of denoising iterations. For example, Lnew=256 with T=256 implies a linear schedule where exactly 1 token is unmasked per step. Content describes the task domains represented in each dataset.
Benchmark
Method
Lp
Avg. Processed Length
Metric ↑
FLOPs Saved (%)↑
Perf. Δ(%)↑
IfEval
FC
87.2
1367.2
0.4801
-
-
SC-0.5
192.1
0.5342
98.47∗∗∗
+11.25∗
SC-0.75
208.0
0.5521
98.05∗∗∗
+14.99∗∗
SC-0.9
222.0
0.5459
97.64∗∗∗
+13.70∗∗
SC-0.95
230.5
0.5450
97.37∗∗∗
+13.50∗∗
SC-0.99
243.8
0.5694
96.92∗∗∗
+18.58∗∗∗
Table 2: Main Results. Performance comparison between native diffusion decoding and our proposed dynamic cropping method across four benchmarks using the LLaDA architecture. Benchmark indicates the evaluation suite. Method distinguishes between the native Full Context ( FC ) baseline and our SmartCrop ( SC ) approach at various length-prediction quantiles ( τ ). Lp represents the average prompt length in tokens (constant per benchmark). Avg. Processed Length denotes the mean number of total tokens processed by the model in each forward pass, comprising both the prompt and the generated sequence. For the FC baseline, this value is constant at Lc=Lp+Lnew , where Lp and Lnew represent the prompt length and the maximum sequence length, respectively. In contrast, for SC , the output length is defined by the predicted total length L^ . Metric denotes the task-specific performance score. FLOPs Saved quantifies the reduction in total floating-point operations relative to the FC baseline (e.g., a 98% saving implies our method requires only 2% of the baseline computation). Perf. Δ reports the relative percentage change in Metric compared to the FC baseline. For all metrics marked with ↑ , higher is better. We determine statistical significance via pairwise bootstrap hypothesis testing (5000 resamples) between FC and SC within each experimental condition. We compute paired differences on samples matched by document ID for both Metric and FLOPs Saved . We estimate two-sided p -values from the resulting bootstrap distributions. Significance levels are denoted as ∗p<0.05,∗∗p<0.01, and ∗∗∗p<0.001 .
Figure 2: Sensitivity of IfEval Performance to Context Length Perturbations. We analyze the robustness of SmartCrop ( τ=0.9 ) by shifting the predicted length L^ by a deviation factor δ∈[−50%,+50%] . The blue curve shows the model performance (mean ± 95% CI) across these varying context lengths. The red line denotes the control baseline, where lengths are sampled from the empirical length distribution of other benchmarks. The green line represents the Full Context baseline performance. While the model is relatively robust to moderate under-estimation (negative δ ), generation quality degrades as superfluous padding is reintroduced (positive δ ), eventually converging toward the baseline.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 3: Predicted Length Invariance ( HumanEval ). Left: kernel density estimate of predicted new tokens for Lnew∈{512,1024,2048,4096} . Right: boxplots of the same values. The bulk of the distribution is comparatively stable across Lnew , with the main visible difference being a stronger right truncation when Lnew=512 , which is expected when the required completion length approaches the canvas limit. Note: Lnew=256 causes the predicted length distribution to be heavily truncated and uninformative compared to larger canvases.
Figure 4: Predicted Length Invariance ( IfEval ). Left: kernel density estimate of predicted new tokens for Lnew∈{256,512,1024,2048,4096} . Right: boxplots of the same values. The central mass of the predicted-length distribution (roughly 50–150 new tokens) is broadly consistent across Lnew , while larger canvases primarily increase the range of rare long-length outliers.
Figure 5: Predicted Length Invariance ( LongFormQA ). Left: kernel density estimate of predicted new tokens for Lnew∈{256,512,1024,2048,4096} . Right: boxplots of the same values. The predicted length is close to invariant across Lnew for the typical range of outputs, with only small shifts in the median and dispersion. This supports the claim in the main text that, for LongFormQA , the model’s inferred length prior is largely insensitive to the particular (potentially conservative) initial canvas size used for the first forward pass.
Figure 6: Predicted Length Invariance ( GSM8K ). Left: kernel density estimate of predicted new tokens for Lnew∈{256,512,1024,2048,4096} . Right: boxplots of the same predicted token counts. We observe a pronounced dependence of the predicted length distribution on Lnew (most visible when setting Lnew at 256 or 4096), indicating that the per-position EoS probabilities used by SmartCrop are not strictly invariant to the amount of masked padding presented to the model at the first denoising step. The Lnew=256 curve also shows a hard truncation consistent with the finite canvas constraint.
Figure 7: Performance Gain vs. Length ( GSM8K ). Per-instance change in Exact Match (grey; Δ∈{−1,0,+1} ) plotted against generated length. The blue curve reports the mean delta within length bins.
Figure 8: Performance Gain vs. Length ( HumanEval ). Per-instance change in Pass@1 (grey; Δ∈{−1,0,+1} ) plotted against generated length. The binned average (blue) is non-monotonic and fluctuates across bins, which is consistent with either weak dependence on length or limited sample counts per bin.
Figure 9: Performance Gain vs. Length ( IfEval ). Per-instance change in prompt-level Strict Accuracy (grey; Δ∈{−1,0,+1} ) plotted against generated length. The binned average (blue) is positive for shorter generations and approaches zero for longer ones, indicating that the net gains from SmartCrop at this operating point are concentrated on prompts that produce relatively short outputs. The negative dip around intermediate lengths should be interpreted with the corresponding bin sample size in mind.
Figure 10: Performance Gain vs. Length ( LongFormQA ). Per-instance change in the evaluation score (grey; continuous delta) plotted against generated length. The binned mean delta (blue) is positive across the typical length range and does not exhibit a clear monotonic trend.
Diffusion language models (DLMs) generate text by iteratively denoising masked sequences, but standard decoding either fixes the sequence length or relies on ad hoc stopping rules, often leading to unnecessary denoising steps. We recast length selection as a discrete-time survival problem over the end-of-sequence token and propose a plug-in, training-free length predictor that can be added to any existing DLM. Across reasoning and code-generation benchmarks, survival-guided length decoding speeds up inference by up to 7 times while preserving task accuracy. We further find that predicted lengths vary widely even within the same dataset, making model performance sensitive to the chosen length.
Ivan Kobyzev, Abbas Ghaddar, Yufei Cui
Huawei Noah’s Ark Lab, Montreal Research Center, Canada
Diffusion-based Large Language Models (D-LLMs) represent a promising frontier in generative AI, offering fully parallel token generation that can lead to significant throughput advantages and superior GPU utilization over the traditional autoregressive paradigm. However, this parallelism is constrained by the requirement of a fixed-size response length prior to generation. This architectural limitation imposes a severe trade-off: oversized response length results in computational waste on semantically meaningless padding tokens, while undersized response length causes output truncation requiring costly re-computations that introduce unpredictable latency spikes. To tackle this issue, we propose Predict-then-Diffuse, a simple and model-agnostic framework that enables compute-budgeted inference per input query by first estimating the response length and then using it to run inference with D-LLM. At its core lies an Adaptive Response Length Predictor (AdaRLP), which estimates the optimal response length given an input query. As a measure against under-estimating the response length and re-running inference with a higher value, we introduce a data-driven safety mechanism based on a small increase of the predicted length. As a whole, our framework avoids wasting computation on padding tokens, at the same time preserving output quality. Experimental validation on multiple datasets demonstrates that Predict-then-Diffuse significantly reduces computational costs (FLOP) compared to the default D-LLM inference mechanism, while being robust to skewed data distributions.
Michael Rottoli, Subhankar Roy, Stefano Paraboschi
Diffusion Language Models (DLMs) have attracted significant attention for their strong reasoning ability. However, under a bidirectional attention mechanism, DLMs operate over an exponentially large exploration space compared to autoregressive models (ARMs), making it challenging to focus on reasoning-guiding tokens under random masking. We define causal shortcuts as token chains that cover the full sequence and provide explicit guidance towards correct reasoning trajectories. We analyze the effects of causal shortcuts on the reasoning accuracy and convergence speed of DLMs, and find that they largely improve answer convergence efficiency and generation accuracy. Motivated by this, we propose a Causal Shortcut Learning (CSL) Framework for DLMs. Specifically, we introduce a step-by-step token extraction procedure to extract causal shortcuts from data, and apply parallel prioritized masking on these tokens during training to enable efficient and accurate convergence to correct answers via causal shortcuts. Extensive experiments across multiple reasoning benchmarks and two base models demonstrate that CSL consistently outperforms existing SFT-variant baselines, achieving an average improvement of 1.92% over SFT-only models, and up to 4.20% on MATH-500. The code is available at the \href{https://github.com/ZJUDianJin/Causal-Shortcuts-Learning}{https://github.com/ZJUDianJin/Causal-Shortcuts-Learning