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

Predict to Skip: Linear Multistep Feature Forecasting for Efficient Diffusion Transformers
Feb 20, 2026Hanshuai Cui, Zhiqing Tang, Qianli Ma +3Diffusion Model Acceleration Diffusion Sampling
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
LinCa: Accelerating Diffusion Models via Learnable Decomposed Feature Caching
Aug 18, 2026Jinshan Liu, Haoran Qin, Xiaobing Tu +9Diffusion Model Inference Acceleration Diffusion Model Caching
Shanghai Jiao Tong University, China · Xi’an Jiaotong University, China · Shandong University, China +3
Beyond Fixed Formulas: Data-Driven Linear Predictor for Efficient Diffusion Models
Apr 29, 2026Zhirong Shen, Rui Huang, Jiacheng Liu +6Diffusion Transformer Diffusion Model Caching
Shanghai Jiao Tong University · University of Electronic Science and Technology of China · Shandong University +2

Feb 20, 2026J/K move · Enter open · S save

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

Hanshuai Cui, Zhiqing Tang, Qianli Ma, Zhi Yao+2

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

LinCa: Accelerating Diffusion Models via Learnable Decomposed Feature Caching

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