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
把多任务自动驾驶 VLA 从逐 token 自回归改造成并行离散扩散,三阶段蒸馏后规划、感知、问答性能持平,解码吞吐从 44.5 提到 673.4 TPS(15.1 倍)。
端到端自动驾驶的 VLA(Vision-Language-Action,把多路摄像头画面、指令和动作输出统一进一个模型的架构)几乎都建在自回归解码上:token 一个一个生成,延迟随输出长度线性涨,对要实时响应的驾驶是硬伤。训练上还有第二个问题:teacher forcing 让模型只见过「前缀全对」的输入,推理时前面错一个 token,后面跟着错下去,误差在长序列里滚雪球,这就是 exposure bias。
另一条路线,专用扩散策略(比如 DiffusionDrive),并行出轨迹、延迟低,但都是从头训练的单任务模型,没有视觉语言推理能力,看不懂场景也答不了问。两条路线各占一头:要么聪明但慢,要么快但窄。这篇要的是把已有的多任务自回归模型整体翻译成扩散模型,聪明和快都要。
核心动作是把因果 attention 掩码换成块因果(block-causal):一个解码块内的 token 互相可见(双向),块与块之间仍然保持因果。自回归模型就是这个结构在块大小 B=1 时的特例,改造就有了连续的路径。生成范式换成 masked 离散扩散:从全掩码序列出发,每步并行预测所有 token,留下置信度最高的一批,其余重新掩码,再迭代。输出长度为 L 的序列从 O(L) 步压到 O(T) 步,T 远小于 L。
跨范式的迁移不能一步到位,论文把它拆成三阶段:
架构基于 Qwen3-VL(8B 用 SigLIP2-SO400M 视觉编码器 + 36 层,2B 用 SigLIP2-Large + 28 层),检测框和轨迹全部走统一文本 tokenizer,没有任何任务专属输出头。系统层面上了两项工程优化:FlashInfer 定制块因果 attention kernel,CUDA Graphs 消掉固定形状去噪步的 CPU 启动开销。
多任务统一评测(单一 2B checkpoint 同时考问答、检测、规划):
| 设置 | DriveBench | LingoQA | COCO mAP | NAVSIM v1 PDMS | 解码吞吐 TPS |
| AR 基线 (B=1) | 51.23 | 68.40 | 39.20 | 88.14 | 44.5 |
| 扩散 B=32 | 48.80 | 65.80 | 36.30 | 87.44 | 124.8 |
| 扩散 B=32 + 系统优化 | — | — | — | — | 673.4 |
规划专项上,扩散版 NAVSIM v1 PDMS 88.3(v2 EPDMS 88.6),配打分选轨后 91.1 / 90.7,超过 ReCogDrive(90.8 / 83.6)与 DriveVLA-W0(90.2 / 86.5)。闭环 Bench2Drive 成功率 49.55%、驾驶分 78.93,驾驶分高于 UniDriveVLA(78.37)等专用系统,但低于自家 AR 基线(80.51);急刹 61.67% 略好于基线的 60.34%,交通标牌 85.26% 低于基线的 89.25%。
消融把三阶段的账算清楚了:直接把 AR 权重改成双向 attention 掉到 84.1 PDMS,逐块蒸馏拉回 87.7,跨尺度蒸馏补到 88.3,与 AR 基线 88.1 持平。exposure bias 的收益有量化:12146 对 NAVSIM 样本上,平均逐路点 L2 误差降 5.8%(0.5935→0.5589),且领先幅度从首个路点的 0.002 单调扩大到第 8 个路点的 0.082,正对自回归误差累积的病灶。速度账:算法层 2.8 倍,FlashInfer 再乘 1.7,CUDA Graphs 再乘 3.1,合计 15.1 倍,单 token 延迟 22.7ms 降到 1.5ms。
路线之争的实用解法。自回归阵营和多任务阵营、扩散阵营和低延迟阵营吵了很久,这篇给出的答案是:不用从头训练扩散模型,现成的 AR 通用模型可以整体翻译过去,语义能力基本不丢。它是一个通用配方,论文明说任何预训练 AR 驾驶 VLA 都能这么转,对车端部署这类延迟硬约束场景是直接可用的工程路径。
对更广的 LLM 社区也有参照价值:AR→扩散的迁移难点(attention 不匹配、优化目标不一致)用课程式块扩展加同范式蒸馏来解,这个思路不只在驾驶任务上成立。块大小 B 是个运行时可调的延迟旋钮,4 到 32 随时切换,性能损失 0.5 个点以内,部署侧很实用。
论文自己列了两条:离散 tokenization 会引入空间量化误差,高精度轨迹的平滑度偶尔受损;扩散学生的能力上限被初始 AR 老师封顶,老师不会的学生也学不来。
读完补充几点。基线是自家 AR 模型,不是该规模下的最强 AR 系统,「持平」是对着自己的起点说的;题目里的 15.1 倍是算法与系统优化的叠乘,纯范式转换只有 2.8 倍,且 FlashInfer / CUDA Graphs 的收益在 CUDA 平台单独测得,主实验全部跑在昇腾 910C 上,两边数字没有交叉验证。视觉检测(COCO mAP 从 39.2 掉到 36.3)和开放问答(LingoQA 68.4 掉到 65.8)的退化是真实存在的,论文归因于坐标序列化对格式约束敏感,但没有给缓解方案。闭环 Bench2Drive 上扩散版驾驶分(78.93)低于自家 AR 基线(80.51),「无实质退化」的结论在闭环主要指标上并不成立,只有急刹一项以 61.67% 对 60.34% 略胜。