End-to-End Conformal Calibration for Optimization Under Uncertainty
Authors: Christopher Yeh
Nicolas Christianson
Alan Wu
Adam Wierman
Yisong Yue
Venue: Accepted by TMLR ()

Abstract: Machine learning can significantly improve performance for decision-making under uncertainty across a wide range of domains. However, ensuring robustness guarantees requires well-calibrated uncertainty estimates, which can be difficult to achieve with neural networks. Moreover, in high-dimensional settings, there may be many valid uncertainty estimates, each with its own performance profile—i.e., not all uncertainty is equally valuable for downstream decision-making. To address this problem, this paper develops an end-to-end framework to _learn_ uncertainty sets for conditional robust optimization in a way that is informed by the downstream decision-making loss, with robustness and calibration guarantees provided by conformal prediction. In addition, we propose to represent general convex uncertainty sets with partially input-convex neural networks, which are learned as part of our framework. Our approach consistently improves upon two-stage estimate-then-optimize baselines on concrete applications in energy storage arbitrage and portfolio optimization.

Published in Transactions on Machine Learning Research (12/2025)
End-to-End Conformal Calibration for Optimization Under Uncertainty
Christopher Yeh"
cyeh edu nchristi@caltech. edu yww.@calleche edu adamw@ealtech, edu yyue@caltech. edu
Nicolas Christianson"
Alan Wu
Adam Wierman
Yisong Yue
Deportment of Computing and Mathematical Sciences
California Institute of Technology
Reviewed on OpenReview: https: // openreview. form?? fd=
Abstract
Machine learning can significantly improve performance for decision-making under uncertainty across wide range of domains. However, ensuring robustness guarantees requires well- calibrated uncertainty estimates, which can be difficult to achieve with neural networks. Moreover, in high-dimensional settings, there may be many valid uncertainty estimates, each with its own performance profile ie., not all uncertainty is equally valuable for downstream decision-making. To address this problem, this paper develops an end-to-end framework to learn uncertainty sets for conditional robust optimization in a way that is informed by the downstream decision-making loss, with robustness and calibration guarantees provided by conformal prediction. In addition, propose to represent general convex uncertainty sets with partially input-convex neural networks, which are learned as part of our framework. baselines on
Our approach consistently improves upon two-stage estimate-then-opiine conerete applications in energy storage arbitrage and portfolio optimization.
1 Introduction
Well-calibrated estimates of forecast uncertainty are vital for risk-aware decision-making in many real- world systems. For instance, grid-scale battery operators forecast electricity prices to schedule battery charging/ hiihng//dsshaaging to maximize profit, while accounting for forecast uncertainty to mitigate financial or operational risk. Similarly, financial investors use forecasts of asset returns with uncertainty estimates maximize portfolio returns while minimizing downside risk.
Traditional approaches to decision-making under uncertainty often follow "estimato--hhnn-ppiimaa (ETO) paradigm (Chenreddy et al., 2022), in which the uncertainty estimation and decision- stages decoupled. In the "estimation" stage, predictive model is trained to forecast a target quantity along with an uncertainty set. In the "optimization" stage, this uncertainty estimate informs a downstream decision. Crucially, the cost or performance the downstrcam decision is typically not fod back into the training of the predictive model.
A recent line of work (Chenreddy al., 2022; Patel al., 2024; Wang et al., 2024; Chenreddy & Delage, 2024) has made steps toward bridging the gap between uncertainty quantification and robust optimization-driven decision- decision-making, where optimization problems take forecast uncertainty set parameter, is common in energy systems (Yan et al., 2020; Parvar & Nazaripouya, 2022) and financial applications (Gregory et al., 2011; Quaranta & Zaffaroni, 2008). However, existing approaches are suboptimal for several reasons:
*These authors contributed equally.

---

