Diffusion-DFL: Decision-focused Diffusion Models for Stochastic Optimization
Zihao Zhao, Christopher Yeh, Lingkai Kong, Kai Wang
Abstract
Decision-focused learning (DFL) integrates predictive modeling and optimization by training predictors to optimize the downstream decision target rather than merely minimizing prediction error. To date, existing DFL methods typically rely on deterministic point predictions, which are often insufficient to capture the intrinsic stochasticity of real-world environments. To address this challenge, we propose the first diffusion-based DFL approach, which trains a diffusion model to represent the distribution of uncertain parameters and optimizes the decision by solving a stochastic optimization with samples drawn from the diffusion model. Our contributions are twofold. First, we formulate diffusion DFL using the reparameterization trick, enabling end-to-end training through diffusion. While effective, it is memory and compute-intensive due to the need to differentiate through the diffusion sampling process. Second, we propose a lightweight score function estimator that uses only several forward diffusion passes and avoids backpropagation through the sampling. This follows from our results that backpropagating through stochastic optimization can be approximated by a weighted score function formulation. We empirically show that our diffusion DFL approach consistently outperforms strong baselines in decision quality. The source code for all experiments is available at https://github.com/GT-KOALA/Diffusion_DFL.
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.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 3a53c198-b113-47a8-be4a-a1fc866a67d7Cited by top-tier papers2
- Neural Network-Based Score Estimation in Diffusion Models: Optimization and GeneralizationYinbin Han, Meisam Razaviyayn, Renyuan XuICLR 2024 · 33 citations
- HONet: Data-Efficient Learning for Exact Cover Tasks via Hypergraph OptimizationPengyang Huang, Zirui Zhuang, Haifeng Sun, Qi Qi et al.ICML 2026
Builds on14
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Diffusion Models Beat GANs on Image SynthesisPrafulla Dhariwal, Alexander Quinn NicholNeurIPS 2021 · 13,211 citations
- Improved Denoising Diffusion Probabilistic ModelsAlexander Quinn Nichol, Prafulla DhariwalICML 2021 · 5,234 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
- CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series ImputationYusuke Tashiro, Jiaming Song, Yang Song, Stefano ErmonNeurIPS 2021 · 1,245 citations
Related papers
- End-to-end Stochastic Optimization with Energy-based ModelLingkai Kong, Jiaming Cui, Yuchen Zhuang, Rui Feng et al.NeurIPS 2022 · 33 citations
- Learning MDPs from Features: Predict-Then-Optimize for Sequential Decision Making by Reinforcement LearningKai Wang, Sanket Shah, Haipeng Chen, Andrew Perrault et al.NeurIPS 2021 · 44 citations
- Online Decision-Focused LearningAymeric Capitaine, Maxime Haddouche, Eric Moulines, Michael I. Jordan et al.ICLR 2026 · 4 citations
- 3D-Learning: Diffusion-Augmented Distributionally Robust Decision-Focused LearningJiaqi Wen, Lei Fan, Jianyi YangINFOCOM 2026 · 1 citation
- From Sequential to Recursive: Enhancing Decision-Focused Learning with Bidirectional FeedbackXinyu Wang, Jinxiao Du, Yiyang Peng, Wei MaAAAI 2026
