LESA: Learnable Stage-Aware Predictors for Diffusion Model Acceleration

Peiliang Cai; Jiacheng Liu; Haowen Xu; Xinyu Wang; Chang Zou; Linfeng Zhang

cs.CVFeb 24, 2026

LESA: Learnable Stage-Aware Predictors for Diffusion Model Acceleration

Authors: Peiliang Cai, Jiacheng Liu, Haowen Xu, Xinyu Wang, Chang Zou, Linfeng Zhang

Organizations: Shanghai Jiao Tong University

Abstract

Diffusion models have achieved remarkable success in image and video generation tasks. However, the high computational demands of Diffusion Transformers (DiTs) pose a significant challenge to their practical deployment. While feature caching is a promising acceleration strategy, existing methods based on simple reusing or training-free forecasting struggle to adapt to the complex, stage-dependent dynamics of the diffusion process, often resulting in quality degradation and failing to maintain consistency with the standard denoising process. To address this, we propose a LEarnable Stage-Aware (LESA) predictor framework based on two-stage training. Our approach leverages a Kolmogorov-Arnold Network (KAN) to accurately learn temporal feature mappings from data. We further introduce a multi-stage, multi-expert architecture that assigns specialized predictors to different noise-level stages, enabling more precise and robust feature forecasting. Extensive experiments show our method achieves significant acceleration while maintaining high-fidelity generation. Experiments demonstrate 5.00x acceleration on FLUX.1-dev with minimal quality degradation (1.0% drop), 6.25x speedup on Qwen-Image with a 20.2% quality improvement over the previous SOTA (TaylorSeer), and 5.00x acceleration on HunyuanVideo with a 24.7% PSNR improvement over TaylorSeer. State-of-the-art performance on both text-to-image and text-to-video synthesis validates the effectiveness and generalization capability of our training-based framework across different models. Our code is available at https://github.com/caipeiliang2004/LESA.

Figures & tables

Figure 1: Images generated by Qwen-image using LESA with a 6.25 $\times$ acceleration.

Figure 2: Cosine Similarity curves and PCA-projected trajectories show that feature evolution differs markedly across diffusion models, with non-smooth, stage-dependent dynamics that challenge the common assumption of simple or continuous temporal change.

Figure 3: Overview of the training-based stage-aware learnable predictor framework. (a) Training progress of the LESA uses outputs from the previous $K$ steps as input and updates the predictor with the DiT activation output. (b) Inference process with the learned predictor skips part of the DiT computations by predicting features. (c) Predictor architecture processes features from previous steps with a linear projection and processes the timestep with a KAN module, whose outputs are multiplied to generate the feature at step $t$ . (d) Stage-aware denoising divides the trajectory into three stages and assigns a dedicated predictor to each stage.

Method	Acceleration				Quality Metrics		Perceptual Metrics
FLUX	Latency(s) $\downarrow$	Speed $\uparrow$	FLOPs(T) $\downarrow$	Speed $\uparrow$	Image Reward $\uparrow$	CLIP Score $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
[dev]: 50 steps	23.24	1.00 $\times$	3726.87	1.00 $\times$	0.99	32.64	$\infty$	1.00	0.00
$50\%$ steps	11.82	1.97 $\times$	1863.44	2.00 $\times$	0.97	32.57	29.56	0.73	0.31
PAB	17.84	1.30 $\times$	3013.13	1.24 $\times$	0.95	32.55	28.84	0.67	0.40
DBCache	16.88	1.38 $\times$	2384.29	1.56 $\times$	1.01	32.53	33.86	0.87	0.12
FORA ( $\mathcal{N}$ =3)	9.06	2.57 $\times$	1267.89	2.94 $\times$	0.99	32.27	30.65	0.77	0.25

Table 1: Quantitative comparison in text-to-image generation for FLUX.1-schnell.

Figure 4: Comparison of LESA and TaylorSeer under different speedup ratios. LESA preserves better quality metrics and perceptual metrics On QwenImage and FLUX.

Figure 5: Comparison of Sampling Process Image and Final Image. On FLUX, LESA demonstrates significant speedup ratio with excellent spatial organization and accurate color representation.

