Spec-AUF: Accept-Until-Fail Training under Train-Inference Misalignment for Masked Block Drafters

Abstract: Speculative decoding accelerates autoregressive generation by drafting a block of tokens that the target model verifies left-to-right, committing only the longest accepted prefix. Block (DLM-style) drafters predict the whole block in parallel, which is fast but trained with a full-block cross-entropy that supervises every position against the gold continuation -- even though inference discards every token after the first rejection. Recent acceptance-aware objectives patch this by reweighting the full-block loss; we instead use teacher-forced learning as a motivation for how supervision should concentrate on the accepted prefix. A mask-only block drafter has no input-side channel for gold-prefix conditioning, so AUF approximates that prefix-sensitive supervision on the loss side by keeping the cross-entropy support only through the drafter's first predicted failure. AUF is a single, detached change to the CE support -- no auxiliary objective, no verifier rollouts, and no change to the inference pipeline or the exactness contract. Within fixed drafter backbones and serving settings on Qwen3-8B, AUF raises the DFlash drafter's average emitted length $τ$, averaged over six benchmarks, from 2.40 to 2.61, with a gain on every benchmark, and transfers to Domino's two-branch head (2.56 to 2.68). Two findings sharpen the picture: the decay-only baseline reaches higher token accuracy on the shared block mask yet decodes worse, and on DFlash, once AUF truncates the support, the standard exponential position-decay weighting becomes empirically inert.

• The paper introduces the Accept-Until-Fail (AUF) training objective that restricts supervision to the accepted prefix and the first failure, resolving train-inference misalignment.
• It demonstrates that AUF improves acceptance length and speculative decoding throughput compared to traditional decayed cross-entropy methods.
• AUF removes the need for manual position decay priors and induces a curriculum-like training dynamic, streamlining model optimization for block-masked drafters.

Accept-Until-Fail Training for Masked Block Drafters: A Technical Assessment of Spec-AUF

Introduction

Spec-AUF introduces a principled training regime, Accept-Until-Fail (AUF), for speculative block drafting in LLMs that operate under the block-masked (non-autoregressive) paradigm. The primary contribution is a training objective that directly resolves the fundamental train-inference misalignment present in state-of-the-art block drafters used for speculative decoding—settings where only the contiguous accepted prefix from the draft model is committed, while suffix tokens after the first verification failure are discarded. Traditional objectives, such as fixed-position-decayed cross-entropy, supervise all suffix tokens, resulting in suboptimal acceptance profiles and unnecessary optimization of tokens irrelevant at inference time.

The AUF approach, motivated by teacher-forced maximum likelihood for autoregressive models but tailored for the block-drafter setting where input-side conditioning is absent, restricts cross-entropy supervision strictly to the accepted prefix and the breaker token (the first failure). This approach contrasts with recent acceptance-aware methods and overcomes the lack of a gold-prefix channel by performing the prefix-sensitive supervision entirely on the loss side.

Background and Objective Mismatch

Block-masked drafters like DFlash or Domino predict the entire block of candidate tokens in parallel, without input-side gold-token conditioning. During inference, speculative decoding operates by accepting only the contiguous prefix agreed on by a much larger target model; tokens after the first failure are ignored. As a result, maximizing overall block accuracy does not align with maximizing throughput or minimizing the number of verifications.

Previous block-level objectives, most notably DFlash’s exponentially decayed cross-entropy, encode a position prior that heuristically emphasizes early positions, operating under the assumption that earlier tokens have a larger impact on accepted length. However, they continue to allocate nonzero credit to suffix positions that would not be verified (or even observed) under the left-to-right speculative contract. This induces a structural gap between the training and deployment objectives.

Figure 1: Teacher-forced training for an autoregressive draft model; gold prefixes enable prefix-conditioned CE, which is impossible for block-masked drafters.

Figure 2: Full-block decayed CE (left) supervises all positions, including those after the first failure. AUF (right) restricts supervision to the accepted prefix plus the first failure.

The Accept-Until-Fail (AUF) Objective

Spec-AUF postulates that the proper training objective for block-masked drafters is not to maximize additive token-level accuracy across the entire block, but rather to condition loss allocation on the first verification failure. The gold cross-entropy loss remains the learning signal, but it is withheld from all positions beyond this first error.

Formally, for a masked candidate block, AUF computes the proposal $\hat y_{1:B}$ and determines the first index $j^*$ where $\hat y_{j^*} \ne y_{j^*}$ (comparing only for valid positions). The supervised set becomes $\{1, ..., j^*\}$, encompassing the longest match plus the breaker token. Supervision for all $i > j^*$ is masked out.

This is a zero-temperature, maximum-a-posteriori plug-in estimator of the prefix acceptance gate present in soft-gated objectives like D-PACE and SpecDiff-2. Notably, AUF introduces no auxiliary objective, requires no verifier rollouts on teacher trajectories, and makes no changes to architecture or inference procedures.

