💻 computer science

Few-Step Diffusion Sampling Through Instance-Aware Discretizations

This paper proposes an instance-aware discretization framework that learns to adapt timestep allocations based on input-dependent priors, thereby overcoming the limitations of globally shared schedules and consistently improving generation quality across diverse diffusion and flow matching models with negligible inference overhead.

Original authors: Liangyu Yuan, Ruoyu Wang, Tong Zhao, Dingwen Fu, Mingkun Lei, Beier Zhu, Chi Zhang

Published 2026-03-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Liangyu Yuan, Ruoyu Wang, Tong Zhao, Dingwen Fu, Mingkun Lei, Beier Zhu, Chi 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 recreate a complex painting, like a detailed landscape or a portrait of a person. You start with a blank canvas covered in static noise (like TV snow). Your goal is to slowly turn that noise into the final picture.

Diffusion models (the AI behind tools like DALL-E 3 or Midjourney) do this by taking tiny steps, cleaning up a little bit of the noise at a time, until the image is clear.

The Problem: The "One-Size-Fits-All" Map

To get from the noise to the picture, the AI follows a mathematical path. To make this fast, the AI has to take "steps" along this path.

The Old Way: Imagine a tour guide leading a group of 100 people through a city. The guide has a single, pre-written itinerary for the whole day.
- For the people who want to visit a simple park, the guide forces them to walk through a complex, winding maze of museums first. It's a waste of time.
- For the people who want to see a massive cathedral, the guide rushes them through the park too quickly, and they miss the details.
- The Result: Everyone gets to the destination, but some are bored, some are confused, and the final photos aren't perfect. In AI terms, this is called a Global Schedule. It uses the same number of steps for every single image, regardless of how hard that specific image is to generate.

The Solution: A Personal GPS for Every Image

The paper introduces a new method called INDIS (Instance-Aware Discretization).

Think of INDIS as giving every single traveler their own personal GPS that looks at their specific starting point and destination before the trip begins.

It Analyzes the Journey: Before the AI starts generating an image, it looks at the specific "noise" (the starting point) and the prompt (e.g., "a cat" vs. "a chaotic explosion").
It Customizes the Steps:
- If the image is simple (like a blue sky), the GPS says, "Take big, fast steps. We can get there in 3 moves!"
- If the image is complex (like a face with intricate details), the GPS says, "Slow down. Take tiny, careful steps here to get the eyes right."
The Result: Every image gets exactly the amount of "effort" it needs. Simple things are generated instantly; complex things get the extra attention they require.

Why This Matters

Speed vs. Quality: Usually, if you want a high-quality image, you have to wait longer (take more steps). If you want it fast, the quality drops. INDIS breaks this trade-off. It gets high-quality results in very few steps (sometimes just 3 or 5) because it doesn't waste time on easy parts of the image.
No Re-Training: The best part? The AI model itself doesn't need to be retrained from scratch. It's like giving the existing tour guide a new, smarter GPS device. It's a lightweight upgrade that works immediately.

The Analogy in Action

Imagine you are baking cookies.

Global Schedule: You put every batch of dough in the oven for exactly 12 minutes. Some cookies are undercooked (gooey), and some are burnt.
INDIS: You have a smart oven that checks each cookie individually. It sees the chocolate chip cookie needs 10 minutes, but the oatmeal raisin needs 14. It adjusts the timer for every single cookie automatically.

Summary

This paper proposes a smart way to tell AI image generators: "Don't treat every picture the same. Look at what you're making, and decide exactly how many steps you need to make it perfect."

The result is faster generation, sharper images, and less wasted computing power, all without needing to rebuild the AI from the ground up.

1. Problem Statement

Diffusion Probabilistic Models (DPMs) and Flow Matching models generate high-fidelity data by simulating trajectories defined by Ordinary Differential Equations (ODEs) or Stochastic Differential Equations (SDEs). While advanced numerical solvers (e.g., DPM-Solver, UniPC) have accelerated sampling, the discretization strategy (the schedule of timesteps used to step from noise to data) remains a critical bottleneck.

Current Limitation: Existing state-of-the-art methods rely on either handcrafted heuristics (e.g., Uniform, LogSNR) or globally optimized timestep schedules. These global schedules are optimized once for the entire dataset and applied uniformly to every sample.
The Core Issue: This "one-size-fits-all" approach fails to account for instance-specific complexity. Different initial noise samples ( $x_T$ ) and conditional inputs ( $c$ ) follow distinct trajectories with varying levels of difficulty. A global schedule represents a suboptimal compromise that cannot adapt to the specific dynamics of a single generation instance, limiting performance in few-step (low NFE) scenarios.

2. Methodology: INDIS (Instance-Specific Discretization)

The authors propose INDIS, a framework that learns to dynamically allocate timesteps based on the specific input instance (initial noise and condition).

A. Core Insight & Toy Experiments