Method	Acceleration				Quality Metrics		Perceptual Metrics
Qwen	Latency(s) $\downarrow$	Speed $\uparrow$	FLOPs(T) $\downarrow$	Speed $\uparrow$	Image Reward $\uparrow$	CLIP Score $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
$50$ steps	127.40	1.00 $\times$	12917.56	1.00 $\times$	1.25	35.59	$\infty$	1.00	0.00
$20\%$ steps †	25.92	4.89 $\times$	2583.51	5.00 $\times$	0.94	34.95	28.59	0.61	0.52
FORA ( $\mathcal{N}$ =4) †	38.43	3.32 $\times$	3359.99	3.84 $\times$	0.93	34.40	28.66	0.59	0.51
ToCa ( $\mathcal{N}$ =8, $\mathcal{R}$ =75%)	61.37	2.08 $\times$	2991.34	4.32 $\times$	1.02	34.96	28.93	0.63	0.44
DuCa ( $\mathcal{N}$ =9, $\mathcal{R}$ =80%) †	34.73	3.67 $\times$	2958.13	4.37 $\times$	0.77	34.62	28.45	0.58	0.55

Table 2: Quantitative comparison in text-to-image generation for Qwen-Image and Qwen-Image-Lightning.

Figure 6: Qualitative Comparison of Generated Images. On Qwen-Image, LESA delivers higher speedup with high image quality.

Method	Acceleration				Perceptual Metrics
HunyuanVideo	Latency(s) $\downarrow$	Speed $\uparrow$	FLOPs(T) $\downarrow$	Speed $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
Original: 50 steps	98.91	1.00 $\times$	29773.0	1.00 $\times$	$\infty$	1.00	0.00
$22\%$ steps	22.98	4.30 $\times$	6550.1	4.55 $\times$	17.65	0.59	0.42
ToCa ( $\mathcal{N}=5$ , $\mathcal{R}=90\%$ )	26.12	3.79 $\times$	7006.2	4.25 $\times$	17.04	0.54	0.44
DuCa ( $\mathcal{N}=5$ , $\mathcal{R}=90\%$ )	23.08	4.29 $\times$	6483.2	4.48 $\times$	17.08	0.54	0.43
TeaCache ( $l=0.4$ )	21.83	4.53 $\times$	6550.1	4.55 $\times$	18.25	0.61	0.38

Table 3: Quantitative comparison of text-to-video generation on HunyuanVideo.

Figure 7: Qualitative Comparison of Early, Middle, and Late Frames From Generated Videos. On HunyuanVideo, LESA constructs continuous dynamic process with excellent subject accuracy under high speedup ratio.

Experiment Settings			Quality Metrics		Perceptual Metrics
Interval	Segmented	Module	Image Reward $\uparrow$	CLIP Score $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
$\mathcal{N}$ = 5	✔	KAN	1.00	32.70	30.96	0.77	0.24
$\mathcal{N}$ = 5	✘	KAN	1.00	32.67	30.77	0.77	0.25
$\mathcal{N}$ = 5	✔	MLP	0.99	32.69	30.76	0.77	0.25
$\mathcal{N}$ = 5	✘	MLP	0.95	32.50	30.29	0.73	0.31
$\mathcal{N}$ = 7	✔	KAN	0.98	32.88	30.17	0.72	0.32

Table 4: Ablation study for Interval, Stage Segmentation and Timestep Modulation Model for LESA on FLUX

Method	Acceleration				Quality Metrics		Perceptual Metrics
FLUX	Latency(s) $\downarrow$	Speed $\uparrow$	FLOPs(T) $\downarrow$	Speed $\uparrow$	Image Reward $\uparrow$	CLIP Score $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
[schnell]: 4 steps	2.13	1.00 $\times$	278.41	1.00 $\times$	0.93	34.09	$\infty$	1.00	0.00
FORA ( $\mathcal{N}$ =3)	1.30	1.64 $\times$	139.24	2.00 $\times$	0.95	34.43	29.77	0.78	0.18
ToCa ( $\mathcal{N}$ =3, $\mathcal{R}$ =80%)	2.20	0.97 $\times$	154.54	1.80 $\times$	0.94	34.46	29.79	0.78	0.18
DuCa ( $\mathcal{N}$ =3, $\mathcal{R}$ =80%)	1.33	1.60 $\times$	145.46	1.91 $\times$	0.95	34.44	29.80	0.78	0.18
TaylorSeer ( $\mathcal{N}$ =3, $\mathcal{O}$ =2)	1.29	1.65 $\times$	139.24	2.00 $\times$	0.95	34.43	29.77	0.78	0.18

