# End-to-End Stochastic Optimization with Energy-Based Model

arXiv: https://arxiv.org/abs/2211.13837 · HTML: https://mlsanity.com/paper/2211.13837/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 63 references that are in the map.
- First posted on arXiv 25 Nov 2022.

## Builds on (most cited first, 30 of 32)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Pointer Networks](https://mlsanity.com/paper/1506.03134/) (2015)
- [CVXPY: A Python-Embedded Modeling Language for Convex Optimization](https://mlsanity.com/paper/1603.00943/) (2016)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Learning Combinatorial Optimization Algorithms over Graphs](https://mlsanity.com/paper/1704.01665/) (2017)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One](https://mlsanity.com/paper/1912.03263/) (2019)
- [Sliced Score Matching: A Scalable Approach to Density and Score Estimation](https://mlsanity.com/paper/1905.07088/) (2019)
- [Smart "Predict, then Optimize"](https://mlsanity.com/paper/1710.08005/) (2017)
- [How to Train Your Energy-Based Models](https://mlsanity.com/paper/2101.03288/) (2021)
- [Melding the Data-Decisions Pipeline: Decision-Focused Learning for Combinatorial Optimization](https://mlsanity.com/paper/1809.05504/) (2018)
- [Learning to Optimize](https://mlsanity.com/paper/1606.01885/) (2016)
- [A Rewriting System for Convex Optimization Problems](https://mlsanity.com/paper/1709.04494/) (2017)
- [Maximum Entropy Deep Inverse Reinforcement Learning](https://mlsanity.com/paper/1507.04888/) (2015)
- [Improved Contrastive Divergence Training of Energy Based Models](https://mlsanity.com/paper/2012.01316/) (2020)
- [Energy-based Generative Adversarial Network](https://mlsanity.com/paper/1609.03126/) (2016)
- [Interior Point Solving for LP-based prediction+optimisation](https://mlsanity.com/paper/2010.13943/) (2020)
- [SDE-Net: Equipping Deep Neural Networks with Uncertainty Estimates](https://mlsanity.com/paper/2008.10546/) (2020)
- [Contrastive Losses and Solution Caching for Predict-and-Optimize](https://mlsanity.com/paper/2011.05354/) (2020)
- [Learning deep kernels for exponential family densities](https://mlsanity.com/paper/1811.08357/) (2018)
- [Learning Discrete Energy-based Models via Auxiliary-variable Local Exploration](https://mlsanity.com/paper/2011.05363/) (2020)
- [MCMC Should Mix: Learning Energy-Based Model with Neural Transport Latent Space MCMC](https://mlsanity.com/paper/2006.06897/) (2020)
- [Automatically Learning Compact Quality-aware Surrogates for Optimization Problems](https://mlsanity.com/paper/2006.10815/) (2020)

## Built on by (most cited first, 19 of 19)

- [A Survey of Contextual Optimization Methods for Decision Making under Uncertainty](https://mlsanity.com/paper/2306.10374/) (2023)
- [Learning Energy-Based Models by Cooperative Diffusion Recovery Likelihood](https://mlsanity.com/paper/2309.05153/) (2023)
- [Gen-DFL: Decision-Focused Generative Learning for Robust Decision Making](https://mlsanity.com/paper/2502.05468/) (2025)
- [Generative Models in Decision Making: A Survey](https://mlsanity.com/paper/2502.17100/) (2025)
- [MUBen: Benchmarking the Uncertainty of Molecular Representation Models](https://mlsanity.com/paper/2306.10060/) (2023)
- [BPQP: A Differentiable Convex Optimization Framework for Efficient End-to-End Learning](https://mlsanity.com/paper/2411.19285/) (2024)
- [Differentiable Distributionally Robust Optimization Layers](https://mlsanity.com/paper/2406.16571/) (2024)
- [DF2: Distribution-Free Decision-Focused Learning](https://mlsanity.com/paper/2308.05889/) (2023)
- [Locally Convex Global Loss Network for Decision-Focused Learning](https://mlsanity.com/paper/2403.01875/) (2024)
- [Combinatorial Optimization Augmented Machine Learning](https://mlsanity.com/paper/2601.10583/) (2026)
- [E2E-AT: A Unified Framework for Tackling Uncertainty in Task-aware End-to-end Learning](https://mlsanity.com/paper/2312.10587/) (2023)
- [Prediction Loss Guided Decision-Focused Learning](https://mlsanity.com/paper/2509.08359/) (2025)
- [Learning Deterministic Surrogates for Robust Convex QCQPs](https://mlsanity.com/paper/2312.12485/) (2023)
- [Differentiation of Multi-objective Data-driven Decision Pipeline](https://mlsanity.com/paper/2406.00779/) (2024)
- [Learning Solutions of Stochastic Optimization Problems with Bayesian Neural Networks](https://mlsanity.com/paper/2406.03082/) (2024)
- [A Pontryagin Method of Model-based Reinforcement Learning via Hamiltonian Actor-Critic](https://mlsanity.com/paper/2603.28971/) (2026)
- [Smart predict-then-robustly-optimize](https://mlsanity.com/paper/2607.21773/) (2026)
- [Decision-Driven Regularization: A Blended Model for Learning and Optimization](https://mlsanity.com/paper/2608.15124/) (2026)
- [RL-PaO: Prediction as Action in Decision Making under Uncertainty](https://mlsanity.com/paper/2609.37065/) (2026)

## Abstract (arXiv)

Decision-focused learning (DFL) was recently proposed for stochastic optimization problems that involve unknown parameters. By integrating predictive modeling with an implicitly differentiable optimization layer, DFL has shown superior performance to the standard two-stage predict-then-optimize pipeline. However, most existing DFL methods are only applicable to convex problems or a subset of nonconvex problems that can be easily relaxed to convex ones. Further, they can be inefficient in training due to the requirement of solving and differentiating through the optimization problem in every training iteration. We propose SO-EBM, a general and efficient DFL method for stochastic optimization using energy-based models. Instead of relying on KKT conditions to induce an implicit optimization layer, SO-EBM explicitly parameterizes the original optimization problem using a differentiable optimization layer based on energy functions. To better approximate the optimization landscape, we propose a coupled training objective that uses a maximum likelihood loss to capture the optimum location and a distribution-based regularizer to capture the overall energy landscape. Finally, we propose an efficient training procedure for SO-EBM with a self-normalized importance sampler based on a Gaussian mixture proposal. We evaluate SO-EBM in three applications: power scheduling, COVID-19 resource allocation, and non-convex adversarial security game, demonstrating the effectiveness and efficiency of SO-EBM.