End-to-end Stochastic Optimization with Energy-based Model
Authors: Lingkai Kong
Jiaming Cui
Yuchen Zhuang
Rui Feng
B. Aditya Prakash
Chao Zhang
Venue: NeurIPS 2022 Accept (2022)
Keywords: end-to-end stochastic optimization
energy-based model
decision-focused learning

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.

End-to-End Stochastic Optimization with Energy-Based Model
Lingkai Kong
Jiaming Cui
Yuchen Zhuang
Rui Feng
B. Aditya Prakash
Chao Zhang
College of Computing
Georgia Institute of Technology
(1kkong, j jimiinmumm1199977 yewhuang,,neceee rfeng, badityap, chaozhang)@gatechh. edu
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 per- formance to the standard two-stage predict-then-opiimiz pipeline. However, most existing DFL methods only applicable convex problems subset of non- convex problems that can be easily relaxed to convex ones. Further, they can be inefficient in training due 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- -ased 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 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 distribution-basee regularizer to capture the overall energy landscape. Finally, we propose an efficient training procedure for SO-EBM with self-normalized importance sampler based on Gaussian mixture proposal. We evaluate SO-EBM in three applications: power scheduling, COVID-19 19 resource allocation, and non-convex adversarial security game, demonstrating the effectiveness and efficiency of SO-EBM.
Introduction
Many real-life decision making tasks stochastic optimization problems, where needs to make decisions to minimize a cost function that involves stochastic parameters. Oftentimes, the involved stochastic parameters unknown and context-dependent, meaning that they need to be predicted from observed features. For example, when allocating clinical resources for COVID-19, necessary consider future cases in different regions, whose distributions unknown and have be predicted from the current state. As two other examples, hedge funds need to continuously adjust their portfolio for maximal expected return, by forecasting future return rates of different stocks; and in supply chain optimization, facility locations need to be decided to minimize long-term operational costs, by accounting for unknown and uncertain regional supplies and customer demands.
With the feasibility of training powerful deep learning predictors from large amounts of data, is increasingly common to solve such stochastic optimization problems using two-stage predict-them optimize pipeline. In the prediction stage, one learns a predictive model for the unknown parameters using certain prediction loss (e.g., negative likelihood). In the optimization stage, the predicted distributions of the parameters are used to parameterize the stochastic optimization problem, which be then solved using off-the-shelf solvers 2 [2]. This two-stage pipeline relies an implicit
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

