WAM-Diff2: Hierarchical AR-to-Diffusion Distillation for Highly Efficient Autonomous Driving VLA
Zhihao Zhu, Hanlin Shang, Mingwang Xu, Feipeng Cai, Zhuolin He, Yaoyi Li, Jianhua Han, Hang Xu, Siyu Zhu
cs.RO, cs.AI, cs.CV
2026-08-02
A three-stage distillation pipeline converts a multi-task autonomous driving VLA from autoregressive decoding to parallel discrete diffusion, keeping planning, perception and QA on par while decoding throughput rises from 44.5 to 673.4 TPS (15.1x).
End-to-end driving VLAs (Vision-Language-Action models that unify multi-camera input, instructions and action output in one network) are almost all built on autoregressive decoding. Tokens come out one at a time, latency scales linearly with output length, and that is a hard constraint for a system that has to react in real time. Training has a second flaw: teacher forcing means the model only ever sees correct prefixes, so one bad token at inference time drags the rest of the sequence with it. That compounding drift is exposure bias.
The other camp, specialized diffusion policies such as DiffusionDrive, decode in parallel and run fast, but they are trained from scratch as single-task planners with no visual-linguistic reasoning. The field has been split between models that are smart but slow and models that are fast but narrow. This paper translates an existing multi-task autoregressive model into a diffusion model wholesale, and keeps both halves.
The core move is replacing the causal attention mask with a block-causal one: tokens inside a decoding block attend to each other bidirectionally, while blocks stay causal across boundaries. Autoregressive decoding is the B=1 special case, which gives the transition a continuous path. Generation switches to masked discrete diffusion: start from a fully masked sequence, predict all tokens in parallel each step, keep the highest-confidence subset, re-mask the rest, iterate. Decoding an L-token sequence drops from O(L) steps to O(T), with T far below L.
The cross-paradigm transfer is staged:
The architecture is Qwen3-VL (8B pairs a SigLIP2-SO400M encoder with 36 layers; 2B pairs SigLIP2-Large with 28), with bounding boxes and trajectories all routed through one text tokenizer and no task-specific heads. Two system-level optimizations stack on top: FlashInfer kernels tuned for block-causal attention, and CUDA Graphs to strip CPU launch overhead from the fixed-shape denoising loop.
Unified multi-task evaluation (one frozen 2B checkpoint across QA, detection and planning):
| Setting | DriveBench | LingoQA | COCO mAP | NAVSIM v1 PDMS | Decode TPS |
| AR baseline (B=1) | 51.23 | 68.40 | 39.20 | 88.14 | 44.5 |
| Diffusion B=32 | 48.80 | 65.80 | 36.30 | 87.44 | 124.8 |
| Diffusion B=32 + system opt | n/a | n/a | n/a | n/a | 673.4 |
On planning-specific NAVSIM, the diffusion model posts 88.3 PDMS (v1) and 88.6 EPDMS (v2), reaching 91.1 / 90.7 with score-based trajectory selection, ahead of ReCogDrive (90.8 / 83.6) and DriveVLA-W0 (90.2 / 86.5). Closed-loop Bench2Drive: 49.55% success rate and 78.93 driving score, with strong sub-scores on emergency braking (61.67%) and traffic signs (85.26%), though the driving score trails stronger specialized systems such as UniDriveVLA.
The ablation accounts for each stage: converting AR weights straight to bidirectional attention crashes to 84.1 PDMS; block-wise distillation recovers 87.7; cross-scale distillation lands at 88.3, on par with the AR baseline's 88.1. The exposure-bias claim is quantified on 12,146 paired NAVSIM samples: average per-waypoint L2 error drops 5.8% (0.5935 to 0.5589), and the margin widens monotonically from 0.002 at the first waypoint to 0.082 at the eighth, exactly where autoregressive drift bites. The speed ledger: 2.8x from the algorithm, another 1.7x from FlashInfer, another 3.1x from CUDA Graphs, 15.1x combined, taking per-token latency from 22.7 ms to 1.5 ms.
A practical settlement for a long-running paradigm fight. Instead of training a diffusion model from scratch, an off-the-shelf AR generalist can be translated into one with semantics largely intact. The recipe is stated as general, applicable to any pretrained AR driving VLA, which makes it a directly usable engineering path for latency-bound vehicle deployment.
There is a lesson for the wider LLM community too: the mismatched-attention, mismatched-objective problem of AR-to-diffusion transfer yields to curriculum block expansion plus same-paradigm distillation, and nothing in that recipe is driving-specific. Block size B doubles as a runtime latency dial, switchable from 4 to 32 at under half a point of planning cost.
The authors name two: discrete tokenization introduces spatial quantization artifacts that occasionally hurt high-precision trajectory smoothness, and the diffusion student's ceiling is set by the initial AR teacher's reasoning.
A few more from the read. The baseline is their own AR model, not the strongest AR system at that scale, so "parity" is parity with its own starting point. The 15.1x headline multiplies algorithmic and system gains; the pure paradigm shift is 2.8x, and the FlashInfer / CUDA Graphs numbers were measured separately on a CUDA platform while all model-quality results ran on Ascend 910C NPUs, so the two sets are never validated together. Detection (COCO mAP 39.2 to 36.3) and open-ended QA (LingoQA 68.4 to 65.8) do degrade, attributed to coordinate serialization being format-sensitive but left unmitigated. And on closed-loop Bench2Drive the diffusion model's driving score (78.93) sits below its own AR baseline (80.51), so the "no meaningful degradation" claim holds for open-loop metrics, not for the closed-loop headline; only emergency braking (61.67% vs 60.34%) edges out the baseline.