🤖 AI

LESA: Learnable Stage-Aware Predictors for Diffusion Model Acceleration

This paper proposes LESA, a learnable stage-aware predictor framework utilizing Kolmogorov-Arnold Networks and a multi-stage, multi-expert architecture to significantly accelerate Diffusion Transformers while maintaining or improving generation quality across diverse text-to-image and text-to-video models.

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

Published 2026-03-17

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to paint a masterpiece, but you have a very strict rule: you must add one tiny brushstroke at a time, and you have to wait for the paint to dry before adding the next one. This is how Diffusion Models (the AI behind tools like DALL-E 3 or Midjourney) currently work. They start with a canvas full of static noise (like TV snow) and slowly, step-by-step, "denoise" it until a clear image appears.

The problem? This process is incredibly slow. It might take 50 steps to paint a picture, and each step requires the AI to do a massive amount of math. It's like walking across a room by taking 50 tiny, deliberate steps instead of just walking normally.

LESA is a new "shortcut" that lets the AI take giant leaps without ruining the painting. Here is how it works, explained simply:

1. The Problem: The "One-Size-Fits-All" Mistake

Previous attempts to speed this up were like trying to drive a car using the same gear for every situation.

The "Reuse" Method: Some methods just said, "Hey, the picture looks almost the same as the last step, so let's just copy the last step's work." This works for a bit, but eventually, the picture gets blurry or weird because the AI did change, just a little.
The "Guessing" Method: Others tried to predict the future by drawing a straight line based on the last few steps (like guessing where a ball will land). This works well when the ball is rolling smoothly, but fails when the ball suddenly bounces or changes direction.

The authors realized that the AI's "painting process" isn't a straight line. It has three distinct phases:

The Chaos Phase (High Noise): At the start, the image is changing wildly. It's like a stormy sea. Predicting the next wave is hard because everything is moving fast.
The Flow Phase (Medium Noise): The image starts to take shape. The changes are smooth and predictable, like a calm river.
The Detail Phase (Low Noise): The image is almost done. The AI is just adding tiny details (like eyelashes or texture). The changes are very small and precise.

Old methods tried to use the same prediction rule for all three phases, which is like trying to surf a tsunami, ride a river, and thread a needle all while wearing the same pair of shoes. It doesn't work well.

2. The Solution: LESA (The "Specialized Team")

LESA stands for Learnable Stage-Aware predictors. Think of LESA not as one smart AI, but as a team of three specialized experts, each with a specific job:

Expert 1 (The Storm Chaser): Specializes in the chaotic beginning. They are good at handling wild, unpredictable changes.
Expert 2 (The River Guide): Specializes in the middle. They are great at predicting smooth, continuous flow.
Expert 3 (The Micro-Scalpel): Specializes in the end. They are incredibly precise for tiny, final details.

Instead of guessing blindly, LESA looks at the current stage of the painting and switches to the right expert. This ensures the prediction is always accurate for that specific moment.

3. The Secret Weapon: The "Mathematical Swiss Army Knife" (KAN)

How do these experts learn to be so good? The paper uses a special type of neural network called a Kolmogorov-Arnold Network (KAN).

Old Way (MLP): Imagine a standard AI is like a factory worker who can only use a hammer. If they need to screw something in, they have to hit it with the hammer until it works. It's rigid.
LESA Way (KAN): Imagine an AI that is a Swiss Army Knife. It can learn to be a screwdriver, a blade, or a corkscrew depending on the task. KANs are flexible; they can "shape-shift" their math to fit the specific curve of the data perfectly. This allows them to learn the complex, wiggly patterns of the AI's painting process much better than rigid methods.

4. The Training: "Practice Makes Perfect"

To teach this team, the researchers used a two-stage training method:

Teacher-Student Phase: First, they let the AI paint the whole picture normally (the "Ground Truth"). The LESA experts watch and learn, trying to copy the teacher's steps perfectly.
The "Drift" Phase: Then, they let the experts try to paint the picture on their own using their predictions. If they make a mistake, they don't just get corrected; they have to keep painting based on their own mistakes. This teaches them to be robust and not panic when things get slightly off-track.

The Result: Super Speed, No Quality Loss

The results are impressive. By using LESA:

FLUX.1 (a popular image model) became 5x faster with almost no loss in quality.
Qwen-Image became 6.25x faster, and the images were actually better than previous speed-up methods.
HunyuanVideo (a video generator) became 5x faster while keeping the video smooth and clear.

In a nutshell:
Imagine you are walking across a field.

Normal AI: Takes 50 tiny, careful steps. (Slow)
Old Speed-up: Tries to run, but trips because it doesn't know the terrain. (Fast but messy)
LESA: Has a map that tells it, "Here is a muddy swamp, walk carefully. Here is a smooth path, run! Here is a glass floor, tiptoe!" It adapts its speed and style to the terrain, getting you across the field in record time without falling down.

LESA makes AI image and video generation fast enough to be used in real-time applications, without sacrificing the beautiful quality we expect.

1. Problem Statement