Table 5: Quantitative comparison in text-to-image generation for FLUX.1-schnell. Best results are highlighted in bold .

Method	Acceleration				Quality Metrics		Perceptual Metrics
Qwen	Latency(s) $\downarrow$	Speed $\uparrow$	FLOPs(T) $\downarrow$	Speed $\uparrow$	Image Reward $\uparrow$	CLIP Score $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
[Lightning]: 8 steps	4.37	1.00 $\times$	560.96	1.00 $\times$	1.29	35.32	$\infty$	1.00	0.00
FORA ( $\mathcal{N}$ =2)	3.08	1.42 $\times$	350.65	1.60 $\times$	1.27	35.20	31.35	0.82	0.13
ToCa ( $\mathcal{N}$ =2, $\mathcal{R}$ =90%)	6.70	0.65 $\times$	393.19	1.43 $\times$	1.27	35.30	32.03	0.81	0.15
DuCa ( $\mathcal{N}$ =2, $\mathcal{R}$ =90%)	3.37	1.30 $\times$	379.26	1.48 $\times$	1.28	35.16	32.74	0.84	0.11
TaylorSeer ( $\mathcal{N}$ =2)	3.23	1.35 $\times$	350.65	1.60 $\times$	1.27	35.45	29.94	0.68	0.28

Table 6: Quantitative comparison in text-to-image generation for Qwen-Image-Lightning. Best results are highlighted in bold .

Method	Quality Metrics		Perceptual Metrics
Stable Diffusion XL	Image Reward $\uparrow$	CLIP Score $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
[xl]: 50 steps	0.58	33.63	$\infty$	1.00	0.00
DeepCache ( $\mathcal{N}$ =3)	0.52	33.42	29.85	0.77	0.26
LESA ( $\mathcal{N}$ =3)	0.56	33.57	31.28	0.80	0.18
DeepCache ( $\mathcal{N}$ =4)	0.40	33.29	29.22	0.73	0.31
LESA ( $\mathcal{N}$ =4)	0.52	33.39	30.47	0.75	0.24

Table 7: Quantitative comparison of LESA and DeepCache in text-to-image generation for Stable Diffusion XL.

Method	VRAM (GB) $\downarrow$	MACs (T) $\downarrow$	Latency (s) $\downarrow$	FLOPs (T) $\downarrow$	Image Reward $\uparrow$	PSNR $\uparrow$
[dev]: 50 steps	0.62	1859.62	23.24	3726.87	0.99	$\infty$
ToCa ( $\mathcal{N}$ =8, $\mathcal{R}$ =75%)	12.31	414.88	12.39	829.86	0.95	29.07
DuCa ( $\mathcal{N}$ =8, $\mathcal{R}$ =70%)	3.65	428.86	9.40	858.27	0.94	29.06
TeaCache ( $l$ =1.0)	0.69	409.43	7.07	820.55	0.84	28.61
TaylorSeer ( $\mathcal{N}$ =6, $O$ =2)	7.66	372.38	6.73	746.28	1.02	28.94
LESA ( $\mathcal{N}$ =7)	0.81	372.25	5.19	745.51	0.98	30.17

Table 8: Comparison of methods in Cache Memory, MACs, Latency, and FLOPS on FLUX-1.dev . Best results are highlighted in bold .

Method	Quality Metrics		Perceptual Metrics
QwenImage	Image Reward $\uparrow$	CLIP Score $\uparrow$	PSNR $\uparrow$	SSIM $\uparrow$	LPIPS $\downarrow$
50 steps	1.25	35.59	$\infty$	1.00	0.00
LESA ( $1$ prompt)	1.14	35.17	30.18	0.78	0.25
LESA ( $5$ prompts)	1.15	35.31	30.36	0.78	0.25
LESA ( $10$ prompts)	1.16	35.33	30.21	0.78	0.26
LESA ( $20$ prompts)	1.15	35.21	29.78	0.77	0.25

Table 9: Quantitative comparison of LESA trained with different number of prompts for QwenImage. Best results are highlighted in bold .

