Shortcut Diffusion Training with Cumulative Consistency Loss: An Optimal Control View
Paribesh Regmi, Sandesh Ghimire, Rui Li
Abstract
Although iterative denoising (i.e., diffusion/flow) methods offer strong generative performance, they suffer from low generation efficiency, requiring hundreds of steps of network forward passes to simulate a single sample. Mitigating this requires taking larger step-sizes during simulation, thereby allowing one- or few-step generation. Recently proposed shortcut model learns larger step-sizes by enforcing alignment between its direction and the path defined by a base many-step flow-matching model through a self-consistency loss. However, its generation quality is significantly lower than the base model. In this paper, we formulate few-step generation as a controlled base generative process, and show that self-consistency loss can be understood through the lens of optimal control. This perspective naturally motivates its generalization to the proposed cumulative self-consistency loss that cumulatively penalizes misalignment along the entire trajectory. This encourages larger step-sizes that not only align with the base model at the current time step but also support alignment in the subsequent steps, facilitating high-quality generation. Furthermore, we draw a connection between our approach and reinforcement learning, potentially opening the door to a new set of approaches for few-step generation. Experiments show that we significantly improve one- and few-step generation quality under the same training budget. Implementation is available at: https://github.com/paribeshregmi/Shortcut-CSL
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on26
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- Scalable Diffusion Models with TransformersWilliam Peebles, Saining XieICCV 2023 · 5,568 citations
- Consistency ModelsYang Song, Prafulla Dhariwal, Mark Chen, Ilya SutskeverICML 2023 · 1,720 citations
- Score-Based Generative Modeling through Stochastic Differential EquationsYang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar et al.ICLR 2021 · 1,270 citations
Related papers
- One Step Diffusion via Shortcut ModelsKevin Frans, Danijar Hafner, Sergey Levine, Pieter AbbeelICLR 2025 · 2 citations
- Align Your Flow: Scaling Continuous-Time Flow Map DistillationAmirmojtaba Sabour, Sanja Fidler, Karsten KreisNeurIPS 2025 · 91 citations
- Shortcutting Pre-trained Flow Matching Diffusion Models is Almost Free LunchXu Cai, Yang Wu, Qianli Chen, Haoran Wu et al.NeurIPS 2025 · 4 citations
- Self-Corrected Flow Distillation for Consistent One-Step and Few-Step Image GenerationQuan Dao, Hao Phung, Trung Tuan Dao, Dimitris N. Metaxas et al.AAAI 2025 · 11 citations
- Truncated Consistency ModelsSangyun Lee, Yilun Xu, Tomas Geffner, Giulia Fanti et al.ICLR 2025