Published in Transactions on Machine Learning Research (12/2025)
1. Estimate
Optimize
Estimate dhea Optimise (ETO)
(opticual) calibeaation
Taput
task aagnostti loss MSE, NLL, pitholl!
Input
ML morbel
uncertaiaty
robust opsimication
decision
maddel
fiz x1)
procedtuee
gradient de;d0
no back back from task loss
Our End- to End
Inpat
ML model 0
momsonformitt eoure function
conformal calibratiio pöck g)
(1-a)-eue set Ra(a) (0 aa(a,
calibrrttd
robrast optiniiontton min max [(2, (ecky()
decisson =g(x)
robust
fiz, 2 xg(a))
Framework (E2E)
s.t. 8(a,
gradient #11,
Figure 1: Whereas prior "estimateethhen-ppiiiaa (ETO, top) approaches separate the model training from the optimization (decision-making) procedure, we propose framework for end-to-end (E2E, bottom) conformal calibration for optimization under uncertainty that directly trains the machine learning model using gradients from the task loss.
1. Lack of decision-aware training: The predictive model is typically not trained with feedback from the reanstream objective. Because the downstream objective often asymmelric with respect to the forecasting model's errors, prediction models trained with decision-aguostar losses may achieve high predictive accuracy but still perform poorly on the decision-making objective.
2. Restricted uncertainty
parameterizations: To preserve tractability of the robust optimization problem, uncertainty sets often constrained to simple parametric forms (e.g., boxes ellipsoids), limiting the expressivity of uncertainty estimates.
3. Calibration challenges: Because neural network models are often poor at estimating their own uncertainty, the forecasts may not be well-calibrated. Recent approaches such as isotonic regression (Kuleshov et al., 2018) and conformal prediction (Shafer & Vovk, 2008) have made progress in providing calibrated uncertainty estimates from deep learning models, but such methods are typically applied post-hoc to trained models and are therefore difficult to incorporate into an end-to-end training procedure.
As such, there is of yet comprehensive methodology for training calibrated uncertainty-aware deep learning models end-to-end with downstream decision-making objectives. In this work, provide the first such methodology. We make three specific contributions corresponding to the three issues identified above:
1. We develop a framework for training prediction models end-to-end with downstream decision-making objectives and conformal-calibratee uncertainty sets the context of the conditional robust optimization problem. This framework illustrated in Figure (bottom). By including differentiable conformal calibration in our model during training, we close the loop and ensure that feedback from the uncertainty's impact the downstream objective is accounted for in the training process, since not all model errors nor uncertainty estimates will result in the same downstream cost. This end-to-end training enables the model to focus its learning capacity on minimizing error and uncertainty on outputs with the largest decision-makinng cost, with more leeway for outputs that have lower costs.
2. We propose using partially input-convex neural networks (PICNNs) as the nonconformity score function for conformal prediction, enabling the approximate parametrization of arbitrary compact, convex uncertainty sets in the conditional robust optimization problem. existing works PICNNs to parametrize such arbitrary convex
To the best of our knowledge, uncertainty sets. Due to the universal convex function approximation property these networks enjoy (Chen el al., 2019), this approach enables training much more general representations of uncertainly than prior works have considered, which in turn yields substantial improvements on downstream decision- making performance. Importantly, PICNNs are well-matched our conditional robust optimization problem: we show that the robust problem resulting from this parametrization can be reformulated as a tractable convex optimization problem.
2

---

Published in Transactions on Machine Learning Research (12/2025)
3. We propose exact and computationally efficient method to differentiate through the conformal prediction procedure during training. Unlike prior work (Stutz et al., 2022), our method gives exact gradients, without using approximate ranking and sorting methods.
Finally, we evaluate the performance of our approach on two applications: an energy storage arbitrage task and a portfolio optimization problem. We demonstrate that the combination of end-to-end training with the flexibility of the PICNN-based uncertainty sets consistently improves over ETO baseline methods. The performance benefit of our end-to-end method is apparent even under distribution shift.
2 Problem Statement and Background
suppose that data (ar, E IRm TR is sampled i.i.d. from an unknown joint distribution P. Upon observing the input (but not the label y), an agent makes a decision R". After the decision is made, the true label y is revealed, and the agent incurs task loss f(z, y, 2), for some known task loss function f R x R x R² + R. In addition, the agent's decision must satisfy set of joint constraints g(2, 2) S 0 coupling and 2.
Our problem defined formally as follows:
As an illustrative example, consider an agent who would like to minimize the costs of charging and discharging a battery over 24 hours in a day. The agent may use weather forecasts and historical observations x to prodict future energy prices y. Based on the predicted prices, the agent docides on the amount to charge or discharge the battery. The task loss f is the cost incurred by the agent, and the constraints g include limits how fast the battery can charge well as the maximum capacity of the battery. This example is explored in more detail in Section 5.
Because the agent does not observe the label prior to making its decision, ensuring good performance and constraint satisfaction requires that the agent makes decisions z that are robust to the various outcomes of y- A common objective is to choose z to robustly minimize the task loss and satisfy the constraints over all realizations of y within a (1 - -coonidienccee region [xx) C R" of the true conditional distribution Ply x), where (0, 1) is a fixed risk level based on operational requirements. In this case, the agent's robust decision can be expressed as the optimal solution to the following conditional robust optimization (CRO) problem (Chenreddy et al., 2022):
2² (x) arg min max f(z, 0.2 s.t. g(2,2) S O.
(1)
ZER" BER(a)
After the agent decides 2* (x), the true label y revealed, and the agent incurs the task loss f(x, y, 2 (x)). Thus, the agent seeks to minimize expected task loss
B( (,,,,,,(
(2)
While the joint distribution P is unknown, we assume that we have a dataset D = (( y:)) if.. samples
from P. Then, our objective is to train a machine learning model to learn an approximate (1 - a)-confidence set (zz) of possible y values for each input æ. Formally, our learned ((z) should satisfy the following maryinal coverage guarantee.
Definition 2.1 (marginal coverage). An uncertainty set 12(x)) for the distribution P provides marginal coverage at level (1 a) if P(ma))> (y E Q(x)) 2 1-a.
Comparison to related work. The problem of constructing data-driven and machine-learred uncertainty sets with probabilistic coverage guarantees for use in robust optimization has been widely explored in prior literature (e.g., Bertsimas & Thiele (2006); Bertsimas et al. (2018); Alexeenko & Bitar (2020); Goerigk & Kurtz (2023)). Chenreddy et al. (2022) first coined the phrase "conditional robust optimization" for the problem (1) and considered learning contexl-dependdet uncertainty sets 2(x) in this setting. However, their approach results in a mixed integer optimization that is intractable to solve for large-scale problems. al., 2023). As shown in Figure (top), the ETO paradigm separates the machine learning model training from the decision optimization. The lack of feedback from the downstream task loss during model training in ETO generally
Morcover, they follow the "estimate then optimize" (ETO) paradigm (Elmachtoub
3

---

Published in Transactions on Machine Learning Research (12/2025)
leads to uncertainty sets Q(z) which yield suboptimal results. Several other recent papers follow the ETO paradigm using homoskedastic ellipsoidal uncertainty sets (Johnstone & Cox, 2021), heteroskedastic box and ellipsoidal uncertainty sets (Sun et al., 2023), and a "union of balls" parametrization of uncertainty (Patel et al., 2024). In our experiments (Section 5), we demonstrate consistent improvements over the methods of Johnstone & Cox (2021) and Sun et al. (2023).
Closest to our work an "end-to-end" formulation of the CRO problem posed by Chenreddy & Delage (2024), which aims to learn conditional uncertainty sets (xx) using a weighted combination of the downstream task loss along with "conditional coverage loss" to promote calibrated uncertainty. However, they focus solely ellipsoidal uncertainty sets, and their conditional coverage loss does not provably ensure coverage for their learned uncertainty sets. In our experiments, we do not compare against Chenreddy & Delage (2024) because we only consider other methods with a provable coverage guarantee."
A concurrent related work by Wang et al. (2024) approaches the problem of learning unconditional uncortainty sets for robust optimization, while achieving robust constraint satisfaction guarantees. While their method also uses end-to-end task loss, we find their use of the same uncertainty set Zt for every problem instance (i.e., Q is independent of z) to be highly restrictive. For example, in the context of our battery control problem, this restriction would disallow the use of weather forecasts and historical price data to estimate uncertainty in future energy prices. 2 They also use restrictive uncertainty set parametrizations such as box, ellipsoidal, and polyhedral uncertainty. Moreover, their guarantee on constraint satisfaction requires an assumption on the convergence of their algorithm to an approximate stationary point; however, this convergence is only guaranteed asymptotically, and not in finitely many samples.
An alternative line of rescarch designs post-hoc conformal prediction sets aimed at minimizing decision costs in robust classification problems (Kiyani et al., 2025; Cortes-Gomez et al., 2024). However, to date, it is unclear how to adapt these methods, which produce discrete uncertainty sets, for regression settings.
In contrast, our work overcomes these limitations: we incorporate diffcrentiable conformal calibration during training to ensure that uncertainly is learned end-to-end a manner that is both calibrated and minimizes task loss. We apply split conformal post-hoc calibration during inference for provable guarantees on coverage. Furthermore, we use partially input-convex neural networks (Amos et al., 2017) to dircctly parameterize the nonconformity score function in conformal prediction, enabling general and expressive representation of arbitrary conditional convex uncertainty regions that can vary with or and be used efficiently in robust. optimization.
Beyond the above closely related work, this paper builds upon and contributes to several different areas in machine learning and robust optimization; see Appendix B for a comprehensive discussion.
3 End-to-End Training of Conformally Calibrated Uncertainty Sets
In this section, we describe our proposed methodological framework for end-to-end task-aware training of predictive models with conformally calibrated uncertainty for the conditional robust optimization problem (1). Our overarching goal is to learn uncertainty sets s(x)) which provide (1 a) coverage for any choice of a E (0, 1), and which offer the lowest possible task loss (2). To this end, we must consider three primary questions:
1. How should the family of uncertainty sets 2(x) be parametrized?
2. How can we guarantee that the uncertainty set M(x) provides coverage at level 1 a?
3. How can the uncertainty set 2(x) be learned to minimize expected task loss?
Figure 1 (bottom) illustrates the key parts of our framework to answer these questions. First, we use a machine learning model to parametrize a nonconformity score function 8g, and we define the uncertainty set (xx) to be a q-sublevel sot of s0(x, Second, we use conformal calibration to pick q to enforce marginal
As of the time of writing, the appreach of Chenreddy & Delage (2024) also suffers from substantial incomsistencies between their code implementation and the equations from their paper. In particular, the conditional coverage loss proposed in their it will (almost surely) yield gradients.
paper not implementable,
2 After the present paper's 's submission, the authors of Wang et al. (2024) updated their preprint and expanded their formulation to accommodatee conditional uncertainty sets.
4

---

Published in Transactions on Machine Learning Research (12/2025)
coverage. Third, we backpropagate gradients through both the robust optimization and conformal calibration steps to update the machine learning model, thereby enabling end-to-end learning. Sections 3.1 to 3.3 describe each of these parts in detail, and Algorithm shows pseudocode for both training and inference.
For the rest of the paper, we make the following assumptions on the functions f and g to ensure tractability of the resulting optimization problem.
Assumption 3.1. We assume the task loss has the form f(x, U, U, y F 2) + f(x, ()) where F(x, is an affine function of z and f(x,z) convex in 2, Furthermore, we assuine that g(x, z) is convex in z.
Technically, the assumption on I can be relaxed to the more general setting where / is a conic-representable saddle function that is convex in & and concave in y (Juditsky & Nemirovaki, 2022), which admits an automated dual representation (Schiele et al., 2023). However, we believe the core ideas for our exposition and proofs are much more straightforward with the stronger assumption in Assumption 3.1.
3.1 Representations of the uncertainty set We consider convex uncertainty sets of the form
Do(x) (0 E R" so 0) 9).
(3)
where so : R" R" R is an arbitrary nonconformity score function that is convex in D. q is a scalar, and 0 collects the parameters of model that we will seek to learn. Note that this representation loses generality; any family of convex sets Q(z) be represented as such collection of sublevel sets of partially input-convex function s(x,, 9). This particular representation is chosen due to the ease of calibrating sets of this form via conformal prediction to ensure marginal coverage, as we will describe in Section 3.2.
In choosing a particular score function 80, one must balance two considerations: first, the generality of the sets (x) that 89 can represent, and second, the tractubility of the resulting robust optimizattion problem (1). We will now show that our representation (3) generalizes commonly-used box and ellipsoidal uncertainty sets, which are known to have tractable robust problems; later, in Section 4, we will propose to approximate more general convex uncertainty sets using partially input-convex neural networks.
Box uncertainty sets. A simple uncertainty representation is box uncertainty where 2(x) [y(x), ((²) is an n-dimensional box whose lower and upper bounds depend on Let he 1 - Ih x IR" be a neural network that estimates lower and upper bounds: ho(x)) = (Nl (x), hai (z)). To represent box uncertainty set in the form (3), we define a nonconformity score function that generalizes scalar conformalized quantile regression (Romano et al., 2019):
so(r,0) max([|n3((] - v vll.. 1ly nb (z)
Then, the uncortainty set (3) becomes
Oo(x)= [x) - q1, hy (x) q1] =: (g(x), v(a)].
Given a box uncertainty set (z),,, we can take the dual of the inner maximization problem (see Appondix C.1) to transform the robust optimization problem (1) into an equivalent form that is convex, and hence tractable, under Assumption 3.1:
(x) = arg min min DERn
(v(x) - y(x)) v + y(x) T F(x, 2) + f(x, 2)
=ERe
(4)
s.t. v 0, V F(z, 2) > 0, g(x, 2) <
Ellipsoidal uncertainty sets. Another common form of uncertainty set is ellipsoidal uncertainty. Suppose a neural network model he R + R² x S predicts mean and covariance parameters he(x) = (ua(x) Me(x)), that P(y) x; 0) = A1 po(x), Ex)) denotes predicted conditional density, where N(: 14, E) is the
5

---

Published in Transactions on Machine Learning Research (12/2025)
Algorithm Endd-o-end conformal calibration for robust decisions under uncertainty
function TRAIN(training data D {(T,, vi) uncertainty level a, initial model parameters 8)
for mini-batch B [L....,N) N} do Randomly split batch: == (Beal, Bpred) Compute 4 QUANTILE( (2i, Mt) hie 1 1-a)
for E Bpred do
Solve for robust decision =x;) using (4), (5), or (8) Compute gradient task loss: d8, = 0f Vi, 20 (z,))/
Update using gradients EieByro d0,
function INFERENCE(modee parameters 0, calibration data Deal ((z,, M uncertainty level a, input x)
Compute q QUANTILE( (sa
(Ø, (2,0)))) (..))) 1-a)
return robust. decision (x) using (4), (5), or (8)
function QUANTILE(SCOres S = ( level B)
((M))))))(()+))) = SORT ASCENDING(S (+00))
return 8([(M+1))81)
multivariate normal density function. In this case, we define the nonconformity score function based on the squared Mahalanobis distance (Johnstone & Cox, 2021; Sun et al., 2023)
sa(z, v) (y us (x)) (Eo(z)) (()))))) uz())),
which yields uncertainty sets (3) that are ellipsoidal:
(x( (0 (0 (o(x)) ((o))))((=))))))))))))) ( S 5t.
Let Lz(x) denote the unique lower-triangular Cholesky factor of [x (i.c., Eo(x) = Lo( Taking the dual of the inner maximization problem and invoking strong duality (see Appendix C.2), we transform the robust optimization problem (1) into an equivalent form that convex, and hence tractable, under Assumption 3.1:
(x)x) arg min 2E3P
Lo(x) F(z, 2(2,2)ll2 " F(z,2) f(x,2)
(5)
s.t. 2x 50
3.2 Conformal uncertainty set calibration
As long as the uncertainty set Qg(x) can be expressed in the form (3), we can use the split conformal prediction procedure at inference time to choose value q that ensures lo(x)) provides marginal coverage (Definition 2.1) at any confidence level 1 - The split conformal procedure assumes access to calibration dataset Deal = {(T, M drawn exchangeably from P. We refer readers to Angelopoulos et al. (2023) for details on this procedure.
Lemma 3.2 (from Angelopoulos et al. (2023), Appendix D). Let Dest = [1x v.) be a calibration dataset drown ezchangeably (e.g., i.i.d.) from P, and let s; = so V:). If = (s. 1 a) (see Algorithm 1) is the empirical [(M+1)([-o]] -quantile of the set {s MM and (z, y) is drawn exchangeably with Deal, then (2) has the marginal coverage guarantee
a Pay.Daa Ro(x)) $ a M+1
We split conformal prediction, rather than full conformal prediction, both for computational tractability and to avoid the problem of nonconvex uncertainty sets that arise from the full conformal approach, noted in Johnstone & Cox (2021). For the rest of this paper, we assume E ( so that QUANTILE(ss,MM, 1 a) oo is finite. Thus, for appropriate choices of the score function se, the uncertainty set Po(x:) is not unbounded.

---

Published in Transactions on Machine Learning Research (12/2025)
While the split conformal prediction procedure in Lemma 3.2 ensures that the uncertainty set Qo(x) satisfies (1 a) coverage at inference time, this process does not address the question of Iraining the uncertainty set (z) (via the score function se) to ensure optimal task performance while maintaining coverage. In Section 3.3, we propose applying a separate differentiable conformal prediction procedure during training to address this challenge.
3.3 End-to-end training and calibration
Thus far, we have discussed how to calibrate an uncertainty set Qz(x)) of the form (3) to ensure coverage, and we described two choices of score function so parametrizing common box and ellipsoidal uncertainty sets. However, to ensure that the uncertainty sets (2(x)) both guarantee coverage and ensure optimal downstream task performance, it is necessary to design an end-to-end training methodology that can incorporate both desiderata in fully differentiable manner. We propose such a methodology in Algorithm 1.
Our end-to-end training approach minimizes the empirical task loss (()) = E 4()) using minibatch gradient descent, where (0)) fa V/, 2g (x;)). This requires differentiating through both the robust well the conformal prediction step. The gradient of the lask loss on a single af (I; Vr, (x;)) 0zg
optimization problem instance is de, = 0z d0
0zz 80 (a:), where 2o (2;) is computed by diffcrentiating through the
Karush Kuhn Tucker (KKT) conditions of the convex reformulation of the optimization problem (1) (i.e., the problems (4), (5)) following the approach of Amos & Kolter (2017), under mild assumptions on the differentiability of and g. Note that the gradient of any convex optimization problem can be computed with respect to ils parameters such (Agrawal et al., 2019, Appendix B).
To include calibration during training, we assume that for every (z, y) in our training set, sø(x, y) is differentiable w.r.t. f almost everywhere; this assumption holds for common nonconformity score functions, including those used in this paper. We then adopt the conformal training approach (Stutz et al., 2022) in which a separate g is chosen in each minibatch, as shown Algorithm 1. The chosen q depends on azū
Øzū 2g 0 (through 80), and (x:) depends on the chosen Therefore 20 involves calculating By 00 where dg requires differentiating through the empirical quantile function. Whereas Stutz et al. (2022) uses smoothed approximate quantile function for calculating 9. we find the smoothing unnecessary, as the gradient of the empirical quantile function is unique and well-defined almost everywhere. Importantly, our exact gradient. is both more computationally efficient and simpler to implement than the smoothed approximate quantile approach. See Appendix for more details.
50
After training has concluded and we have performed the final conformal calibration step, the resulting model enjoys the following theoretical guarantee on performance (ef. Sun et al. (2023, Proposition 1)).
Proposition 3.3. Under the same assumptions Lemma 3.2, the task loss satisfies the following bound with probability at least 1 - (over I,y, y, and the calibration set Deat):
2(x,;) S z0)
((,,,,2(()))
min max ((x s.t. g(0, 2) SERP (Efa(a)
(6)
Proof. When y E [2),,, the definition of (z) (1) guarantees that (6) holds. By Lemma 3.2, F.y.DDall (y e (x(x)) 2 1 a.
We note that in practice, training models from scratch using only the task loss sometimes leads to poor local optima. Instead, in our experiments, we pretrain our models with a standard loss function (e.g., negative ligeliielihood for the mean-variance predictor used in ellipsoidal uncertainty sets), and then we fine-tune the model using the task loss. This is a standard heuristic commonly used in the decision-focused learning and predict- then-optimize literature (Donti al., 2017; Mandi al., 2024). Appendix provides more details on our training procedure.
7

---

Published in Transactions on Machine Learning Research (12/2025)
2
1
box/ellipse z'
(task loss: 0)
PICNN z*
0
(task loss: -0.5)
-1
-1
0
1
Figure 2: Consider a robust portfolio optimization problem with assets, where y E R is a random vector of asset returns, and the decision = IR² represents portfolio weights: max, mingen = - 9 s.t. 2 2 0, 1 1 S 1. Let the distribution of asset returns y be uniform over 3 discrete points (black). The optimal box (blue), ellipse (orange), and PICNN (green) uncertainty sets shown with their robust decision vectors =². The flexibility of the PICNN uncertainty representation allows it to achieve the lowest expected task loss.
4 Representing General Convex Uncertainty Sets via PICNNS
The previous section discussed how to train calibrated box and ellipsoidal uncertainty sets end-to-end to optimize the ream task loss. However, both box and ellipsoidal uncertainly sets have restrictive shapes which may yield suboptimal task performance. If (2)x)) could represent any arbitrary convex uncertainty set, this more expressive class would enable obtaining better task loss. Figure 2 illustrates an example where a general convex uncertainty set representation provides a clear advantage over box and ellipsoid uncertainty.
To this end, we propose directly learn partially-convex nonconformity score function 59 F R" + E that is convex only in the second input vector. Fixing x, any q-sublevel set {Ø R" s0(2, 9) q} of so is convex set, and likewise every family of convex sets can be expressed the q-sublevel sets of some partially-convex function. To implement this approach, we are faced with two questions.
1. How should we parametrize the score function 89 so aa(a)) can approzimate arbitrary conver sets? A natural answer to parametrize sy with a partially input-convex neural network (PICNN) (Amos et al., 2017), which can efficiently approximate any partially-convex function (Chen et al., 2019). We consider a PICNN defined as so(2,)) y) WLoL +
VLD ,y, where
op 0, uo I,
W₁ Wi diag([Wru wi]+)
#4+ 1 = ReLU (Rau + rI),
Vi = Vi diag (Vrun + vt)
0 = ReLU (Wio: + Viy + bi),
by = Biu + bt.
with weights 0 = (Ri, rı. Wi. Wi, wi, Vi, Vi, uz Bi, b
The matrices Wi are constrained to be entrywise nonnegative to ensure convexity of se with respect to For ease of notation, we assume all hidden layers o1....,OL have the same dimension
2. Does the chosen parumetrization of a2(x)) (via PICNNs) yield a tractable reformulation of the CRO problem (1)? Fortunately, we show in the following theorem that the answer is yes.
Thcorem 4.1. Let Oo(x)= {6 R so(x, 0) S q), where so is PICNN as defined in (7). Then, under A Assumplion 3.1, the CRO problem (1) with uncerlainly set (a( is equivalent to the following convex (and hence tractable) minimization problem:
2 arg min min
b(0,q) v 10++fe f(x,2)
ZERP
vERaLd-1
(8)
A(O) [F z)
s.t.
v 2 0, g(x, x) 50
where A(8) ((2 and b(0,g) q) B 2ld+1 are constructed from the weights of the PICNN (7),
and b also depends on q-
8

---

Published in Transactions on Machine Learning Research (12/2025)
We prove Theorem 4.1 in Appendix C.3; the main idea is that when (g(x) is sublevel set of a PICNN, we can equivalently reformulate the inner maximization problem in (1) as linear program and take the dual to yield a tractable minimization problem.
Since the PICNN uncertainty sets are of the same form as (3) and yield a tractable convex reformulation (8) of the CRO problem (1), we can apply the split conformal procedure detailed Section 3.2 to choose & E R and obtain coverage guarantees on Og(z), and we can employ the same end-to-end training methodology from Section 3.3 to train calibrated uncertainties end-to-end using the downstream task loss. In some cases during training, the inner maximization problem of (1) with PICNN- MCCNN-parametrized uncertainty set may be unbounded (if (x(x) is not compact) or infeasible (if the chosen q is too small causing (z) to be empty). This will lead, respectively, to an infeasible or unbounded equivalent problem (8). We can avoid this concern by adjusting the PICNN architecture to ensure its sublevel sets are compact and by suitably increasing q when needed to ensure (x) never empty. Such modifications do not change the general form of the problem (8) and preserve the marginal coverage guarantee for the uncertainty set l2((2); see Appendix C.4 for details.
5 Experiments
In this section, present experimental results for our F2E method against several ETO baselines. Code to reproduce our results is available on GitHub..
5.1 Problem descriptions
We consider two tasks: price forecasting for battery storage operation and portfolio optimization. Their task loss functions and constraints satisfy Assumption 3.1.
Price forccasting for battery storage. This problem comes from Donti et al. (2017), where a grid-scale battery operator predicts electricity prices E Rº over a T-step horizon and uses the predicted prices to decide a battery charge/discharge schedule for price arbitrage. The input features x include the past day's prices and temperature, the next day's energy load forecast and temperature forecast, binary indicators of weekends or holidays, and yearly sinusoidal features. The operator decides how much to charge (2" E R²) or discharge (zout E ) the battery, which changes the battery's 's state of charge (2state {R The battery has capacity B, charging efficiency 7, and maximum charging/ discharging rates eln and cout, The task loss function represents the multiple objectives of maximizing profit, flexibility to participate in other markets by keeping the battery near half its capacity (with weight A), and battery health by discouraging rapid charging/ discharging (with weight c):
T f(y.2) = Σ ye(2" - zut )2 + X
B -1 1
z
state
+
+ e||:
ont
2
1=1
The constraints are, for all b= :,....
20
state
B/2,
2
state 4
21 + 724
5 t < d
0 < z out
< out
0 < 2/
state < B.
Following Donti al. (2017), we sct T 24, 1, 7 0.9, cn 0.5, out 0.2, x 0.1, and = 0.05.
Portfolio optimization. We adopt the portfolio optimization setting and synthetic dataset from Chenreddy & Delage (2024), where the prediction targets y E I=" are the returns of a set of n sccurities, and the decision RD sets portfolio weights. The task loss is f(y =) = -y z, with constraints 2 0, _ == 1. The synthetic dataset consists of (z, y) E __22 drawn from a mixture of three 4-D multivariate normal distributions. We provide data details in Appendix E.2 and experimental results in Appendix A.2. The results are similar to those for battery storage, except that portfolio optimization is a diver-dimeniioaal and casier problem.
hhttpe://github.u com/ www.eeeeereeeeeeeeeeece htt.:/hhww..ccerrrrreeee///eeeee conformmal
9

---

Published in Transactions on Machine Learning Research (12/2025)
Box
Ellipse
PICNN
task laes
ETO
-20
ETO-JC
distribution shift
ETO-SLL
optimal
0.01
0.05
0.1
0.01
0.1 uncertainty level
0.05
0.01
0.1 uncertainty level
0.05
uncertalntyy level
Box
Ellipse
PICNN
task loss
ETO
with distribution shift
ETO-JC
20
ETO-SLL
E2E
optimal
-30
0.01
0.05
0.1
0.2 0.01
0.05
0.1
0.2 0.01
0.05
0.1
0.2
uncertainty level
uncertainty level
uncertainty level
Figure 3: Task loss performance (mean +1 stddev across 10 runs) for the battery storage problem with no distribution shift (top) and with distribution shift (bottom). Lower values are better.
5.2 Baseline methods
We implemented several "estimate-thennoptiiiz (ETO) baselines, listed below, to compare against our end-to-end (E2E) method. These two-stage ETO baselines are trained using lask-agnostic losses such as pinball loss or negative lig-likelihood (NLL). To ensure fair comparison against our E2E method, we also apply conformal calibration to each ETO method after training to satisfy coverage.
ETO denotes models with identical neural network architectures to our E2E models, differing only in the loss function during training. The box uncertainty ETO model is trained with pinball loss to predict the ž and 1 - Σ quantiles. The ellipsoidal uncertainty ETO model is trained with multivariate normal negative log-likelibae likelihoooo loss. For the PICNN ETO model, we train sø using a negative log-likelihood loss by interpreting 89 as an energy function -.e. Poly xr) o exp(-s0(x; y)) yielding the loss NLL(O)= = In sg(2, In Zz()), where Zo(x) SyER= exp(-se(13 ū)) dÿ, following the approach of Lin Ba (2023). More details of the ETO models are given in Appendix E.
ETO-SLL is our implementation of the box and ellipsoid uncertainty ETO methods from Sun et al. (2023). Unlike ETO, ETO-SLL first trains point estimate model (without uncertainty) with mcan-squared error loss. Then, ETO-SLL box and ellipsoidal uncertainty sets are derived from training a separate quantile regressor using pinball loss to predict the (1 a)-quattiles of absolute residuals fz-norm of residuals of the point estimate. Unlike ETO which can learn cllipsoidal uncertainty sets with different covariance matrices for each input the ETO-SLL ellipsoidal uncertainty sets all share the same covariance matrix (up to scale).
ETО-JC our implementation the ellipsoid uncertainty ETO method by Johnstone Cox (2021). Like ETO-SLL, ETO-JC also first trains a point estimate model (without uncertainty) with mean-squared error loss. ETO-JC uses the covariance matrix (with the same scale) for each input
5.3 Battery storage problem results
Figure 3 (top) compares task loss performance for different uncertainty levels (a E (-01, 10,, .2] and the different uncertainty set representations for the ETO baselines against our proposed E2E methodology on the battery storage problem with no distribution shift. Our E2E approach consistently yields improved performance over all ETO basclines, for all three uncertainty set parametrizations, and over all tested uncertainty levels a. Moreover, the PICNN uncertainty representation, when trained end-to-end, provides up to 42% relative improvement in performance the best ETO box uncertainty set and up to 209%
10

---

Published in Transactions on Machine Learning Research (12/2025)
Box
Ellipse
PICNN
1.0
coverage
0.9
ETO
ETO-JC
distribution
0.8
ETO-SLL
shift
0.7
optimal
0.01
0.05
0.1
0.2
0.01
0.1 uncertainty level
0.05
0.2
0.01
0.1 uncertainty level a
0.05
uncertainty level
Box
Ellipse
PICNN
1.0
with distribution shift
coverage
0.8
ETO-JC
ETO
ETO-SLL
0.6
E2E
optimal
0.01
0.1 uncertainty level a
0.05
0.2 0.01
0.05
0.1
0.2 0.01
0.05
0.1 0.2
uncertainty level
uncertainty level a
Figure 4: Coverage (mean +1 stddev across 10 runs) for the battery storage problem with no distribution shift (top) and with distribution shift (bottom). The dotted black line indicates the target coverage level a. Our E2E models achieve similar coverage the ETO baselines, confirming that the lower task loss of our F2E models does not come at the expense of worse coverage.
relative improvement over the best ETO ellipse uncertainty set. We also show the corresponding coverage obtained by the learned uncertainty sets Figure (top); all models obtain coverage close to the target level, confirming that the improvements task loss performance from our E2E approach do not come at the cost of worse coverage.
In Appendix A, we present additional results on the battery storage problem, including value-at-risk (VaR) and conditional value-at-risk (CVaR) metrics which show that the E2E approach consistently outperforms the ETO baselines on tail risk as well. In addition, Table A1 shows the per-epoch computational time required by our F2E approach compared with the two-stage baselines. While E2F requires additional training time, the training time dominated by time spent solving the conditional robust optimization problem (1) for every training example at every epoch. The development of GPU-acoelerated solvers (Blin, 2024) may help significantly speed up the optimization, though we do not explore GPU-accelerated optimization in our implementation.
5.4 Performance under distribution shift
The aforementioned results were produced without distribution shift, where our training and test sets were sampled uniformly at random, thus ensuring exchangeability and guaranteeing marginal coverage. In this section, we evaluate our method on the more realistic setting with distribution shift by splitting our data temporally: our models are trained on the first 80% of days and evaluated on the last 20% of days. Figure A4 plots the underlying time series data and visually highlights the distribution shift. Figures 3 (bottom) and (bottom) mirror Figures (top) and 4 (top), except that there is now distribution shift. We again find that our E2E approach consistently yields improved performance over all ETO basclines, for all three uncertainty set parametrizations, and for all tested uncertainty levels Likewise, the PICNN uncertainties, when trained end-to-end, improve on the performance offered by box and cllipsoidal uncertainty. We unsurprisingly find that, under distribution shift, the models do not. provide the same level of coverage guaranteed in the i.i.d. case, as the exchangeability assumption needed for conformal prediction no longer holds. The ellipsoidal and PICNN models tend to provide worse coverage than the box uncertainty, which we believe reflects how ellipsoidal and PICNN uncertainty sets offer greater representational power, and thus might be fitting too closely to the pre-shift distribution, which impacts robustness under distribution shift. Devising methods to
11

---

Published in Transactions on Machine Learning Research (12/2025)
anticipate distribution shift when training these more expressive models, and in particular the PICNN-basod uncertainly, remains an interesting avenue for future work.
6 Conclusion
In this work, we develop the first end-to-end methodology for training predictive models with uncertainty estimates (with calibration enforced diffcrentiably throughout training) that are utilized in downstream conditional robust optimization ion problems. We demonstrale an approach utilizing partially input-convex neural networks (PICNNs) to represent general convex uncertainty regions, and we perform extensive experiments on battery storage application and a portfolio optimization task. Whereas prior works on two-stage cehmate-then-optinie approaches emphasized "the convenience brought by the disentanglement of the prediction and the uncertainty calibration" (Sun et al., 2023), our results highlight that such "convenience" comes at substantial cost; end-to-end approach, combined with the expressiveness of the PICNN reprosentation, has clear performance gains over the traditional two-stage methods.
A number of interesting directions for future work on learning decision-aware uncertainty in an end-to-end manner remain. First, while our PICNN-based uncertainty set representation allows the parametrization of general convex uncertainty sets, future work may explore noncowver uncertainty regions. Doing so may require eschewing the analytical methods for differentiating through convex optimization problems and instead use, c.g.; policy gradient methods for passing gradients through general stochastic and robust optimization problemss. Second, developing end-to-end methods to larget conditional calibration (as opposed to marginal calibration) remains an active research direction. Finally, one may explore other types of constraints besides uncertainty set-based robustness, such as value-at-risk (VaR) or conditional value-at-risk (CVaR) constraints.
Acknowledgments
We thank Priya Donti for helpful discussions. The authors acknowledge support from an NSF Graduate Research Fellowship (DGF-2139433); NSF Grants CNS-2146814, CPS-2136197, CNS-2106403, and NGSDI- 2105648; gifts from Amazon, OpenAl, and Latitude Al; and the Caltech Resnick Sustainability Institute.
12

---

Published in Transactions on Machine Learning Research (12/2025)
References
Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and J. Zico Kolter. Differentiable Convex Optimization Layers. In Advances in Newral Information Processing Systems, volume 32, Curran Associates, Inc., 2019. URL https://papers: wapp..c/paaee httpppppee.illleppppppee///////e fwwess/aper/e////// Abstract. html.
Polina Alexeenko and Eilyan Bitar. Nonparametric Estimation of Uncertainty Sets for Robust Optimization. In 2020 59th IEEE Conference on Decision and Control (CDC), pp. 1196 1203, Jeju, Korea (South), December 2020. IEEE. ISBN 978-1-72817-447-1. doi:
202099333822 URL https:
//ieeexplore.eeo leee. wrg//bccmeent/9336663/
Brandon Amos and Zico Kolter. OptNet: Differentiable Optimization as a Layer in Neural Networks. Proceedings of the 34th International Conference on Machine Learning, pp. 136 145. PMLR, July 2017. URL attps://procest
mlr. wwwps//0//ammsss//..eppp html. ISSN: 2640-3498.
Brandon Amos, Lei Xu, and J. Zico Kolter. Input Convex Neural Networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 146 155. PMLR, August 2017. URL https::/prroceednngs mlr, SALL..peeas/v000 ht////00//aaa/ amos/mm. html.
Anastasios N. Angelopoulos, Stephen Bates, Adam Fisch, Lihua Lei, and Tal Schuster. Conformal Risk Control, April 2023. URL http///axvvvvv. wwwrxiiv..gggae.o wwwp///////////eeeeeeeeeeee/ 02814. MEXiv:2288.02814 [cs, math, stat].
Aharon Ben-Tal Laurent El Ghaoui, and Arkadi Nemirovski. Robust Optimization. Princeton University Press, 2009. ISBN 978-0-691-14308-2 URL https://www.jjsttt www.staalle..ctt//...... ctt7sk8p.
Dimitris Bertsimas and Aurélie Thiele. Robust and Data-Driven Optimization: Modern Decision Making Under Uncertainty. In Models, Methods, and Applications for Innovative Decision Making, INFORMS TutORials in Operations Research, pp. 95-122. INFORMS, September 2006. ISBN 978-1-877640-20-- doi: 10.1287 1063.c...030022... URL. https://pubsonl ine informs
infform..rg//o//f. wwtpoo/g////oooooooo httoo///a/////// 1ttp://abb//1...//eeeleee 12w7/educ. 0063.0020
Section: 4.
Dimitris Bertsimas, Eugene Litvinov, Xu Andy Sun, Jinye Zhao, and Tongxin Zheng. Adaptive Robust Optimization for the Security Constrained Unit Commitment Problem. IEEE Transactions on Power Systems, 28(1):52 63, February 2013. ISSN 1558-0679. doi: 21109/PPWSS.002.2222
URL
attps://ieeexplore. www.llorreeee..oooo oeg.eeeeeee.oeeeeeeoeo /tt///dccueen//2411... //0.////.......ee.ooooooo//.//.. Conference Name: IEEE Transactions on Power Systems.
Dimitris Bertsimas, Vishal Gupta, and Nathan Kallus. Data-driven robust optimization. Mathematical Programming, 167(2) 2335 292, February 2018. ISSN 1436-4646. doi: 10.1007/110107-017-11255 URL
https://doi.org: 10. 1100/0100000010000000000122
Nicolas Blin.
Accelerate
Large
Linear Programming
Problems with
NVIDIA cuOpt, Law.aarttte-aare--ie----------oe
October 2024.
URL https: / http://evvlloor..vviia..ooo www.iiiiiiii....m.. /
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight Uncertainty in Neural Networks. Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1613 1622. PMLR, May 2015. URL https://proceeee mlr. www.press/vve /v37// v377blumdee115,
Yize Chen, Yuanyuan Shi, and Baosen Zhang, Oplimal Coutrol Via. Neural Networks: A Convex Approach. nwt.:/oruetrr
In International Conference on Learning Representations, 2019. URL id=HLMW72ACK7.
Abhilash Reddy Chenreddy and Erick Delage. End-to-end Conditional Robust Optimization. In Proceedings of the Fortieth Conference on Uncertainty Artificial Intelligence, pp. 736-748. PMLR, September 2024. http:///nn.......ees//////../ http://244//ceeee///44.a.4a html. TSSN: 2640-3498.
URL https:// httpae////pocceeegg...o mlr. mlr,
13

---

Published in Transactions on Machine Learning Research (12/2025)
Abhilash Reddy Chenreddy, Nymisha Bandi, and Erick Delage. Data-Driven Conditional Robust Optimization. In Advances in Newral Information Processing Systems, volume 35, pp. 9525- 9537, December 2022. URL Abstract- html.
httppa/pppeeeirpppppppppeeeeeeee httppppeetileee/pper////// haash/
Nicolas Christianson, Lucien Werner, Adam Wierman, and Steven Low. Dispatch-awarg planning for feasible power system operation. Electric Power Systems Research, 212:108597, November 2022. ISSN 0378- 7796. doi: 90378779622006733.
URL http://www.iiencciireeeeee wcmm/ccince/aatiileeeeeeeee
Santiago Cortes-Gomez, Carlos Miguel Patiño, Yewon Byun, Steven Wu, Fric Horvitz, and Bryan Wilder. Utillty--iieeted Conformal Prediction: A Decision- Aware Framework for Actionable Uncertainty Quan- tification. In The Thirteenth International Conference on Learning Representations, October 2024. URL https://openrevieww wtttforrummtc---00mnnamaaa
Priya L. Donti, Brandon Amos, and Zico Kolter. Task-based End-to-end Model Learning Stochastic Optimization. In Advances in Neural Information Processing Systems, volume 30, Long Beach, CA, USA, December 2017. Curran Associates, Inc. URL http://papers. tas-ttass--------- -as-d--oot-----d-mmooll wwweeeeeannnnnneee htwp:miminnnsstoccacttooooootiio http:/pppprsn.uuuss..ppp cwtuurps.cc/pppeeee
Yury Dvorkin. A
Stochastic Electricity Market. IEEE Transactions on Power Systems, 35 (4):2993 3003, July 2020. ISSN 1558-0679. doi: 10.1109/ N0.1109/TPWWRS TTPPRS.20109 20199296129111..... URL https: //ieee///eee..ee...... ieee. orgo httg/documentt/890 8937824 Conference Name: IEEE Transactions on Power Systems.
Bat-Sheva Einbinder, Yaniv Romano, Matteo Sesia, and Yanfei Zhou. Training Uncertainty-Awer Classifiers with Conformalized Deep Learning. Tn Advances in Neural Information Processing Systems, New Orleans, URL
LA, USA, November 2022. doi: 10.48550 /atpe//riiiv.222 2205.058755 arXiv:2205.05877 [cs, stat].
05878.
Adam N. Elmachtoub and Paul Grigas. Smart "Predict, then Optimize". Management Science, 68(1):9-26, January 2022. ISSN 0025-1909. doi: 22287/mssc.2000.3922 URL sttpss//pubsonliin informs.oo.... org/ doi/10.
2020. 3922. Publisher: INFORMS.
Adam N. Elmachtoub, Henry Lam, Haofeng Zhang, and Yunfan Zhao. Estimate- ThenattteThe---ee ThinpptttTTe---Opimizeeee versus - inntraatte--sttiintiiooon--- Optimizaiinnn versus Sample Average Approximation: A Stochastic Dominance Perspective, August 2023. URL http://arxiv. http://aaivvv..g/bbbbeeee/ arg/ 06833. arXiv:2304.06831 [cs, stat].
Yarin Gal and Zoubin Ghahramani. Dropout Bayesian Approximation: Representing Model Uncertainty in Deep Learning. In 150i: 021022 2 fes, stat], October 2016. URL
02142.
arXiv: 1504..
Mare Goerigk and Jannis Kurtz. Data-driven robust optimization using deep neural networks. Computers & Operations Research, 151:106087, March 2023. ISSN 0305-0548. doi: 10. wtt.//..co...02..11.... URL http://twt.....ccccc.....ooooaaa
Christine Gregory, Ken Darby-Dowman, and Gautam Mitra. Robust oplimization and portfolio selection: The cost of robustness. European Journal of Operational Research, 212(2):417-428, July 2011. ISSN 0377-2217. doi: 10.1016/j.ccor.201 pi1//03377221110004555 2010//..jo.2201..2.0011 URL https: http:///wwwwwsweeeddeeeeee. wtt.:/cceeacdrreee..cccccceeeeee www.sciiece//arecccccc/c// articc//
Chancellor Johnstone and Bruce Cox. Conformal uncertainty sets for robust optimization. In Proceedings of the Tenth Symposium on Conformal and Probabilistic Prediction and Applications, pp. 72-90. PMLR, September 2021. URL httpss//proceeeding/ mlr. htt.p//eeeaaaaa......rreee////aa
ISSN: 2640-3498.
Anatoli Juditsky and Arkadi Nemirovski. On well-structured convex concave saddle point prob- lems and variational inequalities with monotone operators. Optimization Methods and Soft- ware, 1666)116671602, September 2022. ISSN 1055-6788.
1000//000008882221192111 doi: 100000/0056788.0000 2021.1928111.
URL http://22..or//0000000005777/..
10000/11...100010555.111 2021. 1928121.
Publisher: Taylor & Francis _eprint:
14

---

Published in Transactions on Machine Learning Research (12/2025)
Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations, San Diego, CA, USA, May 2015. URL. http://arxiv.org/// htsp//bb/14122/ 6980.
Shayan Kiyani, George Pappas, Aaron Roth, and Hamed Hassani. Decision Theoretic Foundations for Conformal Prediction: Optimal Uncertainty Quantification for Risk-Averse Agents, February 2025. URL http://arxiv. http:///aaxiv..ogaaabba aba /2022.250.... 02561. arXiv:2502.0256 [cs].
Volodymyr Kuleshov, Nathan Fenner, and Stefano Ermon. Accurate Uncertainties for Deep Learning Using Calibrated Regression. In Proceedings of the 35th International Conference on Machine Learning, pp. 2796 2804. PMLR, July 2018. URL https:// httpse//procceednng.. mlr, pprsss/v0/kulest/ii 2640-3498.
html. TSSN:
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and Scalable Predictive Un- certainty Estimation using Deep Ensembles. In Advances in Newral Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedingg. 90720440000220101/fffssataoaooee -aterraaebbee---btraatttt. html.
saarooeednggg.nuuiss...oooo neurips. httppipppsccpppppe//01///////// /aah///hh///////////////
Alexander Lin and Demba E. Ba. How to Train Your FALCON: Learning Log- Cog-Concave Densities with Energy-Based Neural Networks. In Fifth Symposium on Aduances in Approzimate Bayesion Inference, July 2023. URL https:///penreee
net//orum?iddUP11---tt
Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan. Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness. volume 33, pp. 7498 7512. Curran Associates, Inc., 2020. URL httpp::/prooceee//eee Awttrf-bbbtttttttee .tml.
Jayanta Mandi, James Kotary, Senne Berden, Maxime Mulamba, Victor Bucarey, Tias Guns, and Ferdinando Fioretto. Decision-Focusel Learning: Foundations, State of the Art, Benchmark and Future Opportunities. Journal of Artificial Intelligence Research, 80:1623-1701, August 2024. ISSN 1076-9757. doi: 10.1613/jair. 1.15320. URL https: htwp://wwwwwwo jair owgg/ingees httpa//ww.ppppeaaaaaaaaaae/////e
Zachary Nado, Neil Band, Mark Collier, Josip Djolonga, Michael W. Dusenberry, Scbastian Farquhar, Qixuan Feng, Augelos Filos, Marton Havasi, Rodolphe Jenatton, Ghassen Jerfel, Jeremiah Liu, Zelda Mariet, Jeremy Nixon, Shreyas Padhy, Jie Ren, Tim G. Rudner, Faris Sbahi, Yeming Wen, Florian Wenzel, Kevin Murphy, D. Sculley, Balaji Lakshminarayanan, Jasper Snock, Yarin Gal, and Dustin Tran. Uncertainty Baselines: Benchmarks for Uncertainty & Robustness in Deep Learning. January 2022. doi: 200.5500/rXii.2106..00
URL
org.abs//////////////////.. htrg/abs/2106 2206 2106.04015. 04015. arXiv:2106.04015 [es].
Mariola Ndrio, Avinash N. Madavan, and Subhonmesh Bose. Pricing Conditional Value at Risk-Sensitive Economic Dispatch. In 2021 IEEE Power & Energy Society General Meeting (PESGM), pp. 01 05, July 2021. doi: 10.1109/ 10.1109/PESGM4666 9637845. ISSN: 1944-9933.
URL https: bttts://ieeexpllee ieee. org/document//
Arkadi Nemirovski and Alexander Shapiro. Convex Approximations of Chance Constrained Programs. SIAM Journal on Optimization, 17(4):969-99 996, January 2007. ISSN 1052-6234, 1095-7189. doi: 10.1137 2.1117/050022328.
URL http://epubs siam. http//oo/11..117/05000000
Scyed Shahin Parvar and Hamidreza Nazaripouya. Optimal Operation of Battery Energy Storage
Under Uncertainty Using Data-Driven Distributionally Robust Optimization. Electric Power Sys- tems Research, 211:108180, October 2022. ISSN 0378-7796. doi: 10.1016/j.epsr 200....022...111800 URL
http::://www..ciencccc
Yash P. Patel, Sahana Rayan, and Ambuj Tewari. Conformal Contextual Robust Optimization. In Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, pp. 2485-2493. PMLR, April 2024. URL https://proceeddngg mlr...mm.........
http:/eeeddnnee.....eee/e///e/ wtt.:mll.....///eeemmmmmmm /teess//288/ppaaaee22e p-paaee24a.... htmm.. ISSN: 2640-3498.
Bala Kameshwar Poolla, Ashish R. Hota, Saverio Bolognani, Duncan S. Callaway, and Ashish Cherukuri. Wasserstein Distributionally Robust. Look-Ahead Economic Dispatch. IEFE Transactions on Power
15

---

Published in Transactions on Machine Learning Research (12/2025)
Systems, 36(3):2010-2022 May 2021. ISSN 1558-0679. doi: TLL.09/TTWRS2200..022 URL
https:///eeexpplore. ieee. orgo org.doccment/9922230 Conference Name: TEEE Transactions on Power Systems.
Anna Grazia Quaranta and Alberto Zaffaroni. Robust optimization of conditional value at risk and port- folio selection. Journal of Banking & Finance, 32(10) 2046 2056, October 2008. ISSN 0378-4266. doi: 10.10166jj. 2tt../m.n.i.......2222...... 2007.12.025. URL https: //w::///wwwww.. https://www.scienceeee sww.sceeeededrecc....m wwtp//eett.ccccccceeeeeeee htmm//ccennccaatticl/ppiii
$0378426607004281.
Hamed Rahimian and Sanjay Mehrotra. Distributionally Robust Optimization: A Review. arxi::1908.08 056599 [cs, math, stat), August 2019. URL http://aarin http:/awwwii.r..raaaa.ee www./ogg//////eeeeo/oo//ee abs/190s./ 05659, arXiv: 1908.05659.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning. The MIT Press, November 2005. ISBN 978-0-262-25683-4.
doi: mtt.5//mmmmppee/326......
URL
https://doi. org/10.//// wtt.:o..///.......... weti77551mmmmpp 7tt.://tt/mmmttee/////// 3206. 001.0001.000000. 000.0001.
Yaniv Romano, Evan Patterson, and Emmanuel Candes. Conformalized Quantile Regression. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché-Bue, E. Fox, and R. Garnett (eds.), Advances in Newral Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL http:://tt/pttteee http://aappeeee/ea//eeeeeeeeesss ceapperr.neuuipppeee iss. con.:1raacooooommmmmcmcm www.eoormalizddd-ooo -
pdf.
Philipp Schiele, Eric Sager Luxenberg, and Stephen P. Boyd. Disciplined Saddle Programming. Transactions on Machine Learning Research, June 2023. ISSN 2835-8856. URL KhMLfEIoUm.
Glenn Shafer and Vladimir Vovk. A Tutorial on Conformal Prediction. Journal of Machine Learning Research, 9(12):371 421, 2008. URL
html.
Alexander Shapiro, Darinka Dentcheva, and Andrzej Ruszczyísski. Lectures on Stochastic Programming:: Modeling and Theory. Society for Industrial and Applied Mathematics, January 2009. ISBN 978-0-89871- 687-0 978-0-898711-77-11 doi:
URL http://epubs siam www..aa..aaaa............... org/////b//o///0/./
9197/11.97808987187771
David Stutz, Krishnamurthy, Dvijotham, Ali Taylan Cemgil, and Arnaud Doucet. Learning Optimal Conformal Classifiers. In International Conference Learning Representations, May 2022. URL https:
//openreview.o net./tormmidd-t80442.
Chunlin Sun, Linyu Liu, and Xiaocheng Li. Prodict-then-Calibratee A New Perspective of Robust. Contextual LP. In Advances in Newral Information Processing Systems, volume 36, pp. 17713 17741, December 2023. URL https://prooeeiigg..eerrri htppe///eppppeeeleeppppp/p/ee/ee 597777111222222aa777776606622121
hhash/
html.
Stanislav Uryasev and Tyrrell Rockafellar. Conditional Value-at-Risk: Optimization Approach. In Stanislav Uryasev and Panos M. Pardalos (eds.), Stochastic Optimization: Algorithms and Applications, pp. 411 411-435. Springer US, Boston, MA, 2001. ISBN 978-1-4757-65946. doi: 11--007/978--4757-666 17. URL
https://doi.. orr...rg//10 1trg/10..000 100019007777781 1-78-1-4757-5544
Trina Wang, Cole Becker, Bart Van Parys, and Bartolomeo Stellato. Learning Decision-Focused Uncertainty 1922/ 19225. arXiv:2305.19225
Sets in Robust Optimization, July 2024. URL [math].
J. Warrington, P. J. Goulart, S. Mariéthoz, and M. Morari. Robust reserve operation in power systems using affine policies. In 2012 IEEE 51st IEEE Conference on Decision and Control (CDC), pp. 1111- 1117, December 2012. doi: 10.1109 CDC22012..422511. URL https://eeeepploree leee. www.:eeeeeeee........e wtg.:bstra///a/////////o document/6425913. ISSN: 0743-1546.
Bryan Wilder, Bistra Dilkina, and Milind Tambe. Melding the Data-Decisions Pipeline: Learning for Combinatorial Optimization. Proceedings of the AAAT Conference on Artificial Intelligence, 33(01):1658-1665 July 2019. ISSN 2374-3468. doi: 121609/aaai 1300..33800.333116665 URL https:/ //jjs. aaai. http://www./pp...aaaaaaaaa/////a
Number: 01.
16

---

Published in Transactions on Machine Learning Research (12/2025)
Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P Xing. Deep Kernel Learning. In Arthur Gretton and Christian C. Robert (eds.), Proceedings of the 19th International Conference Artificial Intelligence and Statistics (AISTATS), pp. 370 378, Cadiz, Spain, May 2016. PMLR. URL html. ISSN: 1938-7228.
http://procee http://oroceddnnn.ll.ec wom.m.l.ppee. www....rre..///wnaeee...eeeeeeee
Xiaohe Yan, Chenghong Gu, Xin Zhang, and Furong Li. Robust Optimization Liiiiizzatto-BBBedd Based Energy Storage Operation for System Congestion Management. IEEE Systems Journal, 14(2):2694 2702, June 2020. ISSN 1932-8184, 1937-9234, 2373-7816. doi: 101109//SSSTT.219999222222 URL https://ieeexplore ieee. wrg/doccument/8802555//
Qipeng P. Zheng, Jianhui Wang, and Andrew L. Liu. Stochastic Optimization for Unit Commitment A Review. IEEE Trunsactions Power Systems, 30(4):1913 1924, July 2015. ISSN 1558-0679. doi:
10.1109/TPWRS. NPW.S.2014.2355200 URL https://ieeexplorrs ieee. org/document/8912008 Conference Name: IEEE Transactions on Power Systems.
Weifeng Zhong, Kan Xie, Yi Liu, Shengli Xie, and Lihua Xie. Chance Constrained Scheduling and Pricing for Multi-Service Battery Energy Storage. IEEE Transactions Smart Grid, 12(6):5030 5042, November 2021. ISSN 1949-3061. doi: 10.1109/TSG.20 2021.1022110091400 URL https: ieee. org/ org/document/ 9525453. Conference Name: IEFE Transaction ions on Smart Grid.
17

---

Published in Transactions on Machine Learning Research (12/2025)
Table A1:
Time per epoch of pretraining, optimization, and end-to-end training, measured in seconds (wall-clock time). Values are mean #1 stddev across 10 epochs. Lower values are better.
pretrain
optimize (train+val)
E2E training
Box
0.03 = 0.01
4.29 = 0.07
6.34 = 0.32
Ellipse
0.03 + 0.01
7.79 + 0.06
7.12 + 1.28
PICNN
4.20 = 0.20
50.49 + 0.19
33.24 = 0.25
A Appendix: Additional experimental results
A.1 Experimental results: Battery storage
In Figures A1 and A2, we plot the (1 - a ----le---a--- (VaR) and (1 - co-odiiiionnal value-at-risk (CVaR) (Uryasev & Rockafellar, 2001) of the task loss obtained via our method on the battery storage problem to evaluate model robustness. In particular, for each target coverage level 1 a, we evaluate VaR2-*[task loss] and CVaR CVaR!'"[tase ftask loss] for both the ETO baseline(s) and our E2E methodologies across box, ellipsoid, and PICNN uncertainty sets. In general, it is clear that our E2F methodology uniformly improves over ETO, and while the performance of E2E ellipsoid close to that of E2E PICNN, the latter appears to perform better for certain a. These results highlight that our methodology improves not only the average task loss, but robustness more generally, even when compared with the robust ETO baselines.
The optimal task losses shown in black dotted lines in Figures 3, A1 and A2 are the lowest achievable task loss on the lest set given perfect knowledge of the target y. The optimal task loss is calculated for each example (z, y) in the test set as f(z, y, ppp)) where
opt = arg min (2,,1,2) s.1. g(r, S O.
SERP
Performance. In Table A1, we report the wall-clock time for pretraining, optimization, and end-to-end training. These times were mcasured on machine with 2x AMD EPYC 7513 32-Core Processors, 1TiB RAM, and 4 NVIDIA A100 GPUs (although only 1 of the GPUs was used in these experiments). The "pretrain" column gives the time per epoch of training with the standard loss function (pinball loss for Box, negative log-likelihood loss for Ellipsc and PICNN). The "optimize" column gives the time needed to compute the decision z (x) given a pretrained model. The "E2E" column gives the time per epoch of E2E training, which requires computing (x) for each training and calibration example. Evidently, the bulk of the time spent on E2E training comes from computing the decision = =z); the additional overhead from computing the gradient through the KKT conditions of the optimization problem is much less than solving the optimization problem itself.
We note that in theory, E2E times should always be higher than the optimization time. The optimization time accounts for the time required to compute 25 (z), and E2E training additional accounts for the time required to compute the gradient 30 (x). However, practice, the optimization time depends heavily on the choice of numerical solver. In our implementation, we use the default cvxpy solver (Clarabel) for the optimization step in ETO, whereas we use the default cvxpylayers solver (SCS) during E2E training, In the case of the Ellipse and PICNN uncertainty sets, the SCS solver is generally faster than Clarabel; hence, Table A1 counterintuitively reports lower E2E training times for Ellipse and PICNN than optimization times.
We further note that the optimization and E2E training times reflect our particular implementation, but significantly faster times are likely possible. In particular, we solve all of the optimization problems (in both the ETO optimization step and during E2E training) using CPU-basod solvers; however, recently developed CPU-accelerated solvers may be able to reduce the optimization time required by orders of magnitude, especially when solving batch of problems with the same structure (Blin, 2024).
18

---

Published in Transactions on Machine Learning Research (12/2025)
Ellipse
PICNN
VaR "[task loss]
ETO
-10
ETO-JC
distribution shift
ETO-SLL
-20
optimal
0.01
0.05
0.1
0.01
0.1 uncertainty level
0.05
0.01
0.1 uncertainty level
0.05
uncertainty level
Box
Ellipse
PICNN
ETO
with distribution shift
VaR -atask
-10
ETO-JC ETO-SLL
E2E optimal
-20
0.01
0.05
0.1
0.2 0.01
0.05
0.1
0.2 0.01
0.05
0.1
0.2
uncertainty level
uncertainty level
uncertainty level a
Figure A1: Similar to Figure 3, except that the value plotted the VaR loss] (mean +1 stddev across runs) on the test set for the battery storage problem with no distribution shift (top) and with distribution shift (bottom). Lower values are better.
Box
Ellipse
PICNN
CVaR -ttas
ETO
no distribution shift
-10
ETO-JC
ETO-SLL
E2E
-20
optimal
0.01
0.05
0.1
0.2
0.01
0.05
0.1
0.2 0.01
0.1 uncertainty level a
0.05
0.2
uncertainty level
uncertainty level a Ellipse
Box
PICNN
CVaR1 (task loss]
with distribution shift
ETO-JC ETO-SLL
optimal
-20
0.01
0.05
0.05
uncertainty level
uncertainty level
uncertainty level
Figure A2: Similar Figure except that the value plotted is the CVaR loss] (mean +1 stddev across 10 runs) the test set for the battery storage problem with no distribution shift (top) and with distribution shift (bottom). Lower values are better.
A.2 Experimental results: Portfolio optimization
Tables A2 and A3 show the task loss and coverage results for the portfolio optimization problem. We again find that our E2E approach generally improves upon the ETO basclines at all uncertainty levels a, with the exception of box uncertainly, where all the methods achieve similar performance. Our PICNN-based uncertainty representation, when learned end-to-end, performs better than box uncertainty and comparably with ellipse uncortainty. The similarity in performance between E2E ellipsoidal uncertainty and E2E PICNN uncertainty is likely due the underlying alcatoric uncertainty (i.e., the uncertainty in P x)) generally taking ellipsoidal shape the conditional distribution P(y x) is Gaussian mixture model, and it tends
19

---

Published in Transactions on Machine Learning Research (12/2025)
Table A2: Task loss performance (mean + 1 stddev across 10 runs) for the portfolio optimization problem. Lower values are better, and the best performance for each uncertainty-level highlighted. The results show that our E2E methods consistently outperform the ETO baselines.
uncerlainty level
0.01
0.05
0.1
0.2
ETO
Box
-1.16 + 0.42
-1.37 + 0.12
-1.39 = 0.13
-1.41 + 0.12
ETO
Ellipse
-1.09 + 0.12
-1.24 + 0.11
-1.29 = 0.10
-1.33 + 0.10
ETO
PICNN
-0.95 - 0.24
-1.11 + 0.24
-1.20 = 0..2
-1.31 0.16
ETO-SLL
Box
-1.41 = 0.13
-1.42 + 0.12
-1.42 = 0.12
-1,44 0.11
ETO-SLL
Ellipse
-1.12 0.22
-1.37 = 0.12
-1.40 = 0.12
+1.33
ETO-JC
Ellipse
-1.16 + 0.17
-1.40 + 0.11
-1.42 01.1
+1.44 0.11
E2E
Box
-1.21 + 0.44
-1.40 + 0.14
-1.43 = 0.11
-1.43 + 0.10
E2E
Ellipse
-1.48 0.12
-1.47 + 0.11
-1.48
0.11
-1.47 0.11
E2E
PICNN
-1.45 + 0.14
-1.48
0.10
-1.48
0.10
-1.47 0.11
Table A3: Coverage (mean +1 stddev across 10 runs) for the portfolio optimization problem. Our E2E models achieve similar coverage to the ETO baselines, confirming that the lower task loss of our E2E models does not come at the expense of worse coverage.
uncertainty level
0.01
0.05
0.1
0.2
ETO
Box
.984
.007
.947
.017
.902 .017
.786 .020
ETO
Ellipse
.988
.004
.944 .020
.894 .022
.794 .027
ETO
PICNN
.989
.006
.949
.014
.901 + .019
.801 + .034
ETO-SLL
Box
.985 + .012
.945 = .021
.885 + .030
.796 + .029
ETO-SLL
Ellipse
.989 = .011
.945 = .024
.885 + .039
.795 + .030
ECO-JC
Ellipse
.991 - .006
.953 + .017
.902 + .226
.796 + .026
E2E
Box
.989
.006
.949
.012
.903 .016
.785 .019
E2E
Ellipse
.992
.006
.954
.010
.903 + .02
.788 + .022
E2E
PICNN
.993 + .002
.953 + .010
.912 + .017
.798 + .024
to have dominant mode (e.g., see Figure A3). In terms of coverage, we find that all the models and training methodologies obtain coverage very close to the target level, confirming that the improvements in task loss performance from our E2E approach do not come at the cost worse coverage.
Because the conditional distribution Ply ax) for the portfolio optimization problem 2-dimensional, we can visualize the conditional distribution well the uncertainty sets estimated by our models. Figure A3 plots the conditional density for input x [11.167 0.024] T along with the a = 0.1 uncertainty sets 2t (a)) and the resulting decision vectors (x) for each uncertainty set parametrization. Uncertainty sets and decision vectors from both ЕTO and E2E models are shown in different colors. The key takeaway from this figure that smaller uncertainty sets (which is what ETO training tends to produce) do not always result in lower task loss. Furthermore, the more flexible parametrization of the PICNN allows it to learn uncertainty set shapes that may be more amenable to the downstream robust decision task than box ellipsoidal uncertainty, even if the resulting uncertainty set has larger odder shape.
B Appendix: Related Work
"Task-basodl or "decisionnfocuusod" learning. The notion of "task-basod" end-to-end model Icarning was introduced by Donti al. (2017), which proposed train machine learning models in an end-to- end fashion to minimize downstream stochastic optimization objective. To achieve this, the authors
20

---

Published in Transactions on Machine Learning Research (12/2025)
Box
Ellipse ETO task loss -0.07 E2E task loss -0.08
PICNN ETO task loss 0.03 E2E task loss -1.06
ETO task loss 0.67 E2E task loss -1.06
5.0
ETO 9a(x)
2.5
ETO (x)
E2E Ro(x)
0.0
E2E z (x)
2.5
true y
-5.0
-7.5
-5
0
-5
0
5
-5
0
5
V1
Y1
y1
Figure A3: This figure plots the density of the conditional distribution P(y a) for = [-..167 0.024] from the portfolio optimization problem, with darker colors indicating higher density. Also plotted are the a = 0.1 uncertainty sets (2) (dashed lines) and the resulting decision vectors =3 (x) (arrows) for each uncertainty set parametrization. Results for ETO models are shown in blue, whereas results for E2E are shown in orange. The "truc" y sampled from Ply x) drawn green, and the task loss for this example is computed using this y- The decision vectors have been artificially scaled larger to be easier to see.
backpropagate gradients through a stochastic optimization problem, which made possible for various types of convex optimization problems via the implicit function theorem (Donti et al., 2017; Amos & Kolter, 2017; Agrawal et al., 2019). However, Donti al. (2017) does not train the model to estimate uncertainty and thereby does not provide any explicit guarantecs on robustness on their decisions to uncertainty. Our framework improves upon this baseline by yielding calibrated uncertainty sets which can then be used to obtain robust decisions.
Another line of related works relies on surrogate loss functions approximate the gradient of the downstream lask loss, typically in settings where the exact gradient does not exist is otherwise hard to compute. For example, the "Smart Predict, then Optimize" approach (Elmachtoub & Grigas, 2022) specifically considers differentiating the solution of linear programs, whereas Wilder et al. (2019) differentiates the solution of discrete (combinatorial) optimization problems. As with Donti et al. (2017), though, these approaches focus only on average task loss without any estimation of model uncertainty, thereby lacking explicit guarantees on the robustness of the resulting decisions. A recent survey by Mandi et al. (2024) provides a comprehensive review of existing decision-focusel learning approaches.
Uncertainty Quantification. Various designs for deep learning regression models that provide uncertainty estinates have been proposed in the literature, including Bayesian neural networks (Blundell et al., 2015; Gal & Ghahramani, 2016), Gaussian process regression and deep kernel learning (Rasmussen & Williams, 2005; Wilson et al., 2016; Liu et al., 2020), ensembles of models (Lakshminarayanan et al., 2017), and quantile regression (Romano et al., 2019), among other techniques. These methods typically only provide heuristic uncertainty estimates that are not necessarily well-calibrated (Nado et al., 2022).
Post-hoc methods such as isotonic regression (Kuleshov et al., 2018) or conformal prediction (Shafer & Vovk, 2008) may be used to calibrate the uncertainty outputs of deep learning models. These calibration methods generally treat the model as a black box and scale predicted uncertainty levels so that they are calibrated on a held-out calibration set. Isotonic regression guarantees calibrated outputs in the limit of infinite data, whereas conformal methods provide probabilistic, finite-sample calibration guarantees when the calibration set is exchangeable (e.g., drawn i.id. from the same distribution) with test data. These calibration methods are generally not included in the model training procedure because they involve non-differentiable operators, such as sorting. However, recent works have proposed differentiable losses (Einbinder et al., 2022; Stutz et al., 2022) that approximate the conformal prediction procedure during training and thus allow end-to-end training of models to output more calibrated uncertainty. As approximations, these methods lose the marginal coverage guarantees that true conformal methods provide. However, such guarantees can be recovered test time by replacing the approximations with true conformal prediction.
21

---

Published in Transactions on Machine Learning Research (12/2025)
Robust and stochastic optimization. The optimization community has proposed number of techniques over the years to improve robust decision-making under uncertainty, including stochastic, risk-sensitive, chance-constrainee
distributionally robust, and robust optimization (e.g., Ben-TTll et al. (2009); Shapiro et al. (2009); Nemirovski & Shapiro (2007); Rahimian & Mchrotra (2019)). These techniques have been applied to a wide range of applications, including energy systems operation (Zheng et al., 2015; Ndrio el. al., 2021; Dvorkin, 2020; Zhong et al., 2021; Poolla al., 2021; Warrington al., 2012; Bertsimas et al., 2013; Christianson al., 2022) and portfolio optimization (Gregory et al., 2011; Quaranta & Zaffaroni, 2008; Bertsimas et al., 2018). In these works, the robust and stochastic optimization methods enable selecting decisions (grid resource dispatches or portfolio allocations) in a manner that is aware of uncertainty, e-g., so an energy system operator can ensure that sufficient generation is available to meet demand even on a cloudy day without much solar generation. Typically, however, the construction of uncertainly sels, estimated probability distributions over uncertain parameters, or ambiguity sets over distributions takes place offline and is unconnected to the eventual docision-making task. Thus, our proposed end-to-end approach allows for simulltnnous calibration of uncertainly sets with optimal decision-making.
C Appendix: Maximizing over the uncertainty set
We consider robust optimization problems of the form
min max y Fzz + 2( s.t. zER BER"
g(0,2) S 0. g E 1(x),
For fixed 2, the inner maximization problem
max 9 Fz
s.t.
i e (()))
@ER"
which we analyze in the more abstract form
max T y s.t. yEΩ EER*
for arbitrary e E R" \ (0). The subsections of this appendix derive the dual form of this maximization problem for specific representations of the uncertainty set Ω.
Suppose y is standardized or whitened by an affine transformation with E R" and invertible matrix W e Rax
Utransfrmeed Wy - 11)
so that Q an uncertainty set on the transformed Ycransformed- Then, the original primal objective can be recovered
y c T (W ytransformed + (a) (We) Yrransformed + c μ.. In our experiments, we use element-win wise standardization of y by setting W diag(Hsta), where yetd E R" is the element-wise standard-deviation of y-
C.1 Maximizing over box constraint
Let [g C R. be box uncertainty set for R.. Then, for any vector R",, the primal linear program
max c VER"
S.L.
gSysy
has dual problem
min vER2
T 7 8.6.
In -In] =c 420,
which can also be cquivalently written as
min VER
( v-
s.t. V w-c20 0.
Since strong duality always holds for lincar programs, the optimal values of the primal and dual problems will be equal so long as one of the problems feasible, e.g., so long as the box [y. yl is nonempty. We can thus incorporate this dual problem into the outer minimization of (1) to yield the non-robust form (4).
22

---

Published in Transactions on Machine Learning Research (12/2025)
C.2 Maximizing ellipsoid
For any E 1 (0), E st and q > O, the primal quadratically constrained linear program (QCLP)
max c'y y VER"
s.t. (y x)
= =
has dual problem
min -c c s.t. V20 PER
4
By Slater's condition, strong duality holds by virtue of the assumption that > 0 (which implies strict feasibility of the primal problem), and thus the primal and dual problems have the same optimal value. Morcover, since E is positive definite and q > 0, this problem has a unique optimal solution at v* = 2// clll. where L the unique lower- twerrttranggllr Cholesky factor of (i.e., = L.. Substituting v* into the dual problem yields
+ CCc. c.c.
Plugging this into (1) yields the non-robust form (5).
We write the dual objective terms of the Cholesky factor because our predictive models for ellipsoidal uncertainty directly output the entries of L (see Appendix E). Note, however, that the dual problem solution can be equivalently written in terms of the squarc-root of Σ, because
L |= - Se T / 1//
|
C.3 Proof of Theorem 4.1: Maximizing over the sublevel set of PICNN
Let Rm R" + R be partially input-convex neural network (PICNN) with ReLU activations described in (7), so that sn(x, y) is convex in Suppose that all the hidden layers have the same dimension d (i.e., VI = o,...,.- L-.,,-L 1: Wi E idxd Vi E Rdx"" b E R²), and the final layer L has Wt. E glxd Vz E R1 by E R. Let c IR" be any vector. Then, the optimization problem
max c y yER"
s.t.
so s6(2,8) 9) 5
(9)
can be equivalently written as
c y
(10a)
yERn di o ERd
σi 2 0s
1/ 1, 1,.., 1 - 0 1-1
(10b)
01+1 2 Wio: + Vig + b:
(10c)
WLOL + VLY + by
(10d)
To see that this is the case, first note that (10) is a relaxed form (9), obtained by replacing the equalities 01+11 = ReLU (Wia + Vay + by) in the definition of the PICNN (7) with the two separate inequalities 01+1 2 0₂ and 01-1 2 Wio: Viy + by for each / - 0, ... L- 1. As such, the optimal value of (10) is no less than that of (9). However, given optimal solution y, ,.. or to (10), it is possible to obtain another feasible solution y,āi...., with the same optimal objective value by iteratively decreasing each component of until one of the two incquality constraints (10b), (10c) is tight, beginning at = 1 and incrementing 1 once all entries of a cannot be decreased further. This procedure of decreasing the entries in each will maintain problem feasibility, since the weight matrices W; are all assumed to be entrywise nonnegative in the PICNN construction; in particular, this procedure will not increase the left-hand side of (10d). Morcover, since one of the two constraints (10b), (10c) will hold for each entry of each o this immediately implies that y is feasible for the unrelaxed problem (9), and (9) and (10) must have the same optimal value.
23

---

Published in Transactions on Machine Learning Research (12/2025)
Having shown that may replace the convex program (9) with linear equivalent (10), can write this latter problem in the matrix form
01
cTy
max
s.t.
S
BER",
BLER
0L
where
-aa
0g
la
Od
A
Vo
- Id
E IR 22 (2
b=
-
E [R.22d+11
(11)
W1
Ia Wr.
bL-1 bb
VE.
By strong duality, if this linear program has an optimal solution, its optimal value is equal to the optimal value of its dual problem:
min dT V PERSLA+1
s.t. A = Ora
2 O.
(12)
We can incorporate this dual problem (12) into the outer minimization of (1) to yield the non-robust form (8). For a more interpretable form of this dual problem, let y(0 denote the portion of the dual
vector corresponding to the i-th block-row matrix A, indexed from That is, plt Vid+ia(i+11d for
0.. 1. Furthermore, let μ P2Ld+1 be the last entry of v. Written out, the dual problem (12)
becomes
min eRRd, PER
plg- - br.)
Σ by (((+))
(00) (2
1=0
s.t.
[V=T
Vi VLd+1: =0
WI,p((++41] ((++]) - -1()) Oa
VI O....,
V 2 0.
C.4 Ensuring feasibility of the PICNN maximization problem
As noted at the end of Section 4, it may sometimes be the case that the inner maximization problem of (1) is unbounded or infeasible when 24 (x) is parametrized by PICNN, since in general, the sublevel sets of the PICNN might be unbounded, or the q selected by the split conformal procedure detailed Section 3.2 may be sufficiently small that (z)x) {0 R" s0 9) S 9} empty for certain inputs We can address each of these concerns using separate techniques.
Ensuring compact sublevel sets. To ensure that the PICNN- parammttreer score function sø(x, y) has compact sublevel sets in y, we can redefinc the output layer by setting Vz 01xn and adding small goo norm term penalizing growth in y::
s0(2, y) = Wy.,, + ellvll + br,
(13)
where e 2 is a small penalty term, and where all the remaining parameters and layers remain identical to their definition (7). This modification ensures that, for any fixed x, sp y) has compact sublevel sets, since σg 2 by construction (7) and the penalty term elly|loo will grow unboundedly large as y goes infinity in any direction, long as > 0. Moreover, long is chosen sufficiently small and the PICNN
24

---

Published in Transactions on Machine Learning Research (12/2025)
is sufficiently deep, this modification should not negatively impact the ability of the PICNN to represent general compact convex uncertainty sets.
Using this modified PICNN (13), the maximization problem (9) can be written as an equivalent linear program
max
c y
(14a)
yER", 1 KER
8.1.
σi 0ad
14 = 1, ,..., 1A =
(14b)
Ø41 Wioi Viy + bj
(14c)
K 2 Vi, K N Vii
14
(14d)
W Wror + EK + by S
(14e)
where the equivalence between (9) and (14) follows the same argument as that employed in the previous section when showing the equivalence of (9) and (10). We can thus likewise apply strong duality to obtain an equivalent minimization form of the problem (14) and incorporate this into the outer minimization of (1) to yield a non-robust problem the general form (8).
Ensuring (2)x)) is nonempty. If the q chosen by the split conformal procedure is too small such that (z) is empty, i.c., q S Jwin where
Amin
min s9(2, 0),
(15)
then we simply increase y to Jmnn that (z) is guaranteed to be nonemply. That is, for input H, set
q max min so 9). QUANTILE((SO(FT ()) 1 1- -a))
VER"
This preserves the marginal coverage guarantee, as increasing q can only result in larger uncertainty set (z)..
In theory, Qmin varies as a function of 0 and it is possible to differentiate through the optimization problem (15) using the methods from Agrawal et al. (2019) since the problem convex and 88 is assumed to differentiable w.r.t. almost everywhere. However, to avoid this added complexity, practice, we treat quin constant. In other words, on inputs where have increase g to Jmin, treat 0g =0
00
D Appendix: Exact differentiable conformal prediction
In this section, we prove how to exactly differentiate through the conformal prodiction procedure, unlike the approximate derivative first introduced in Stutz al. (2022).
Theorem D.1. Let a E (0,1) 1) be a risk level, and let 8, == se (o(x,, vi) denole the scores computed by a score function 8q m. x I - IR over data points ((x,, W.) M Suppose so s0(x,,30 Ws) differentiablee w.r.t. for all = M.
Define SM+1 == co. Let o: (l...., M M+ + ascending order, such that (o( I 8o(i) for all i < For simplicity of notation, we may write S(i) = 8o())-
(1,.... .M....MMM..... + 1 denote the permutation that sorts the scores in
Let QUANTILE((s:)M I 1- 1-a) a) where the QUANTILE function is as defined in Algorithm 1. That is, = S(k), where & := [M + 1)(1 - a)l (I,...,,A,, M, M + +1.. If ((K) unique, then
dq
4 (T fa2 गम
d0 0,
otherwise.
Proof. First, when a E (0, M we have k M + 1, so q = oo is constant regardless of the choice of 0.
Thus, dg = 0.
25

---

Published in Transactions on Machine Learning Research (12/2025)
Now, suppose 2 M+1 The QUANTILE function returns the k-th largest value of {ss)== foo). Since we assume S(k) is unique, we have de
1l kl. Finally, we have
()
dq M dq ds
M dg ds()) ds()) d so( (((a))))Bo()))))
d0
ds, d0
ds (s) d0 d0 d0
iml
The two key assumptions in this theorem are that (1) sq is differentiable w.r.t. 0, and (2) s(e)) is unique. When Sg is a neural network with a common activation function (c.g., ReLU), (1) holds for inputs (x, E R x R" almost everywhere and 0 almost everywhere. Regarding (2), in practice, just as the gradient of the max function is typically implemented without checking whether its inputs have ties, we do not check whether S(k) is unique.
E
Appendix: Experiment details
Our experiments were conducted across a variely of machines, including private servers and Amazon AWS EC2 instances, ranging from 12-core to 128-core machines. Our ETO experiments benefited from GPU acceleration across combination of NVIDIA GeForce GTX 1080 Ti, Titan RTX, T4, and A100 GPUs. Our E2E experiments did not GPU acceleration, due the lack of GPU support in the evxpylayers Python package (Agrawal et al., 2019).
In all experiments, we use batch size of 256 and the Adam optimizer (Kingma & Ba, 2015). Models were trained for up to 100 epochs with early stopping if there was no improvement in validation loss for 10 consecutive epochs.
For box and ellipsoid ETO baseline models, we performed hyperparameter grid search over learning rates
(10-4.5 10 10-4, 10 10-3.5, 10-3 3, 10 2.5 2.5, 10 10-1.5) and L2 weight decay values (0, 10 -4-4, 10-3, 10 1002) 2). For
PICNN ETO models performed hyperparameter grid search over learning rates (10-1, 10-3 10-2) and L2 weight decay values (10-4, 10-3, 10-2).
E.1
Uncertainty representation
Box uncertainty. Our box uncertainty model uses a neural network he with 3 hidden layers of 256 units each and ReLU activations with batch-normalivation. The output layer has dimension 2n, where dimensions predict the lower bound. Output dimensions n n 1 2n, after passing through softplus to ensure positivily, represents the difference between the upper and lower bounds. That is,
The (x) [ (x)
ha(2)mm [ho()) coftplus(ho(())44...)
This architecture ensures that hz)()) > (x))
In the two-stage ETO baseline, we first train hg to estimate the a/22- and (1 - a/2)-quantillss, so that (x), (x)) represents the centered (1 c)-confidencc region. Quantile regression a common method for generating uncertainty sets for scalar predictions by estimating quantiles of the conditional distribution P(y x) (Romano al., 2019). For scalar true label y, quantile regression models commonly trained to minimize pinball loss (a.k.a quantile loss) where is the quantile level being estimated:
BB (y 9). if y 0 (1- v). if y 5 g.
pinbally
To generalize the pinball loss to our setting of multi-dimmnnionall y E R. we sum the pinball loss across the dimensions of y: pinball vf. Σ pinballl yt).
26

