Flash-VAED: Plug-and-Play VAE Decoders for Efficient Video Generation

Lunjie Zhu¹, Yushi Huang¹, Xingtong Ge¹, Yufei Xue¹, Zhening Liu¹, Yumeng Zhang¹, Zehong Lin², Jun Zhang¹*

¹iComAI Lab, The Hong Kong University of Science and Technology    ²School of Data Science, Lingnan University

The Forty-Third International Conference on Machine Learning (ICML), 2026

[Paper] [Code] [Model]

* Corresponding author: [email protected]

---

📝 Abstract

Latent diffusion models have enabled high-quality video synthesis, yet their inference remains costly and time-consuming. As diffusion transformers become increasingly efficient, the latency bottleneck inevitably shifts to VAE decoders. To reduce their latency while maintaining quality, we propose a universal acceleration framework for VAE decoders that preserves full alignment with the original latent distribution. Specifically, we propose (1) an independence-aware channel pruning method to effectively mitigate severe channel redundancy, and (2) a stage-wise dominant operator optimization strategy to address the high inference cost of the widely used causal 3D convolutions in VAE decoders. Based on these innovations, we construct a Flash-VAED family. Moreover, we design a three-phase dynamic distillation framework that efficiently transfers the capabilities of the original VAE decoder to Flash-VAED. Extensive experiments on Wan and LTX-Video VAE decoders demonstrate that our method outperforms baselines in both quality and speed, achieving approximately a 6× speedup while maintaining the reconstruction performance up to 96.9%. Notably, Flash-VAED accelerates the end-to-end generation pipeline by up to 36% with negligible quality drops on VBench-2.0.

Qualitative and quantitative comparisons of video reconstruction. Flash-VAED (bottom) offers the fastest decoding speed with minimal fidelity loss vs. the original VAE decoder (top) and prior baselines (middle).

Overview of the Flash-VAED architecture. Stage-wise dominant operator optimization (left) substitutes CausalConv3D with stage-specific efficient operators; independence-aware channel pruning (right) reduces channels to 12.5%–25% of the original with minimal quality loss.

🤗 Model Weights

Weights are hosted on Hugging Face: Aoko955/Flash-VAED.

Download:

pip install -U "huggingface_hub[cli]"
huggingface-cli download Aoko955/Flash-VAED \
  --local-dir . \
  --include "models/wan/*.pth" \
  --include "models/ltx/*.pth" \
  --include "models/ltx/teacher/*"

Expected files after download:

models/wan/Flash_VAED_Wan.pth
models/wan/Wan_VAE_Teacher.pth
models/ltx/Flash_VAED_LTX.pth
models/ltx/teacher/config.json
models/ltx/teacher/diffusion_pytorch_model.safetensors

💻 Installation

conda create -n flashvaed python=3.10 -y
conda activate flashvaed
pip install -r requirements.txt

🚀 Inference

Reconstruct a video (teacher encode → student decode)

python infer.py --model wan --input demo.mp4 --output outs/wan
python infer.py --model ltx --input demo.mp4 --output outs/ltx

With more options:

python infer.py --model wan --input demo.mp4 --output outs/wan \
  --num_frames 81 --img_h 480 --img_w 832 --fps 8 --device cuda:0

Outputs: frames/*.png and optional recon.mp4 (disable with --no_mp4).

Decode from a latent

python infer.py --model wan --mode decode --latent z.pt --output outs/decode

z.pt should be a torch.Tensor, or a dict with key latent / latents / z.

📈 Results

Video reconstruction (UCF-101)

Flash-VAED vs. original VAE decoders and competitive baselines on RTX 5090D / Jetson Orin.

Visual comparison of video generation results. Flash-VAED (bottom) matches the original Wan 2.1 VAE (top) in fidelity and texture, while LightVAE (middle) shows severe artifacts.

📂 Repository Layout

Flash-VAED/
  infer.py
  requirements.txt
  assets/
  models/
    wan/
      model_hybrid_aggressive.py   # Wan student
      model_original.py            # Wan teacher (encode)
    ltx/
      ltx_prune_1_4.py              # LTX student
      vendor/                       # minimal local helpers (no pip diffusers)
      teacher/config.json           # LTX teacher config (weights on HF)

📚 Citation

If you find this work useful, please cite:

@inproceedings{zhu2026flashvaed,
  title={Flash-{VAED}: Plug-and-Play {VAE} Decoders for Efficient Video Generation},
  author={Lunjie Zhu and Yushi Huang and Xingtong Ge and Yufei Xue and Zhening Liu and Yumeng Zhang and Zehong Lin and Jun Zhang},
  booktitle={Forty-third International Conference on Machine Learning},
  year={2026},
  url={https://openreview.net/forum?id=PDBLtVDb0d}
}

⚖️ Acknowledgements & License

This work builds on:

• Wan2.1
• LTX-Video
• Turbo-VAED

Please respect upstream Wan / LTX licenses when redistributing weights.