Source-linked AI summary
Accelerating Diffusion Transformers with Gaussian Process Rectified Feature Cache
Zhirong Shen, Rui Huang, Chang Zou, Shikang Zheng, Jiacheng Liu, Peiliang Cai, Zhengyi Shi, Yaosong Du, Liang Feng, Xiaobing Tu, Jinkui Ren, Xiantao Zhang, Linfeng Zhang
TL;DR
Feature caching accelerates diffusion-transformer inference, but accumulated prediction mismatch creates a label-data problem for online regression correction. GP-Refiner uses Gaussian-process correction and uncertainty-triggered full-computation calibration; on FLUX, it improves image fidelity while maintaining approximately 3× latency speedup and 5.8× FLOPs reduction.
Problem
Accumulated feature-cache errors make accelerated features differ from the reference trajectory, leaving online regression without accurate real features for unbiased training samples.
Source: §3.2 · p. 8Method
GP-Refiner treats cached full-computation features as Gaussian-noisy observations of reference features, uses GPR to correct baseline predictions, and triggers full-computation calibration from posterior variance.
Source: §1 · pp. 2–3, §5 · p. 15Results
Approximately 3× latency speedup and 5.8× FLOPs reduction accompany enhanced image fidelity across FLUX cache frameworks.
Source: §4.2 · p. 12Takeaways & Limitations
The framework is designed to plug into baseline cache forecasters with nearly no additional computation overhead.
Source: §1 · p. 3The regression assumes observed outputs relate to latent function values through additive i.i.d. Gaussian noise, with residuals modeled as zero-mean Gaussian noise.
Source: §3.1 · p. 7, §3.3 · p. 9
Abstract
from arXiv · showhide
Diffusion Transformers have become the dominant paradigm in generative AI, but their high computational costs severely hinder real-time applications. Prediction-based feature caching is widely used to accelerate diffusion transformers; however, as the number of steps increases, the deviation between its predictions and the reference full-compute trajectory gradually grows. An intuitive idea is to use an online regression model to dynamically correct this deviation, but it faces the issue of label data being unavailable during the acceleration process. This paper presents a statistical observation that the residuals between the features of full computation steps using caching methods and reference full-compute trajectory locally exhibit a zero-mean Gaussian distribution. By treating the features of full computation steps as noisy observations of reference features, the data acquisition problem is resolved. Based on this observation, a plug-and-play GP-Refiner correction framework is proposed. This method utilizes Gaussian Process Regression for correction and, leveraging the properties of GPR, introduces an uncertainty-adaptive computation strategy that triggers necessary full-computation calibration by monitoring the posterior variance in real time. Experiments demonstrate significant improvements across different models when combined with various state-of-the-art methods. Integrating the proposed framework with TaylorSeer reduces the computational load by 19.3% while improving PSNR by 0.9 dB and reducing LPIPS from 0.46 to 0.29. Code is available in https://github.com/Aredstone/GP-Refiner.
1 Introduction
Forecasting-based caching reduces Diffusion Transformer computation, but prediction errors accumulate across skipped steps and drift from the reference trajectory. GP-Refiner treats full-compute features obtained under caching as noisy reference observations, uses Gaussian Process Regression for correction, and triggers recalibration when posterior uncertainty exceeds a threshold.
1 Introduction
- Forecasting-based caching reduces computation, but prediction errors accumulate across skipped steps and drift from the reference full-compute trajectory.After steps have been skipped, a full forward computation operates on an already drifted state, so its features are not unbiased supervision for rectification.Source: §1 · p. 2
- The residuals are approximately zero-mean Gaussian, allowing full-compute features obtained under caching to serve as noisy observations of reference features.The Q–Q plots of randomly projected residuals follow a standard normal line with R2 ≈0.97–0.99.Source: §1 · p. 2, Figure 2 · p. 3
- GP-Refiner applies Gaussian Process Regression to correct baseline cache predictions, bringing the accelerated trajectory closer to reference features with almost no additional computation overhead.The approach uses full-compute-step features as noisy observations to address the unavailable-label problem during accelerated inference.Source: §1 · p. 3
- Posterior variance quantifies correction confidence; exceeding a predefined threshold triggers a full forward computation to recalibrate the state.This uncertainty-adaptive mechanism is intended to balance inference speed and generation quality.Source: §1 · pp. 3–4
2 Related Work
Diffusion Transformer acceleration targets both costly per-step computation and many sequential denoising iterations. This section contrasts model and scheduler changes with cache-based methods, whose accumulated drift motivates corrective mechanisms.
2 Related Work
- DiT computational burden comes from heavy Transformer operations on high-resolution inputs and the many sequential denoising iterations needed to preserve quality.Most existing acceleration efforts target one or both of these bottlenecks.Source: §2 · p. 4
2.1 Acceleration via Model Modification and Retraining
- Model modification and trajectory compression can reduce inference cost, but often require retraining or fine-tuning, making them less plug-and-play for pretrained DiTs.This motivates training-free inference-time acceleration.Source: §2.1 · p. 4
2.2 Sampling Schedulers and Numerical Solvers
- Higher-order numerical solvers permit larger denoising steps, but still require a full denoising-network forward pass at every step.They reduce total steps while leaving per-step DiT computation largely unchanged.Source: §2.2 · p. 5
2.3 Cache-Based Inference Acceleration
- Cache-based methods reuse or forecast intermediate features to skip denoising-network computation, but forecasting errors can still accumulate over long horizons.Feedback-driven verification or probing can trigger recomputation to suppress drift, though existing corrective designs add overhead.Source: Cache-then-Forecast · p. 5, §2.3 · p. 5, Cache-then-Reuse · p. 5
3 Method
The method treats full-computation features obtained under caching as noisy observations of reference features, then uses Gaussian Process Regression to correct cached predictions. It also uses posterior variance to decide when to trigger full computation.
3.1 Preliminary
- 3.1 Preliminary: Forecasting caches predict skipped-step features from temporal differences, while GPR provides predictive means and variances from noisy observations.TaylorSeer uses an anchor feature and temporal finite differences in a truncated Taylor expansion.Source: §3.1 · pp. 6–7
3.2 Observation
- 3.2 Observation: Residuals are approximately zero-mean Gaussian; 100 random projections averaged correlation above 0.98, with a minimum of 0.9424.Single-vector elements also closely fit a standard Gaussian bell curve.Source: §3.2 · p. 8
3.3 Feature Correction Modeling via Gaussian Processes
- 3.3 Feature Correction Modeling via Gaussian Processes: GPR maps baseline predicted features to reference features using historical full-computation features as noisy observations.The observation pairs use a baseline prediction for the previous step and a full-computation feature; pairing adjacent steps focuses correction on local residuals.Source: §3.3 · pp. 9–10
3.4 Adaptive Computation via Uncertainty
- 3.4 Adaptive Computation via Uncertainty: The system triggers full computation when posterior variance exceeds a threshold; otherwise, it uses the corrected mean.The authors report that threshold-triggered full-computation steps cluster where denoising nonlinearity and error sensitivity are strongest.Source: §3.4 · p. 11
4 Experiments
Experiments evaluate GP-Refiner with multiple cache methods on FLUX and Qwen-Image, measuring image fidelity and computational efficiency. The reported results include improvements in fidelity and an ablation of rectification and uncertainty-guided scheduling.
4.1 Experiment Settings
- Experiments use Qwen-Image and FLUX.1-dev at 50 inference steps, plus Qwen-Image-Lightning at 8 steps, integrating GP-Refiner with TaylorSeer, ClusCa, or HiCache.Evaluation uses 200 Draw-Bench prompts and measures SSIM, PSNR, LPIPS, latency, FLOPs, and speedup against full-step baselines.Source: §4.1 · pp. 11–12
4.2 Main Results
- On FLUX, ClusCa + ours records the best PSNR (29.823) and lowest LPIPS (0.3334); on Qwen-Image, TaylorSeer + ours raises PSNR from 28.20 to 29.48 and lowers LPIPS from 0.65 to 0.29 versus N = 7.The FLUX comparison reports about 3× latency speedup and 5.8× FLOPs reduction; HiCache and TaylorSeer exceed 6.2× FLOPs speedup on Qwen-Image.Source: Table 1 · p. 12, §4.2 · p. 12, Table 2 · p. 13
4.3 Ablation Study
- At fixed TaylorSeer skip intervals N = 6, 7, and 8, GP rectification improves fidelity; at N = 7, PSNR rises from 28.671 to 29.266 and LPIPS falls from 0.4542 to 0.4151.The authors report that correction compensates for Taylor-expansion feature residuals across caching frequencies.Source: Table 3 · p. 13, §4.3 · pp. 13–14
- Taylor-based dynamic scheduling alone reaches 5.72× speedup but PSNR 28.930; the full TaylorSeer + ours (O = 2) configuration reaches PSNR 29.550 and LPIPS 0.3497.The full configuration uses GP-derived uncertainty for both scheduling and rectification.Source: Table 3 · p. 13, §4.3 · p. 14
- Across 200 samples, dynamic scheduling's reported NFE counts cluster at 8 or 9, and images at 9 versus 10 NFEs differ only marginally in visual quality.The paper says this allocation prunes redundant computations while preserving image fidelity.Source: §4.3 · pp. 14–15, Figure 5 · p. 15
5 Conclusion
GP-Refiner is a plug-and-play feature rectification framework for diffusion transformers, using a local zero-mean Gaussian residual observation and posterior variance to adaptively trigger full-computation calibration. Experiments report enhanced image fidelity alongside reduced computational complexity and faster inference.
5 Conclusion
- Experiments on large-scale generative models report enhanced image fidelity, reduced computational complexity, and faster inference.The paper presents this as improving the trade-off between computational efficiency and generation quality.Source: §5 · p. 15
- GP-Refiner uses posterior variance to trigger full-computation calibration when inference confidence falls below a threshold.The framework is described as plug-and-play and transforms acceleration into baseline estimation with rectification.Source: §5 · p. 15
- The framework is motivated by a statistical observation that baseline prediction residuals exhibit a local zero-mean Gaussian distribution.It uses this observation to rectify cumulative feature-cache errors during accelerated inference.Source: §5 · p. 15