Forecasting When to Forecast: Accelerating Diffusion Models with Confidence-Gated Taylor

TL;DR

This work tackles the slow inference of Diffusion Transformers by proposing a training-free acceleration framework that combines a Last Block Forecast (LBF) with Prediction Confidence Gating (PCG). By shifting Taylor-based prediction to the last block and using the first-block prediction error as a reliability proxy, the method enables a dynamic caching strategy that selectively skips intermediate block computations. The approach delivers substantial speedups across multiple models and modalities (up to 4.14x on Wan Video and 3.17x on FLUX) while maintaining high visual fidelity, outperforming prior training-free methods like TaylorSeer in both speed and quality. The result is a practical, adaptive acceleration mechanism for diffusion-based generation that preserves output quality and reduces memory and compute overhead.

Abstract

Diffusion Transformers (DiTs) have demonstrated remarkable performance in visual generation tasks. However, their low inference speed limits their deployment in low-resource applications. Recent training-free approaches exploit the redundancy of features across timesteps by caching and reusing past representations to accelerate inference. Building on this idea, TaylorSeer instead uses cached features to predict future ones via Taylor expansion. However, its module-level prediction across all transformer blocks (e.g., attention or feedforward modules) requires storing fine-grained intermediate features, leading to notable memory and computation overhead. Moreover, it adopts a fixed caching schedule without considering the varying accuracy of predictions across timesteps, which can lead to degraded outputs when prediction fails. To address these limitations, we propose a novel approach to better leverage Taylor-based acceleration. First, we shift the Taylor prediction target from the module level to the last block level, significantly reducing the number of cached features. Furthermore, observing strong sequential dependencies among Transformer blocks, we propose to use the error between the Taylor-estimated and actual outputs of the first block as an indicator of prediction reliability. If the error is small, we trust the Taylor prediction for the last block; otherwise, we fall back to full computation, thereby enabling a dynamic caching mechanism. Empirical results show that our method achieves a better balance between speed and quality, achieving a 3.17x acceleration on FLUX, 2.36x on DiT, and 4.14x on Wan Video with negligible quality drop. The Project Page is \href{https://cg-taylor-acce.github.io/CG-Taylor/}{here.}

Figures (6)

• Figure 1: Comparison of our method and baselines under different speedup ratios.
• Figure 2: The expanded architecture of the DiT and the difference between different methods. (a) DiTs follow a hierarchical architecture which are composed of multiple transformer blocks, where each block typically consists of three modules. (b) In the original method, every module is computed at each timestep, resulting in a high computational cost. TaylorSeer reduced computation by calculating a module at one timestep and using Taylor expansions to predict its features for the remaining timesteps within a given interval $N$. Our method uses a confidence-gated method: if the first block’s prediction error at timestep t-1 is small, we directly predict the last block’s feature at this timestep, skipping computation of the $B-1$ blocks to achieve acceleration. Otherwise, we compute all blocks normally.
• Figure 3: The framework of our method. At timestep $t-1$, we first compute the actual output of the first block and simultaneously predict it using Taylor expansion with the cached feature at timestep $t$. The prediction error is then evaluated. If the error is below a threshold $\epsilon$, the Taylor prediction is considered reliable. In this case, we use it to approximate the last block feature, thereby skipping the computation of the remaining $(B-1)$ blocks to accelerate inference. Otherwise, we fall back to full computation of all blocks. This process is repeated for each subsequent timestep.
• Figure 4: Visualization of different acceleration methods on FLUX. Our method achieves higher visual quality and greater similarity to the original image while operating at a faster speed.
• Figure 5: Number of denoising steps below and above the $\epsilon$ threshold for different $\epsilon$ values.