The authors first validated their hypothesis using synthetic 2D data. They compared:

Uniform Schedule: Fixed spacing.
Globally Optimized: A single set of timesteps optimized for the whole dataset.
Instance-Adaptive (Oracle): A unique set of timesteps optimized for each specific noise sample.

Finding: The instance-adaptive oracle significantly outperformed the global schedule, proving that a single global schedule is mathematically suboptimal for individual trajectories.

B. The INDIS Framework

To scale this to high-dimensional image and video generation, INDIS introduces a lightweight neural network, $\phi(\cdot)$ , that predicts the discretization schedule.

Input: The network takes the initial noise sample $x_T$ and the conditional guidance $c$ (e.g., class label, text prompt) as input.
Output: It outputs a tailored discretization set $\xi_\phi = \{\tau_n, \Delta\tau_n, \gamma_n\}_{n=1}^N$ $ξ_{ϕ} = {τ_{n}, Δ τ_{n}, γ_{n}}_{n = 1}^{N}$ , which includes:
- $\tau_n$ : The primary timestep locations.
- $\Delta\tau_n$ : Time shift factors to mitigate exposure bias (the mismatch between training noise schedules and sampling trajectories).
- $\gamma_n$ : Scale factors to adjust the magnitude of the predicted noise.
Architecture:
- For noise $x_T$ , the network uses Singular Value Decomposition (SVD) to extract features, followed by MLPs.
- For conditions $c$ , it processes embeddings (e.g., CLIP for text, linear layers for class labels).
- The outputs are normalized via Softmax (for timesteps) and Tanh (for shifts/scales) to ensure stability and monotonicity.
Training Objective:
- The network is trained via gradient-based optimization to minimize the distance (e.g., LPIPS or MSE) between the student sample (generated with $\xi_\phi$ ) and a teacher sample (generated with a high-step, high-quality solver).
- Crucially, the teacher is a pre-computed "oracle" trajectory, and the student learns to approximate it using fewer steps by adapting the schedule per instance.

3. Key Contributions

Identification of Global Suboptimality: The paper provides quantitative evidence that globally shared timestep schedules are inherently suboptimal for individual samples, establishing a theoretical upper bound for instance-specific methods.
Instance-Aware Discretization Paradigm: Proposes INDIS, the first method to generalize gradient-based discretization search to the conditional generative setting, allowing the timestep schedule to be a function of the input noise and condition.
Handling Exposure Bias: Integrates learnable time-shift and scale factors directly into the discretization search, offering a novel way to alleviate the exposure bias problem common in few-step sampling without heavy model retraining.
Scalability: Demonstrates that this lightweight adaptation can be applied to diverse models (Pixel-space, Latent-space, Video) with negligible inference overhead.

4. Experimental Results

The authors evaluated INDIS across multiple datasets and model architectures, comparing it against best-in-class heuristics, globally optimized schedules (GOD), and other discretization methods (DMN, GITS, LD3).

Pixel-Space Diffusion (CIFAR-10, ImageNet, FFHQ):
- INDIS consistently achieved the lowest Fréchet Inception Distance (FID).
- At NFE=3 (3 steps), INDIS improved FID by 35.33% on average compared to the strongest baseline.
- Performance gains were most pronounced in low-step regimes (3–6 steps).
Latent-Space Diffusion (Stable Diffusion on LSUN-Bedroom):
- Achieved a 14.12% average FID improvement over baselines.
Flow Matching & Video (FLUX.1-dev, LTX-Video):
- On FLUX.1-dev (text-to-image), INDIS outperformed both global heuristics and globally optimized schedules in FID, CLIP scores, and CMMD metrics.
- On LTX-Video, it improved aesthetic quality, imaging quality, and subject consistency as measured by VBench.
Efficiency:
- Inference Overhead: The lightweight network $\phi$ adds negligible cost (e.g., ~2.3% extra time for FLUX.1-dev, ~2.5% for CIFAR-10).
- Training Cost: Requires only a one-time tuning phase using a pre-generated teacher dataset; it does not require retraining the massive base diffusion model.

5. Significance

Paradigm Shift: INDIS moves the field away from static, global solver configurations toward dynamic, instance-aware sampling. It acknowledges that "one size does not fit all" in generative trajectories.
Practical Impact: It offers a highly efficient "plug-and-play" acceleration method. Unlike model distillation (which requires expensive retraining) or complex solver redesigns, INDIS only requires training a small auxiliary network to predict better timesteps.
Generalizability: The method is solver-agnostic (works with iPNDM, DPM-Solver, etc.) and modality-agnostic (works for images and videos), making it a versatile tool for improving the efficiency of modern generative AI systems.

In conclusion, INDIS demonstrates that by tailoring the discretization strategy to the specific characteristics of each generation instance, it is possible to achieve significantly higher quality in few-step diffusion sampling with minimal computational cost.