token-decoding acceleration
Preprint · 2026
WAM-Diff2
Hierarchical AR-to-Diffusion Distillation
for Highly Efficient Autonomous Driving VLA
1 Fudan University
2 Yinwang Intelligent Technology Co., Ltd.
* Equal contribution † Corresponding author
01 · Overview
Abstract
Vision-Language-Action models unify scene understanding, perception, and motion planning, but sequential autoregressive decoding makes them slow and susceptible to exposure bias. Specialized diffusion policies execute in parallel, yet typically sacrifice the broad visual-linguistic reasoning learned by large autoregressive models.
WAM-Diff2 bridges this divide with a multi-task discrete diffusion VLA and a three-stage hierarchical distillation strategy. Progressive block-wise adaptation, block-wise distillation, and model-wise cross-scale distillation preserve the semantic foundations of a pre-trained generalist while converting it into a parallel refinement architecture. Across driving understanding, perception, and planning benchmarks, WAM-Diff2 matches the autoregressive baseline while mitigating long-horizon error accumulation. The transition yields a 2.8× decoding speedup, scaling to 15.1× with FlashInfer and CUDA Graphs.
optimized decode tokens / second
lower average waypoint L2 error
02 · Method
From sequential generation
to parallel refinement
WAM-Diff2 replaces token-by-token causal decoding with block-causal attention: tokens refine each other bidirectionally within a block, while causality is preserved across blocks. A staged curriculum protects the model's multi-task knowledge throughout the architectural shift.
01
Stage I
Progressive Block-Wise Adaptation
Gradually expand the decoding block from B = 1 → 4 → 8 → 16 → 32, smoothly relaxing causal attention without destabilizing the pre-trained model.
02
Stage II
Block-Wise Distillation
A stable small-block diffusion teacher guides progressively larger-block students on intermediate noisy states, recovering quality and reducing exposure bias.
03
Stage III
Model-Wise Cross-Scale Distillation
An advanced 8B diffusion teacher transfers holistic semantic capabilities into an efficient 2B student sharing the same generation paradigm.
03 · Experiments
A unified model for understanding, perception, and planning.
Preserving performance at higher decoding throughput.
A single frozen checkpoint is evaluated across DriveBench, LingoQA, COCO, NAVSIM v1, and NAVSIM v2. Larger blocks increase parallelism while retaining competitive understanding, perception, and planning performance.
Unified multi-task evaluation
Performance & decoding efficiency
WAM-Diff2
| Method | DriveBench ↑ | LingoQA ↑ | COCO ↑ | NAVSIM v1 ↑ | NAVSIM v2 ↑ | Decode TPS ↑ |
|---|---|---|---|---|---|---|
| Qwen3-VL-2B | 47.68 | 48.00 | 34.20 | — | — | 44.5 |
| Ours-2B B=1 · AR | 51.23 | 68.40 | 39.20 | 88.14 | 88.32 | 44.5 |
| WAM-Diff2 B=4 | 48.92 | 68.20 | 37.80 | 87.87 | 87.99 | 68.3 / 401.4 |
| WAM-Diff2 B=8 | 48.81 | 68.00 | 37.30 | 87.57 | 87.78 | 89.5 / 561.6 |
| WAM-Diff2 B=16 | 48.77 | 66.00 | 36.50 | 87.32 | 87.30 | 108.4 / 635.3 |
| WAM-Diff2 B=32 | 48.80 | 65.80 | 36.30 | 87.44 | 87.50 | 124.8 / 673.4 |
Dual TPS values denote standard / FlashInfer + CUDA Graph optimized execution. NAVSIM v1 values omit score-based candidate selection for a direct architectural comparison.
Long-horizon robustness
Less drift,
step after step.
Autoregressive prediction commits permanently to early mistakes. Bidirectional refinement lets WAM-Diff2 revise uncertain tokens within each block, reducing the average per-waypoint L2 error by 5.8% across 12,146 paired NAVSIM samples.
0.5935→0.5589average L2 error
Qualitative results
Driving understanding, visual perception,
and motion planning.
04 · Reference
Citation
If you find WAM-Diff2 useful in your research, please consider citing our paper.
@article{zhu2026wamdiff2,
title = {WAM-Diff2: Hierarchical AR-to-Diffusion
Distillation for Highly Efficient Autonomous Driving VLA},
author = {Zhu, Zhihao and Shang, Hanlin and Xu, Mingwang and
Cai, Feipeng and He, Zhuolin and Li, Yaoyi and
Han, Jianhua and Xu, Hang and Zhu, Siyu},
journal = {arXiv preprint arXiv:2608.01035},
year = {2026}
}