arXiv preprint ·
Cumulative Marginal Matching for One-Step Diffusion
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
Abstract
The sampling procedures of diffusion and flow-matching models involve hundreds of denoising steps incurring substantial inference time, which has motivated a growing body of work on training schemes for few-step and one-step generation. Existing approaches primarily seek to align the generation trajectory of few-step and base models, yet their training objectives only involve discrepancies at a single point in the trajectory, encouraging only local \emph{point-level} alignment. While recent methods such as \emph{Shortcut-CSL} introduce cumulative objective that enforce alignment over the entire trajectory, this is suboptimal as they do not directly enforce alignment between the distributions induced by the two models throughout the generation process. We propose \emph{Cumulative Marginal Matching} (CMM), a principled objective for few-step generation: by viewing flow-matching as a sequence of intermediate marginal distributions connecting the prior and data distributions, CMM directly optimizes the cumulative discrepancy between the base and few-step marginal distributions over the entire sequence. Specifically, we formulate this discrepancy using the Wasserstein-2 distance. By exploiting the pushforward representation of flow-matching, we transform an intractable cumulative Wasserstein objective into a tractable surrogate that can be optimized using samples from the intermediate marginal, which are readily generated through the conditional interpolation process of flow-matching. Furthermore, we show that the existing \emph{Shortcut-CSL} objective arises as a first-order approximation of that of the CMM, establishing a theoretical connection between the two objectives. Experiments on standard benchmarks demonstrate improvement in generation quality over \emph{Shortcut-CSL} and other state-of-the-art one- and few-step baselines under identical training budgets.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
Continual Learning Exam: Benchmarking Long-Horizon RSI and Experiential Learning in LLM Agents
Dream to Assist: Intent-Conditioned World Models for Shared Vehicle Control
A Nonconvex Framework for Joint Sparse Sample Weighting and Parameter Estimation
AdaTutor: Adaptive Tutoring through Pedagogical Strategy Selection
How Do Small Coding Agents Use Code Graphs and Embeddings?
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.