Explore similar work

Feb 20, 2026

Predict to Skip: Linear Multistep Feature Forecasting for Efficient Diffusion Transformers

Diffusion Transformers (DiT) have emerged as a widely adopted backbone for high-fidelity image and video generation, yet their iterative denoising process incurs high computational costs. Existing training-free acceleration methods rely on feature caching and reuse under the assumption of temporal stability. However, reusing features for multiple steps may lead to latent drift and visual degradation. We observe that model outputs evolve smoothly along much of the diffusion trajectory, enabling principled predictions rather than naive reuse. Based on this insight, we propose PrediT, a training-free acceleration framework that formulates feature prediction as a linear multistep problem. We employ classical linear multistep methods to forecast future model outputs from historical information, combined with a corrector that activates in high-dynamics regions to prevent error accumulation. A dynamic step modulation mechanism adaptively adjusts the prediction horizon by monitoring the feature change rate. Together, these components enable substantial acceleration while preserving generation fidelity. Extensive experiments validate that our method achieves up to

5.54\times

latency reduction across various DiT-based image and video generation models, while incurring negligible quality degradation.

Hanshuai Cui, Zhiqing Tang, Qianli Ma +3

School of Artificial Intelligence, Beijing Normal University, Beijing, China · Institute of Artificial Intelligence and Future Networks, Beijing Normal University, Zhuhai, China · Shenzhen University of Advanced Technology, Shenzhen, China

Diffusion Model Acceleration Diffusion Sampling Diffusion-Based Image Generation Diffusion Transformer Video Diffusion Models

Ask PDF

Aug 18, 2026

LinCa: Accelerating Diffusion Models via Learnable Decomposed Feature Caching

Diffusion models have achieved remarkable success in image and video generation, yet the high computational cost of iterative sampling remains a critical bottleneck for practical deployment. Feature caching has emerged as a promising acceleration paradigm by reusing or predicting intermediate features across timesteps. However, existing training-free methods apply uniform prediction strategies that cannot adapt to the heterogeneous feature dynamics, causing significant quality degradation under high acceleration ratios. We propose LinCa, a feature caching framework based on learnable invertible networks. LinCa decomposes cached features into sub-components with distinct continuity properties via a lightweight invertible network and applies differentiated prediction orders matched to each component. The strict invertibility guarantees lossless reconstruction back to the original feature space, forming a unified Decompose-Predict-Reconstruct pipeline. By training separate predictors for different models and timestep segments, LinCa adapts to heterogeneous feature dynamics. Experiments on FLUX, Qwen-Image, and HunyuanVideo demonstrate that LinCa, with less than 0.2% additional parameters, significantly outperforms existing methods and maintains near-lossless quality at 5-7x speedup. Code: https://github.com/QHR69/LinCa

Jinshan Liu, Haoran Qin, Xiaobing Tu +9

Shanghai Jiao Tong University, China · Xi’an Jiaotong University, China · Shandong University, China +3

Diffusion Model Inference Acceleration Diffusion Model Caching Diffusion Models

Ask PDF

Apr 29, 2026

Beyond Fixed Formulas: Data-Driven Linear Predictor for Efficient Diffusion Models

To address the high sampling cost of Diffusion Transformers (DiTs), feature caching offers a training-free acceleration method. However, existing methods rely on hand-crafted forecasting formulas that fail under aggressive skipping. We propose L2P (Learnable Linear Predictor), a simple data-driven caching framework that replaces fixed coefficients with learnable per-timestep weights. Rapidly trained in ~20 seconds on a single GPU, L2P accurately reconstructs current features from past trajectories. L2P significantly outperforms existing baselines: it achieves a 4.55x FLOPs reduction and 4.15x latency speedup on FLUX.1-dev, and maintains high visual fidelity under up to 7.18x acceleration on Qwen-Image models, where prior methods show noticeable quality degradation. Our results show learning linear predictors is highly effective for efficient DiT inference. Code is available at https://github.com/Aredstone/L2P-Cache.

Zhirong Shen, Rui Huang, Jiacheng Liu +6

Shanghai Jiao Tong University · University of Electronic Science and Technology of China · Shandong University +2

Diffusion Transformer Diffusion Model Caching Diffusion Model Inference Acceleration Diffusion Models

Ask PDF