---

Published in Transactions on Machine Learning Research (12/2025)
Our end-to-end (E2E) box uncertainty models use the same architecture above, initialized with weights from the trained ETO model. We found it helpful to use weighted combination of the task loss and pinball loss during training of the E2E models to improve training stability. In our experiments, we used weight of 0.9 on the task loss and 0.1 on the pinball loss. The E2E models used the best L2 weight decay from the ETO models, and the learning rate was luned across 10-2, 10-3, and 10-1
Ellipsoidal uncertainty. Our ellipsoidal uncertainty model uses a neural network ho with 3 hidden layers of 256 units each and ReLU activations with batch-normallizttion The output layer has dimension n 1++n(n+1))) where dimensions 1 : n predict the mean (а( and the remaining output dimensions are used to construct a lower-trianuulat Cholesky factor (x(x) of the covariance matrix [x = Lo(x))o(x)). We
pass the diagonal entries of Lo (x) through a softplus function to ensure strict positivity, which then ensures Ex)x is positive definite.
For the ETO baseline, we trained the model using the negative log-likelihood (NLL) loss
NL N
1
Σ
WI [o(x),
(a.))ED
where (: | p4, (1, denotes the density of a multivariate normal distribution with mean and covariance matrix Σ.
Our end-to-end (E2F) ellipsoidal uncertainty models use the same architecture above, initialized with weights from the the trained ETO model. We found it helpful to weighted combination of the task loss and NLL loss during training of the E2E models to improve training stability. In our experiments, we used a weight of 0.9 on the task loss and 0.1 on the NLL loss. The E2E models used the best L2 weight decay from the ETO models, and the learning rate was tuned across 10-2, 10-3, and 10-4.
PICNN uncertainty. Our PICNN has 2 hidden layers with ReLU activations.
For the battery storage problem, we used 64 units per hidden layer. We did not run into any feasibility issues for the PICNN maximization problem, so we did not restrict V2 as described in Appendix C.4, and we set c=0.
For the portfolio optimization problem, we tried 32, 64, and 128 units per hidden layer, finding that 32 units worked best. We did run into feasibility issues for the PICNN maximization problem, which we resolved by sctting Vz 01xm as described in Appendix C.4. This change alone was sufficient, and we set c= 0.
For the ETO baseline, we take inspiration from the approach by Lin & Ba (2023) to give probabilistic interpretation to a PICNN model Sg via the energy-based model Pe(y x) = Zo(z) exp(-sa (z, v)) where 7o(x) = SSER= exp(-so(x; ū)) dÿ is the normalizing constant. We train our ETO PICNN models with an approximation to the true NLL loss based on samples from the Metropolis-Adiiuee (MALA), Markov Chain Monte Carlo (MCMC) method. We refer readers to our code for the specific hyperparameters and implementation details we used.
Langevin Algorithm
Note that under this energy-based model, adding scalar constant e to the PICNN (i.e., sa(2, y) + +c) does not change the probability distribution. That is, exp(-sa(x.)) y)) & exp(-se(z, y) + c). To regularize the PICNN model, which has a bias term in its output layer, we therefore introduce a regularization loss of Wzero s0 2 where Wxero is a regularization weight. This regularization loss encourages sø(x, y) to be close to 0, for all examples in the training set. In our experiments, we set Wsero 1.
Our end-to-end (E2E) PICNN uncertainty models use the same architecture as above, initialized with weights from the the trained ETO model. Unlike for box and ellipsoidal uncertainty which used weighted combination of lask loss and NLL loss, our E2E PICNN uncertainty models are trained only with the task loss. Similar to the ETO PICNN model, we also regularize the E2E PICNN. Here, we add a regularization loss of Wy wg. 92, where Wy is a regularization weight and q is the conformal prediction threshold computed in cach minibatch of E2E training. This regularization loss term aims to keep q near O: without this regularization, we found that q tended to grow dramatically over training epochs with poor task loss. In our experiments, we set wy = 0.01.
27

---

Published in Transactions on Machine Learning Research (12/2025)
The E2E models used the best L2 weight decay from the ETO models. For the battery storage problem, we tested learning rates of 10-3 and 10-1. For the portfolio optimization problem, we used learning rate of 5 x 10 10-3
E.2 Data
Price forecasting for battery storage. We use the same dataset Donti al. (2017) in price forecasting for battery storage problem. In this dataset, the target y E R is the hourly PJM day-ahead system energy price for 2011-2016, for a total of 2189 days. Unlike Donti et al. (2017), though, we do not exclude any days whose electricity prices are too high (>500s/ (/500. Whereas Donti et al. (2017) treated these days as outliers, our conditional robust optimization problem is designed to output robust decisions. For predicting target for a given day, the inputs I E R00 include the previous day's log- poggpprces, the given day's hourly load forecast, the previous day's hourly temperature, the given day's hourly temperature, and several calendar-based features such as whether the given day is a weekend or US holiday.
For the setting without distribution shift, we take a random 20% subset of the dataset as the test set; because the test set is selected randomly, it considered exchangcable with the rest of the dataset. For the sctting with distribution shift, we take the chronologically last 20% of the dataset as the test set; because load, electricity prices, and temperature all have distribution shifts over time, the set is not exchangeable with the rest of the dataset. Figure A4 illustrates the data from the distribution shift setting, where the electricity price tends to be lower and has lower variability in the test set compared to the training/ training:validdation validattin splits. For each seed, we further use 80/20 random split of the remaining data for training and calibration.
Portfolio optimization. For the portfolio optimization task, we used synthctically generated data. We sample x E RR25v E RR2 from a mixture of three 4-D multivariate Gaussian distributions as used in Chenreddy & Delage (2024). Formally,
I
~ Pa PaN(aa Ea) + ps No.6. Ex) + Pe Pe Ee)
y
where Po + Ph + Pc = Specifically,
=,
Pb 0GMM + (1 Φ),
0GGM
(1 Φ))
0GMM +
μ4b [0 5 0] 0
P4 04,
He Pb:
0.37
0
1.5
0
0
Σ,
Σ, CGMM
E,
0.37
0
2
0.73
QGMM
0
0
0.73
3
for some & [0, and QGMM E [0, 1]. In our experiments, we used 6 = 0.7 and CGMM 0.9. (Chenreddy & Delage (2024) do not disclose the values of & and 0GMM chosen for their experiments.) For each random seed, we generate 2000 samples and use (train, calibration, test) split of (600, 400, 1000).
28

---

Published in Transactions on Machine Learning Research (12/2025)
200
split
day-ahead price (S/MWh)
train/val
150
test
125
100
75
50
25
0
2
7
10
12
month
(a) Month of year electricity price
Q1 (Jan-Mar)
(Apr-jun)
Q3 (Jul-Sep)
Q4 (Oct-Dec)
day saa-ahead price (S-WWh)
200
year
2011
2013
100
2014
2015
50
train/val
12
18
12 hour of day
test
hour day
hour
hour
(b) Hour of day electricity price (shaded region shows + 1 std dev)
daily peak price (S/MWh)
(Jan-Mar)
(Apr-jun)
Q3 (Jul-Sep)
Q4 (Oct-Dec)
2011
150
2012
2013
100
2014
2015
2016
train/val
20
40
60
80
100
20
40
60
80
100
20
40
60
80
100
20
40
60
80
100
test
daily high temp (*F)
daily high temp
daily high temp ("F)
daily high temp
(c) Daily high temperature (rounded to nearest 5'F)
daily peak electricity price (shaded region shows 1 std dev
Figure A4: Visualization of data from the battery storage problem under distribution shift. The train/ val splits are sampled randomly from the range 2011-01-04 to 2015-10-20, whereas the test set comprises 2015-10- 21 to 2016-12-31. These plots evidently show that the electricity price is generally lower in the test set and also has lower variability by both hour of day and the daily maximum temperature.
29


Reviews:
[{"review_id": "kvbkXEJ0AG", "reviewer": "Reviewer_A4ot", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- **[S1]**: Novelty and methodological contributions.  Overall, this paper makes two major contributions.  The first is the integration of differentiating through the conformal prediction within an end-to-end learning context.  The second is the learned uncertainty sets within the context of conditional robust optimization.  These two contributions are both well motivated and a major strength of the paper.   \n- **[S2]**: Experimental results.  Overall, the numerical results of the paper are promising, although I'll discuss some limitations and suggestions for improvement in the weaknesses.  That said, the end-to-end approach consistently achieves not only good task loss, but also good coverage.  In addition, the authors present some results on distributional shifts that are favorable for the proposed methods.\n- **[S3]**:  Clarity and presentation. The paper is well-written.  Contributions are clearly stated, theory is appropriate, and numerical results are presented cleanly.  \n\n### Weaknesses\n- **[W1]**: Evaluation metrics.  Given the focus on conditional robust optimization (CRO), standard CRO metrics should be reported, such as VaR/CVaR, rather than just the task loss, as these metrics prioritize robustness, i.e., what decision makers are often concerned with.  \n- **[W2]**: Baselines.  While the authors evaluate and compare their approach with existing ETO approaches, they do not compare it with the other end-to-end approach presented in [1].  The authors state that they only compare against approaches that maintain guaranteed marginal coverage.  However, even the proposed approach only guarantees marginal coverage for the calibration set and not the test set.  For this reason, I believe it would be reasonable to compare to [1], especially given the marginal coverage %'s reported are close to the target marginal coverage.   \n- **[W3]**: Distribution shifts experiments.  These experiments are not particularly convincing.  The authors split the data temporally.  However, there is no discussion of how the distribution of data differs temporally.  Insight into this, or perhaps providing an ablation on varying levels of shifts through perturbation, would be more convincing.  \n- **[W4]**: No reporting on times to give insight into scalability.  Given the context of optimization, it would be helpful to report the time required for training between the different variants of methods, e.g., if box/ellipsoid are more efficient in training, they may be preferable in specific contexts.  \n- **[W5]**:  Assumption 3.1, in particular $g$ not being a function of $y$.  While most existing approaches make the same assumption, this should be stated earlier and as a separate assumption from the assumption on the objective, given that it may easily be missed.  Alternatively, framing the approach as focusing on objective uncertainty and dropping $y$ from the constraints would make the notation more consistent with the approach's limitations.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an end-to-end approach for condition robust optimization. The authors propose an approach based on minimizing the decision-loss while differentiating through conformal prediction to ensure a desired marginal coverage is met.  In addition, the authors propose partially input-convex neural networks (PCINNs) as a means to parameterize uncertainty sets.  This approach is then evaluated on a price forecasting problem for battery storage and a portfolio optimization, where their approach, along with the uncertainty sets learning through PCINNs, achieves both a good task loss and coverage.  In addition, they evaluate their approach within the context of distribution shifts and demonstrate favorable results.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**W1/2** would strengthen the contribution and provide a more comprehensive evaluation that helps me lean toward acceptance.  **W4/5** issues need to be addressed, but they are minor.  Address.  For **W3**, the authors should either tone down their claims, unless a more robust evaluation is done.", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "dEjrHptyol", "reviewer": "Reviewer_iVud", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper is quite well written and walks the reader coherently through the problem statement and proposed solution, highlighting the arising research questions and proposals to address them. I also appreciate the mix of visualizations and theoretical backing, and the proper explanations and related work review in the appendix.\n- The problem of rendering these models more valuable for down-stream decision-making is important. Many papers end with the conformal procedure and motivate obtained sets as useful for decision-making, but do not go the extra steps of actually modeling decision-making and addressing this aspect. So this direction is good for furthering the field.\n\nWeaknesses\n- The paper does not talk much about the weaknesses or limitations of the approach, both in terms of imposed convexity restrictions (Ass. 3.1), and the limitations of proposing an end-to-end framework restricted to the particular input-convex model (PICNN). In essence, this approach loses most of the fundamental appeals of conformal prediction, i.e. it’s post-hoc, modelling-free, training-free nature. The appeals of the particular modelling-heavy pipeline should be motivated further (especially in light of limited experimentation), or the scope properly set to the types of problems where this does not constitute any major limitations. I’d be happy to learn more about the range of use cases the authors have in mind.\n- Experimental validation, while claimed to “extensively evaluate” and “demonstrate conclusively” only contains one synthetic experiment and one tabular regression task, with two external baselines (ETO-SLL, ETO-JC). Furthermore, baseline comparisons are limited to particular settings, e.g. PICNN only compares internal baselines, and the only full comparison is on ellipsoidal sets. Two different groups of baselines seem relevant: 1) when it comes to multivariate prediction sets, the literature offers a substantially wider range of shapes beyond box/ellipse, e.g. consider recent comparison (Dheur et al. 2025). This includes similarly flexible but non-convex approaches; 2) as internalization of CP into the training objective is performed, it seems that a fair comparison should similarly incorporate conformal training (Stutz et al. ) or related uncertainty “tuning” methods like conformal boosting (Ran et al. 2024) across baselines. When it comes to performed experiments, while interesting they are mainly limited to the iid data case, and I find the shift results not very convincing especially in light of employing actual time series data. More experiments highlighting realistic use cases should be done to help highlight the usefulness of going the modelling-specific route. For example, in Tab. 1 all results are well within standard deviations and obtained task loss benefits (and visualized uncertainty sets) seem questionable. Overall, more experiments are needed and the claims on empirical validation/advantage should be softened.\n\nMethod Clarifications / Questions\n- Perhaps I am misunderstanding but does the PICNN include both prediction and actually modeling the score function? Is any distinction in that matter made? When we update the model parameters, should I interpret this as underlying “finetuning” or rather updates to the conformal calibration procedure only? I was confused about this looking at Fig. 1.\n- ETO is invalidated as lacking task loss feedback. But instead of doing gradient-based model updates, could we not simply feed task loss information into the conformal calibration procedure by designing smart nonconformity scores accounting for task loss, or doing some other form of post-hoc quantile adjustments?\n- Looking at Algo. 1, do I understand correctly that a conformal quantile and solution of the optimization problem needs to be recomputed for every test input at inference time? If so, any comments on computational costs? And if so, why do we separate a train and inference step but not simply continuously update the model parameters by computing gradients (as we already solve the optimization) to adapt the model even further, potentially robustifying to slow shifts?\n- In sec. 3.3., is there a reason to do the end-to-end training with ERM? Can we not directly go for some more robust optimization procedure here? I perhaps misunderstood how this relates to the “robust” optimization objective Eq. 1.\n- Fig. 4 displays marginal coverage levels and states “confirming that the improvements in task loss performance from our E2E approach do not come at the\ncost of worse coverage” . From my understanding of the method after the training/updating step, split CP is simply added on top during inference, so the coverage rate by the conformal quantile does not directly relate to the task loss and is trivially satisfied (when exchangeable). Am I missing the connection to the task here?\n\nReferences\n- Dheur, Victor, et al. \"Multi-Output Conformal Regression: A Unified Comparative Study with New Conformity Scores.\" arXiv preprint arXiv:2501.10533 (2025).\n- Xie, Ran, Rina Barber, and Emmanuel Candes. \"Boosted conformal prediction intervals.\" Advances in Neural Information Processing Systems 37 (2024): 71868-71899.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper focuses on optimizing uncertainty-aware decision making in a more holistic, trainable fashion. The paper combines conformal prediction (to generate uncertainty sets) with a decision-based task loss to obtain a min-max optimization objective that incorporates conformal and task information, and obtains subsequent loss gradients to update the underlying model. The paper makes convexity assumptions on the objectives and employs a partial-convex neural network (PICNN) to render the optimization tractable. The PICNN models the conformal scoring function in a flexible, convex fashion, permitting generating multi-variate convex uncertainty sets that improve upon simpler box or ellipsoidal sets. A synthetic 2D portfolio experiment is used to visually motivate improvements over two baselines, and tabular regression on energy pricing time series (re-shuffled to ensure exchangeability) further motivates improvements on the task loss (suggesting better decision-making) while keeping conformal coverage.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- The proof of Prop. 3.3 is rather a comment or description highlighting where the results come from, but not a proper proof. A more explicit proof should be included in the Appendix for completeness. Also, more details on the conformal training approach taken from Stutz et al. should be added to the Appendix, as this forms an important part of the pipeline.\n- Visualizations of the obtained conformal uncertainty sets/intervals for the battery experiment would be good to have.\n\nOther comments\n- Nonconformity score uniqueness could also be ensured with some extra jitter in practice (instead of simply trusting in it).\n- What is the intuition behind permitting negative task losses?\n- Is my understanding correct that the task loss reduces (Fig. 3) with higher uncertainty level alpha because the optimization becomes less restrictive, i.e. the effect of the inner max. is lessened?", "broader_impact_concerns": "I have no ethical concerns regarding the paper.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "cPJpbCMp00", "reviewer": "Reviewer_qLpG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n* **Novelty and Significance**: The work presents an interesting approach to bridge the gap between uncertainty quantification and downstream optimization tasks. The proposed end-to-end framework, which directly incorporates the task loss into the learning of uncertainty sets, offers a compelling alternative to the prevailing estimate-then-optimize (ETO) paradigm.\n\n* **Theoretical Guarantees**: The framework is theoretically sounding. It provides formal, finite-sample marginal coverage guarantees for its uncertainty sets by leveraging split conformal prediction. \n\n* **Methodological Soundness**: The proposed use of Partially Input-Convex Neural Networks (PICNNs) to parameterize general convex uncertainty sets is methodologically sound. \n\n* **Rigorous Empirical Validation**: The claims are substantiated by empirical validation on two relevant applications: energy storage arbitrage and portfolio optimization. The framework's performance is superior to multiple ETO baselines across various uncertainty models and risk levels.\n\n\n**Weaknesses and Areas for Improvement**\n\n* **Performance Degradation Under Distribution Shift**: The empirical results show the framework's performance degradation under distribution shift, where the exchangeability assumption essential for conformal prediction is violated. The observation that more expressive models exhibit a more significant drop in coverage is important. However, the analysis could be deepened by investigating the mechanisms behind this vulnerability rather than primarily noting it as an avenue for future work\n\n* **Practical Complexity and Training Stability**: From the appendix, it seems that the practical implementation of the end-to-end training procedure is non-trivial and sensitive to the optimization setup. For example, as detailed in the appendix, stable training necessitated the use of specific regularization of the conformal quantile for the PICNN model to prevent divergence. This suggests that the decision-focused loss landscape is challenging to optimize directly and that the reported performance may depend on these carefully chosen heuristics.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a novel methodology for learning decision-aware uncertainty representations for conditional robust optimization. The primary contributions are threefold:\n\n* The paper develops a novel end-to-end (E2E) training framework that directly accounts for the downstream task loss. By backpropagating through both the robust optimization problem and a conformal calibration procedure, the framework learns uncertainty sets that are explicitly shaped by their utility in the decision-making context. This presents a significant departure from conventional estimate-then-optimize (ETO) paradigms, which decouple the prediction and optimization stages.\n\n* The authors propose the use of Partially Input-Convex Neural Networks (PICNNs) to parameterize general convex uncertainty sets. This approach enhances the expressiveness of the uncertainty representation beyond standard approaches e.g., boxes or ellipsoids. \n\n* The work introduces an efficient method for differentiating through the quantile function. This technique avoids the need for smoothed approximations used in prior work, enabling precise gradient-based learning while preserving the finite-sample coverage guarantees of conformal prediction.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* **Dataset description**: I am not familiar with the specific datasets used by authors. However, they mention that the price forecast for the battery storage problem is conditioned on the price on the past day. I wonder whether that violates the exchangeability assumption given that the dataset seems to comprise time series.  \n\n* **Deepen the Discussion on Distribution Shift**: In Section 5.4, the authors may expand the analysis of why the more expressive models exhibit worse coverage under distribution shift. A plausible hypothesis is that their flexibility allows them to overfit to spurious correlations in the training distribution that do not generalize. Suggesting potential avenues for mitigation (e.g., specific regularization techniques, domain adaptation methods) would strengthen this section.\n\n* **Reproducibility**: authors may consider to move crucial details about the hybrid loss functions and regularization from the appendix into the main methodology (Section 3.3 or 4). That would provide readers with a better understanding of what is required to successfully train these models and strengthens the paper's contribution to practical application.\n\n* **Computational Requirements**: (related to the previous point) - the authors may add more details concerning scalability limitation of their approach. E.g., concerning the optimization library used. This is an important practical consideration for other researchers looking to build on this work.", "broader_impact_concerns": "I have no concerns regarding the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)