One critical consequence is that AUF eliminates the need for manually set position priors—decay hyperparameters—that were previously entrenched in block drafter pipelines. Empirically, when AUF is applied, reintroducing exponential position decay within the accepted region is inert: acceptance lengths do not increase further.

Comparative Analysis with Other Objectives

AUF differs from competing acceptance-aware objectives both in mechanism and effect:

• GRIFFIN: Drops the loss past Top-$K$ mismatches but does not perform first-match truncation and uses a coarser gating criterion.
• D-PACE: Softens the prefix acceptance gate, assigning fractional credit for token correctness in a prefix-product, but keeps full-block support and requires additional surrogate-based dynamic weighting.
• SpecDiff-2: Optimizes a soft prefix-product reward via verifier rollouts and distillation, at a substantial computational and implementation cost.

Spec-AUF positions itself as the minimal intervention for aligning training and inference mechanics. It removes a source of hand-designed inductive bias (decay scheduling) and places position selection entirely under the control of the model’s actual predictive ability.

Empirical Results

Evaluations are performed using the Qwen3-8B backbone, trained on ShareGPT, over six standard benchmarks (GSM8K, MATH-500, HumanEval, MBPP, MT-Bench, Alpaca). The principal metric is acceptance length $\tau$, which quantifies the mean number of drafted tokens accepted per verification pass plus the correction token.

Key numerical findings:

• For DFlash, the mean acceptance length $\tau$ improves from $2.40$ (decayed CE) to $2.61$ under AUF, with every benchmark showing a positive gain.
• Transfer to Domino’s two-branch setup also shows an increase ($j^*$0).
• Reintroducing exponential decay inside AUF’s support provides no further improvement—for DFlash, AUF-only and AUF+decay are virtually identical.
• The decay-only baseline achieves higher block-wide token accuracy (measured on the mask of natively-valid positions) but underperforms for throughput, validating the hypothesis that block-wide accuracy is insufficient for speculative decoding.

  Figure 3: Two-dimensional trajectory for active token ratio and accuracy on the AUF support; AUF enables expansion of support as the model improves, contrasting with fixed-support methods.

  

  Figure 4: DFlash training dynamics show that AUF’s support expands with training, and AUF achieves higher acceptance length $j^*$1 with lower common-token accuracy but higher targeted support accuracy.

  

  Figure 5: Per-position conditional acceptance rates $j^*$2 under DFlash and Domino; AUF variants consistently lift per-position acceptance relative to baselines.

In Domino, variations of AUF applied to only the base branch, both branches, or by sharing across branches all yield consistent gains over decayed CE, though the margin is smaller due to the stronger baseline performance of Domino.

Theoretical and Practical Implications

Spec-AUF’s findings have important implications for both practice and theory in draft model training for speculative decoding:

• Train-Inference Consistency: The work provides strong evidence that objectives aligned with downstream inference mechanics—here, the left-to-right speculative contract—yield better practical system throughput than heuristically aligned surrogates.
• Removal of Inductive Bias: By eliminating the necessity for position decay priors, Spec-AUF reduces manual hyperparameter tuning and architectural constraints, providing a cleaner design space for future block drafters.
• Curriculum-Like Training Dynamics: AUF implicitly induces an easy-to-hard curriculum, with training focus expanding rightward as the drafter’s predictions improve; early training operates on short prefixes, later training on longer ones, with expansion governed by the model’s own first-error behavior.

Despite the appeal, limitations exist: all comparisons are at fixed block size and scale, head-to-head with only some acceptance-aware baselines, and further interpretability work is needed to explain causal pathways. Whether AUF’s effects generalize to substantially larger models or alternative serving settings, or whether hybrid methods (e.g., partial reweighting inside AUF support) would recapture the negligible gains seen in some architectures, remains an open research problem.

Future Directions

Given AUF’s simplicity and empirical efficacy, future work should include:

• Extending AUF across broader block-size and model-scale regimes to characterize its scaling properties.
• Incorporating AUF into training procedures for fully diffusion-based or hybrid drafters, where the alignment between training and left-to-right decoding contracts is even more pronounced.
• Hybridizing AUF with sequence-level or RL-style objectives, evaluating whether graded credit or adaptive support expansion can provide additional gains.
• Deep analysis of the interaction between AUF-induced curricula and convergence properties, particularly in regimes where block sizes are dynamically scheduled during training.

Conclusion

Spec-AUF delivers a formally motivated, computationally efficient resolution to the core train-inference misalignment in speculative block drafting. By restricting supervision to the prefix up to the first verification failure, block-masked drafters gain measurable throughput improvements and shed unnecessary inductive biases. The demonstration that position decay becomes entirely inert under AUF in DFlash establishes AUF as a strong candidate for a new default in speculative-drafter training objectives. Future work remains to validate this setting under scale, in diverse model families, and against stronger baselines using more sophisticated alignment techniques.