GeoShrink speeds up diffusion transformers for better image and video generation
GeoShrink: Accelerating Diffusion Transformers with Two Lines of Code
Computer Vision and Pattern Recognition
Summary
Diffusion transformers are powerful but slow because they need to run many steps to create images, videos, or sounds. The authors propose GeoShrink, a simple way to skip some steps without losing quality by predicting outputs based on previous ones. This method works without extra training and keeps the same overall schedule, making generation much faster. Tests show it improves quality measures while running about five times faster on various types of data.
What this means in practice
- •For computer vision engineers: Accelerate image and video generation pipelines while improving output quality using fewer model evaluations per sample.
- •For audio processing developers: Speed up audio and music generation models by predicting intermediate outputs instead of fully recomputing all steps.
Authors
Haosen Li, Wenshuo Chen, Shaofeng Liang, Lei Wang, Bowen Tian, Yutao Yue
Abstract
Diffusion transformers incur substantial inference cost through repeated model evaluations along a sampling trajectory. We introduce GeoShrink, a training-free acceleration method that retains the original solver grid while evaluating the model only at a prescribed set of anchors. At skipped stages, GeoShrink predicts the solver-facing output by adding a geometrically retained fraction of the latest observed innovation to the most recent exact output. We derive this rule from chordal tangent transport and round-trip line projection, and establish a geometric anchor-spacing principle that minimizes the largest adjacent gap expansion under fixed coverage and first span. The analysis characterizes the geometric closure and propagation of prediction errors without assuming access to future model outputs. Experiments cover image, video, motion, and audio generation, together with adapted 3D backends. At approximately $5\times$ acceleration, GeoShrink improves FLUX PSNR by 3.10 dB over the strongest listed baseline. On HunyuanVideo, it achieves a reported $4.99\times$ speedup and improves ChronoMagic-Bench-150 PSNR by 5.44 dB over the strongest listed fidelity baseline. Comparisons at fixed evaluation budgets further show substantial gains on motion, audio, music, and 3D generation.