Original: End-to-End Stochastic Optimization with Energy-Based Model
Revised: End-to-End Stochastic Optimization with Energy-Based Model

Original: Abstract
Revised: Abstract

Original: Decision-focused learning (DFL) was recently proposed for stochastic optimization problems that involve unknown parameters.
Revised: Decision-focused learning (DFL) was recently proposed for stochastic optimization problems that involve unknown parameters.

Original: By integrating predictive modeling with an implicitly differentiable optimization layer, DFL has shown superior performance to the standard two-stage predict-then-optimize pipeline.
Revised: By integrating predictive modeling with an implicitly differentiable optimization layer, DFL has shown superior performance to the standard two-stage predict-then-optimize pipeline.

Original: 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.
Revised: 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.

Original: Further, they can be inefficient in training due to the requirement of solving and differentiating through the optimization problem in every training iteration.
Revised: Further, they can be inefficient in training due to the requirement of solving and differentiating through the optimization problem in every training iteration.

Original: We propose S O -E BM , a general and efficient DFL method for stochastic optimization using energy-based models.
Revised: We propose S O -E BM , a general and efficient DFL method for stochastic optimization using energy-based models.

Original: Instead of relying on KKT conditions to induce an implicit optimization layer, S O -E BM explicitly parameterizes the original optimization problem using a differentiable optimization layer based on energy functions.
Revised: Instead of relying on KKT conditions to induce an implicit optimization layer, S O -E BM explicitly parameterizes the original optimization problem using a differentiable optimization layer based on energy functions.

Original: 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.
Revised: 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.

Original: Finally, we propose an efficient training procedure for S O -E BM with a self-normalized importance sampler based on a Gaussian mixture proposal.
Revised: Finally, we propose an efficient training procedure for S O -E BM with a self-normalized importance sampler based on a Gaussian mixture proposal.

Original: We evaluate S O -E BM in three applications: power scheduling, COVID-19 resource allocation, and non-convex adversarial security game, demonstrating the effectiveness and efficiency of S O -E BM .
Revised: We evaluate S O -E BM in three applications: power scheduling, COVID-19 resource allocation, and non-convex adversarial security game, demonstrating the effectiveness and efficiency of S O -E BM .

Original:  Introduction
Revised:  Introduction

Original: Many real-life decision making tasks are stochastic optimization problems, where one needs to make decisions to minimize a cost function that involves stochastic parameters.
Revised: Many real-life decision making tasks are stochastic optimization problems, where one needs to make decisions to minimize a cost function that involves stochastic parameters.

Original: Oftentimes, the involved stochastic parameters are unknown and context-dependent , meaning that they need to be predicted from observed features.
Revised: Oftentimes, the involved stochastic parameters are unknown and context-dependent , meaning that they need to be predicted from observed features.

Original: For example, when allocating clinical resources for COVID-19, it is necessary to consider future cases in different regions, whose distributions are unknown and have to be predicted from the current state.
Revised: For example, when allocating clinical resources for COVID-19, it is necessary to consider future cases in different regions, whose distributions are unknown and have to be predicted from the current state.

Original: 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.
Revised: 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.

Original: With the feasibility of training powerful deep learning predictors from large amounts of data, it is increasingly common to solve such stochastic optimization problems using a two-stage predict- thenoptimize pipeline.
Revised: With the feasibility of training powerful deep learning predictors from large amounts of data, it is increasingly common to solve such stochastic optimization problems using a two-stage predict- thenoptimize pipeline.

Original: In the prediction stage, one learns a predictive model for the unknown parameters using certain prediction loss ( e .
Revised: In the prediction stage, one learns a predictive model for the unknown parameters using certain prediction loss ( e .

Original: g ., negative likelihood).
Revised: g ., negative likelihood).

Original: In the optimization stage, the predicted distributions of the parameters are used to parameterize the stochastic optimization problem, which can be then solved using off-the-shelf solvers [20, 2, 12].
Revised: In the optimization stage, the predicted distributions of the parameters are used to parameterize the stochastic optimization problem, which can be then solved using off-the-shelf solvers [21, 2, 12].

Original: This two-stage pipeline relies on an implicit36th Conference
Revised: This two-stage pipeline relies on an implicit

Original: Neural and commonly-accepted assumption: improvements in parameter prediction in term of the prediction loss will always translate to better optimization outcomes .
Revised: Conference and commonly-accepted assumption: improvements in parameter prediction in terms of the prediction loss will always translate to better optimization outcomes .

Original: However, this is not the case: machine learning models make errors and the impact of prediction errors is not uniform w.r.t.
Revised: However, this is not the case: machine learning models make errors and the impact of prediction errors is not uniform w.r.t.

Original: the optimization loss.
Revised: the optimization loss.

Original: Thus, a smaller predictive loss does not necessarily lead to a smaller decision regret.
Revised: Thus, a smaller predictive loss does not necessarily lead to a smaller decision regret.

Original: A better approach is decision-focused learning (DFL) [13, 1, 4], which integrates prediction and optimization layers into a unified model to learn them end-to-end.
Revised: A better approach is decision-focused learning (DFL) [13, 1, 4], which integrates prediction and optimization layers into a unified model to learn them end-to-end.

Original: Most DFL methods implement the optimization procedure as an implicitly differentiable layer and develop techniques ( e . g ., using
Revised: Most DFL methods implement the optimization procedure as an implicitly differentiable layer and develop techniques ( e . g ., using

Original: KKT conditions) to compute the gradients w.r.t.
Revised: KKT conditions) to compute the gradients w.r.t.

Original: the decision variables and enable back-propagating through it.
Revised: the decision variables and enable back-propagating through it.

Original: Compared to the two-stage pipeline, the prediction layer so learned is tailored for the optimization problem and better in the sense that it can yield decisions with smaller regrets.
Revised: Compared to the two-stage pipeline, the prediction layer so learned is tailored for the optimization problem and better in the sense that it can yield decisions with smaller regrets.

Original: Besides
Revised: Besides

Original: DFL, there are also approaches where a policy network is trained to directly map from the input to the solution of the optimization problem using supervised or reinforcement learning [45, 25, 32].
Revised: DFL, there are also approaches where a policy network is trained to directly map from the input to the solution of the optimization problem using supervised or reinforcement learning [47, 26, 34].

Original: 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.
Revised: 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.

Original: Despite their promising results, existing DFL methods [13, 4, 1, 39, 46] suffers from two drawbacks.
Revised: Despite their promising results, existing DFL methods [13, 4, 1, 41, 48] suffer from two drawbacks.

Original: (1) Generality .
Revised: (1) Generality .

Original: To leverage the KKT condition, they are mostly applicable to only convex optimization objectives [13, 4, 1].
Revised: To leverage the KKT condition, they are mostly applicable to only convex optimization objectives [13, 4, 1].

Original: Though a few works approximate nonconvex objectives by quadratic functions [39, 46], their applicability is still limited to easy-to-relax nonconvex problems and can suffer from poor gradient estimates caused by the relaxation.
Revised: Though a few works approximate nonconvex objectives by quadratic functions [41, 48], their applicability is still limited to easy-to-relax nonconvex problems and can suffer from poor gradient estimates caused by the relaxation.

Original: (2) Scalability .
Revised: (2) Scalability .

Original: Due to the reliance of using the
Revised: Due to the reliance of using the

Original: KKT condition to compute derivatives, they require repeatedly solving the optimization program and back-propagating through it during training, which makes them unscalable.
Revised: KKT condition to compute derivatives, they require repeatedly solving the optimization program and back-propagating through it during training, which makes them unscalable.

Original: The problem is more severe when the expectation of the cost cannot be computed analytically.
Revised: The problem is more severe when the expectation of the cost cannot be computed analytically.

Original: We propose a new end-to-end stochastic optimization method using an energy-based model (EBM), named S O -E BM .
Revised: We propose a new end-to-end stochastic optimization method using an energy-based model (EBM), named S O -E BM .

Original: Similar to DFL, S O -E BM models the algorithmic predict- and -optimize structure by stacking a differentiable optimization layer on top of a neural predictor.
Revised: Similar to DFL, S O -E BM models the algorithmic predict- and -optimize structure by stacking a differentiable optimization layer on top of a neural predictor.

Original: Different from DFL, S OE BM eliminates the need of using KKT conditions to induce an implicit differentiable optimization layer.
Revised: Different from DFL, S OE BM eliminates the need of using KKT conditions to induce an implicit differentiable optimization layer.

Original: Instead, it leverages expressive EBMs [31] to directly model the conditional probability of the decisions and parameterizes an explicit energy-based differentiable layer as a surrogate to the original optimization problem.
Revised: Instead, it leverages expressive EBMs [33] to directly model the conditional probability of the decisions and parameterizes an explicit energy-based differentiable layer as a surrogate to the original optimization problem.

Original: To better approximate the optimization landscape with the EBM surrogate, we design two complementary learning objectives in S O -E BM : 1) a local matching objective that maximizes the likelihood of the optimal decisions; and 2) a global matching objective that minimizes the distribution distance between the posterior distribution of the decision variables and the surrogate-based decision distribution.
Revised: To better approximate the optimization landscape with the EBM surrogate, we design two complementary learning objectives in S O -E BM : 1) a local matching objective that maximizes the likelihood of the optimal decisions; and 2) a global matching objective that minimizes the distribution distance between the posterior distribution of the decision variables and the surrogate-based decision distribution.

Original: Due to its flexibility, S O -E BM is not constrained to convex objectives, but can be applied to a wide class of stochastic optimization problems.
Revised: Due to its flexibility, S O -E BM is not constrained to convex objectives, but can be applied to a wide class of stochastic optimization problems.

Original: Another key advantage of S O -E BM is its computational efficiency.
Revised: Another key advantage of S O -E BM is its computational efficiency.

Original: As the optimization layer is parameterized by an energy-based model, S O -E BM eliminates the need of solving and differentiating through the optimization procedure at each training iteration.
Revised: As the optimization layer is parameterized by an energy-based model, S O -E BM eliminates the need of solving and differentiating through the optimization procedure at each training iteration.

Original: Rather, S O -E BM estimates the derivatives of the energy-based surrogate model using a self-normalized importance sampler.
Revised: Rather, S O -E BM estimates the derivatives of the energy-based surrogate model using a self-normalized importance sampler.

Original: We design a Gaussian-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 S O -E BM .
Revised: We design a Gaussian-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 S O -E BM .

Original: The main contributions of this work are: (1) We propose a new end-to-end stochastic optimization method based on energy-based model.
Revised: The main contributions of this work are: (1) We propose a new end-to-end stochastic optimization method based on energy-based model.

Original: It avoids the needs of solving and differentiating through the optimization problem during training and can be applied to a wide range of stochastic optimization problems beyond convex objectives.
Revised: It avoids the needs of solving and differentiating through the optimization problem during training and can be applied to a wide range of stochastic optimization problems beyond convex objectives.

Original: (2) We propose a coupled training objective to encourage the energy-based surrogate to well approximate the optimization landscape and an efficient selfnormalized importance sampler based on a mixture-of-Gaussian proposal.
Revised: (2) We propose a coupled training objective to encourage the energy-based surrogate to well approximate the optimization landscape and an efficient selfnormalized importance sampler based on a mixture-of-Gaussian proposal.

Original: (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.
Revised: (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.

Original:  Preliminaries
Revised:  Preliminaries

Original: Problem Formulation.
Revised: Problem Formulation.

Original: Consider a stochastic optimization problem:
Revised: Consider a stochastic optimization problem:

Original: [Equation]
Revised: [Equation]

Original: [Table][Table] [Figure] [Table]
Revised: [Table]

Original: [Figure] [Table] [Figure] [Table]
Revised: [Table]

Original: [Figure]

Original: [Table]
Revised: [Table]

Original: +P62Tc0Vqb6KxUZ9hIYahlLAeSkfr0NF+VudRQSXlEJZphIXfVLCKXmJ1WYAI2iqPzkpKxZXE//cXlVHYY0QoKdKm7+F/eKPdjp1xt+y5+APyfBjDTYDO1+/Wc4SEWeoCahwNpu4GfUc15ICoVVLcwtZiAuYYhdRzUkaHvl5PKKbztlwOPUuKeJT9SHEyUk1hZJ5DrHNu3T2lh8qdbNKf7YK6XOckItpoviXHFK+ThGPpAGBanCERBGOq9cjMCAIBd2zYUQPD35OTndbwXvW/tfDhrHB7M4ltgWe8eaLGAf2DH7zNqswwT7zn6zv+zG++H98W69u2nrnDebecsewbv/B0fHruo=</latexit> Generic Loss e.g., NLL Backpropagation Decision-focused where y denotes the parameters of the optimization problem, a denotes the decision variables within a feasible space C , and f is the cost function to be optimized.
Revised: +/vkv+R8o+Zu1jZOt8r1rWEcBbJMVkiFuGSb1MkROSENQsk9eSKv5M16sF6sd+tjUDpiDXuWyA9Yn185p6tG</latexit> where y denotes the parameters of the optimization problem, a denotes the decision variables within a feasible space C , and f is the cost function to be optimized.

Original: In many applications, the parameters y are unknown and stochastic , which must be inferred from some correlated features x .
Revised: In many applications, the parameters y are unknown and stochastic , which must be inferred from some correlated features x .

Original: We assume a dataset D
Revised: We assume a dataset D

Original: = { x i , y i } N i =1 drawn from the joint distribution of the features and problem parameters.
Revised: = { x i , y i } N i =1 drawn from the joint distribution of the features and problem parameters.

Original: Our task is to learn a decision-making model M parameterized by θ , which takes the features x as input and outputs the optimal decisions a ∗ ( x ; θ ) .
Revised: Our task is to learn a decision-making model M parameterized by θ , which takes the features x as input and outputs the optimal decisions a ∗ ( x ; θ ) .

Original: The model should be learned such that its output optimal decisions minimize the expected decision cost under the joint distribution of ( x , y ) , namely: 
Revised: 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 min θ E ( x ,y )

Original: [Equation]
Revised: ∼D f ( y, a ∗ ( x ; θ )) .

Original: For example, in supply chain optimization, y can be regional customer demands, a can be pre-ordered products for each region, f 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 .
Revised: For example, in supply chain optimization, y can be regional customer demands, a can be pre-ordered products for each region, f 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 .

Original: Herein, the actual customer demands y are usually unknown and must be predicted from features x such as customer purchase history and regional economy indices.
Revised: Herein, the actual customer demands y are usually unknown and must be predicted from features x such as customer purchase history and regional economy indices.

Original: Two-stage learning v.s. Decision Focused Learning (DFL).
Revised: Two-stage learning v.s. Decision Focused Learning (DFL).

Original: A common practice to the above stochastic optimization problem is the two-stage predict-then-optimize framework.
Revised: A common practice to the above stochastic optimization problem is the two-stage predict-then-optimize framework.

Original: It first learns an probabilistic predictive model p ( y | x ; θ ) and then uses existing stochastic optimization solvers to obtain the optimal action that minimizes the expected cost: a∗ ( x
Revised: It first learns an probabilistic predictive model p ( y | x ; θ ) and then uses existing stochastic optimization solvers to obtain the optimal action that minimizes the expected cost: a

Original: [Equation]
Revised: [Equation]

Original: Though the two-stage approach is simple and efficient, it can suffer from suboptimal performance due to the misalignment of the prediction loss and the optimization loss.
Revised: Though the two-stage approach is simple and efficient, it can suffer from suboptimal performance due to the misalignment of the prediction loss and the optimization loss.

Original: In contrast, DFL integrates prediction and optimization into an end-to-end model, thus tailoring the predictive model for the optimization task, as shown in Figure 1.
Revised: In contrast, DFL integrates prediction and optimization into an end-to-end model, thus tailoring the predictive model for the optimization task, as shown in Figure 1.

Original: By directly optimizing the task loss, the gradient of the model parameters can be obtained with the following chain rule:
Revised: By directly optimizing the task loss, the gradient of the model parameters can be obtained with the following chain rule:

Original: [Equation]
Revised: [Equation]

Original: The key challenge here is to compute the Jacobian ∂a ∗ ( x ; θ 
Revised: The key challenge here is to compute the Jacobian ∂a ∗ ( x ; θ )

Original: ) ∂y : it is needed to apply the chain rule to learn the model using gradient decent methods.
Revised: [Equation] : it is needed to apply the chain rule to learn the model using gradient decent methods.

Original: This is nontrivial, because a ∗ is the solution of a stochastic optimization solver and not directly differentiable.
Revised: This is nontrivial, because a ∗ is the solution of a stochastic optimization solver and not directly differentiable.

Original: To address this challenge, OptNet [13, 4] assumes quadratic optimization objectives and differentiates through the KKT conditions using the implicit function theorem.
Revised: To address this challenge, OptNet [13, 4] assumes quadratic optimization objectives and differentiates through the KKT conditions using the implicit function theorem.

Original: This way, OptNet can obtain the Jacobian by solving the optimization problem along with a set of linear equations in each training iteration.
Revised: This way, OptNet can obtain the Jacobian by solving the optimization problem along with a set of linear equations in each training iteration.

Original: Several works [1, 39, 46] extend this technique to more general cases.
Revised: Several works [1, 41, 48] extend this technique to more general cases.

Original: For example, cvxpylayers [1] extends it to more general cases of convex optimization using disciplined parametrized programming (DPP) grammar.
Revised: For example, cvxpylayers [1] extends it to more general cases of convex optimization using disciplined parametrized programming (DPP) grammar.

Original: Although existing DFL approaches can achieve better decisions compared to two-stage learning, they have several drawbacks.
Revised: Although existing DFL approaches can achieve better decisions compared to two-stage learning, they have several drawbacks.

Original: First, they are often constrained to convex optimization objectives as they rely on the KKT conditions to compute derivatives.
Revised: First, they are often constrained to convex optimization objectives as they rely on the KKT conditions to compute derivatives.

Original: Though [39, 46] propose to 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 to poor decision quality.
Revised: Though [41, 48] propose to 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 to poor decision quality.

Original: Second, they suffer from high computational complexity because the computation of the Jacobian ∂a ∗ ( x ; θ ) ∂y requires repeatedly solving the optimization program and back-propagating through it.
Revised: Second, they suffer from high computational complexity because the computation of the Jacobian ∂a ∗ ( x ; θ ) ∂y requires repeatedly solving the optimization program and back-propagating through it.

Original: The problem is more severe when the expectation of the cost cannot be computed analytically.
Revised: The problem is more severe when the expectation of the cost cannot be computed analytically.

Original: In that case, we need to use sample average approximation [26, 44, 29] and draw multiple IID samples from the predictive distribution, which makes the objective much more complicated and expensive.
Revised: In that case, we need to 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.

Original:  Proposed Method 3.1 Energy-based Model for End-to-end Stochastic Programming
Revised:  Proposed Method 3.1 Energy-based Model for End-to-end Stochastic Programming

Original: Our task is to learn an optimal decision-making model M θ , such that its output optimal decision a ∗ ( x ; θ ) for input x minimizes the expected optimization cost: E ( x ,y )
Revised: Our task is to learn an optimal decision-making model M θ , such that its output optimal decision a ∗ ( x ; θ ) for input x minimizes the expected optimization cost: E ( x ,y )

Original: ∼D f ( y, a ∗ ( x ; θ )) .
Revised: ∼D f ( y, a ∗ ( x ; θ )) .

Original: The core idea of our method is to directly model M θ ’s probability distribution over the decisions conditioned on the features, denoted as q ( a | x ; θ ) , using energy-based parameterization [31]:
Revised: The core idea of our method is to directly model M θ ’s probability distribution over the decisions conditioned on the features, denoted as q ( a | x ; θ ) , using energy-based parameterization [33]:

Original: [Equation]
Revised: [Equation]

Original: To parameterize the energy function E ( x , a ; θ ) , a natural option is to directly use a deep neural network which takes a feature-decision pair ( x , a ) as input and output a scalar value.
Revised: To parameterize the energy function E ( x , a ; θ ) , a natural option is to directly use a deep neural network which takes a feature-decision pair ( x , a ) as input and output a scalar value.

Original: However, such a design ignores the algorithmic structure of the optimization problem and thus can be data-inefficient during learning.
Revised: However, such a design ignores the algorithmic structure of the optimization problem and thus can be data-inefficient during learning.

Original: Instead, we propose to explicitly model the problem structure by using the expected task loss as the energy function:
Revised: Instead, we propose to explicitly model the problem structure by using the expected task loss as the energy function:

Original: [Equation]
Revised: [Equation]

Original: where p ( y | x ; θ ) is the predictive distribution of an uncertainty-aware neural network.
Revised: where p ( y | x ; θ ) is the predictive distribution of an uncertainty-aware neural network [19, 31].

Original: Eq.
Revised: Eq.

Original: 2 creates an explicit energy-based differentiable layer as a surrogate to the original optimization problem.
Revised: 2 creates an explicit energy-based differentiable layer as a surrogate to the original optimization problem.

Original: Compared with the two-stage model, it builds a direct connection between the input features x and decision variable a and thus is more tailored to the downstream task.
Revised: Compared with the two-stage model, it builds a direct connection between the input features x and decision variable a and thus is more tailored to the downstream task.

Original: 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 a lot of learning.
Revised: 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 a lot of learning.

Original: This energy-based surrogate function also has a intuitive interpretation.
Revised: This energy-based surrogate function also has a intuitive interpretation.

Original: When the decision a is in a region with smaller expected task loss, it has lower energy and higher probability; when it leads to higher expected task loss, it has higher energy and lower probability.
Revised: When the decision a is in a region with smaller expected task loss, it has lower energy and higher probability; when it leads to higher expected task loss, it has higher energy and lower probability.

Original: Since we have the ground truth for the parameters y in the training data, the feature-decision pairs
Revised: Since we have the ground truth for the parameters y in the training data, the feature-decision pairs

Original: D a = { ( x i , a ∗ i ) } Ni =1 can be easily constructed by solving a ∗ i = arg min
Revised: D a = { ( x i , a ∗ i ) } Ni =1 can be easily constructed by solving a ∗ i = arg min

Original: a i ∈ C
Revised: a i ∈ C

Original: f
Revised: f

Original: ( y i , a i ) for each ( x i , y i ) using any off-the-shelf optimization solvers.
Revised: ( y i , a i ) for each ( x i , y i ) using any off-the-shelf optimization solvers.

Original: Note that the construction of such feature-action training pairs needs to be done only once during preprocessing (Fig. 3).
Revised: Note that the construction of such feature-action training pairs needs to be done only once during preprocessing (Fig. 3).

Original: Then, we minimize the negative log-likelihood (NLL) of all the feature-decision pairs:
Revised: Then, we minimize the negative log-likelihood (NLL) of all the feature-decision pairs:

Original: [Equation]
Revised: [Equation]

Original: By minimizing the NLL, we are essentially minimizing the energy of the optimal actions while maximizing the energy of other points.
Revised: By minimizing the NLL, we are essentially minimizing the energy of the optimal actions while maximizing the energy of other points.

Original: Thus the end-to-end stochastic programming problem is translated to learning a neural network that outputs the smallest energy for the optimal actions.
Revised: Thus the end-to-end stochastic programming problem is translated to learning a neural network that outputs the smallest energy for the optimal actions.

Original: 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.
Revised: 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.

Original: Our energy-based formulation can also be interpreted as a non-sequential maximum entropy inverse reinforcement learning (MaxEN-IRL) model [51, 50].
Revised: Our energy-based formulation can also be interpreted as a non-sequential maximum entropy inverse reinforcement learning (MaxEN-IRL) model [53, 52].

Original: From the IRL perspective, the input features { x i } Ni =1 can be interpreted as environment states, the feature-decision pairs { ( x i , a i ) } Ni =1 as expert demonstrations, and the negative expected task loss − E p ( y | x ; θ ) f ( y, a ) as reward.
Revised: From the IRL perspective, the input features { x i } Ni =1 can be interpreted as environment states, the feature-decision pairs { ( x i , a i ) } Ni =1 as expert demonstrations, and the negative expected task loss − E p ( y | x ; θ ) f ( y, a ) as reward.

Original: Eq.
Revised: Eq.

Original: 3 is then equivalent to maximizing the likelihood of the expert demonstrations (optimal decisions) to recover the hidden reward function.
Revised: 3 is then equivalent to maximizing the likelihood of the expert demonstrations (optimal decisions) to recover the hidden reward function.

Original: 3.2 Augmenting Energy-Based Objective with Distribution Regularizer
Revised: 3.2 Augmenting Energy-Based Objective with Distribution Regularizer

Original: One limitation of learning with EBM-based likelihood for the optimization problem is that we model only one data point a ∗ for each conditional distribution q ( a | x ; θ ) .
Revised: One limitation of learning with EBM-based likelihood for the optimization problem is that we model only one data point a ∗ for each conditional distribution q ( a | x ; θ ) .

Original: Minimizing the NLL for a single data point ignores the overall matchness between the landscape of original optimization problem and that of the EBM-based probability density, which can easily cause overfitting.
Revised: Minimizing the NLL for a single data point ignores the overall matchness between the landscape of original optimization problem and that of the EBM-based probability density, which can easily cause overfitting.

Original: To learn the overall shape of the energy function better, we propose a distribution-based regularizer which augments the energy-based objective from a global training perspective.
Revised: To learn the overall shape of the energy function better, we propose a distribution-based regularizer which augments the energy-based objective from a global training perspective.

Original: The distribution-based regularizer is based on minimizing the distance between the model distribution and an oracle posterior distribution.
Revised: The distribution-based regularizer is based on minimizing the distance between the model distribution and an oracle posterior distribution.

Original: Specifically, we assume that a follows a prior distribution p ( a ) .
Revised: Specifically, we assume that a follows a prior distribution p ( a ) .

Original: With the ground-truth label y , the posterior distribution of a is then given by p ( a | y )
Revised: With the ground-truth label y , the posterior distribution of a is then given by p ( a | y )

Original: =
Revised: =

Original: p ( a ) exp ( − f ( y,a ))
Revised: p ( a ) exp ( − f ( y,a ))

Original: Z ( y ) when we define the task loss as the negative log-likelihood, i .
Revised: Z ( y ) when we define the task loss as the negative log-likelihood, i .

Original: e ., − log p ( y | a ) =
Revised: e ., − log p ( y | a ) =

Original: f ( y, a ) .
Revised: f ( y, a ) .

Original: However, such an expression depends on the ground-truth label, which cannot be obtained at test time.
Revised: However, such an expression depends on the ground-truth label, which cannot be obtained at test time.

Original: We let the model distribution q ( a | x ; θ ) mimic this oracle posterior p ( a | y ) distribution by minimizing their
Revised: We let the model distribution q ( a | x ; θ ) mimic this oracle posterior p ( a | y ) distribution by minimizing their

Original: KL-divergence:
Revised: KL-divergence:

Original: [Equation]
Revised: [Equation]

Original: where H ( · ) denotes entropy of the probability distribution.
Revised: where H ( · ) denotes entropy of the probability distribution.

Original: Our final training loss is a weighted combination of the NLL and the distribution based regularizer:
Revised: Our final training loss is a weighted combination of the NLL and the distribution based regularizer:

Original: [Equation]
Revised: [Equation]

Original: where λ is a hyper-parameter.
Revised: where λ is a hyper-parameter.

Original: In this final objective, the KL-divergence loss and the NLL loss complement each other (Fig. 2).
Revised: In this final objective, the KL-divergence loss and the NLL loss complement each other (Fig. 2).

Original: The NLL loss is a local training strategy while the KL-divergence loss learns the energy function from a global level.
Revised: The NLL loss is a local training strategy while the KL-divergence loss learns the energy function from a global level.

Original: Specifically, the NLL loss helps the model capture the location of the optimal decision better but ignores the overall energy shape.
Revised: Specifically, the NLL loss helps the model capture the location of the optimal decision better but ignores the overall energy shape.

Original: The distribution-based regularizer essentially adds more anchor points besides the optimal decisions and thus helps fit the overall energy shape.
Revised: The distribution-based regularizer essentially adds more anchor points besides the optimal decisions and thus helps fit the overall energy shape.

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: /T0xppPUkCmxnRM1Ir3pz8T+vm5rwxp9ymaQGJVsuClNBTEzmX5MBV8iMmFhCmeL2VsJGVFFmbDYFG4K3+vI6aVXK3lW50qiWatUsjjycwTlcggfXUIN7qEMTGCA8wyu8OY/
Revised: /T0xppPUkCmxnRM1Ir3pz8T+vm5rwxp9ymaQGJVsuClNBTEzmX5MBV8iMmFhCmeL2VsJGVFFmbDYFG4K3+vI6aVXK3lW50qiWatUsjjycwTlcggfXUIN7qEMTGCA8wyu8OY/

Original: Oi/PufCxbc042cwp/4Hz+AJYFjL [Figure] </latexit
Revised: Oi/PufCxbc042cwp/4Hz+AJYFjL8=</latexit

Original: [Figure]
Revised: [Figure]

Original: 3.3 Training with Contrastive Divergence and Self-normalized Importance Sampling
Revised: 3.3 Training with Contrastive Divergence and Self-normalized Importance Sampling

Original: Optimizing Eq.
Revised: Optimizing Eq.

Original: 5 requires evaluating the partition function and computing the KL-divergence between two continuous distributions which typically involves intractable integrals.
Revised: 5 requires evaluating the partition function and computing the KL-divergence between two continuous distributions which typically involves intractable integrals.

Original: In this subsection, we propose to use a self-normalized importance sampler based on mixture of Gaussians to estimate the gradient of the model parameters efficiently.
Revised: In this subsection, we propose to use a self-normalized importance sampler based on mixture of Gaussians to estimate the gradient of the model parameters efficiently.

Original: First, we derive (see supplementary for details)
Revised: First, we derive (see supplementary for details)

Original: the gradient of the training loss with respect to the model parameters θ as:
Revised: the gradient of the training loss with respect to the model parameters θ as:

Original: [Equation]
Revised: [Equation]

Original: As can be seen, the gradient can be estimated by sampling from the model distribution q ( a | x ; θ ) and oracle distribution p ( a | y ) .
Revised: As can be seen, the gradient can be estimated by sampling from the model distribution q ( a | x ; θ ) and oracle distribution p ( a | y ) .

Original: Unfortunately, we cannot easily draw samples because of the unnormalized constant.
Revised: Unfortunately, we cannot easily draw samples because of the unnormalized constant.

Original: Existing methods usually resort to MCMC methods ( e .
Revised: Existing methods usually resort to MCMC methods ( e .

Original: g ., Langevin dynamics) to use this gradient estimator.
Revised: g ., Langevin dynamics) to use this gradient estimator.

Original: However, MCMC is an iterative process and can be time consuming.
Revised: However, MCMC is an iterative process and can be time consuming.

Original: To improve the training efficiency, we propose to use a self-normalized importance sampler based on a Gaussian mixture proposal to estimate the gradient.
Revised: To improve the training efficiency, we propose to use a self-normalized importance sampler based on a Gaussian mixture proposal to estimate the gradient.

Original: Specifically, for each x , we first sample a set of M candidates { a m } Mm =1 from a proposal distribution π ( a | x ) , and then sample ˜ a from the empirical distribution located at each a m and weighted proportionally to exp ( − E ( a | x ; θ )) /π
Revised: Specifically, for each x , we first sample a set of M candidates { a m } Mm =1 from a proposal distribution π ( a | x ) , and then sample ˜ a from the empirical distribution located at each a m and weighted proportionally to exp ( − E ( a | x ; θ )) /π

Original: ( a | x ) .
Revised: ( a | x ) .

Original: To reduce the variance of the self-normalized importance sampler, we propose to use a mixture of K Gaussians which are centered at the location of the corresponding optimal decision as the proposal distribution: π ( a | x ) = 1 K (cid:80) Ki =1 N ( a ∗ ; σ k ) , where K and
Revised: To reduce the variance of the self-normalized importance sampler, we propose to use a mixture of K Gaussians which are centered at the location of the corresponding optimal decision as the proposal distribution: π ( a | x ) = 1 K (cid:80) Ki =1 N ( a ∗ ; σ k ) , where K and

Original: [Table]
Revised: [Table]

Original: + ( )
Revised: [Equation] [Table] [Equation]

Original: Preprocessing
Revised: [Table] [Equation] [Table] [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: OptimizationSolver
Revised: [Equation]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: +LfVSwxPKxnTIu5YqGnHjZ/Njp+TMKgMSxtqWQjJXf09kNDJmEgW2M6I4MsveTPzP66YYXvuZUEmKXLHFojCVBGMy+5wMhOYM5cQSyrSwtxI2opoytPmUbAje8surpF2repfV2l290qjncRThBE7hHDy4ggbcQhNawEDAM7zCm6OcF+fd+Vi0Fpx85hj+wPn8AZ38joM=</latexit>
Revised: +lfMQ042iDytsQvOWXV0mjXPKuSuX7SrFayeLIkVNyRi6IR65JldyRGqkTTjR5Jq/kzXlyXpx352PRuuZkMyfkD5zPH0/gklc=</latexit>

Original: { σ k } Kk =1 are hyper-parameters.
Revised: { σ k } Kk =1 are hyper-parameters.

Original: This mixture of Gaussian based self-normalized importance sampler enjoys great computational efficiency.
Revised: This mixture of Gaussian based self-normalized importance sampler enjoys great computational efficiency.

Original: We only need to estimate the gradient of the energy-based surrogate layer by drawing samples from a simple mixture of Gaussians, instead of solving and differentiating through the optimization problem as in DFL.
Revised: We only need to estimate the gradient of the energy-based surrogate layer by drawing samples from a simple mixture of Gaussians, instead of solving and differentiating through the optimization problem as in DFL.

Original: This significantly reduces the training time compared with DFL as shown in Section 4.
Revised: This significantly reduces the training time compared with DFL as shown in Section 4.

Original: Further, locating the proposal distribution at the optimal decisions mimics the contrastive divergence method [21, 15] used in Langevin dynamics where the MCMC chain starts from the training data.
Revised: Further, locating the proposal distribution at the optimal decisions mimics the contrastive divergence method [22, 15] used in Langevin dynamics where the MCMC chain starts from the training data.

Original: This makes the proposal distribution close to the model distribution and thus has better sample efficiency.
Revised: This makes the proposal distribution close to the model distribution and thus has better sample efficiency.

Original: The sampling procedure from the oracle distribution p ( a | y ) is similar with q ( a | x ; θ ) .
Revised: The sampling procedure from the oracle distribution p ( a | y ) is similar with q ( a | x ; θ ) .

Original: When the expected task loss has no analytical expression, we can draw multiple samples from p ( y | x ; θ ) to estimate the expectation and use reparameterization trick [28, 22, 34, 40] to make it differentiable.
Revised: When the expected task loss has no analytical expression, we can draw multiple samples from p ( y | x ; θ ) to estimate the expectation and use reparameterization trick [29, 23, 36, 42] to make it differentiable.

Original: Finally, the model can be efficiently trained via gradient-based method, such as Adam [27].
Revised: Finally, the model can be efficiently trained via gradient-based method, such as Adam [28].

Original: The detailed training procedure is given in Alg.
Revised: The detailed training procedure is given in Alg.

Original: 1 in the supplementary.
Revised: 1 in the supplementary.

Original: Model inference.
Revised: Model inference.

Original: With the optimal model parameter θ opt , we draw samples from q ( a | x ∗ ; θ opt ) for an unseen test instance x ∗ .
Revised: With the optimal model parameter θ opt , we draw samples from q ( a | x test ; θ opt ) for an unseen test instance x test .

Original: However, in the real-world applications, we usually only need the most optimal decision.
Revised: However, in the real-world applications, we usually only need the most optimal decision.

Original: This can be obtained by solving arg min a ∗  ∈ C E p ( y | x  ; θ opt ) f
Revised: This can be obtained by solving arg min a ∗ test ∈ C E p ( y | x test ; θ opt ) f

Original: ( y, a ) with any existing black-box optimization solver such as CVXPY [12, 2] and Pyomo [5, 20] (Fig. 3).
Revised: ( y, a ) with any existing black-box optimization solver such as CVXPY [12, 2] and Pyomo [5, 21] (Fig. 3).

Original:  Experiments
Revised:  Experiments

Original: In this section, we empirically evaluate S O -E BM .
Revised: In this section, we empirically evaluate S O -E BM .

Original: We conduct experiments in three applications: (1)
Revised: We conduct experiments in three applications: (1)

Original: 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 a non-convex optimization objective.
Revised: 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 a non-convex optimization objective.

Original: Finally, we do ablation studies to show the effect of each component in S O -E BM .
Revised: Finally, we do ablation studies to show the effect of each component in S O -E BM .

Original: 4.1 Load Forecasting and Generator Scheduling
Revised: 4.1 Load Forecasting and Generator Scheduling

Original: In this task, a power system operator needs to decide how much electricity a ∈ R 24 to schedule for each hour in the next 24 hours to meet the actual electricity demands.
Revised: In this task, a power system operator needs to decide how much electricity a ∈ R 24 to schedule for each hour in the next 24 hours to meet the actual electricity demands.

Original: The optimization objective is a combination of an under-generation penalty, an over-generation penalty, and a mean squared loss between supplies and demands:
Revised: The optimization objective is a combination of an under-generation penalty, an over-generation penalty, and a mean squared loss between supplies and demands:

Original: [Equation]
Revised: [Equation]

Original: subject to
Revised: subject to

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Figure]

Original: [Figure]

Original: [Figure]

Original: d

Original: [Figure]
Revised: [Figure]

Original: True landscape 1.

Original: 1.0 d

Original: Recovered landscape 1.

Original: 1.0 d 1.0 d where [ · ] + = max { v, 0 } .
Revised: where [ · ] + = max { v, 0 } .

Original: Usually, the penalty coefficients satisfy γ e ≫ γ s since under-generation is more serious than over-generation.
Revised: Usually, the penalty coefficients satisfy γ e ≫ γ s since under-generation is more serious than over-generation.

Original: The quadratic regularization term indicates the preference for generation schedules that closely match actual demands.
Revised: The quadratic regularization term indicates the preference for generation schedules that closely match actual demands.

Original: The ramping constraint c r restricts the change in generation between consecutive timepoints, which reflects physical limitationsassociated with quick changes in electricity output levels.
Revised: The ramping constraint c r restricts the change in generation between consecutive timepoints, which reflects physical limitations with quick changes in electricity output levels.

Original: Experiment Setup .
Revised: Experiment Setup .

Original: We forecast the electricity demands y ∈ R 24 over all 24 hours of the next day using a 2-hidden layer neural network.
Revised: We forecast the electricity demands y ∈ R 24 over all 24 hours of the next day using a 2-hidden layer neural network.

Original: We assume y
Revised: We assume y

Original: i is a Gaussian with mean µ i and variance σ 

Revised: i is a Gaussian with mean µ i and variance σ

Original: i ; as such, the expectation in the optimization objective can be computed analytically.
Revised: i ; as such, the expectation in the optimization objective can be computed analytically.

Original: 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.
Revised: 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.

Original: Following [13], we set γ s = 0 .
Revised: Following [13], we set γ s = 0 .

Original: 4 , γ e = 50 and c = 0 .
Revised: 4 , γ e = 50 and c = 0 .

Original: We compare with the following baselines on this task: (1) A two-stage predict-then-optimize model trained with negative likelihood loss for the prediction task.
Revised: We compare with the following baselines on this task: (1) A two-stage predict-then-optimize model trained with negative likelihood loss for the prediction task.

Original: (2) Decision-focused learning with the
Revised: (2) Decision-focused learning with the

Original: QPTH solver [13].
Revised: QPTH solver [13].

Original: It uses sequential quadratic programming (SQP) to iteratively approximate the resultant convex objective as a quadratic objective, iterates until convergence, and then computes the necessary Jacobians using the quadratic approximation at the solution.
Revised: It uses sequential quadratic programming (SQP) to iteratively approximate the resultant convex objective as a quadratic objective, iterates until convergence, and then computes the necessary Jacobians using the quadratic approximation at the solution.

Original: (3) DFL with cvxpylayers [1]  which provides a differentiable layer with disciplined parameterized programming (DPP) grammar.
Revised: (3) DFL with cvxpylayers [1] (DFL-CVX) which provides a differentiable layer with disciplined parameterized programming (DPP) grammar.

Original: Since the analytical expectation of Eq.
Revised: Since the analytical expectation of Eq.

Original: 7 cannot be written in DPP, we use Monte Carlo sampling to estimate the expectation for this baseline.
Revised: 7 cannot be written in DPP, we use Monte Carlo sampling to estimate the expectation for this baseline.

Original: (4) Policy-net.
Revised: (4) Policy-net.

Original: It direct maps from the input features to the decision variables by minimizing the task loss using supervised learning [13].
Revised: It direct maps from the input features to the decision variables by minimizing the task loss using supervised learning [13].

Original: Our supplementary provides more details of the model parameters.
Revised: Our supplementary provides more details of the model parameters.

Original: [Figure]
Revised: [Figure]

Original: Results.
Revised: Results.

Original: Fig.
Revised: Fig.

Original: 5 shows the end task loss for all the methods.
Revised: 5 shows the end task loss for all the methods.

Original: Our method S O -E BM outperforms all the baselines with a significant reduction of training time.
Revised: Our method S O -E BM outperforms all the baselines with a significant reduction of training time.

Original: Compared with the strongest baseline
Revised: Compared with the strongest baseline

Original: DFL-QPTH, S O -E BM improves the task loss by 7 . 3% .
Revised: DFL-QPTH, S O -E BM improves the task loss by 7 . 3% .

Original: The improvement is because DFL-QPTH needs to use SQP to iteratively obtain the solutions for non-quadratic optimization problems.
Revised: The improvement is because DFL-QPTH needs to use SQP to iteratively obtain the solutions for non-quadratic optimization problems.

Original: Differentiating through all the steps of SQP is prohibitively expensive in memory and time.
Revised: Differentiating through all the steps of SQP is prohibitively expensive in memory and time.

Original: To address this issue, existing works [13] differentiate through just the last step of SQP to obtain approximate gradients.
Revised: To address this issue, existing works [13] differentiate through just the last step of SQP to obtain approximate gradients.

Original: The inaccurate gradients may accrue during training and thus hurt decision quality.
Revised: The inaccurate gradients may accrue during training and thus hurt decision quality.

Original: This shows that our explicit differentiable energy-based function is an effective surrogate for the original implicit optimization layer.
Revised: This shows that our explicit differentiable energy-based function is an effective surrogate for the original implicit optimization layer.

Original: In term of efficiency, S O -E BM is more than 136 times faster than DFL-QPTH (0.68 second/epoch v.s. 93.
Revised: In term of efficiency, S O -E BM is more than 136 times faster than DFL-QPTH (0.68 second/epoch v.s. 93.

Original: second/epoch) in training.
Revised: second/epoch) in training.

Original: This is because we only need to 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.
Revised: This is because we only need to 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.

Original: DFL-cvxpylayers is even slower than DFL-QPTH since it needs to use sample average estimation to draw multiple samples from p ( y | x ; θ ) , which results in a much more complicated optimization objective.
Revised: DFL-CVX is even slower than DFL-QPTH since it needs to use sample average estimation to draw multiple samples from p ( y | x ; θ ) , which results in a much more complicated optimization objective.

Original: This is also likely the reason why DFL-cvxpylayers underperforms DFL-QPTH in terms of task loss here.
Revised: This is also likely the reason why DFL-CVX underperforms DFL-QPTH in terms of task loss here.

Original: 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).
Revised: 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).

Original: This is not surprising: as aforementioned, Policy-Net is a pure end-to-end model that needs a large amount of data to rediscover the algorithmic structure of the optimization task.
Revised: This is not surprising: as aforementioned, Policy-Net is a pure end-to-end model that needs a large amount of data to rediscover the algorithmic structure of the optimization task.

Original: In contrast, DFL and S O -E BM model the predict-and-optimize structure in the model design, which save learning and are more data-efficient.
Revised: In contrast, DFL and S O -E BM model the predict-and-optimize structure in the model design, which save learning and are more data-efficient.

Original: Fig.
Revised: Fig.

Original: 4 shows the ground-truth and S O -E BM learned landscapes of the energy function.
Revised: 4 shows the ground-truth and S O -E BM learned landscapes of the energy function.

Original: As we can see, S O -E BM can recover the landscape of the original optimization objective effectively though with a small discrepancy.
Revised: As we can see, S O -E BM can recover the landscape of the original optimization objective effectively though with a small discrepancy.

Original: The small discrepancy is expected since the ground-truth landscape is computed by directly using the ground-truth label, while S O -E BM uses the uncertainty-aware 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.
Revised: The small discrepancy is expected since the ground-truth landscape is computed by directly using the ground-truth label, while S O -E BM uses the uncertainty-aware 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.

Original: 4.2 Resource Allocation for COVID- 
Revised: 4.2 Resource Allocation for COVID-

Original: As seen in the COVID-19 pandemic, an increasing number of infected patients leads to increasing demand for medical resources; which makes it challenging for policymakers and epidemiologists to plan ahead.
Revised: As seen in the COVID-19 pandemic, an increasing number of infected patients leads to increasing demand for medical resources; which makes it challenging for policymakers and epidemiologists to plan ahead.

Original: Mechanistic epidemiological models based on Ordinary Differential Equations (ODEs) are often used to capture and forecast the dynamics of epidemics including for the COVIDpandemic [33, 23, 19, 10].
Revised: Mechanistic epidemiological models based on Ordinary Differential Equations (ODEs) are often used to capture and forecast the dynamics of epidemics including for the COVIDpandemic [35, 24, 20, 10].

Original: Such forecasts are typically then used as guidance to help plan for future resource allocation [36, 18, 41, 24].
Revised: Such forecasts are typically then used as guidance to help plan for future resource allocation [38, 18, 43, 25].

Original: In this task, we study the optimization problem of hospital bed preparation for COVID-19, one of the most common and important tasks epidemiologists focus on during the pandemic [3, 7].
Revised: In this task, we study the optimization problem of hospital bed preparation for COVID-19, one of the most common and important tasks epidemiologists focus on during the pandemic [3, 7].

Original: We need to decide how many beds a ∈ R 7 are needed to prepare for the next week based on the forecasted number of hospitalized patients y ∈ R 7 .
Revised: We need to decide how many beds a ∈ R 7 are needed to prepare for the next week based on the forecasted number of hospitalized patients y ∈ R 7 .

Original: The optimization objective is a combination of linear and quadratic costs, which accounts for over-preparations [ a − y ]
Revised: The optimization objective is a combination of linear and quadratic costs, which accounts for over-preparations [ a − y ]

Original: + and under-preparations [ y − a ] + in the next 7 days over ODE-derived dynamics: minimize [Equation]
Revised: + and under-preparations [ y − a ] + in the next 7 days over ODE-derived dynamics: minimize a ∈ R7 (cid:88) i

Original: [Equation]
Revised: [Equation]

Original: Experiment Setup .
Revised: Experiment Setup .

Original: We use the SEIR+HD ODE model proposed in [23] to capture the dynamics of the COVID-19 pandemic, which has been used in policy studies for interventions.
Revised: 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.

Original: The model is driven by a key parameter: the transmission rate β that reflects the probability of disease transmission.
Revised: The model is driven by a key parameter: the transmission rate β that reflects the probability of disease transmission.

Original: The forecasting model is a two-layer gated recurrent unit (GRU) [6] 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 β for the next 7 days.
Revised: The forecasting model is a two-layer gated recurrent unit (GRU) [6] 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 β for the next 7 days.

Original: We assume that β follows a Gaussian distribution.
Revised: We assume that β follows a Gaussian distribution.

Original: Since SEIR-HD is a complex non-linear ODE system, there is no closed expression for the distribution of the number of hospitalized patients y .
Revised: Since SEIR-HD is a complex non-linear ODE system, there is no closed expression for the distribution of the number of hospitalized patients y .

Original: Hence, we choose to sample from the forecasted distribution of βtimes and then simulate the SEIR-HD model to obtain the empirical distribution of the number of hospitalized patients y .
Revised: Hence, we choose to sample from the forecasted distribution of βtimes and then simulate the SEIR-HD model to obtain the empirical distribution of the number of hospitalized patients y .

Original: [Figure]
Revised: [Figure]

Original: Results .
Revised: Results .

Original: Fig.
Revised: Fig.

Original: 6 shows the results on the resource allocation for COVID-19 task.
Revised: 6 shows the results on the resource allocation for COVID-19 task.

Original: Both S O -E BM and DFL significantly outperforms the two-stage model in terms of the task loss, yielding 10 . 9%+ improvement.
Revised: Both S O -E BM and DFL significantly outperforms the two-stage model in terms of the task loss, yielding 10 . 9%+ improvement.

Original: DFL-QPTH is not included here because it needs to use SQP to iteratively approximate the objective and is much slower than
Revised: DFL-QPTH is not included here because it needs to use SQP to iteratively approximate the objective and is much slower than

Original: DFL-CVX in this task.
Revised: DFL-CVX in this task.

Original: The performance of S OE BM and DFL are similar, however, the training time of S O -E BM is 4.3 times faster than DFL.
Revised: The performance of S OE BM and DFL are similar, however, the training time of S O -E BM is 4.3 times faster than DFL.

Original: The time improvement of S O -E BM is not as large as on the energy scheduling task because all the methods need to compute the complex
Revised: The time improvement of S O -E BM is not as large as on the energy scheduling task because all the methods need to compute the complex

Original: ODE equations in the forward pass which is a time consuming part.
Revised: ODE equations in the forward pass which is a time consuming part.

Original: Therefore, considering the overall training efficiency, we argue that S O -E BM is a better option than DFL on this task in practice.
Revised: Therefore, considering the overall training efficiency, we argue that S O -E BM is a better option than DFL on this task in practice.

Original: 4.3 Network Security Game
Revised: 4.3 Network Security Game

Original: In this task, we study stochastic optimization in network security games [47].
Revised: In this task, we study stochastic optimization in network security games [49].

Original: Given a network
Revised: Given a network

Original: G =
Revised: G =

Original: ( V, E ) , a source node s ∈ V , and a set of target nodes T ∈ V , a network security game (NSG) [48, 17, 42] defines the min-max game where the attacker attempts to travel from s to any t ∈ T while the defender places checkpoints on certain number of edges in the graph.
Revised: ( V, E ) , a source node s ∈ V , and a set of target nodes T ∈ V , a network security game (NSG) [50, 17, 44] defines the min-max game where the attacker attempts to travel from s to any t ∈ T while the defender places checkpoints on certain number of edges in the graph.

Original: Each of the target node has a reward u ( t ) should the attacker reach them.
Revised: Each of the target node has a reward u ( t ) should the attacker reach them.

Original: The defender first chooses a mixed strategy.
Revised: The defender first chooses a mixed strategy.

Original: Having observed the defender’s mixed strategy but not the sampled pure strategy, the attacker attempts to choose a path.
Revised: Having observed the defender’s mixed strategy but not the sampled pure strategy, the attacker attempts to choose a path.

Original: 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.
Revised: 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.

Original: Suppose that a defines the placement of checkpoints by the defender where a e is the probability that edge e is covered.
Revised: Suppose that a defines the placement of checkpoints by the defender where a e is the probability that edge e is covered.

Original: The attackers will perform a random walk on the graph, generating a path r , stopping at either a target or a checkpoint.
Revised: The attackers will perform a random walk on the graph, generating a path r , stopping at either a target or a checkpoint.

Original: The transition between two nodes is determined by the defender’s strategy a and node-level parameters y , where y u defines the “attractiveness” of node u , representing the attacker’s idiosyncratic preferences.
Revised: The transition between two nodes is determined by the defender’s strategy a and node-level parameters y , where y u defines the “attractiveness” of node u , representing the attacker’s idiosyncratic preferences.

Original: The defender wants to choose a to maximize its expected utility: max a E T ∼ a E r ∼ p ( r | a , y ) − g ( r ) , where T is covered edges sampled according to a , g ( r ) =
Revised: The defender wants to choose a to maximize its expected utility: max a E T ∼ a E r ∼ p ( r | a , y ) − g ( r ) , where T is covered edges sampled according to a , g ( r ) =

Original: u ( t ) if the attacker reaches target t with path r and 0 if the attacker is stopped at a checkpoint.
Revised: u ( t ) if the attacker reaches target t with path r and 0 if the attacker is stopped at a checkpoint.

Original: The defender’s objective comes with the constraint (cid:80) a ≤ k , where k =represents the resources available to the defender.
Revised: The defender’s objective comes with the constraint (cid:80) a ≤ k , where k =represents the resources available to the defender.

Original: The optimization objective is non-convex due to the irrational strategy of the attacker [46].
Revised: The optimization objective is non-convex due to the irrational strategy of the attacker [48].

Original: Setup.
Revised: Setup.

Original: We follow the setup of [47], see supplementary for details.
Revised: We follow the setup of [49], see supplementary for details.

Original: [Figure]
Revised: [Figure]

Original: Results .
Revised: Results .

Original: Fig.
Revised: Fig.

Original: 7 shows the results on the adversarial network security game task.
Revised: 7 shows the results on the adversarial network security game task.

Original: Since the original DFL fails on this large scale problem, we compare a block-variable sampling approach specialized to this task (DFL-Block)[46].
Revised: Since the original DFL fails on this large scale problem, we compare a block-variable sampling approach specialized to this task (DFL-Block)[48].

Original: Surrogate [47] performs a linear dimension reduction for DFL to speed up the training time and improve the performance by smoothing the training landscape.
Revised: Surrogate [49] performs a linear dimension reduction for DFL to speed up the training time and improve the performance by smoothing the training landscape.

Original: We did not present the results of Surrogate in the previous two tasks because the numbers of decision variables there are relatively small and it has even worse performance than the original DFL approach.
Revised: We did not present the results of Surrogate in the previous two tasks because the numbers of decision variables there are relatively small and it has even worse performance than the original DFL approach.

Original: As we can see,
Revised: As we can see,

Original: S O -E BM can achieve competitive task loss with Surrogate and outperforms DFL-block and the twostage method.
Revised: S O -E BM can achieve competitive task loss with Surrogate and outperforms DFL-block and the twostage method.

Original: Meanwhile, S O -E BM is 1 . 8 times faster than Surrogate.
Revised: Meanwhile, S O -E BM is 1 . 8 times faster than Surrogate.

Original: 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 a larger matrix to inverse.
Revised: 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 a larger matrix to inverse.

Original: However, this issue can be mitigated by using some advanced matrix inverse algorithm [9, 30].
Revised: However, this issue can be mitigated by using some advanced matrix inverse algorithm [9, 32].

Original: DFL-Block is even worse than the two-stage model because it back-propagates through randomly sampled variables which results in inaccruate gradient estimation.
Revised: DFL-Block is even worse than the two-stage model because it back-propagates through randomly sampled variables which results in inaccruate gradient estimation.

Original: 4.4 Ablation Study
Revised: 4.4 Ablation Study

Original: [Table]
Revised: [Table]

Original: We investigate the effectiveness of the coupled training objective and alternative EBM training algorithms via ablation studies on the load forecasting and generator scheduling task.
Revised: We investigate the effectiveness of the coupled training objective and alternative EBM training algorithms via ablation studies on the load forecasting and generator scheduling task.

Original: Tableshows the results.
Revised: Tableshows the results.

Original: Our findings can be summarized as follows: (1) Both the MLE and KLD training objectives work and outperform the twostage model.
Revised: Our findings can be summarized as follows: (1) Both the MLE and KLD training objectives work and outperform the twostage model.

Original: This verifies the effectiveness of using the energy-based surrogate function to approximate the optimization landscape.
Revised: This verifies the effectiveness of using the energy-based surrogate function to approximate the optimization landscape.

Original: (2) By dropping either the MLE or KLD term, we observe a performance degradation in our method.
Revised: (2) By dropping either the MLE or KLD term, we observe a performance degradation in our method.

Original: Specifically, without the MLE term, the task loss increases by 13% ; without KLD term, the task loss increases by 2% .
Revised: Specifically, without the MLE term, the task loss increases by 13% ; without KLD term, the task loss increases by 2% .

Original: The larger degradation when simply minimizing the KLD term is possibly because accurately recover the entire optimization landscape is too difficult and it needs the MLE term to help capture the location of the optimal decision.
Revised: The larger degradation when simply minimizing the KLD term is possibly because accurately recover the entire optimization landscape is too difficult and it needs the MLE term to help capture the location of the optimal decision.

Original: (3) Training S O -E BM with Contrastive Divergence based
Revised: (3) Training S O -E BM with Contrastive Divergence based

Original: Langevin Dynamics (CD-LD) [15] achieves similar performance but incurs longer training time (1.
Revised: Langevin Dynamics (CD-LD) [15] achieves similar performance but incurs longer training time (1.

Original: second/epoch
Revised: second/epoch

Original: v.s. 0.68 second/epoch).
Revised: v.s. 0.68 second/epoch).

Original: This phenomenon is likely because the optimization problem is relatively low-dimensional, thus importance sampling works well and enjoys better efficiency in this regime.
Revised: This phenomenon is likely because the optimization problem is relatively low-dimensional, thus importance sampling works well and enjoys better efficiency in this regime.

Original: However, even with CD-LD, S O -E BM is still 63 times faster than DFL-QPTH (1.
Revised: However, even with CD-LD, S O -E BM is still 63 times faster than DFL-QPTH (1.

Original: second/epoch v.s. 93.12 second/epoch).
Revised: second/epoch v.s. 93.12 second/epoch).

Original: [Figure]
Revised: [Figure]

Original: We also investigate the impact of training data size for each method.
Revised: We also investigate the impact of training data size for each method.

Original: Fig.
Revised: Fig.

Original: 8 provides the task losses and training time under different ratios of training data on the load forecasting and generator scheduling task.
Revised: 8 provides the task losses and training time under different ratios of training data on the load forecasting and generator scheduling task.

Original: As we can see, our method outperforms the baselines constantly except for the ratio of 0 . 05 .
Revised: As we can see, our method outperforms the baselines constantly except for the ratio of 0 . 05 .

Original: When the ratio is below 0 .
Revised: When the ratio is below 0 .

Original: 05 , all the methods cannot work properly due to the extremely low resource.
Revised: 05 , all the methods cannot work properly due to the extremely low resource.

Original: The superior performance of our method is because we design the energy function as the expected task loss, which leverages the algorithmic structure inherent in the optimization problem.
Revised: The superior performance of our method is because we design the energy function as the expected task loss, which leverages the algorithmic structure inherent in the optimization problem.

Original: Hence, our method is not very data demanding.
Revised: Hence, our method is not very data demanding.

Original: In terms of efficiency, our method reduces the training time significantly compared with DFL-QPTH across different amounts of training data.
Revised: In terms of efficiency, our method reduces the training time significantly compared with DFL-QPTH across different amounts of training data.

Original:  Additional Related Work
Revised:  Additional Related Work

Original: The DFL works described in Section 2 focus on continuous optimization.
Revised: The DFL works described in Section 2 focus on continuous optimization.

Original: There are also studies that extend DFL for combinatorial optimization.
Revised: There are also studies that extend DFL for combinatorial optimization.

Original: [49] relaxes the discrete decision into its continuous counterpart and adds a quadratic regularization term to avoid vanishing gradients.
Revised: [51] relaxes the discrete decision into its continuous counterpart and adds a quadratic regularization term to avoid vanishing gradients.

Original: [35] proposes a log barrier regularizer and differentiates through the homogeneous self-dual embedding.
Revised: [37] proposes a log barrier regularizer and differentiates through the homogeneous self-dual embedding.

Original: [37] proposes a noise contrastive objective by maximizing the distance between the optimal solution and noisy samples.
Revised: [39] proposes a noise contrastive objective by maximizing the distance between the optimal solution and noisy samples.

Original: Our framework may be also extended for combinatorial problems by using discrete energy-based model [11].
Revised: Our framework may be also extended for combinatorial problems by using discrete energy-based model [11].

Original: The SPO+ loss [16] has been proposed to measure the prediction errors against optimization objectives, but it is only applicable to linear programming.
Revised: The SPO+ loss [16] has been proposed to measure the prediction errors against optimization objectives, but it is only applicable to linear programming.

Original: ProjectNet [8] approximately solves the linear programming problems using a differentiable projection architecture.
Revised: ProjectNet [8] approximately solves the linear programming problems using a differentiable projection architecture.

Original: Our work is also related to learning to optimize [45, 32, 25], which learns a policy network that solves optimization problems using supervised or reinforcement learning.
Revised: Our work is also related to learning to optimize [47, 34, 26], which learns a policy network that solves optimization problems using supervised or reinforcement learning.

Original: However, these pure approaches need to rediscover the structure of the optimization problem and thus data-inefficient.
Revised: However, these pure approaches need to rediscover the structure of the optimization problem and thus data-inefficient.

Original:  Limitations and Discussion
Revised:  Limitations and Discussion

Original: We focused on addressing the scalability and generality of existing decision-focused learning (DFL) for end-to-end stochastic optimization.
Revised: We focused on addressing the scalability and generality of existing decision-focused learning (DFL) for end-to-end stochastic optimization.

Original: We argue that these deficiencies stem from the reliance of implicitly differentiable optimization layers based on KKT conditions.
Revised: We argue that these deficiencies stem from the reliance of implicitly differentiable optimization layers based on KKT conditions.

Original: As a remedy, we circumvent such deficiencies by replacing the implicit optimization layer with a newly parameterized energybased surrogate function.
Revised: As a remedy, we circumvent such deficiencies by replacing the implicit optimization layer with a newly parameterized energybased surrogate function.

Original: We proposed a coupled training objective to encourage the energy-based surrogate well approximate the optimization landscape, as well as an efficient training procedure based on self-normalized importance sampling.
Revised: We proposed a coupled training objective to encourage the energy-based surrogate well approximate the optimization landscape, as well as an efficient training procedure based on self-normalized importance sampling.

Original: Empirically, we demonstrated that our energy-based model is effective in a wide range of stochastic optimization problems with either convex or nonconvex objectives.
Revised: Empirically, we demonstrated that our energy-based model is effective in a wide range of stochastic optimization problems with either convex or nonconvex objectives.

Original: It can achieve better or comparable performance than state-of-the-art DFL methods for stochastic optimization, while being several times or even orders of magnitude faster.
Revised: It can achieve better or comparable performance than state-of-the-art DFL methods for stochastic optimization, while being several times or even orders of magnitude faster.

Original: We discuss limitations and possible extensions of S O -E BM : (1) Handling more complex constraints .
Revised: We discuss limitations and possible extensions of S O -E BM : (1) Handling more complex constraints .

Original: Our method handles the constraints implicitly through the pre-processing step and explicitly through the inference step.
Revised: Our method handles the constraints implicitly through the pre-processing step and explicitly through the inference step.

Original: However, when the feasible space is extremely small, training the constrained EBM may have more benefits.
Revised: However, when the feasible space is extremely small, training the constrained EBM may have more benefits.

Original: To train EBMs in the constrained space, one direction is to project samples into the feasibility space.
Revised: To train EBMs in the constrained space, one direction is to project samples into the feasibility space.

Original: Another direction is to explore adding soft constraints to the energy function during training, e .
Revised: Another direction is to explore adding soft constraints to the energy function during training, e .

Original: g ., Augmented Lagrangian penalty and Barrier penalty.
Revised: g ., Augmented Lagrangian penalty and Barrier penalty.

Original: (2) More effective training methods .
Revised: (2) More effective training methods .

Original: Our method is a general framework for end-to-end stochastic programming problem based on EBM.
Revised: Our method is a general framework for end-to-end stochastic programming problem based on EBM.

Original: There are a number of training techniques  that can be applied [43] and we can further improve the task performance using the advanced EBM training algorithms [14, 38].
Revised: There are a number of training techniques [45] that can be plugged into our framework and we can further improve the task performance using the advanced EBM training algorithms [14, 40].

Original: Acknowledgments: 
Revised: Acknowledgments: We thank the anonymous reviewers for their helpful comments.

Original: This work was supported in part by the NSF (Expeditions CCF-1918770, 
Revised: This work was supported in part by the NSF (Expeditions CCF-1918770, CAREER IIS-2028586, IIS-2027862,

Original: CAREER IIS-2028586, IIS-2027862, IIS-1955883, IIS-2106961, IIS-2008334, CAREER IIS2144338, PIPP CCF-2200269), CDC MInD program, faculty research award from Facebook and funds/computing re-sources from Georgia Tech. 
Revised: IIS-1955883, IIS-2106961, IIS-2008334, CAREER IIS-2144338, PIPP CCF-2200269), CDC MInD program, faculty research award from Facebook and funds/computing resources from Georgia Tech. Checklist

Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Revised: [Yes] (b) Did you describe the limitations of your work?

Revised: [Yes] (c) Did you discuss any potential negative societal impacts of your work?

Revised: [Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Revised: [Yes] 2.

Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Revised: [N/A] (b) Did you include complete proofs of all theoretical results?

Revised: [N/A] 3.

Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Revised: [Yes] (b) Did you specify all the training details (e

Revised: g., data splits, hyperparameters, how they were chosen)?

Revised: [Yes] (c) Did you report error bars (e

Revised: g., with respect to the random seed after running experiments multiple times)?

Revised: [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)?

Revised: [Yes] 4.

Revised: If you are using existing assets (e

Revised: g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Revised: [Yes] (b) Did you mention the license of the assets?

Revised: [Yes] (c) Did you include any new assets either in the supplemental material or as a URL?

Revised: [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Revised: [N/A] 5.

Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Revised: [N/A]