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

Abstract

Verbatim: "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)."

Key Contributions

• Accept-Until-Fail (AUF) — truncates the cross-entropy support at the drafter's first greedy token error (the "breaker" token is kept in support), rather than supervising the full block. Credit indicator c_i = 1[i ≤ j*], where j* is the first exact mismatch among valid positions. Support is computed from detached greedy predictions (no gradient flow through the selection).
• Train-inference alignment — directly targets the structural mismatch: full-block CE supervises positions that inference would never reach after an earlier rejection.
• Removes a hyperparameter — replaces the fixed exponential position-decay weighting (γ) with a model-dependent acceptance boundary.
• No infrastructure changes — "a single, detached change to the CE support—no auxiliary objective, no verifier rollouts, and no change to the inference pipeline."

Methodology

Motivated by teacher forcing: gold-prefix conditioning should decide which positions receive supervision. A mask-only drafter has no input-side channel for that, so AUF enforces it on the loss side by truncating CE support at the first predicted failure. Contrasted against D-PACE's graded/soft reweighting — AUF is a hard plug-in gate. Two optimization axes are tracked during training: active-token ratio and support accuracy.

Results

DFlash (Qwen3-8B, block size B=16)

Per-benchmark (T=0): GSM8K 2.29→2.50-2.51 · MATH-500 2.39→2.64 · HumanEval 2.86→3.09 · MBPP 2.85→3.08-3.09 · MT-Bench 2.03→2.18-2.19 · Alpaca 2.01→2.14-2.15. Gain on every benchmark.

Critical finding: AUF-only and AUF+decay are indistinguishable at both temperatures — "decay is fully redundant once the support is selected." I.e., once AUF truncates support, the standard exponential position-decay weighting becomes empirically inert (contradicts the assumption that decay and prefix-awareness are complementary).

Second finding: the decay-only baseline achieves higher common-token accuracy on the fixed block mask yet decodes worst — per-position token accuracy is not the right proxy for accepted length.

Domino (Qwen3-8B, same backbone, two-branch head)

All seven AUF variants beat the baseline; the best variant differs by temperature.

Protocol

Six benchmarks (GSM8K 200, MATH-500 200, HumanEval 164, MBPP 200, MT-Bench 80, Alpaca 200). SGLang backend, batch size 1, single GPU, max_new_tokens=512, greedy (T=0) and sampling (T=1). Epoch-6 checkpoints reported. Drafters: DFlash (5-layer block-diffusion, masked CE, target layers {1,9,17,25,33}, B=16) and Domino (same backbone + two-branch base/causal-correction head). Training data: ShareGPT (6 epochs, thinking disabled).

Verification note (against prior KB record)

• Validated model scale — "reportedly 8B": CONFIRMED, and it is the only scale. The paper's own limitation: "All results are at a single configuration: block size B=16, target Qwen3-8B, and ShareGPT training data... We do not vary B or model scale." So "validated at 8B" is exactly right; nothing above or below 8B is tested.
• τ figures 2.40→2.61 (DFlash) and 2.56→2.68 (Domino), and "no inference-pipeline change," all confirmed verbatim.

Limitations (author-stated)

Single model scale (Qwen3-8B only); fixed block size (B=16 only); no head-to-head retraining against D-PACE or other baselines under matched compute; uncertainty whether AUF's stable SFT-side update is preferable to a tuned on-policy RL objective or merely cheaper; the causal mechanism behind the gains is not definitively established. Preprint, not peer-reviewed.

Source: arXiv:2607.01893 — Spec-AUF: Accept-Until-Fail Training under Train-Inference Misalignment for Masked Block Drafters, Yang & Li, 2 July 2026. Full text retrieved via arXiv HTML5 on 2026-07-23.