[Paper Review] End-to-End Stochastic Optimization with Energy-Based Model

Lingkai Kong, Jiaming Cui+4 more|arXiv (Cornell University)|Nov 25, 2022

Advanced Neural Network Applications7 citations

TL;DR

This paper proposes So-Ebm, a general and efficient end-to-end stochastic optimization framework using energy-based models (EBMs) that bypasses reliance on KKT conditions. By parameterizing the optimization problem via a differentiable EBM surrogate and employing a coupled training objective with likelihood and distribution matching, So-Ebm achieves superior performance and up to 63x faster training than state-of-the-art DFL methods on convex and nonconvex problems including power scheduling and adversarial security games.

ABSTRACT

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.

Motivation & Objective

• To address the limitations of existing decision-focused learning (DFL) methods, which are restricted to convex problems and suffer from high training cost due to repeated optimization solves.
• To develop a generalizable and scalable DFL framework applicable to both convex and nonconvex stochastic optimization problems.
• To eliminate dependency on KKT conditions by replacing implicit differentiable optimization layers with an explicit energy-based surrogate model.
• To improve training efficiency through a self-normalized importance sampling procedure and a coupled objective that captures both optimal decision locations and the global energy landscape.
• To demonstrate effectiveness and efficiency across diverse real-world applications, including power scheduling, pandemic resource allocation, and nonconvex security games.

Proposed method

• So-Ebm uses an energy-based model (EBM) to explicitly parameterize the optimization landscape, replacing implicit KKT-based differentiable layers.
• The method introduces a coupled training objective: a maximum likelihood loss to match the optimal decisions and a distribution-based regularizer to align the posterior of decisions with the EBM's predictive distribution.
• A self-normalized importance sampler with a Gaussian mixture proposal is used to efficiently estimate expectations during training, reducing computational cost.
• Constraints are handled via pre-processing and inference-time projection, with potential extensions using soft constraints like barrier or augmented Lagrangian penalties.
• The EBM surrogate is trained end-to-end with the predictive model, enabling backpropagation through the optimization layer without solving the original problem at each step.
• The framework is general and can be extended to combinatorial problems using discrete EBMs or advanced training techniques such as MCMC or improved EBM training algorithms.

Experimental results

Research questions

• RQ1 Can a general and efficient decision-focused learning framework be developed that works beyond convex problems?
• RQ2 How can the optimization landscape be effectively approximated using an energy-based surrogate without relying on KKT conditions?
• RQ3 Can training efficiency be significantly improved by replacing iterative optimization solves with a differentiable EBM-based surrogate?
• RQ4 How does the coupled training objective—combining local and global matching—enhance decision quality and generalization?
• RQ5 To what extent can So-Ebm reduce data dependency and maintain performance under low-data regimes?

Key findings

• So-Ebm achieves up to 63 times faster training than DFL-QPTH, reducing training time from 93.12 seconds per epoch to 1.47 seconds per epoch on the power generator scheduling task.
• In nonconvex adversarial security games, So-Ebm outperforms all baselines in decision regret, demonstrating its effectiveness on complex, non-relaxable problems.
• On the COVID-19 resource allocation task, So-Ebm maintains strong performance even with limited training data, showing low data dependency due to its algorithmic structure-aware energy function.
• The coupled training objective—combining likelihood of optimal decisions and distribution matching—leads to better approximation of the true optimization landscape than single-objective alternatives.
• So-Ebm maintains consistent performance across varying training data ratios, outperforming baselines except at the extreme low-data ratio of 0.05, where all methods fail due to insufficient data.
• The method generalizes well to both convex and nonconvex problems, including those that are difficult to relax, and scales efficiently even when analytical expectations are unavailable.