Few-Step Diffusion Sampling Through Instance-Aware Discretizations

• URL: http://arxiv.org/abs/2603.17671v1
• Date: Wed, 18 Mar 2026 12:49:38 GMT
• Title: Few-Step Diffusion Sampling Through Instance-Aware Discretizations
• Abstract summary: Diffusion and flow matching models generate high-fidelity data by simulating paths defined by Ordinary or Differential Equations (ODEs/SDEs)<n>We propose an instance-aware discretization framework, which learns to adapt timestep allocations based on input-dependent priors.<n>Our method consistently improves generation quality with marginal tuning cost compared to training and negligible inference overhead.
• Score: 16.93105640007355
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Diffusion and flow matching models generate high-fidelity data by simulating paths defined by Ordinary or Stochastic Differential Equations (ODEs/SDEs), starting from a tractable prior distribution. The probability flow ODE formulation enables the use of advanced numerical solvers to accelerate sampling. Orthogonal yet vital to solver design is the discretization strategy. While early approaches employed handcrafted heuristics and recent methods adopt optimization-based techniques, most existing strategies enforce a globally shared timestep schedule across all samples. This uniform treatment fails to account for instance-specific complexity in the generative process, potentially limiting performance. Motivated by controlled experiments on synthetic data, which reveals the suboptimality of global schedules under instance-specific dynamics, we propose an instance-aware discretization framework. Our method learns to adapt timestep allocations based on input-dependent priors, extending gradient-based discretization search to the conditional generative setting. Empirical results across diverse settings, including synthetic data, pixel-space diffusion, latent-space images and video flow matching models, demonstrate that our method consistently improves generation quality with marginal tuning cost compared to training and negligible inference overhead.

Related papers

• Integration Matters: Rollout-Based Training for Constrained Diffusion Models [16.46255955344152]
  Constrained generative models aim to produce samples that satisfy complex feasibility constraints while remaining faithful to the data distribution.<n>Existing constrained generation methods enforce constraints either through training-time optimization or sampling-time correction.<n>We propose a fine-tuning framework that incorporates constraint guidance obtained through online rollout into the training process.
  arXiv  Detail & Related papers  (2026-07-15T22:25:59Z)
• Synergizing Transport-Based Generative Models and Latent Geometry for Stochastic Closure Modeling [1.665466637453776]
  We show that flow matching in a lower-dimensional latent space is suited for fast sampling of closure models.<n>We control the latent space distortion and thus ensure the physical fidelity of the sampled closure term.
  arXiv  Detail & Related papers  (2026-02-19T05:24:00Z)
• Formalizing the Sampling Design Space of Diffusion-Based Generative Models via Adaptive Solvers and Wasserstein-Bounded Timesteps [4.397130429878499]
  Diffusion-based generative models have achieved remarkable performance across various domains, yet their practical deployment is often limited by high sampling costs.<n>We propose SDM, a principled framework that aligns the numerical solver with the intrinsic properties of the diffusion trajectory.<n>By analyzing the ODE dynamics, we show that efficient low-order solvers suffice in early high-noise stages while higher-order solvers can be progressively deployed to handle the increasing non-linearity of later stages.
  arXiv  Detail & Related papers  (2026-02-13T05:02:07Z)
• An Elementary Approach to Scheduling in Generative Diffusion Models [55.171367482496755]
  An elementary approach to characterizing the impact of noise scheduling and time discretization in generative diffusion models is developed.<n> Experiments across different datasets and pretrained models demonstrate that the time discretization strategy selected by our approach consistently outperforms baseline and search-based strategies.
  arXiv  Detail & Related papers  (2026-01-20T05:06:26Z)
• Quantizing Diffusion Models from a Sampling-Aware Perspective [43.95032520555463]
  We propose a sampling-aware quantization strategy, wherein a Mixed-Order Trajectory Alignment technique is devised.<n>Experiments on sparse-step fast sampling across multiple datasets demonstrate that our approach preserves the rapid convergence characteristics of high-speed samplers.
  arXiv  Detail & Related papers  (2025-05-04T20:50:44Z)
• On the Trajectory Regularity of ODE-based Diffusion Sampling [79.17334230868693]
  Diffusion-based generative models use differential equations to establish a smooth connection between a complex data distribution and a tractable prior distribution. In this paper, we identify several intriguing trajectory properties in the ODE-based sampling process of diffusion models.
  arXiv  Detail & Related papers  (2024-05-18T15:59:41Z)
• Align Your Steps: Optimizing Sampling Schedules in Diffusion Models [63.927438959502226]
  Diffusion models (DMs) have established themselves as the state-of-the-art generative modeling approach in the visual domain and beyond. A crucial drawback of DMs is their slow sampling speed, relying on many sequential function evaluations through large neural networks. We propose a general and principled approach to optimizing the sampling schedules of DMs for high-quality outputs.
  arXiv  Detail & Related papers  (2024-04-22T18:18:41Z)
• Diffusion Models as Constrained Samplers for Optimization with Unknown Constraints [55.39203337683045]
  We propose to perform optimization within the data manifold using diffusion models.<n>Depending on the differentiability of the objective function, we propose two different sampling methods.<n>Our method achieves better or comparable performance with previous state-of-the-art baselines.
  arXiv  Detail & Related papers  (2024-02-28T03:09:12Z)
• A Geometric Perspective on Diffusion Models [57.27857591493788]
  We inspect the ODE-based sampling of a popular variance-exploding SDE. We establish a theoretical relationship between the optimal ODE-based sampling and the classic mean-shift (mode-seeking) algorithm.
  arXiv  Detail & Related papers  (2023-05-31T15:33:16Z)
• Reflected Diffusion Models [93.26107023470979]
  We present Reflected Diffusion Models, which reverse a reflected differential equation evolving on the support of the data. Our approach learns the score function through a generalized score matching loss and extends key components of standard diffusion models.
  arXiv  Detail & Related papers  (2023-04-10T17:54:38Z)
• Learning Sampling Distributions for Model Predictive Control [36.82905770866734]
  Sampling-based approaches to Model Predictive Control (MPC) have become a cornerstone of contemporary approaches to MPC. We propose to carry out all operations in the latent space, allowing us to take full advantage of the learned distribution. Specifically, we frame the learning problem as bi-level optimization and show how to train the controller with backpropagation-through-time.
  arXiv  Detail & Related papers  (2022-12-05T20:35:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.