Geometry‑Aware, Training‑Free Acceleration of Diffusion Transformer Sampling (CVPR 2026 Highlight)

Introduction

Diffusion models have become the dominant foundation for image and video generation, with systems such as DiT, FLUX, and HunyuanVideo relying on diffusion Transformers to iteratively denoise from noise to content. However, the generation process is slow because many sequential denoising steps are required, leading to high latency and computational cost.

Motivation

The key challenge is that accelerating diffusion models by simply reducing the number of sampling steps often degrades quality: image structures distort, semantics weaken, textures become inconsistent, and video frames flicker. The underlying reason is that diffusion sampling does not move in a flat Euclidean space but follows a curved feature manifold learned by the model. Large‑step sampling that ignores this geometry causes "manifold drift".

Observations on Feature Geometry

Analysis of intermediate activations in DiT‑XL/2 and FLUX.1‑dev shows that most variance (>99%) is captured by the top 64 principal directions, indicating that denoising evolves on a low‑dimensional, curved manifold rather than the full high‑dimensional space.

Method: GeoRK2

GeoRK2 is a training‑free, plug‑and‑play acceleration framework that integrates geometry awareness into diffusion sampling. It consists of three core modules:

Geometry‑aware RK2 prediction : Instead of naïvely extrapolating the current direction, GeoRK2 performs a second‑order Runge‑Kutta (RK2) step, projecting the intermediate state onto the dominant feature subspace.

Low‑rank metric correction : A local covariance matrix G_t = covariance(H_t) + epsilon * I is estimated from intermediate activations. Using truncated SVD (rank = 64) yields a low‑rank metric that guides a preconditioned correction geometry_correction = -lambda * G_inverse * prediction_error.

Adaptive stability : The method monitors the variance of the acceleration a_t. If a sudden increase is detected, it falls back to a conservative two‑point extrapolation and applies momentum mixing h_out = rho * corrected_prediction + (1 - rho) * h_t to smooth the output.

Implementation Details

GeoRK2 does not require retraining or modifying the model architecture. It intercepts intermediate activations during inference, computes the low‑rank metric, and replaces the original sampler update with the geometry‑aware RK2 step. The additional FLOPs are modest: projection (≈0.014 TFLOPs), metric‑preconditioned inversion (≈0.022 TFLOPs), and amortized truncated SVD (≈0.011 TFLOPs), totaling about 5.1 % extra computation and ~3.8 % wall‑clock overhead.

Experimental Results

GeoRK2 was evaluated on three representative tasks:

ImageNet‑256 (class‑conditional image generation with DiT‑XL/2) : Achieved 1.95× speedup (latency 4.42 s) with FID 2.41, 2.70× speedup with FID 2.67, and 4.92× speedup with FID 3.32, outperforming other accelerators that suffered larger FID increases.

FLUX.1‑dev (text‑to‑image on DrawBench) : At NFE = 50, GeoRK2(N=5) gave 3.52× speedup with ImageReward 0.9889 and CLIP Score 34.96; higher N maintained >0.97 ImageReward while still delivering >4× speedup.

HunyuanVideo (text‑to‑video on VBench) : GeoRK2(N=8) reduced latency from 323.89 s to 69.44 s (4.66× speedup), FLOPs speedup 6.77×, and achieved VBench Score 80.73, preserving temporal consistency where other methods exhibited flickering.

Ablation Study

Controlled ablations on DiT‑XL/2 (N=3) showed that removing any component degrades quality: without geometry correction FID rises to 3.02, without RK2 (using Euler) FID 2.87, and without Euclidean‑related design FID 3.41. This confirms that the three components act synergistically.

Rank Sensitivity

Increasing the truncation rank improves quality up to rank 64, after which gains saturate, matching the spectral analysis that the top 64 directions capture >99 % variance.

Conclusion

GeoRK2 demonstrates that the primary obstacle to aggressive diffusion acceleration is manifold drift, not merely low‑order numerical integration. By explicitly modeling the curved feature geometry with RK2 integration, low‑rank metric correction, and adaptive stability, GeoRK2 delivers 4‑5× faster generation across image, text‑to‑image, and video tasks without retraining, offering a practical plug‑and‑play solution for high‑performance diffusion inference.