Diffusion Models (DMs), particularly Diffusion Transformers (DiTs), have achieved state-of-the-art results in image and video generation but suffer from high computational costs due to iterative denoising processes. To address this, Feature Caching has emerged as a popular acceleration strategy, aiming to reuse or predict features from previous timesteps to skip redundant computations.

However, existing methods face two critical limitations:

Static Assumptions: Current "Cache-Then-Reuse" or "Cache-Then-Forecast" (e.g., Taylor expansion) methods assume feature evolution is either minimal or smooth and continuous over time.
Stage-Dependent Dynamics: The paper argues that diffusion processes are non-uniform. They exhibit distinct dynamics across stages:
- High-noise stage: Drastic, unstable, and rapid feature changes.
- Mid-noise stage: Stable and continuous evolution.
- Low-noise stage: Fine-grained detail refinement.
- Consequence: A single fixed prediction scheme (like a global Taylor expansion) fails to adapt to these varying dynamics, leading to quality degradation, temporal inconsistency, and "drift" during accelerated inference.

2. Methodology: The LESA Framework

The authors propose LESA (Learnable Stage-Aware Predictors), a training-based framework designed to explicitly model the temporal characteristics of the diffusion process.

A. Core Architecture

Kolmogorov–Arnold Network (KAN): Instead of standard Multi-Layer Perceptrons (MLPs), LESA uses KANs as the foundational predictor. KANs utilize learnable spline-based univariate functions at each neuron, offering superior parameter efficiency and the ability to approximate complex, non-linear temporal dynamics compared to fixed-activation MLPs.
Prediction Mechanism: The predictor takes cached features from previous $K$ steps and the current timestep embedding. It uses a linear projection for features and a KAN module to generate a scalar temporal modulation factor ( $\alpha$ ). The predicted feature is computed as a residual update: $\hat{h}_{t-1} = h_t + \alpha \cdot z$ .

B. Stage-Aware Multi-Expert Architecture

To handle the non-uniform dynamics, LESA partitions the denoising trajectory into distinct stages based on noise levels, assigning a dedicated "expert" predictor to each:

High-Noise Expert: Handles the initial stage with rapid shifts. Uses a shorter temporal window ( $K=4$ ) to capture weak correlations.
Mid-Noise Expert: Manages the stable, continuous middle phase. Uses a longer window ( $K=8$ ).
Low-Noise Expert: Responsible for final detail refinement.
This design ensures that the specific statistical properties of each generation phase are modeled by a specialized network.

C. Two-Stage Training Strategy

LESA employs a specialized training pipeline to ensure robustness against prediction drift:

Ground-Truth (GT) Guided Training: The predictor is trained using actual features from the base model to learn the underlying denoising dynamics.
Closed-Loop Autoregressive (CL-AR) Training: The predictor is fine-tuned by conditioning on its own historical predictions mixed with actual model outputs. This exposes the model to accumulated errors, teaching it to remain robust during fast, multi-step inference where drift is inevitable.

3. Key Contributions

Learnable Temporal Prediction: Introduction of a KAN-based predictor that learns temporal dynamics directly from data, overcoming the limitations of training-free, heuristic forecasting methods.
Stage-Aware Multi-Expert Design: A novel architecture that segments the diffusion process and assigns specialized predictors to different noise regimes, significantly improving prediction fidelity.
Robust Training Strategy: A two-stage training approach (GT-guided + CL-AR) that mitigates error accumulation, enabling stable high-ratio acceleration.
Generalization: The framework is model-agnostic and effective across both text-to-image and text-to-video tasks.

4. Experimental Results

The authors evaluated LESA on FLUX.1-dev/schnell, Qwen-Image, and HunyuanVideo, comparing it against SOTA methods like TaylorSeer, TeaCache, and FORA.

FLUX.1-dev: Achieved 5.00× acceleration with only a 1.0% drop in quality metrics (ImageReward/CLIP) compared to the baseline.
Qwen-Image: Achieved 6.25× speedup. Notably, it showed a 20.2% quality improvement over the previous SOTA (TaylorSeer) at similar acceleration levels.
HunyuanVideo: Achieved 5.00× acceleration with a 24.7% PSNR improvement over TaylorSeer.
Efficiency: LESA demonstrated lower latency and fewer FLOPs compared to competitors while maintaining superior visual fidelity (higher PSNR, SSIM, and lower LPIPS).
Qualitative Analysis: Visual comparisons showed LESA preserved spatial structures, color consistency, and object semantics better than baselines, which often suffered from structural distortion or object drift at high acceleration ratios.

5. Significance

LESA represents a paradigm shift in diffusion model acceleration:

From Heuristic to Learnable: It moves away from fixed mathematical assumptions (like Taylor series) toward data-driven, learnable predictors that adapt to the specific physics of the diffusion process.
From Uniform to Stage-Aware: It acknowledges and exploits the non-stationary nature of diffusion, proving that "one size fits all" caching strategies are suboptimal.
Practical Deployment: By maintaining high fidelity even at 6.25× acceleration, LESA makes high-quality generative AI significantly more accessible and efficient for real-world applications, reducing hardware costs and inference latency without sacrificing output quality.

The code is released in supplementary materials, facilitating further research and adoption in the community.