Learn2PD: Adaptive Parallel Decoding

• Learn2PD is an adaptive parallel decoding framework that dynamically finalizes tokens in diffusion-based LLMs to reduce iterative denoising steps.
• It employs a lightweight filter model trained to approximate an oracle Extremely Greedy Parallel strategy, ensuring context-sensitive token unmasking.
• Incorporating an End-of-Text Prediction mechanism, Learn2PD prevents unnecessary computation, significantly boosting throughput on large generation tasks.

Learn2PD refers to a set of adaptive parallel decoding methodologies for diffusion-based LLMs (dLLMs) that accelerate inference by learning when to unmask and finalize token predictions during iterative denoising. This approach employs a lightweight filter model trained to approximate an oracle “Extremely Greedy Parallel” (EGP) strategy, enabling dynamic, input-specific acceleration of generation without degrading output quality. Additionally, an End-of-Text Prediction (EoTP) mechanism is introduced to prevent unnecessary computation for padding tokens past the end-of-sequence marker.

1. Motivation and Scope

Autoregressive decoding in conventional LLMs requires $\mathcal{O}(n)$ sequential steps for $n$ output tokens, which fundamentally limits throughput during inference. Recent dLLMs enable parallel token generation via iterative denoising, but typical parallel decoding strategies rely on fixed or input-agnostic heuristics (e.g., static confidence thresholds). These heuristics lead to suboptimal trade-offs between speed and quality, since they cannot adapt to the tokenwise stability or uncertainty variance present in diverse NLP tasks and inputs (Bao et al., 29 Sep 2025).

Learn2PD provides a flexible, learned alternative: for each token at each step in the denoising process, a filter model predicts whether the token’s present value matches the final reference and thus can be “unmasked” (finalized) for all subsequent iterations. The purpose is to mimic an oracle that always unmasks tokens at the earliest safe moment, minimizing redundant decoding. The filter enables dynamic and context-sensitive parallelism, boosting throughput on various generation tasks.

2. Adaptive Filter Model

The filter model $f_\theta$ is trained as a binary classifier over token decisions. For each token position, it receives features (generally logit-based model confidence scores from the denoising trajectory) and outputs the probability that the token should be unmasked, i.e., its current prediction matches the final output.

Model architecture:

• Input: Block-wise confidence features for tokens, extracted from the denoising sequence. Block sizes such as 32 are typical.
• Layers: Two-layer multilayer perceptron (MLP), total parameter count $\approx$2,112 for block size 32 (as reported).
• Output: Sigmoid activation for a binary decision per token.

Training objective:

• Binary Cross Entropy (BCE) Loss:

$$\mathcal{L}_{BCE} = -\frac{1}{m} \sum_{i=1}^m[y_i \log \sigma(z_i) + (1-y_i)\log (1-\sigma(z_i))]$$

where $z_i$ is the logit output for token $i$, $\sigma$ is the sigmoid, and $y_i$ is the oracle label (1: unmask, 0: remask).

The filter is trained post-hoc, after dLLM model fine-tuning, using reference answers to supervise the correct unmasking pattern. The oracle EGP strategy serves as ground truth.

Characteristically, training is very compute-efficient: data collection (on 4 GPUs) and filter training (on a single T4 GPU) require only minutes, with convergence achieved in approximately 5,000 epochs at learning rate $0.001$ using AdamW (Bao et al., 29 Sep 2025).

3. End-of-Text Prediction (EoTP)

EoTP supplements Learn2PD by eliminating inefficiency in handling padding tokens after an End-of-Text ([EoT]) marker. Conventional dLLMs may continue denoising for all positions up to the preallocated sequence length, even after [EoT] appears, consuming substantial computation—up to 90% of decoding cost for long outputs.

EoTP intervenes as follows:

• Once [EoT] is confidently generated in a decoding block, all subsequent token positions are immediately assigned [EoT].
• Decoding terminates for those positions, and further iterations are halted.

This mechanism is particularly beneficial for large generation lengths (e.g., 1024 tokens), further multiplying throughput gains, as evidenced in benchmarks.

4. Oracle Parallel Decoding and Training Paradigm

The oracle EGP strategy provides the upper bound for parallel decoding—every token is unmasked at the earliest step when its prediction matches ground truth. Learn2PD’s filter is trained to approximate this oracle. Practical implementation utilizes reference completions to assign labels for the filter. During inference, the filter’s binary output determines, per token, whether to keep remasking or finalize its value.

The learning paradigm thus marries efficiency (reduced iterations) with error minimization: since the filter is trained using reference answers rather than heuristics, it better respects per-token uncertainty and adapts to input content.

5. Performance Results and Speed-Quality Trade-Offs

Comprehensive evaluation is presented on the LLaDA benchmark, notably on LLaDA-8B-Instruct:

• On GSM8K (5-shot) with generation length 1024, Learn2PD plus EoTP raises throughput from 0.54 tokens/second (TPS) to 12.26 TPS, for a $22.58\times$ speedup.
• When combined with KV-Cache optimizations (“Dual Cache” method), speedup reaches $57.51\times$.
• Accuracy and quality metrics remain stable, with only 1–2 points deviation from baseline scores.
• Performance is robust across diverse tasks (GSM8K, Math, HumanEval, MBPP) and generation lengths.

The table below summarizes throughput improvements per task and decoding length:

Task	Gen Length	Baseline TPS	Learn2PD+EoTP TPS	Speedup
GSM8K	1024	0.54	12.26	22.58×
Math	1024	0.55	12.53	22.78×
MBPP	1024	0.56	12.66	22.65×
HumanEval	1024	0.56	12.71	22.70×

This direct empirical evidence affirms that substantial acceleration is possible with negligible loss of generation quality (Bao et al., 29 Sep 2025).

6. Comparison with Prior Parallel Decoding Methods

Fixed-threshold parallel decoding unmask tokens based solely on confidence values exceeding a common preset threshold, leading to the same criterion for all inputs and positions. Such settings cannot adapt to positional uncertainty; tokens requiring more denoising may be prematurely finalized or inefficiently deferred.

Learn2PD’s adaptive filter:

• Learns context-sensitive patterns in block-wise confidence.
• Tailors unmasking decisions per token, per input, yielding closer approximation to oracle EGP.
• Orthogonally composes with other optimizations (KV-caching), suggesting modularity.

A plausible implication is that the learned approach is preferable whenever input variability or tokenwise semantic complexity is significant.

7. Applications and Limitations

Learn2PD is suitable for any setting where fast inference in dLLMs is desired, including:

• Real-time conversational agents and chatbots.
• Large-scale question answering and summarization.
• Code generation and program synthesis (HumanEval, MBPP benchmarks).
• Industry deployments with cost-sensitive throughput requirements.

Since the filter model is lightweight and trained post-hoc, integration does not require modification of the underlying dLLM parameters. The principal limitation, as reported, is the necessity for an extra training phase for filter optimization, but given minute-level GPU requirements this overhead is minimal (Bao et al., 29 Sep 2025).

In summary, Learn2PD provides a dynamic, inference-side acceleration framework for diffusion-based LLMs, leveraging adaptive parallel decoding with negligible compromise to output quality and offering substantial gains in real-world throughput.