WAM-Diff2: Hierarchical AR-to-Diffusion Distillation for Highly Efficient Autonomous Driving VLA

Zhihao Zhu^1,* · Hanlin Shang^1,* · Mingwang Xu^1,* · Feipeng Cai^2,* · Zhuolin He¹ · Yaoyi Li²
Jianhua Han² · Hang Xu² · Siyu Zhu^1,†

¹Fudan University    ²Yinwang Intelligent Technology Co., Ltd.

^*Equal contribution    ^†Corresponding author

News

• 2026-08-16 — Released the initial training and inference code for WAM-Diff2.
• 2026-08-02 — Released WAM-Diff2 on arXiv.

Roadmap

• Release progressive block-wise adaptation code.
• Release hierarchical distillation code and configurations.
• Release GPU and NPU inference entry points.
• Release the WAM-Diff2 Block-32 checkpoint on Hugging Face.
• Release benchmark evaluation code and reproducible paper results.
• Release optimized FlashInfer and CUDA Graph inference.

Quick Start

Installation

Python 3.11–3.12 is supported. The reference environment uses PyTorch 2.8 and CUDA 12.x.

git clone https://github.com/fudan-generative-vision/WAM-Diff2.git
cd WAM-Diff2

conda create -n wam-diff2 python=3.12 -y
conda activate wam-diff2
pip install -r environment/requirements_cuda.txt
pip install -e .

The default configuration uses PyTorch SDPA. FlashAttention is optional and must match the installed PyTorch and CUDA versions.

Model checkpoint

Download the WAM-Diff2 checkpoint from Hugging Face and place it at:

checkpoints/WAM-Diff2-B32-2B/

Inference

Inference defaults to Block-32 decoding with 32 denoising steps and dynamic low-confidence remasking at a confidence threshold of 0.9.

Run inference on NVIDIA GPUs:

scripts/infer_gpu.sh \
  --model_id checkpoints/WAM-Diff2-B32-2B \
  --input_file /path/to/eval.json \
  --output_file outputs/predictions.json

Run inference on Ascend NPUs with a compatible torch-npu installation:

scripts/infer_npu.sh \
  --model_id checkpoints/WAM-Diff2-B32-2B \
  --input_file /path/to/eval.json \
  --output_file outputs/predictions.json

Set NUM_GPUS or NUM_DEVICES to launch multiple workers. The decoding defaults can be overridden with --block_size, --denoising_steps, --remasking_strategy, and --confidence_threshold.

Training

Data

Training data sample

A training JSON file is a list of samples in the following format:

[
  {
    "datasource": "Navsim",
    "id": "43ebecce459052f2",
    "image": [
      "data/opensource/navsim_dataset/sensor_blobs/test/2021.05.25.14.16.10_veh-35_01690_02183/CAM_F0/d28200c994c15229.jpg"
    ],
    "conversations": [
      {
        "from": "human",
        "value": "Here is a front-view images from a driving vehicle: <image>\nThe navigation information is: right\nThe current position is (0.00,0.00)\nThe current velocity is: (4.66,-0.09) and current accelerate is: (0.17,-1.53)\nInstruction: Based on the visual motion cues from the image (such as the relative speed of other vehicles and the changing distance to the intersection) and the provided telemetry, predict the optimal driving action for the next 4 seconds with 8 new waypoints."
      },
      {
        "from": "gpt",
        "value": "2.26,-0.17,4.43,-0.70,6.52,-1.55,8.48,-2.64,10.28,-3.88,11.99,-5.25,13.76,-6.83,15.59,-8.59"
      }
    ]
  }
]

Progressive block-wise adaptation

Run Block-4 adaptation from the prepared Qwen3-VL checkpoint on one GPU:

CONFIG_PATH=configs/training/block4.yaml scripts/train_block.sh

For multi-GPU training with explicit data paths:

NUM_GPUS=8 \
CONFIG_PATH=configs/training/block4.yaml \
scripts/train_block.sh \
  --dataset.path_or_dataset /path/to/train.json

The tutorial recipe is defined in configs/training/block4.yaml. Any --section.key value argument overrides the corresponding YAML field.

Continue the progressive adaptation with block8.yaml, block16.yaml, and block32.yaml.

Hierarchical distillation

Launch block-wise/model-wise distillation with Accelerate and DeepSpeed:

NUM_GPUS=8 scripts/train_distil.sh

The experiment and distributed-training settings are defined in configs/distillation/distil.yaml and configs/accelerate/, respectively. They can also be selected explicitly:

NUM_GPUS=8 \
EXPERIMENT_CONFIG=/path/to/distil.yaml \
ACCELERATE_CONFIG=/path/to/deepspeed_zero2.yaml \
scripts/train_distil.sh

Citation

If WAM-Diff2 is useful for your research, please cite:

@article{zhu2026wam,
  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}
}

Acknowledgements

We sincerely thank the teams behind NVIDIA NeMo AutoModel and Bard-VL. This project builds upon their excellent open-source work.