and commonly-accepied assumption: improvements in parameter prediction in terms of the prediction loss will always translate to better optimization outcomes. However, this is not the case: machine learning models make errors and the impact of prediction errors is not uniform w.E.t the optimization loss. Thus, a smaller predictive loss does not necessarily lead to a smaller decision regret.
A better approach is decision-focused learning (DFL) [3 optimization layers into a unified model to learn them end-to-end. Most DFL methods implement the optimization procedure as an implicitly differentiable layer and develop techniques (e.g., using KKT conditions) to compute the gradients w.r.f. the decision variables and enable back-proper mack-propagating through it. Compared to the two- wwo-stage pipeline, the prediction layer so learned is tailored for the optimization problem and better the sense that can yield decisions with smaller regrets. Besides DFL, there are also approaches where policy network is trained to directly map from the input to the solution of the optimization problem using supervised or reinforcement learning [47 26 [4] However, their performance are often inferior to DFL. as they ignore the algorithmic structure of the problem and typically require a large amount of data to rediscover the algorithmic structure.
which integrates prediction and
Despite their promising results, existing DFL methods [33 (1) Generality. To leverage the KKT condition, they are mostly applicable to only convex optimization objectives 13, 4 I. Though a few works approximate nonconvex objectives by quadratic functions their applicability is still limited to easy-to-relax nonconvex problems and can suffer from poor gradient estimates caused by the relaxation. (2) Scalability. Due to the reliance of using the KKT condition to compute derivatives, they require repeatedly solving the optimization program and back-propagating through it during training, which makes them unscalable. The problem is more severe when the expectation of the cost cannot be computed analytically.
48) suffer from two drawbacks
41 48
We propose a new end-to-end stochastic optimization method using an energy-based model (EBM), named SO-EBM. Similar to DFL, SO-EBM models the algorithmic predict-and-optimize structure by stacking a differentiable optimization layer on top of neural predictor. Different from DFL, So- EBM eliminates the need of using KKT conditions to induce an implicit differentiable optimization layer. Instead, it leverages expressive EBMs 33 to directly model the conditional probability of the decisions and parameterizes an explicit energy-based based differentiable layer as surrogate to the original optimization problem. To better approximate the optimization landscape with the EBM surrogate, we design two complementary learning objectives in SO-EBM: 1) a local matching objective that maximizes the likelihood of the optimal decisions; and 2) global matching objective that minimizes the distribution distance between the posterior distribution of the decision variables and the surrogate-based decision distribution.
Due to its flexibility, SO-EBM not constrained convex objectives, but can be applied to wide class of stochastic optimization problems. Another key advantage of SO-EBM is its computational efficiency. As the optimization layer is parameterized by an energy-based model, SO-EBM eliminates the need of solving and differentiating through the optimization procedure at each training iteration. Rather, SO-EBM estimates the derivatives of the energy-based surrogate model using solf-normalized importance sampler. We design a coussian-mixture proposal distribution for the sampler, which not only reduces sampling variance, but also borrows the idea of contrastive divergence for MCMC methods to speed up the training of SO-EBM.
The main contributions of this work are: (1) We propose new end-to-end stochastic optimization method based on energy-based model. It avoids the needs of solving and differentiating through the optimization problem during training and can be applied problems beyond convex objectives.
wide range of stochastic optimization (2) We propose coupled training objective to encourage the energy-based surrogate well approximate the optimization landscape and an efficient self- normalized importance sampler based on a mixture-of- -f-uussia- proposal. (3) Experiments on power scheduling, COVID-19 resource allocation and adversarial network security game show that our method can achieve better or comparable performance than existing DFL methods while being much more efficient.
2 Preliminaries
Problem Formulation. Consider a stochastic optimization problem:
arg min f(y
2

---

Two-stage learning:
Predictive Distribution
Generic Loss e.g., NLL
Stochastic Program
Features
p xx 0)
argmin. plllisco) ((r.
Backpropagation
Decision-focused learning:
Features
Predictive Distribution
p(y x d)
Stochastic Program
a Ppttlkk)))))) [(y.
Optimization B ( (xx: 6))
Backpropagation
Backpropagation
Figure 1:
Two-stage learning back-propagates from predictive loss to the model, ignoring the latter effect of the optimization problem. decision-focused learning directly optimizes the task loss but needs to solve and back-propagate through the optimization problem at every training iteration.
where y denotes the parameters of the optimization problem, a denotes the decision variables within feasible space C, and is the cost function to be optimized. In many applications, the parameters y are unknown and stochastic, which must be inferred from some correlated features x. We assume a dataset D {x,, Vi 1==1 drawn from the joint distribution of the features and problem parameters. Our task is to learn a decision-making model M parameterized by O, which takes the features x as input and outputs the optimal decisions a" [x; 0). The model should be learned such that its output optimal decisions minimize the expected decision cost under the joint distribution of (x, y). namely:
arg ming ((x.)) (y, a" (x: 0))-
For example, in supply chain optimization, y can be regional customer demands, a can pre-ordered products for each region, measures the gap between actual customer demands and pre-ordered supplies, and the problem is to decide the best a to minimize the cost f. Herein, actual customer demands y usually unknown and must be predicted from features such as customer purchase history and regional economy indices.
Two-stage learning v.s. Decision Focused Learning (DFL). A common practice to the above framework. first learns an
stochastic optimization problem the two-stage probabilistic predictive model p(y|x; 0) and then uses existing stochastic optimization solvers to ob- tain the optimal action that minimizes the expected cost: a (x; 8) = arg min.ec Ey--psi (x;)) f (y. a). Though the two-stage approach is simple and efficient, it can suffer from suboptimal performance due to the misalignment the prediction loss and the optimization loss. In contrast, DFL integrates prediction and optimization into end-to-end model, thus tailoring the predictive model for the optimization task, as shown in Figure [ By directly optimizing the task loss, the gradient of the model parameters can be obtained with the following chain rule:
0f a" (x; Of(y, a" (x: da (x; 0) Oy Oa" (x;
00
Oy
00
The key challenge here to compute the Jacobian Oa" (x:0)) it is needed to apply the chain rule By to learn the model using gradient decent methods. This is nontrivial, because a" is the solution of stochastic optimization solver and not directly differentiable. To address this challenge, OptNet assumes quadratic optimization objectives and differentiates through the KKT conditions using the implicit function theorem. This way, OptNet can obtain the Jacobian by solving the optimization problem along with a set of linear equations in each training iteration. Several works 1 41 48
B
extend this technique more general cases. For example, cvxpylayers ① extends it to more general cases of convex optimization using disciplined parametrized programming (DPP) grammar.
Although existing DFL approaches can achieve better decisions compared to two-stage learning, they have several drawbacks. First, they are often constrained to convex optimization objectives as they rely on the KKT conditions to compute derivatives. Though [4] 48 propose approximate some non-convex objectives by a quadratic function around a local minimum, the inaccurate gradients may be aggregated during the training iterations and thus lead poor decision quality. Second, they suffer from high computational complexity because the computation of the Jacobian du (x:0) requires repeatedly solving the optimization program and back- peeeppropggtting through it. The problem is more severe when the expectation of the cost cannot be computed analytically. In that case, we need to
By
3

---

use sample average approximation [27 46 30 and draw multiple IID samples from the predictive
distribution, which makes the objective much more complicated and expensive.
Proposed Method
3.1 Energy-based Model for End-to-end Stochastic Programming
Our task is to learn an optimal decision-making model Ma. such that its output optimal decision a" [x;0) () for input minimizes the expected optimization cost: [x (y, (x; (x: The idea of our method is to directly model M&'s 's's probability distribution over the decisions conditioned on the features, denoted as q(a|x: 0), using energy-based parameterization [33
(a(x;;) 0) = xxp(-F(xx,;0)) Z(x,0) 0)
0) = expl exe(p-EE - a; fad
(1)
To parameterize the energy function Exx, a: f). a natural option is to directly use a deep neural network which takes feature-decision pair (x, a) input and output scalar value. However, such a design ignores the algorithmic structure of the optimization problem and thus can be data-inefficient during learning. Instead, propose explicitly model the problem structure by using the expected task loss as the energy function:
E(x, a; 9) = Ep( ((. a),
where p(y x; 0) is the predictive distribution of an uncertainty-aware neural network 19 31
Eq. creates an explicit energy- based differentiable layer surrogate the original optimization problem. Compared with the two-stage model, it builds direct connection between the input features and decision variable and thus is more tailored to the downstream task. Compared with pure end-to-end architectures, the task-based surrogate energy function explicitly leverages the algorithmic structure of the optimization problem and thus saves lot of learning. This energy-based surrogate function also has a intuitive interpretation. When the decision a is in a region with smaller expected task loss, it has lower energy and higher probability; when it leads higher expected task loss, has higher energy and lower probability.
Since we have the ground truth for the parameters in the training data, the feature-decision pairs Da = {(x, at ) can be easily constructed by solving ot = arg min,, f (Jt, a,) for each (x,, V/) using any off-the-sheef optimization solvers. Note that the construction of such feature-action training pairs needs to be done only once during preprocessing (Fig. 3). Then, we minimize the negative likeliidl (NLL) of all the feature-decision pairs:
EILE = (x, g(a" xx: Exx,-)),, (Eplulx,0)
log( 0)))
(3)
By minimizing the NLL., we are essentially minimizing the energy of the optimal actions while maximizing the energy of other points. Thus the end-to-end stochastic programming problem translated to learning a neural network that outputs the smallest energy for the optimal actions. This new perspective avoids the need of solving and differentiating through the optimization problem at every training iteration, but meanwhile tailors the model for the downstream decision making task.
Our energy- formulation can also be interpreted as a non-sequential maximum entropy inverse From the IRL perspective, the input features [x can be interpreted as environment states, the feature-decision pairs {(x;, a as expert demonstrations, and the negative expected task loss - (p(p) f a) as reward. Eq. 3 is then equivalent to maximizing the likelihood of the expert demonstrations (optimal decisions) to recover the hidden reward function.
reinforcement learning (MaxEN-IRe) IRL) model [33 52
3.2 Augmenting Bhergy-Based Objective with Distribution Regularizer
One limitation of learning with EBM-based likelihood for the optimization problem that we model only one data point a² for each conditional distribution q(a|x; xx 0). Minimizing the NLL for a single data point ignores the overall matchness between the landscape of original optimization problem and that of the basedased probability density, which can casily cause overfitting. To learn the overall

---

shape of the energy function better, we propose distribution- based regularizer which augments the energy- objective from global training perspective.
The distribution-basee regularizer is based on minimizing the distance hetween the model distribution and oracle posterior distribution. Specifically, we assume that follows a prior distribution p(a). With the ground-truie truth label y, the posterior distribution of a is then given by p(aly)
Z(v)
when we define the task loss as the negative log-likelihood, i.e.- log p(yla) fy, a). However, such an expression depends on the ground-truth label, which cannot be obtained at test time. We let the model distribution ((a x; 0) mimic this oracle posterior p(aly) distribution by minimizing their KL-divergence:
kKL = KLLaaa-
0) (xx))~ ( (-Ep(alo) log (()xx 0) H(p(aly))), (4)
where H-()) denotes entropy of the probability distribution. Our final training loss is a weighted combination of the NLL and the distribution based regularizer:
&oal MMLE + ACKL
where a complement each other (Fig. 2 The NLL loss is local training strategy while the KL-divergence loss learns the energy function from a global level. Specifically, the NLL loss helps the model capture the location of the optimal decision better but ignores the overall energy shape. The distribution-base regularizer essentially adds more anchor points besides the optimal decisions and thus helps fit the overall energy shape.
In this final objective, the KL-divergence loss and the NLL loss
True energy
E
MLE training
KL-ddivergence training
E
MLE=-LL-divrrencc training
(b)
(c)
Different training objectives of energy-based optimization. (a) the ground-truth energy- based landscape: (b) MLE training captures the location of the optimum but ignores the overall energy shape; (c) KL-divergence training learns the overall energy shape but captures a blurry optimum location; (d) Our MLE+KL-divergence training wins the best of both worlds.
Figure 2:
3.3 Training with Contrastive Divergence and Seff-normalized Importance Sampling
Optimizing Eq. rerquires evaluating the partition function and computing the KL-divergence between two continuous distributions which typically involves intractable integrals. In this subsection, we propose to use self-normalized importance sampler based on mixture of Gaussians to estimate gradient of the model parameters efficiently. First, we derive (see supplementary for details) the gradient of the training loss with respect to the model parameters as:
@Lrotll 00
(E'aa", 0)
@E(ā, x; 0)
-E(xaa")~z JND.
00
00
(6)
0E x; 0)
@E(ā, x;
(
Ep(alu)
E [
00
00
As can be seen, the gradient can be estimated by sampling from the model distribution q (x: 6) and oracle distribution p(a y). Unfortunately, we cannot easily draw samples because of the unnormalized constant. Existing methods usually resort to MCMC methods (e.g., Langevin dynamics) use this gradient estimator. However, MCMC is an iterative process and can be time consuming. To improve the training efficiency, we propose to use a nelf-normalized importance sampler based on a Gaussian mixture proposal to estimate the gradient.
JM Specifically, for each x, we first sample set of M candidates {a" Jm=1 from a proposal distribution (a)x), and then sample à from the empirical distribution located at each at and weighted pro- portionally to exp(- E( x; ())/=(a) xx). To reduce the variance of the self-normalized importance sampler, we propose to use mixture of K Gaussians which are centered at the location of the corre- sponding optimal decision as the proposal distribution: π (aa & N aK), where K and
5

---

Preprocessing
Inference
Training data
Optimization Solver
Feature-action pairs Da {x., )=
Test Instance Xtest
Optimization Solver
D {x,, d6 J.
Optimal Energy Function
Optimal Action
Pope)x.Pop1) (v.
Model Training
E
E NLL loss
Energy Function
Push up
D(x, a; 0) Bptlidy /y. (y.a)
Push down
+
DNN
CD-SNIS Sampler
Vo
Ax KL-divergence loss
F
E*A
Training data D 36
Feature-action pairs
D, (x,,,]]=
Figure 3: Illustration of So-Enr 50-EBM: end-to-end stochastic optimization with energy- energy-based based model.
{ox) hyper-parametert This mixture of Gaussian based self-normalized importance sampler enjoys great computational efficiency. We only need to estimate the gradient of the mnergy-based surrogate layer by drawing samples from a simple mixture of Gaussians, instead of solving and differentiating through the optimization problem as in DFL. This significantly reduces the training time compared with DFL as shown in Section 4. Further, locating the proposal distribution the optimal decisions mimics the contrastive divergence method [22] 15] used in Langevin dynamics where the MCMC chain starts from the training data. This makes the proposal distribution close to the model distribution and thus has better sample efficiency. The sampling procedure from the oracle distribution p(aly) is similar with x()xx;0 0).
When the expected task loss has no analytical expression, we can draw multiple samples from p(y]x; f) to estimate the expectation and use reparameterization trick (29 23
36 42 to make it
differentiable. Finally, the model can be efficiently trained via gradient-baat based method, such as Adam The detailed training procedure is given in Alg. in the supplementary.
28
Model inference. With the optimal model parameter oopt. we draw samples from q(a Xrest oopt) for an unseen test instance Xtest However, in the real-world applications, we usually only need the most optimal decision. This can be obtained by solving arg min, a EC Epprlx f a) with
any existing black-box optimization solver such as CVXPY 12
and Pyomo 15 (Fig.) 3).
Experiments
In this section, we empirically evaluate SO-EBM. We conduct experiments in three applications: (1) Load forecasting and generator scheduling where the expected task loss has a closed-form expression; (2) Resource allocation for COVID-19 where the expected task loss has no closed-form expression;
(3) Adversarial behavior learning in network security with non-convex optimization objective.
Finally, we do ablation studies to show the effect of each component in SO-EBM.
4.1 Load Forecasting and Generator Scheduling
In this task, a power system operator needs decide how much electricity R4 schedule for each hour in the next 24 hours to meet the actual electricity demands. The optimization objective is combination of an auder-generation penalty, an over-generation penalty, and mean squared loss between supplies and demands:
24
minimize,gra4
E Ey-plule:0 - as]+ + 7ela; w/]+ + =((: )
subject to
la; - a; S Vi,

---

True landscape
Recovered landscape
True landscape
Recovered landscape
16-
03-
-10
-15
-10
0.6 11
-20
-55
05
-13 -15
66 20 -13
-03
65
13
01
d1
Figure 4: Ground-truth and SO-EBM recovered landscapes of the energy function in the power generator scheduling task. Darker colors represent lower energy in the heat maps. For test sample Xrest, we choose the corresponding optimal action deest the center point and select two random direction vectors 1 and U2 to plot the energy landscape, i.e,E' (d1, E(x, a" diui i + d2u2).
where [+ = ma 0). Usually, the penalty coefficients satisfy "ee > "/s since under-generation is more serious than over-generation. The quadratic regularization term indicates the preference for generation schedules that closely match actual demands. The ramping constraint e, restricts the change in generation between consecutive timepoints, which reflects physical limitations associated with quick changes in electricity output levels.
Experiment Setup. We forecast the electricity demands IR over all 24 hours of the next day using a 2-hidden layer neural network. We assume y/ is a Gaussian with mean μi and variance of: as such, the expectation in the optimization objective can be computed analytically. The input features x is a 150-dimensional vector including the past day's electrical load and temperature, the next day' s temperature forecast, non-linear functions of the temperatures, binary indicators of weekends or holidays, and yearly sinusoidal features. Following 13 we set 7 = 0.4, Ye = 50 and c = 0.4.
We compare with the following baselines on this task: (1) A two-stage phadictthen-oppimizee model trained with negative likelihood loss for the prediction task. (2) Decision-focusse learning with the QPTH solver [13] It uses sequential quadratic programming (SQP) to iteratively approximate the resultant convex objective as quadratic objective, iterates until convergence, and then computes necessary Jacobians using the quadratic approximation at the solution. (3) DFL with evxpylayers
[ (DFL-CVX) which provides a differentiable layer with disciplined parameterized programming (DPP) grammar. Since the analytical expectation of Eq. 7 cannot be written in DPP, we use Monte Carlo sampling to estimate the expectation for this baseline. (4) Policy-net. direct maps from the input features the decision variables by minimizing the task loss using supervised learning [13 Our supplementary provides more details of the model parameters.
Results. Fig. Jisows the end task loss for all the methods. Our method SO-EBM outperforms the baselines with a significant reduction of train- ing time. Compared with the strongest baseline DFL-QPTH. SO-EBM improves the task loss by 7.3%. The improvement is because DFL-QPTH needs to use SQP to iteratively obtain the solu- tions for non-quadratic optimization problems. Differentiating through all the steps of SQP is prohibitively expensive in memory and time. To address this issue, existing works [13 differen- tiate through just the last step of SQP to obtain approximate gradients. The inaccurate gradients may accrue during training and thus hurt decision quality. This shows that our explicit differentiable energy- function an effective surrogate for the original implicit optimization layer. In term of efficiency. SO-EBM is more than 136 times faster than DFL-QPTH (0.68 second/epoch v.s. 93.12 second/epoch) training. This is because only need draw samples from a simple mixture of Gaussians to estimate the model parameters instead of solving and differentiating through the optimization problem at every training iteration. DFL-CVX is even slower than DFL-QPTH since needs to use sample average estimation to draw multiple samples from p(y|x; 0), which results in much more complicated optimization objective. This is also likely the reason why DFL-CVX under- performs DFL-QPTH in terms of task loss here. For Policy-Net, we have tuned its hyper-parameters extensively but still cannot achieve good performance on this task (2x-3x larger task loss than the two-stage baseline). This is not surprising: aforementioned, Policy-Net pure end-to-end model
Frening Time (seconduslepach)
10a
10.
10
"Noo-Stage Stage
Policy Net
DFL- OPTH
DFL-CVX CVX
DFL-
So EBM
Two-Stagg
Policy Net
DFL DFL-QPTH LL-QTTH
SO EBM DFL-CVA LL-CVX
Figure 5: Results on power generator scheduling.
7

---

that needs a large amount of data to rediscover the algorithmic structure of the optimization task. In contrast, DFL and SO-EBM model the peddict-and-optimin structure in the model design, which save learning and are more data-efficient.
Fig. showws the ground-truth and SO-EBM learned landscapes of the energy function. As can see, SO-EBM can recover the landscape of the original optimization objective effectively though with small discrepancy. The small discrepancy is expected since the ground-trul landscape is computed by directly using the ground-truth label, while SO-EBM uses the uncertainty-awar neural network to first forecast the distribution of the label and then uses the predictive distribution to compute the expected task loss as the energy function.
4.2 Resource Allocation for COVID-19
As seen in the COVID-19 pandemic, an increasing number of infected patients leads to increasing demand for medical resources; which makes challenging for policymakers and epidemiologists to plan ahead. Mechanistic cpidemiological models based on Ordinary Differential Equations (ODEs) are often used to capture and forecast the dynamies of epidemics including for the COVID-19 pandemic 33 24 20 10 Such forecasts are typically then used as guidance to help plan for future In this task, we study the optimization problem hospital bed preparation for COVID- 19, one of the most common and important tasks epidemiologists focus on during the pandemic [3] [7] We need to decide how many beds R are needed to prepare for the next week based on the forecasted number of hospitalized patients y E R". The optimization objective a combination linear and quadratic costs, which accounts for over-preparations (a y]+ and under-preparations
resource allocation [8
43
25
[y a|+ in the next 7 days over ODE-derived dynamics:
minimiZe
coly. a.l+ + cala; vol+ + (lu o,)) + 9n (l (()
(8)
Experiment Setup. We use the SEIR+HD ODE model proposed in [24] to capture the dynamics of the COVID-19 pandemic, which has been used in policy studies for interventions. The model is driven by a key parameter: the transmission rate 3 that reflects the probability of disease transmission. The forecasting model is a two-layer gated recurrent unit (GRU) [ which takes the number of people in each state of the SEIR+HD model in the last 21 days as input features and outputs the transmission rate 3 for the next 7 days. We assume that ,3 follows a Gaussian distribution. Since SEIR-HD is complex non-lineat ODE system, there is no closed expression for the distribution of the number of hospitalized patients y. Hence, we choose to sample from the forecasted distribution of 3 100 times and then simulate the SEIR-HD model to obtaín the empirical distribution of the number of hospitalized patients y.
Results. Fig. shows the results on the resource allocation for COVID-19 task. Both So-EBM
10000
Training Time (seconds/epoch)
and DFL significantly outperforms the two-stage model in terms of the task loss, yielding 10.9%+ improvement. DFL-QPTH is not included here because it needs to use SQP to iteratively ap- proximate the objective and is much slower than DFL-CVX in this task. The performance of So- EBM and DFL are similar, however, the training time of SO-EBM is 4.3 times faster than DFL. The time improvement of SO-EBM is not as large as on the energy scheduling task because all the methods need to compute the complex ODE equations in the forward pass which is time consuming part. Therefore, considering the overall training efficiency, we argue that S0-EBM is a better option than DFL on this task in practice.
8000
300
6000
200
4000
2000
100
Two-Stage
DF DFL-CVX
So- EBM
Two-Stage Stage
DFL-CVX
50-EBM
Figure 6: Results on Covid-19 resource allocation.
4.3 Network Security Game
In this task, we study stochastic optimization in network security games (4). Given a network G' (V. E), source node V. and a set of target nodes T V, network security game (NSG) 50 17 44 defines the min-max game where the attacker attempts to travel from a to any t T while the defender places checkpoints on certain number of edges in the graph. Each of the
8

---

target node has a reward u(t) should the attacker reach them. The defender first chooses mixed strategy. Having observed the defender's mixed strategy but not the sampled pure strategy, the attacker attempts to choose a path. To minimize the attacker's scores, the defender can try to predict the attacker's path decisions with node features and their past decisions, since the attacker is not perfectly rational in reality. Suppose that defines the placement of checkpoints by the defender where ag is the probability that edge e is covered. The attackers will perform a random walk on the graph, generating a path r, stopping at either target or checkpoint. The transition between two nodes is determined by the defender's strategy a and node-level parameters y, where Yu. defines the "attractiveness" of node u. representing the attacker's s idiosyncratic preferences. The defender wants to g(r), where is covered edges sampled according to a, g(r) u(t) if the attacker reaches target with path r and O if the attacker is stopped at a checkpoint. The defender's objective comes with the constraint Σ k, where k 3 represents the resources available to the defender. The optimization objective is non-conver due to the irrational strategy of the attacker 48).
choose a to maximize its expected utility: maxa Er Er-aEr-ptrppp,
Setup. We follow the setup of 49 see supplementary for details.
Results. Fig. 7 shows the results on the adver-
sarial network security game task. Since the original DFL fails on this large scale problem, compare a block-variable sampling approach specialized to this task (DFL.-Block)[45 Surro- gate [49 performs linear dimension reduction for DFL to speed up the training time and im- prove the performance by smoothing the train- ing landscape. We did not present the results of Surrogate in the previous two tasks because the numbers of decision variables there are rel- atively small and has even worse performance than the original DFI. approach. As we can see, SO-EBM can achieve competitive task loss with Surrogate and outperforms DFI.-block and the two-stage method. The increase of training time in our method is mainly because every evaluation of the defender's utility requires a matrix inverse; our method needs to evaluate the utility for all the samples drawn from the proposal distribution which results larger matrix inverse. However, this issue can be mitigated by using some advanced matrix inverse algorithm g 32]. DFL-Block is even worse than the two-stage model because it back proppprpsss through randomly sampled variables which results in inaccruate gradient estimation.
0.8
[secondsiepach)
400
300
Task
0.4
200
200
maining
0.0 T.O Two--ttge
DFL-Block
Surrogate
S0-EBM
Two--Sthe Stage
DFL-Bloce Block
Surragate
S0. EB-
Figure 7: Results on network security game.
4.4 Ablation Study
We investigate the effectiveness of the coupled training objective and alternative EBM training algorithms via ablation studies on the load fore- casting and generator scheduling task. Table shows the results. Our findings can be summa- rized as follows: (1) Both the MLE and KL.D training objectives work and outperform the two- stage model. This verifies the effecctiveness of using the energy-based surrogate function to ap- proximate the optimization landscape. (2) By dropping either the MLE or KLD term, we observe performance degradation in our method. Specif- ically, without the MLE term, the task loss increases by 13%; without KLD term, the task loss increases by 2%. The larger degradation when simply minimizing the KLD term is possibly because accurately recover the entire optimization landscape is too difficult and needs the MLE term to help capture the location of the optimal decision. (3) Training SO-EBM with Contrastive Divergence based Langevin Dynamics (CD-LD) (15) achieves similar performance but incurs longer training time (1.47 second/epoch v.s. 0.68 second/epoch). This phenomenon is likely because the optimization problem relatively low-dimensional, thus importance sampling works well and enjoys better efficiency in this regime. However, even with CD-LD, SO-EBM is still 63 times faster than DFL-QPTH (1.47 second/epoch v.s. 93. 12 second/epoch).
Method
Task loss
Two-stage
4.52 + 0.33
SO-EBM w/o MLE
4.31 + 0.17
SO-EBM w/o KLD
3.89 + 0.18
SO-EBM w/ CD-LD
3.85 + 0.13
SO-EBM
3.83 + 0.15
Table 1: Ablation study on power generator
scheduling.

---

We also investigate the impact of training data size for each method. Fig. X provides the task losses and training time under different ratios of training data on the load forecasting and genera- tor scheduling task. As we can see, our method outperforms the baselines constantly except for the ratio of 0.05. When the ratio is below 0.05,
6.0
Tan DFL OFLH
Training Time (seconds//ppod
EB
107
Task Loss
5.0
10
10-1
all the methods cannot work properly due to the extremely low resource. The superior perfor- mance of our method is because we design the energy function the expected task loss, which leverages the algorithmic structure inherent in the optimization problem. Hence, method is not very data demanding. In terms of effi- ciency, our method reduces the training time significantly compared with DFL-QPTH across different amounts of training data.
4.0
10-
0,50
1.00
Ratio training data
Ratio training data
Figure 8: Task losses and training time under dif- ferent ratios of training data on power generator scheduling.
5 Additional Related Work
The DFL works deseribed in Section 2 focus on continuous optimization. There are also studies that extend DFL for combinatorial optimization. [51] relaxes the discrete decision into its continuous counterpart and adds a quadratic regularization term to avoid vanishing gradients. [37] proposes log barrier regularizer and differentiates through the homogeneous self-dual embedding. [39 proposes a noise contrastive objective by maximizing the distance between the optimal solution and noisy samples. Our framework may be also extended for combinatorial problems by using discrete energy-based model [T]. The SPO+ loss [6] has been proposed to measure the prediction errors against optimization objectives, but it is only applicable to linear programming. ProjectNet [8] approximately solves the linear programming problems using a differentiable projection architecture. Our work is also related learning to optimize [4 34 26). which learns a policy network that solves optimization problems using supervised or reinforcement learning. However, these pure approaches need rediscover the structure of the optimization problem and thus data-inefficient.
6 Limitations and Discussion
We focused on addressing the scalability and generality of existing srcision-focusee learning (DFL) for end-to-end stochastic optimization. We argue that these deficiencies stem from the reliance of implicitly differentiable optimization layers based on KKT conditions. As a remedy, we circumvent such deficiencies by replacing the implicit optimization layer with a newly parameterized energy- based surrogate function. We proposed coupled training objective to encourage the energy- -baed surrogate well approximate the optimization landscape, as well as an efficient training procedure based on self-normalized importance sampling. Empirically, we demonstrated that our mnergy-based model effective in a wide range stochastic optimization problems with either convex or nonconvex objectives. It can achieve better comparable performance than state-of offte-of-the DFL methods for stochastic optimization, while being several times or even orders of magnitude faster.
We discuss limitations and possible extensions of SO-EBM: (1) Handling more complex constraints. Our method handles the constraints implicitly through the pre-processing step and explicitly through the inference step. However, when the feasible space is extremely small, training the constrained EBM may have more benefits. To train EBMs in the constrained space, one direction is to project samples into the feasibility space. Another direction is to explore adding soft constraints to the energy function during training, e.g., Augmented Lagrangian penalty and Barrier penalty. (2) More effective training methods. Our method is general framework for end-to-end stochastic programming problem based on EBM. There are number of training techniques (45) that can be plugged into our framework and can further improve the task performance using the advanced EBM training algorithms 14 20).
Acknovledgments: We thank the anonymous reviewers for their helpful comments. This work was supported in part by the NSF (Expeditions 19F-1918770, CAREER IIS-2028586. IIS-2027862, 115-1955883, IIS-21 116611 IIS-2008334, CAREER 114-2144338, PIPP CCF-2200269). CDC MInD
program, faculty research award from Facebook and funds/computing resources from Georgia Tech.
10

---

References
[1] Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J Zico Kolter. Differentiable convex optimization layers. Advances in neural information processing systems, 32, 2019.
[2] Akshay Agrawal, Robin Verschueren, Steven Diamond, and Stephen Boyd. A rewriting system for convex optimization problems. Journal of Control and Decision, 5(1):42-60, 2018.
[3] Nick Altieri, Rebecca L Barter, James Duncan, Raaz Dwivedi, Karl Kumbier, Xiao Li, Robert Netzorg, Briton Park, Chandan Singh, Yan Shuo Tan, et Curating a covid-19 data repository and forecasting county-level death counts in the united states, arXiv preprint arXiv:2005.07882, 2020.
[4] Brandon Amos and J Zico Kolter. Optnet: Differentiable optimization as a layer in neural
networks. In International Conference on Machine Learning, pages 136-145. PMLR, 2017.
[5] Michael L. Bynum, Gabriel A. Hackebeil, William E. Hart, Carl D. Laird, Bethany L. Nicholson, John D. Siirola, Jean-Paul Watson, and David L.. Woodruff. python, volume 67. Springer Science & Business Media, third edition, 2021.
modeling in
[6] Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555 2014.
[7] Estee Y Cramer, Evan L Ray, Velma K Lopez, Johannes Bracher, Andrea Brennen, Alvaro J Castro Rivadeneira, Aaron Gerding, Tilmann Gneiting, Katie H House, Yuxin Huang, al. Evaluation of individual and ensemble probabilistic forecasts of covid- 19 mortality in the united states. Proceedings of the National Academy of Sciences, (19(15):e21 13561 199, 2022.
[8] Rares Cristian, Pavithra Harsha, Georgia Perakis, Brian L Quanz, and Ioannis Spantidakis. End- to-end learning via constraint ceastrinnt-enforing approximators for linear programs with applications to supply chains. In Al for Decision Optimization Workshop of the AAAI Conference on Artificial Intelligence, 2022.
[9] Laszlo Csanky. Fast parallel matrix inversion algorithms. In 16th Annual Symposium on Foundations of Computer Science (sfcs 1975), pages 11-12. IEEE, 1975.
[10] Jiaming Cui, Arash Haddadan, ASM Ahsan-UI Haque, Bijaya Adhikari, Anil Vullikanti, and B Aditya Prakash. Information theoretic model selection for accurately estimating unreported covid-19 infections. medRxiv, pages 2021-09, 2021.
[11] Hanjun Dai, Rishabh Singh, Bo Dai, Charles Sutton, and Dale Schuurmans. Learning discrete energy-based models via auxiliary-variable local exploration. Advances in Neural Information Processing Systems, 14114004114555 2020.
[12] Steven Diamond and Stephen Boyd. CVXPY: A Python-embeddeed modeling language for convex optimization. Journal of Machine Learning Research, 17(83)11-5, 2016.
[13] Priya Donti, Brandon Amos, and Zico Kolter. Task-based end-to-end model learning in stochastic optimization. Advances neural information processing systems, 30, 2017.
[14] Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy-based models. In International Conference on Machine Learning, pages 2837-2848. PMLR, 2021.
[15] Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, 32, 2019.
[16] Adam N Elmachtoub and Paul Grigas. Smart "predict, then optimize". Management Science, 68(1):9-26, 2022.
[17] Matteo Fischetti, Ivana Ljubić, Michele Monaci, and Markus Sinnl. Interdiction games and monotonicity, with application to knapsack problems. INFORMS Journal on Computing, 3102):390-410. 2019.
11

---

[18] IHME COVID-19 forecasting team. Modeling covid- 19 scenarios for the united states, Nature medicine, 2020.
[19] Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty deep learning. international conference on machine learning, pages 1050-1059. PMLR, 2016.
[20] Xingjie Hao, Shanshan Cheng, Degang Wu, Tangehun Wu, Xihong Lin, and Chaolong Wang. Reconstruction of the full transmission dynamics of covid-19 in wuhan. Nature, 584(7821):420- 424, 2020.
[21] William E Hart, Jean-Paul Watson, and David L Woodruff. Pyomo: modeling and solving mathematical programs in python. Mathematical Programming Computation, 3(3):219-260, 2011.
[22] Geoffrey Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):1771-1800, 2002.
[23] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv: 1X11001611.014.. 2016.
[24] Morgan P Kain, Marissa L Childs, Alexander D Becker, and Erin A Mordecai. Chopping the tail: How preventing superspreading can help maintain covid-19 control. Epidemics, 34:100430, 2021.
[25] Harshavardhan Kamarthi, Lingkai Kong, Alexander Rodríguez, Chao Zhang, and Aditya Prakash. When in doubt: Neural non-parametric uncertainty quantification for epidemic forecasting. Advances in Neural Information Processing Systems, 34, 2021.
[26] Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial optimization algorithms over graphs. Advances in neural information processing systems, 30, 2017.
[27] Sujin Kim, Raghu Pasupathy, and Shane Henderson, guide to sample average approximation. Handbook of simulation optimization, pages 207-243, 2015.
[28] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference Representation Learning, 2015.
[29] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
[30] Anton Kleywegt, Alexander Shapiro, and Tito Homem-de Mello. The sample average approximation method for stochastic discrete optimization. SIAM Journal Optimization, 12(2):479-502, 2002.
[31] Lingkai Kong, Jimeng Sun, and Chao Zhang. Sde-net: Equipping deep neural networks with uncertainty estimates. In International Conference on Machine Learning, pages 5405-5415. PMLR, 2020.
[32] Aravindh Krishnamoorthy and Deepak Menon. Matrix inversion using cholesky decomposition. In 2013 signal processing: Algorithms, architectures, arrangements, and applications (SPA). pages 70-72. IEEE, 2013.
[33] Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and F Huang. A tutorial on mergy-based learning. Predicting structured data, 1(0), 2006.
[34] Ke Li and Jitendra Malik, Learning to optimize. arXiv preprint arXiv: 1606.008888 2016.
[35] Ruiyun Li, Sen Pei, Bin Chen, Yimeng Song, Tao Zhang, Wan Yang, and Jeffrey Shaman. Substantial undocumented infection facilitates the rapid dissemination of novel coronavirus (sars-cov-2). Science, (686649) 2020.
12

---

[36] C Maddison, A Mnih, and Y Teh. The concrete distribution: A continuous relaxation of discrete random variables. In Proceedings of the international conference on learning Representations. International Conference on Learning Representations, 2017.
[37] Jayanta Mandi and Tias Guns. Interior point solving for Ip-based prediction+ optimisation. Advances in Neurul Information Processing Systems, 33:7272-7282. 2020.
[38] Seyed M Moghadas, Affan Shoukat, Meagan C Fitzpatrick, Chad R Wells, Pratha Sah, Abhishek Pandey, Jeffrey D Sachs, Zheng Wang, Lauren A Meyers, Burton H Singer, et al. Projecting hospital utilization during the covid-19 outbreaks in the united states. Proceedings of the National Academy of Sciences, 117(16):912 9126, 2020.
[39] Maxime Mulamba, Jayanta Mandi, Michelangelo Diligenti, Michele Lombardi, Victor Bucarey, and Tias Guns. Contrastive losses and solution caching for predict-and aappit-and-optimize. arXiv preprint arXiv:2011.05354 2020.
[40] Erik Nijkamp. Ruiqi Gao, Pavel Sountsov, Srinivas Vasudevan, Bo Pang, Song-Chur Zhu, and Ying Nian Wu. Mcmc should mix: Learning energy-based model with neural transport latent space meme. In International Conference on Learning Representations, 2021.
[41] Andrew Perrault. Bryan Wilder, Eric Ewing, Aditya Mate, Bistra Dilkina, and Milind Tambe. End-to-end focusfo learning of adversary behavior in security games, In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 1378-1386, 2020.
[42] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference on machine learning, pages 1278-1288 PMLR, 2014.
[43] Alexander Rodríguez. Anika Tabassum, Jiaming Cui, Jiajia Xie, Javen Ho, Pulak Agarwal, Bijaya Adhikari. and B. Aditya Prakash. Deepcovid: An operational deep learning- driven framework for explainable real-time covid-19 forecasting. Proceedings of the AAAI Conference May 2021.
on Artificial Intelligence, 35(17)115333-144)
[44] Sankardas Roy, Charles Ellis, Sajjan Shiva, Dipankar Dasgupta, Vivek Shandilya, and Qishi Wu. A survey of game theory as applied to network security. In 2010 43rd Hawaii International Conference System Sciences, pages 1-10. IEEE, 2010.
[45] Yang Song and Diederik Kingma. How to train your energy- models. arXiv preprint arXiv:2101.032 2021.
[46] Bram Verweij, Shabbir Ahmed, Anton J Kleywegt, George Nemhauser, and Alexander Shapiro. The sample average approximation method applied to stochastic routing problems: a computa- tional study. Computational optimization and applications, 24(2):289-33 -33,, 2003.
[47] Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. Advances in neural information processing systems, 28, 2015.
[48] Kai Wang, Andrew Perrault, Aditya Mate, and Milind Tambe. Scalable game-focused learning of adversary models: Data-to-decision in network security games. In AAMAS, pages 1449-1457, 2020.
[49] Kai Wang, Bryan Wilder, Andrew Perrault, and Milind Tambe. Automatically learning compact quality-aware surrogates for optimization problems. Advances in Neural Information Processing Systems, 33:9586-9596, 2020.
[50] Alan Washburn and Kevin Wood. Two-person zero-sum games for network interdiction. Operations research, 43(2):243-251, 1995.
[51] Bryan Wilder, Bistra Dilkina, and Milind Tambe. Melding the data-decisions pipeline: Decision- focused learning for combinatorial optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 1658-166660 1665, 2019.
[52] Markus Wulfmeier, Peter Ondruska, and Ingmar Posner. Maximum entropy deep inverse reinforcement learning. arXiv preprint arXiv: 1507.10877.4888 2015.
13

---

[53] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence. volume 8, pages 1433-1438 Chicago, II., USA, 2008
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [No]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "nAG3QymYlnp", "reviewer": "Reviewer_zDEe", "summary": "This paper tackles stochastic optimization problems of the form $\\max_{a \\in C} f(y, a)$, where $y$ is a stochastic variable and $a$ is an action taken by the learner given an observation $x$. \n\nRecent work have introduced solutions that involve Deep Neural Networks. One body of work uses DNNs to make a prediction $\\hat{y}$ given $x$, and then running an off-the-shelf optimizer for the inferred problem $f(\\hat{y}, \\cdot)$. This can lead to suboptimal solutions if $\\hat{y}$ is not accurate. Another body of work can overcome this in certain scenarios by optimizing $(\\hat{y}, a)$ jointly. They do so via end-to-end architectures that use so-called optimization-layers. These layers run a differentiable convex optimizer, and hence the incurred cost $f(\\hat{y}, \\hat{a} )$ can be optimized by jointly through backpropagation. A limitation of this approach is that the optimization layers must backpropagate through a convex optimiser; hence if it does not have a closed form solution, the computational cost can be steep.\n\nThe authors propose a novel architecture for joint optimization of $(\\hat{y}, \\hat{a})$ that avoids introducing optimization layers by taking an EBM approach. They construct an EBM that takes the expected cost $f(\\hat{y}, \\hat{a})$ under the model as the energy. To optimize the energy function, the authors operate a two-stage process. First, given a dataset $D = \\\\{ (x_i, y_i) \\\\}\\_{i=1}^K$, they compute expert actions $a^i$ for each $y^i$ to construct an auxiliary dataset $D' = \\\\{ (x_i, a_i) \\\\}\\_{i=1}^K$. They then use MLE to optimise the energy model under $D'$. \n\nHowever, a main complication of this approach is that it requires stochastic outputs from the DNN. To this end, the authors rely on a Gaussian parameterization and self-normalized importance sampling. This means that the main computational difference between optimization-layer approaches and the author's approach is that they rely on stochastic relaxation instead of exact optimization. \n\nThe authors provide three synthetic experiments to evaluate their architecture: load balancing in power systems, healthcare resource allocation, and network security. In all cases, they show that their proposed method is competitive in terms of final performance while being computationally more efficient than optimization-layer based approaches.   ", "questions": "- The ability to generate expert actions is critical for this method to work. The authors suggest that previous method are limited in that they rely on convex solver - but it seems the proposed method suffer from the same limitation in that the solver is required to obtain expert actions?\n\n- The authors opted for an EBM approach. Another equally valid method would have been expert distillation, which would have the benefit of not requiring estimating the energy function. Could the authors comment on why they opted against this simpler approach?", "limitations": "The authors discuss some limitations of the EBM approach, mainly the risk of overfitting. They provide an ablation to study the severity of the problem. The authors do not discuss other limitations, sensitivity to the choice of sampler and sensitivity to the choice of solver for expert actions.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper is generally well written. The main idea, its motivation, mechanics, and intuition for design choices are all clearly spelled out. The experiments are well described with uncertainty estimates.\n- The proposed architecture is quite clean and follows naturally from the EBM perspective. \n- The experiments are well executed; clearly described, with thorough analysis of results. \n\nWeaknesses\n- The main contribution stated on page 2 does not seem to be well supported. First, the authors claim that the method is end-to-end, but this is not true since to construct the dataset, an off-the-shelf solver is needed to generate expert actions. Second, it does require solving the optimization problem (though not backpropagating through the solver), and while the EBM objective does not require convexity, imputing expert actions does. Overall, I think the paper can be strengthen significantly by a more nuanced discussion of the relative merits of the proposed method.\n- The experiments are somewhat underwhelming in that they are all synthetic. It would have been significantly more interesting to test the proposed architecture on a real world problem with well-developed benchmarks. Moreover, the results themselves suggest that the gains are marginal; in most cases, the two-stage approach of pretraining a predictor $\\hat{y}$ is faster and the loss in performance is not statistically significant.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "PkSi1caTof", "reviewer": "Reviewer_gqgA", "summary": "This paper proposes a method to mitigate some of the computational challenges associated with the decision-focused learning paradigm under stochastic optimization-based decisions with unknown parameters. In particular, the authors construct (a) an energy-based surrogate function for the stochastic optimization-based decisions, (b) a loss function encouraging both good predictive accuracy and good global behavior of the energy-based surrogate, and (c) an efficient training procedure based on a mixture-of-Gaussians proposal. This enables them to avoid expensive implicit KKT differentiation steps that are needed when representing the decision procedure exactly as an optimization problem, as well as better enabling the representation of non-convex decision-making processes. The authors provide extensive experiments on several settings, and show empirically that their method achieves comparable or improved performance over baseline methods while drastically improving on computation time.", "questions": "* Section 4.1: Why does SO-EBM improve over DFL-QPTH? It would be useful to give some intuition for this, as at first glance it seems strange that a method using a surrogate loss would outperform a method using a more “exact” version of the loss.\n* What are some of the limitations and potential negative impacts of the method?", "limitations": "In the checklist, the authors mention that they describe limitations and potential negative societal impacts of their work, but I cannot find this discussion. This discussion should be made more explicit, and should acknowledge that while the applications presented in the paper are social good applications (which I appreciate), the method is general enough that it may also have negative uses.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n* The proposed method is well-motivated, nicely integrates concepts from energy-based learning into the decision-focused learning literature, and mitigates a number of existing challenges with decision-focused learning paradigms.\n* The experimental evaluation is thorough and well-done, demonstrating performance on different kinds of settings with respect to the form of the task loss, and showing insightful ablations. The method strongly outperforms other baselines (and the authors both include error bars *and* transparently leave in one setting where a baseline outperforms their method, both of which I appreciate).\n* Everything is very well-explained, making the paper an absolute pleasure to read.\n\nWeaknesses: \n* In the checklist, the authors mention that they describe limitations and potential negative societal impacts of their work, but I cannot find this discussion.\n\nQuestions:\n* Section 4.1: Why does SO-EBM improve over DFL-QPTH? It would be useful to give some intuition for this, as at first glance it seems strange that a method using a surrogate loss would outperform a method using a more “exact” version of the loss.\n\nMinor points:\n* Section 2, Problem Formulation paragraph: The terms $M$ and $a^*(x;\\theta)$ should be more formally defined in the text (not just in Figure 1), as the rest of the problem formulation and method hinge on this. In particular, the current statement about learning a model that “takes the features $x$ as input and outputs the optimal decisions $a^*(x;\\theta)$” could be misunderstood as a model that maps directly from data to decisions, without the intermediate prediction of the distribution.\n* Section 2, Problem Formulation paragraph: As previously defined, $\\mathcal{D}$ is the dataset. However, the optimization loss is introduced as “the expected decision cost under the joint distribution.” The wording here should be clarified. The goal is to learn the expected decision cost under the joint test-time distribution, *as approximated by* the dataset at training time, which is what the loss shows.\n* Line 92: The left side should be the total derivative, rather than the partial derivative. Similarly, $\\partial y/\\partial \\theta$ should be $\\mathrm{d}y/\\mathrm{d}\\theta$. Similarly, in Equation (6), some of the partial derivatives should be total derivatives.\n\nTypos:\n* Line 84: “DFL method” -> “DFL methods”\n* Problem Formulation paragraph in Section 2: “optimal decisions minimizes” -> “optimal decisions minimize”\n* Line 99: “CVXlayers” -> “cvxpylayers”\n* Line 100: “grammer” -> “grammar”\n* Line 132: The optimization problem is missing $\\mathbf{x}_i$ \n* Line 305: “target node” -> “target nodes”\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "CM01dno-gul", "reviewer": "Reviewer_pnpr", "summary": "The authors propose an approach for \"decision-focused learning\" (DFL).  This is where there is a prediction stage to predict stochastic parameters from some input features, that are then used to solve an optimization problem to determine an optimal action.  For example, an optimization of electricity generator scheduling based on forecasted electricity load / demand.  The old two-stage approach would be to separately fit a forecast model, and then use its predictions in an optimization problem solved by optimization solvers.  The DFL approach is to fit the forecast model in conjunction with the optimization problem, so the forecasts are optimal for the down-stream optimization problems themselves.\n\nTo do this for modern models using neural networks for the forecasters would generally require computing and back-propagating the gradient through the optimization loss and problem to the model parameters.  However, the gradient of the optimal action with respect to the forecast / prediction values is not always differentiable and the optimal action is the solution to an optimization problem, so computing gradient updates can be challenging and time consuming.\n\nMost past work tended to either make restrictive assumptions or approximations, or used computational complex and inefficient implicitly differentiable layers.   This work proposes a workaround by modeling the optimal actions / optimization problem solution as a function of the features directly, but ties it to the forecast model as the energy term of the energy model is given by the expected optimization loss of under the forecast model (the forecast gives the distribution over the predicted parameters for the optimization problem).  In this way the energy model can be trained directly on historical examples with optimal solutions previously derived, by minimizing the negative log likelihood (which in turn optimizes the forecast model for the optimization problem) along with a KL-divergence regularizer added to the objective as well.  \n\nThe authors compare their proposed approach to past methods on 3 datasets / tasks from different domains, and show it consistently has as good or better performance to other past DFL approaches, while also having as good or better run time. ", "questions": "Please see questions and suggestions raised under weaknesses.  \n\nIt would be good to mention and contrast with the other related work mentioned above.\n\nHow does the proposed approach handle constraints?  What is the impact of not including constraints?  What is the impact on large numbers of constraints / more complex optimization problems?\n\nHow much data does the proposed approach require?  How does the amount of data affect performance?\n\n-Figures 4, 5, and 6 - labels are not lined up with the bars so it's hard to match the methods to the bars", "limitations": "There doesn't seem to be discussion of what are the limitations of this approach - it could benefit from some thought and discussion around what the limitations are.  Perhaps this would include dealing with constraints, or it could be the case where a large number of solved problems are needed for it to work well, as this should also be investigated ideally.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n-The authors point out key limitations of past work and aim to tackle a significant and important problem.  I.e., most practical applications of ML involve combining predictions and optimization and the DFL approach is arguably the way things should be done for all these applications / ML in general.\n\n-The authors propose a novel solution that is interesting and shows promising results.\n\n-The authors provide experimental results on a variety of different tasks and compare to multiple past DFL methods.  \n\n-The authors pledge to make the code publicly available, which will help a lot with reproducibility and understanding the approach.\n\n\nWeaknesses:\n\n-One piece of the problem, in particular the optimization part, that seems overlooked and not mentioned is the constraints.  Constraints are not included at all in the model of the end-to-end problem (mapping the features to the action) - they could only be included implicitly since the optimal actions, a, the model is trained on must satisfy the constraints.  It would help to have discussion about why it's not necessary to consider the constraints and what kind of impact it would have.  My speculation is that this can actually have a negative impact, because it puts more burden on the forecast neural network part of the model to account for the constraints (this model then needs to be more complex) - as what naturally fits the data better would have to be abandoned to predict a solution that leads to the best action when dropping the constraints.  I feel the problems chosen had very few and basic constraints which may be why the approach still worked in the experiments, but I would be curious to see how more constraints, and more complex constraints, might impact the performance.  It would also be useful to see how the learned forecast model compares to the separately trained model, and what cases it's making predictions farther off from the true values.\n\n-There is one closely-related work overlooked here that appeared in the AAAI-22 Workshop on AI for Decision Optimization - this pointed out the same limitations and similarly targeted a somewhat similar approach.  However, that approach was different in that it approximated the optimization algorithm itself with a neural net, so the solution could be easily computed and differentiated, as opposed to effectively approximating the entire end-to-end problem as was done here (though here essentially only using the neural net to approximate the forecast piece of the model).  It would be good to mention this alternative approach, and it would be great if the two could be compared, as I'm very curious to see how they compare.  Paper:\n\"End-to-End Learning via Constraint-Enforcing Approximators for Linear Programs with Applications to Supply Chains\"\nhttps://research.ibm.com/haifa/Workshops/AAAI-22-AI4DO/PDF/End-to-End%20Learning%20via%20Constraint-Enforcing%20Approximators%20for%20LinearPrograms%20with%20Applications%20to%20Supply%20Chains.pdf\n\n-The clarity of the presentation and description could be improved.  It's hard to follow the details of the method and the experiments.  For instance, the method description is not easy to follow, especially the KL-divergence regularization and the rationale behind it.  Another example, the experiment procedure is not clearly explained - i.e., for all methods, is the model fit and then just the forecast part used in conjunction with an optimization solver to get the final action / decision?  If so, what solver is used?  How are various hyper parameters set for different methods?  \n\n-Further experimental study would be useful.  In particular it would be good to see the impact on training data size on the performance of the different approaches.  In particular, with enough training data one would not expect the DFL approach to out-perform the two-stage one.  I also wonder if this approach may need more training data points to work well as it directly models the action given the inputs. \n Additionally run times will also be affected by amount of training data.  It would be good to understand these trade-offs.  Additionally it would be best to see the impact of varying lambda.   Additionally, what is the impact of using different network architectures / sizes?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "-CVtYWzeW5", "reviewer": "Reviewer_kQia", "summary": "This paper proposes to address the problem of end-to-end stochastic optimization. That is; finding argmin_a f(y, a) for some cost function f, input variables y, and actions a. The key caveat here is that y is not observed and instead, we observe x, a variable correlated with y (but the relationship between x, y is not provided). \n\nStandard approaches to this problem learn a model of p(y|x), then use an off-the-shelf stochastic optimizer to find argmin_a E_{p(y|x)}[f(y, a)]. This approach of course can accrue errors due to the p(y|x) model not being perfect leading to subpar results. \n\nThis work proposes to learn the p(y|x) model, not to predict y, but such that the induced optimal action a performs best. To do this, they parametrize an energy-based model by taking the expected value of the cost function over the p(y|x) model, i.e. E(a; x) = E_{p(y|x)}[f(y, a)] and train this EBM to maximize likelihood of obtaining the ground truth optimal actions a, given the input x. \n\nThe model is trained using contrastive divergence but an importance sampling estimator is used to deal with the partition function. \n\nGiven that we also have knowledge of the true y for each x in the dataset, we also have access to a noiseless version of the loss function, not available at test time. This information is utilized by adding a regularizer to encourage low KL divergence. Between the noisy model and the noise-free model. \n\nThe authors present results on 3 stochastic optimization problems and show that on average their approach either outperforms the baselines or performs comparable to the strongest baselines while having a considerably lower training cost. ", "questions": "\nAfter reading this paper, I am left with a few questions. First, did you explore alternative methods for EBM training such as contrastive divergence, PCD, noise contrastive estimation, or score matching? As I mentioned above, at the scale of the presented experiments, I think it is likely that a number of these techniques could be successful. \n\nNext, given that your training procedure relies on self-normalized importance sampling, did you perform any verification that the SNIS estimator is accurate? One way to do this would be to show the effective sample size of the importance samples. This can also be used to help tune many parts of the sampler such as the proposal distribution.", "limitations": "The work discusses some of the limitations of EBM training and evaluation but does not spend much time talking about the specific drawbacks of the proposed method. I think the work would benefit from a limitations section.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThis is an interesting application of the theory of energy-based models. It is always nice to see work embracing the methods that have been developed to work with and train unnormalized models, rather than replacing them with normalized surrogates. The proposed method is simple and enables end-to-end optimization of the target lost function rather than a multi-stage approach. Unfortunately, I am not very familiar with this problem-space and the baselines, so I cannot speak much to the quality of the experimental evaluation. \n\nWeaknesses:\n\nI am somewhat concerned by the procedure used to train the models and I feel like the authors could provide more evidence to verify that the method is working as expected. Partition function estimation and training in energy-based models is notoriously difficult and, in my experience, importance sampling estimators only tend to work in *very* small scale settings. In larger scale settings, MCMC-based training such as PCD or CD tend to scale much better. I would expect to see at least 1 MCMC-based training procedure as a baseline. But, given that all of the experiments are in relatively low-dimensional problems, we are in a regime where alternative EBM training procedures may also such such as score-matching, noise contrastive estimation, or conditional noise contrastive estimation. I believe the work would be improved if some investigation was done into the method chosen to train the model.\n\nHaving said that, I would love to see a larger-scale application -- something maybe with a few hundred or thousand variables. EBM training has been successfully scaled to very high-dimensional data so I expect the method should work. Again, I am not very familiar with this space so I am not sure if any such benchmark problems exist, but if they do, that would be nice to include. \n\nBesides this, I found the experimental details to be very lacking. There is no mention of the number of importance samples used in training and the scales of the importance weighting distribution. The neural networks used in the model are not described very well and I had to re-read a few times to understand that p(y|x) = N(y; neural_net1(x), neural_net2(x)). Is that right? The types of neural networks should be explained.\n\nThere is also no discussion of how long the models are trained, how they are optimized, and how model selection is performed. This is very difficult for EBMs, and therefore people typically use proxies for model fit instead of test-set likelihood. For that reason, I feel like it would be important for the authors to comment on this in the paper. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept