Task-based End-to-end Model Learning in Stochastic Optimization

1 Introduction

• The paper targets the mismatch between generic predictive losses and the ultimate closed-loop task objective in larger processes.Examples include using classifications for planning or forecasts for power-grid scheduling, where different errors can have different task costs.Source: §1 · p. 1
• The motivating formulation allows a model with worse overall likelihood to be preferable when it makes more useful predictions for the downstream task.This reflects task-specific error tradeoffs rather than a single generic predictive criterion.Source: §1 · p. 1
• It proposes training probabilistic models to capture the objective of the stochastic-programming task in which they will later be used.The approach replaces training solely for predictive accuracy with learning aimed at minimizing the ultimate task-based loss.Source: §1 · p. 2
• The method is evaluated on a classical inventory stock problem, real-world electrical-grid scheduling, and real-world energy-storage arbitrage.The paper compares the proposed approach with traditional modeling and purely black-box policy optimization approaches.Source: §1 · p. 2

Stochastic programming

• Stochastic programming makes decisions under uncertainty by optimizing objectives governed by random processes, often using estimated probability distributions.Applications include energy, finance, and manufacturing.Source: §2 · p. 2

End-to-end training

• The paper uses end-to-end learning to mean adapting a model to the complete task the user ultimately wants to accomplish, rather than only predicting an output.This emphasizes closed-loop performance on the real task at hand.Source: §2 · p. 2
• Related work includes end-to-end policy learning, value-iteration-based optimization, and multi-objective optimization, but these approaches differ methodologically from the paper’s stochastic-programming focus.The cited lines of work more closely match the paper’s pure function-approximation alternative.Source: §2 · p. 2

Optimizing alternative loss functions

• Alternative-loss research has optimized criteria differing from standard prediction or inference objectives, including structured prediction, auxiliary losses, control performance, and meta-learning objectives.These studies motivate changing the optimized loss while addressing different settings.Source: §2 · pp. 2–3
• Bengio is identified as the closest philosophical predecessor because it tunes a predictive financial model using returns from a differentiable hedging strategy.The paper differs by coupling model tuning to stochastic programming and an explicit task-cost objective.Source: §2 · p. 3

3 End-to-end model learning in stochastic programming

• The method chooses parameters for p(y|x; θ) so that stochastic-programming decisions perform well under the true distribution, rather than merely producing accurate predictions.The framework supports both unconstrained and constrained decisions, with task loss incorporating expected cost and constraint deviations.Source: §3 · pp. 3–4

3.3 Differentiating the optimization solution to a stochastic programming problem

• The differentiation procedure is practical for many cases, including strongly convex functions and constraints, but cannot handle every optimization problem because argmin mappings may be discontinuous.The implementation uses sequential quadratic programming and quadratic-program argmin differentiation to compute derivatives at the optimum.Source: §3 · p. 4, §3.3 · p. 5

3.1 Discussion and alternative approaches

• Traditional log-likelihood training can select opaque error tradeoffs that perform poorly when the true conditional distribution is outside the model class.This model-bias setting may make the predictive model unsuitable for a later optimization task.Source: §3.1 · p. 4
• Black-box policy optimization learns actions directly without a model, but can be data-inefficient because its policy class must represent complex policies without an underlying model.The proposed approach retains a surrogate model while adapting it using task loss, and can combine log-likelihood with task loss.Source: §3.1 · p. 4

3.2 Optimizing task loss

• Task-loss optimization solves a proxy stochastic program, evaluates the resulting decision on observed outcomes, and updates parameters using gradients of violated constraints or the objective.Probabilistic constraints may require minibatches or penalty terms, while deterministic policy constraints remain enforced by the optimizer.Source: §3.2 · pp. 4–5, Figure 1 · p. 5

3.3 Differentiating the optimization solution to a stochastic programming problem

• The central technical challenge is differentiating the optimal decision z⋆(x; θ) with respect to model parameters θ.The approach differentiates KKT conditions with the implicit function theorem to obtain linear equations for the required Jacobians.Source: §3.3 · p. 5

4 Experiments

• The task-based approach substantially improves upon alternative modeling and policy methods across synthetic inventory, real-world grid scheduling, and battery arbitrage tasks.The experiments evaluate a synthetic inventory problem, a real-world energy scheduling task, and a real-world battery arbitrage task.Source: §4 · p. 5

4.1 Inventory stock problem

• In the inventory problem, task-based models outperform MLE and policy counterparts except when a linear MLE matches the true linear model.Under nonlinear truth, task-based learning also improves misspecified linear models by learning task-appropriate error tradeoffs rather than predictive-fit tradeoffs.Source: §4.1 · pp. 6–7, Figure 2 · p. 7

4.2 Load forecasting and generator scheduling

• 38.6% lower task loss than the RMSE net and 8.6% lower task loss than cost-weighted RMSE show the task net improves the grid-scheduling objective despite RMSE favoring the RMSE net.The approach retains computational tractability by using Gaussian assumptions and closed-form or sequential-quadratic-programming approximations while adjusting model parameters through task loss.Source: §4.2 · pp. 7–8, Figure 4 · p. 8

4.3 Price forecasting and battery storage

• On battery arbitrage data, the task-based model has better average task-loss performance and lower variability than RMSE minimization, although improvements are not statistically significant.The task-based method is described as hedging against price spikes and producing more reliable performance when predictions are inaccurate under high stochasticity.Source: §4.3 · pp. 8–9, Table 1 · p. 9

• 38.6% performance improvement over a highly optimized real-world stochastic-programming algorithm was achieved for electricity-generation scheduling based on predicted load.Source: §5 · p. 9
• The task-based model substantially outperformed MLE and policy-optimizing approaches except when MLE perfectly characterized the underlying distribution.Source: §5 · p. 9
• For energy-price prediction, the task-based approach produced more reliable task performance than a traditional predictive method under high inherent stochasticity.Source: §5 · p. 9
• Future work extends the approach to multiple-round stochastic learning, model predictive control, and full reinforcement learning settings.Source: §5 · p. 9

• The appendix presents computational and architectural details for the proposed task-based learning model and its experiments.Source: A Appendix · p. 12

• The chief technical challenge is computing the gradient of an objective that depends on the optimizer’s argmin solution z⋆(x; θ).Source: §A.1 · p. 12
• The method formulates convex stochastic programs through KKT conditions involving modeled-distribution expectations, linear equality constraints, and inequality constraints.Source: §A.1 · p. 12
• Differentiating the KKT optimality conditions and applying the implicit function theorem yields linear equations for the required Jacobians.Source: §A.1 · p. 12
• Sequential quadratic programming finds the optimal policy, while fast quadratic-program argmin differentiation solves the necessary linear equations for derivatives.Source: §A.1 · p. 12

• The inventory objective assigns linear and quadratic costs to ordering, over-ordering, and under-ordering relative to stochastic demand.Source: §A.2 · p. 12
• Discrete demand values with conditional probabilities transform the stochastic inventory problem into a joint quadratic program.Source: §A.2 · p. 13
• The quadratic program is expressed in inequality form, enabling computation of the optimal primal-dual solution and its Jacobian with respect to the probability vector.Source: §A.2 · p. 13
• The overall task-loss gradient is obtained by applying the chain rule after computing the optimizer Jacobians.Source: §A.2 · p. 13
• Model-probability Jacobians can be computed efficiently by backpropagation using a differentiable batch QP solver.Source: §A.2 · p. 13

• The objective penalizes under-generation and over-generation relative to stochastic demand while constraining generator ramping between periods.Here, z is the generator schedule, y is demand affected by features x, γs and γe are penalty parameters, and cr is the ramping limit.Source: §A.3 · p. 13
• Gaussian demand assumptions yield a closed-form expectation obtained by analytically integrating the Gaussian probability density.The closed form uses the Gaussian PDF and CDF evaluated with the demand mean and variance.Source: §A.3 · p. 13
• The resulting objective is convex in the generator schedule because it is an expectation of a convex function.Convexity enables efficient optimization and computation of the required Jacobians.Source: §A.3 · pp. 13–14
• Sequential quadratic programming repeatedly approximates the convex objective with a quadratic objective until convergence.At the solution, the quadratic approximation supplies the Jacobians, while differentiation with respect to μ and σ2 uses a differentiable batch QP solver [31].Source: §A.3 · p. 14

• The battery objective jointly represents charging, discharging, and battery state decisions under stochastic electricity prices.The formulation includes battery capacity, charging efficiency, charge and discharge limits, and flexibility and battery-health hyperparameters.Source: §A.4 · p. 14
• The expectation in the battery objective has a closed form when each stochastic price is modeled with a specified mean.Source: §A.4 · p. 14
• For this experiment, prices are modeled as lognormal, so the predictor estimates the mean of log(y), equivalently log(μ).The predicted means are inserted into the optimization problem before differentiation.Source: §A.4 · p. 15
• The method solves the storage optimization, computes Jacobians at the solution, and updates the underlying price-mean parameter through backpropagation.The update uses the differentiable batch QP solver [31].Source: §A.4 · p. 15

• Linear models use one-layer linear neural networks with appropriate input and output dimensions, whereas nonlinear models use two hidden layers combining linear, batch normalization, ReLU, and dropout with p = 0.2.Source: §A.5 · p. 15
• Models predicting probability distributions add a softmax output layer.Source: §A.5 · p. 15
• All models are implemented in PyTorch and optimized with Adam, while QPs and their Jacobians use a differentiable batch QP solver [31].Jacobian computation is performed automatically through backpropagation.Source: §A.5 · p. 15
• Source code for all experiments is publicly available in the authors’ repository.Source: §A.5 · p. 15