Astraea: A Token-wise Acceleration Framework for Video Diffusion Transformers

TL;DR

Astraea is a framework that accelerates video diffusion transformers by optimizing token selection and attention strategies, achieving significant speedups with minimal quality loss.

Contribution

Astraea introduces a novel token selection and sparse attention method combined with an evolutionary search for optimal configurations in vDiT-based video generation.

Findings

Up to 2.4× inference speedup on a single GPU

Achieves up to 13.2× speedup on 8 GPUs

Over 10 dB improvement in video quality compared to state-of-the-art

Abstract

Video diffusion transformers (vDiTs) have made tremendous progress in text-to-video generation, but their high compute demands pose a major challenge for practical deployment. While studies propose acceleration methods to reduce workload at various granularities, they often rely on heuristics, limiting their applicability. We introduce Astraea, a framework that searches for near-optimal configurations for vDiT-based video generation under a performance target. At its core, Astraea proposes a lightweight token selection mechanism and a memory-efficient, GPU-friendly sparse attention strategy, enabling linear savings on execution time with minimal impact on generation quality. Meanwhile, to determine optimal token reduction for different timesteps, we further design a search framework that leverages a classic evolutionary algorithm to automatically determine the distribution of the token budget effectively. Together, Astraea achieves up to 2.4$\times$ inference speedup on a single GPU with great scalability (up to 13.2$\times$ speedup on 8 GPUs) while achieving up to over 10~dB video quality compared to the state-of-the-art methods ($<$0.5\% loss on VBench compared to baselines).