SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chunlin Sun*,1 Linyu Liu*,2 Xiaocheng Li3
$^{1}$  Institute for Computational and Mathematical Engineering, Stanford University
$^{2}$  Department of Automation, Tsinghua University
<sup>3</sup> Imperial College Business School, Imperial College London
Contextual optimization, also known as predict-then-optimize or prescriptive analytics, considers an optimization problem with the presence of covariates (context or side information). The goal is to learn a prediction model (from the training data) that predicts the objective function from the covariates, and then in the test phase, solve the optimization problem with the covariates but without the observation of the objective function. In this paper, we consider a risk-sensitive version of the problem and propose a generic algorithm design paradigm called predict-then-calibrate. The idea is to first develop a prediction model without concern for the downstream risk profile or robustness guarantee, and then utilize calibration (or recalibration) methods to quantify the uncertainty of the prediction. While the existing methods suffer from either a restricted choice of the prediction model or strong assumptions on the underlying data, we show the disentangling of the prediction model and the calibration/uncertainty quantification has several advantages. First, it imposes no restriction on the prediction model and thus fully unleashes the potential of off-the-shelf machine learning methods. Second, the derivation of the risk and robustness guarantee can be made independent of the choice of the prediction model through a data-splitting idea. Third, our paradigm of predict-then-calibrate applies to both (risk-sensitive) robust and (risk-neutral) distributionally robust optimization (DRO) formulations. Theoretically, it gives new generalization bounds for the contextual LP problem and sheds light on the existing results of DRO for contextual LP. Numerical experiments further reinforce the advantage of the predict-then-calibrate paradigm in that an improvement on either the prediction model or the calibration model will lead to a better final performance.
Contextual linear programming (LP) considers the linear programming problem with the presence of covariates, and it can be viewed as a prediction problem under a constrained optimization context where the output of the prediction model serves as the objective of a downstream LP. The goal is to develop a model (trained from past data) that prescribes a decision/solution for the downstream LP problem using directly the covariates but without observation of the objective function. A similar problem formulation was also studied as prescriptive analytics [Bertsimas and Kallus, 2020] and predict-then-optimize Elmachtoub and Grigas [2022]. The existing literature mainly develops two approaches for the problem: (i) end-to-end algorithms that map the covariates directly to a recommended decision for the LP or (ii) two-step algorithms that first predict the unobserved objective using the covariates and then solve the LP with the predicted objective. While the literature has been largely focused on the risk-neutral objective, our paper considers the risk-sensitive/robust
formulation of the problem. The robust formulation of the contextual LP can be viewed as both a risk-sensitive version of the risk-neutral objective in contextual LP and an adaptive/contextual version of the standard (context-free) robust optimization problem. In comparison with the robust optimization literature, the contextualized robust optimization problem allows a contextual and data-driven way to design the uncertainty set and ideally can utilize the contextual information to make less conservative decisions. The contextual/data-driven uncertainty set has been considered by several existing works on transportation (Guo et al. [2023]), portfolio management (Wang et al. [2022]), and healthcare (Gupta et al. [2020]). While these works focus on applications with special structures or distributions, our paper considers a generic setup and aims for a general solution to the problem.
In our paper, we make the following contribution:
First, we formulate an algorithm design paradigm of predict-then-calibrate and introduce the notion of uncertainty calibration/quantification to the problem of contextual optimization. The idea is to first develop a prediction model without worrying about the downstream robustness, and then tailor the uncertainty calibration methods to quantify the prediction uncertainty towards the robust objective.
Second, we develop two algorithms for robust contextual LP under the paradigm of predict-then-calibrate. The derivation of theoretical guarantees shows the advantage of the disentanglement of the prediction phase and the uncertainty calibration phase. Moreover, we quantify the advantage of a better prediction and/or a better uncertainty calibration both theoretically and empirically.
Third, we also draw a connection with the related problem of distributionally robust optimization. We utilize the same idea of predict-then-calibrate and derive a new generalization bound for the distributional robust version of contextual LP utilizing tools from the nonparametric regression.
In the following, we review the several streams of literature that are related to our work.
Conditional robust optimization. The problem studied in our paper can be viewed as a conditional version of the classic robust optimization problem. The robust optimization stems from the replacement of the risk-neutral expectation objective with a risk-sensitive robust objective such as Value-at-Risk(VaR)/quantile. The literature on robust optimization develops methods to approximate VaR by specifying the shape and size of an uncertainty set to guarantee that a VaR constraint would be satisfied Ghaoui et al. [2003], Chen et al. [2007], Natarajan et al. [2008]. Bertsimas et al. [2018] designs the uncertainty set with a finite-sample probabilistic guarantee on optimal solutions. Goerigk and Kurtz [2020] constructs the uncertainty set as the union of  $K$  convex sets by clustering historical data. While all these works concern the context-free setting, the contextual (or conditional) setting, or the problem of robust contextual LP, considers the presence of covariates in robust optimization. Along this line, Ohmori [2021] proposes a volume-minimized ellipsoid covering  $k$ -nearest samples in the contextual space, leading to a nonlinear programming problem. Chenreddy et al. [2022] extends the idea in Goerigk and Kurtz [2020] and novelly adopts a deep learning approach to the problem which identifies which clusters the uncertain parameters belong to based on the covariates. Compared to the end-to-end formulations Ohmori [2021], Chenreddy et al. [2022], we do not impose restrictions on the prediction model and have a larger flexibility in choosing the uncertainty quantification method.
Contextual LP/Predict-then-optimize. Our work complements the existing literature on contextual LP by considering a risk-sensitive objective. The existing works Donti et al. [2017], Elmachtoub and Grigas [2022], Ho-Nguyen and Klinç-Karzan [2022], Hu et al. [2022] mainly study the risk-neutral objective for the problem and aim to develop prediction methods that account for the downstream optimization tasks. To model the remaining uncertainty after prediction in optimization, Sen and Deng [2018] considers training a parameterized point-wise estimation model first, and then approximating the conditional distribution by imposing randomness to the learned parameters empirically. Ban et al. [2019] proposes to forecast the distribution of new product demand by adopting a linear regression model to fit the demands of similar products with covariates, and then using the empirical distribution of covariate-independent residuals to build a sample average approximation (SAA) model. Some recent works Kannan et al. [2020, 2022, 2021] further generalize it to a residuals-based SAA framework to include more machine learning methods and adopt distributionally robust optimization to mitigate the overfitting under limited data. A key difference between our work and this line of works is that we consider a risk-sensitive objective, while these works consider a risk-neutral objective. In this light, our framework of predict-then-calibrate is compatible with any prediction model developed along this line, and it tailors the prediction model for a robustness task.
Consider a linear program (LP) that takes the following standard form
$$
\mathrm {L P} (c, A, b) := \min  c ^ {\top} x, \tag {1}
$$
$$
\begin{array}{l} \text {s . t .} A x = b, x \geq 0, \end{array}
$$
where  $c \in \mathbb{R}^n$ ,  $A \in \mathbb{R}^{m \times n}$ , and  $b \in \mathbb{R}^m$  are the inputs of the LP. In addition, there is an available feature vector  $z \in \mathbb{R}^d$  that encodes the covariates (side information) associated with the LP.
The tuple  $(c, A, b, z)$  is drawn from an unknown probability distribution  $\mathcal{P}$ . The problem of contextual LP (predict-then-optimize, or prescriptive analytics) usually assumes the availability of a learning (training) dataset sampled from the distribution  $\mathcal{P}$ ,
$$
\mathcal {D} _ {l} := \left\{\left(c _ {t} ^ {(l)}, A _ {t} ^ {(l)}, b _ {t} ^ {(l)}, z _ {t} ^ {(l)}\right) \right\} _ {t = 1} ^ {T _ {l}}.
$$
During the test phase, one needs to recommend a feasible solution  $x_{\mathrm{new}}$  to a new LP problem using the observation of  $(A_{\mathrm{new}}, b_{\mathrm{new}}, z_{\mathrm{new}})$  but without the observation of the objective vector  $c_{\mathrm{new}}$ :
$$
\left(A _ {\text {n e w}}, b _ {\text {n e w}}, z _ {\text {n e w}}\right)\rightarrow x _ {\text {n e w}}.
$$
The focus of the existing literature has been mainly on the risk-neutral objective
$$
\min  _ {x} \mathbb {E} [ c | z ] ^ {\top} x, \tag {2}
$$
$$
\begin{array}{l} \text {s . t .} A x = b, x \geq 0 \end{array}
$$
where the recommended decision  $x$  can be viewed as a function of  $(A, b, z)$ .
Alternatively, we consider a risk-sensitive objective
$$
\min  _ {x} \operatorname {V a R} _ {\alpha} \left[ c ^ {\top} x | z \right], \tag {3}
$$
where  $\alpha \in (0,1)$ . Here  $\mathrm{VaR}_{\alpha}(U)$  denotes the  $\alpha$ -quantile/value-at-risk (VaR) of a random variable  $U$ ; specifically,  $\mathrm{VaR}_{\alpha}(U) \coloneqq F_U^{-1}(\alpha)$  with  $F_U^{-1}(\cdot)$  be the inverse cumulative distribution function of  $U$ .
We can interpret the problem (3) in two ways: First, it can be viewed as a risk-sensitive version of the contextual LP problem (2). From the prediction viewpoint, the standard problem (2) concerns only the prediction of  $\mathbb{E}[c|z]$ , while solving (3) may involve a distributional prediction of the conditional distribution  $c|z$ . Second, the classic robust LP problem considers a setting where there is no covariate, and the VaR is taken with respect to the (marginal) distribution of  $c$ . Comparatively, the problem (3) can be viewed as a conditionally robust version where the VaR is taken with respect to the conditional distribution  $c|z$ . The presence of the covariates  $z$  reduces our uncertainty on the objective  $c$  and thus will induce less conservative decisions with the same level of risk guarantee.
Throughout the paper, we work with a well-trained ML model  $\hat{f}:\mathbb{R}^d\to \mathbb{R}^n$  that predicts the objective with the covariates and is learned from the training data  $\mathcal{D}_l$ . We emphasize that we do not impose any assumption on the model, but assume the availability of a validation set  $\mathcal{D}_{\mathrm{val}}\coloneqq \{(c_t,A_t,b_t,z_t)\}_{t = 1}^T$  that can be used to quantify the uncertainty of  $\hat{f}$ .
3 Robust Contextual LP
The problem (3) has been extensively studied under the context-free setting in the literature of robust optimization [Ben-Tal et al., 2009], and the key challenge is the non-convexity of the objective function. Equivalently, the problem (3) can be written as the following optimization problem
$$
\min  _ {x, \mathcal {U}} \max  _ {c \in \mathcal {U}} c ^ {\top} x, \tag {4}
$$
$$
\mathrm {s . t .} A x = b, \quad x \geq 0, \mathbb {P} _ {c | z} (c \in \mathcal {U}) \geq \alpha
$$
where the decision variables are  $x$ ,  $c$  and a measurable set  $\mathcal{U}$ . The last constraint is with respect to the set  $\mathcal{U}$ , and the probability is taken with respect to the conditional distribution  $c|z$ . One issue to
solve this problem is the intractability in optimization over the uncertainty set. Particularly, even if the conditional distribution  $c|z$  is discrete, (4) is equivalent to a mixed integer linear programming problem, which is generally NP-hard and computationally intensive. Similarly, when solving a sample average approximation form of (4) (Bertsimas and Kallus [2020]), the problem becomes similar to the discrete case, and the intractability issue is still inevitable. To resolve this issue, one can solve the following approximation problem with a fixed uncertainty set  $\mathcal{U}$  satisfying  $\mathbb{P}_{c|z}(c\in \mathcal{U})\geq \alpha$
$$
\mathrm {L P} (\mathcal {U}) := \min  _ {x} \max  _ {c \in \mathcal {U}} c ^ {\top} x, \tag {5}
$$
$$
\begin{array}{l} \text {s . t .} A x = b, x \geq 0. \end{array}
$$
A box- or ellipsoid-shaped uncertainty set enables a tractable optimization of the problem (5). The uncertainty set hopes to cover the high-density region so that the approximation to (3) is tighter.
Now we present two algorithms that implement this approximation idea in the contextual setting. The key is to quantify the uncertainty of the prediction model. Both algorithms construct contextual uncertainty sets for the prediction model; that is, the size of the uncertainty set changes with respect to the covariates  $z$ . In addition, importantly, the construction of these uncertainty sets accounts for the tractability of the downstream robust optimization problem (5) by outputting the box and ellipsoid shapes, respectively. In both algorithms, we first split the available validation data into two sets  $\mathcal{D}_1$  and  $\mathcal{D}_2$ , and use the first set  $\mathcal{D}_1$  for a preliminary calibration, and the second set  $\mathcal{D}_2$  for an additional adjustment. The output from both algorithms gives a contextual uncertainty set  $\mathcal{U}(z)$  that works as the input of (5) and thus solves the downstream robust optimization problem.
Algorithm 1 Box Uncertainty Quantification (BUQ)
1: Input: Dataset  $\mathcal{D}_{val}$ , ML model  $\hat{f}$ , parameter  $\alpha$
2: Randomly split the validation data into two index sets  $\mathcal{D}_1\cup \mathcal{D}_2 = \{1,\dots,T\}$  and  $\mathcal{D}_1\cap \mathcal{D}_2 = \emptyset$
3:  $\% \%$  Preliminary calibration of the uncertainty sets
4: for  $t \in \mathcal{D}_1$  do
5: Calculate the residual vector on the  $t$ -th validation sample
$$
r _ {t} := c _ {t} - \hat {f} \left(z _ {t}\right) \tag {6}
$$
and denote  $r_t = (r_{t1},\dots,r_{tn})^\top$
7: Train a quantile regression model  $\hat{h}(z): \mathbb{R}^d \to \mathbb{R}^n$  that minimizes
$$
\sum_ {t \in \mathcal {D} _ {1}} \sum_ {i = 1} ^ {n} \rho_ {\alpha} \left(\left| r _ {t i} \right| - \hat {h} \left(z _ {t}\right) _ {i}\right) \tag {7}
$$
where  $\rho_{\alpha}(\cdot)\coloneqq \alpha (\cdot)^{+} + (1 - \alpha)(\cdot)^{+}$  denotes the pinball loss
8:  $\% \%$  Additional adjustment of the size of the uncertainty sets
9: for  $t \in \mathcal{D}_2$  do
10: Let
$$
\bar {c} _ {t} (\eta) := \hat {f} (z _ {t}) + \eta \hat {h} (z _ {t}) \in \mathbb {R} ^ {n}, \underline {{c}} _ {t} (\eta) := \hat {f} (z _ {t}) - \eta \hat {h} (z _ {t}) \in \mathbb {R} ^ {n}
$$
11: end for
12: Choose a minimal  $\eta > 0$  such that
$$
\sum_ {t \in \mathcal {D} _ {2}} 1 \left\{\underline {{c}} _ {t} (\eta) \leq c _ {t} \leq \bar {c} _ {t} (\eta) \right\} \geq \min  \left\{| \mathcal {D} _ {2} |, \alpha (| \mathcal {D} _ {2} | + 1) \right\} \tag {8}
$$
where  $1\{\cdot \}$  is the indicator function and the inequality is required to hold element-wise
13: Output:  $h$ ,  $\eta$
In Algorithm 1, we first compute the residual/ error vector  $r_t \in \mathbb{R}^n$  for samples in dataset  $\mathcal{D}_1$ . Then for each coordinate  $i = 1, \dots, n$ , we fit a quantile regression model  $\hat{h}$  that predicts the quantile of the absolute error by minimizing (7). Specifically, the model  $\hat{h}$  aims to predict the quantiles of each coordinate of the residual vector  $r_t$ . Then, in the second step of Algorithm 1, we pretend the quantile model  $\hat{h}$  is "correct" and construct the uncertainty sets accordingly for samples in dataset
$\mathcal{D}_2$ . Importantly, the uncertainty set is parameterized by a scalar  $\eta$ . We use the parameter  $\eta$  to perform an additional adjustment of the uncertainty set so that the coverage probability  $\alpha$  is met on  $\mathcal{D}_2$  by (8). The algorithm outputs the function  $\hat{h}$  and  $\eta$ , and it will output a box-shaped uncertainty set  $\mathcal{U}_{\alpha}^{(1)}(z) = [\hat{f}(z) - \eta \hat{h}(z), \hat{f}(z) + \eta \hat{h}(z)]$  for a new sample with covariates  $z$ .
Algorithm 2 Ellipsoid Uncertainty Quantification (EUQ)
$$
r _ {t} := c _ {t} - \hat {f} (z _ {t})
$$
7: Use  $\mathcal{D}_1$  to fit a model  $\hat{g}:\mathbb{R}^d\to \mathbb{R}$  that predicts the  $\alpha$ -quantile of  $\| r_t\| _2$  using the covariates  $z_{t}$
8: Fit a zero-mean normal distribution for the scaled residual vector
$$
\left\{r _ {t} / \hat {g} \left(z _ {t}\right): t \in \mathcal {D} _ {1} \right\}
$$
and denote the covariance matrix as  $\hat{\Sigma}$
9: %%Additional adjustment of the size of the uncertainty sets
10: Choose a minimal  $\eta > 0$  such that
$$
\sum_ {t \in \mathcal {D} _ {2}} 1 \left\{\sqrt {\left(c _ {t} - \hat {f} (z _ {t})\right) ^ {\top} \hat {\Sigma} ^ {- 1} \left(c _ {t} - \hat {f} (z _ {t})\right)} \leq \eta \hat {g} (z _ {t}) \right\} \geq \min  \left\{| \mathcal {D} _ {2} |, \alpha (| \mathcal {D} _ {2} | + 1) \right\} \tag {9}
$$
where  $1\{\cdot \}$  is the indicator function and the inequality is required to hold element-wise.
11: Output:  $\hat{g},\hat{\Sigma},\eta$
Algorithm 2 follows a similar spirit as Algorithm 1. It first computes the residual vector  $r_t$  for samples in dataset  $\mathcal{D}_1$ , but differently from Algorithm 1, it fits a regression model for the norm  $\| r_t \|_2$ . Then it fits a zero-mean normal distribution for the scaled residual vector. The motivation is that the residual vector (up to a scale  $\hat{g}(z)$ ) follows a multivariate normal distribution. We remark that although the algorithm design is motivated by such a setting of normal distribution, its theoretical guarantee does not rely on this normal structure. The key is the second step of additional calibration: it pretends the uncertainty set is "correct" and uses an additional scalar  $\eta$  to make further adjustments such that the coverage probability  $\alpha$  is met on  $\mathcal{D}_2$  by (9). For a new sample with covariates  $z$ , the algorithm outputs an ellipsoid-shaped uncertainty set  $\mathcal{U}_{\alpha}^{(2)}(z) = \left\{c : \sqrt{(c - \hat{f}(z))^{\top} \hat{\Sigma}^{-1}(c - \hat{f}(z))} \leq \eta \hat{g}(z)\right\}$ .
In comparison, for the first step of preliminary calibration, Algorithm 1 quantifies the uncertainty of each coordinate of the residual vector individually, while Algorithm 2 presumes a Gaussian structure. We remark that the two algorithms provide two examples for the design of the preliminary calibration step, but these are not the only two options. As we will see shortly, the second step of additional adjustment is in charge of the coverage guarantee, so the first step can be further modified based on other prior knowledge of the underlying data/optimization problem.
3.1 Algorithm Analysis
The following proposition provides a coverage guarantee for both Algorithm 1 and 2 in that the output uncertainty set will cover a new sample with a probability of approximately  $\alpha$ .
Proposition 1. For a new sample  $(c, A, b, z)$  from distribution  $\mathcal{P}$ , denote the uncertainty sets output from Algorithm 1 and Algorithm 2 by  $\mathcal{U}_{\alpha}^{(1)}(z)$  and  $\mathcal{U}_{\alpha}^{(2)}(z)$ , respectively. The following inequalities hold for  $k = 1, 2$ ,
$$
\alpha \leq \mathbb {P} \left(c \in \mathcal {U} _ {\alpha} ^ {(k)} (z)\right) \leq \alpha + \frac {1}{| \mathcal {D} _ {2} | + 1} \tag {10}
$$
where the probability is with respect to the new sample  $(c,A,b,z)$  and dataset  $\mathcal{D}_2$ .
The proposition gives the following insights into the two algorithms and the general problem of robust contextual LP. First, the second step (in both algorithms) of the adjustment via the scalar  $\eta$  is the key to the coverage guarantee (10), with several advantages. First, as the concentration argument is made with respect to only this second step, the bound is rather tight and free of unnecessary conservation in the uncertainty set construction. Second, it allows full flexibility in the choice of the prediction model  $\hat{f}$  and the preliminary calibration model  $\hat{h}$  (Algorithm 1) and  $\hat{g}$  (Algorithm 2). This reinforces the notion of predict-then-calibrate. The design of the prediction model  $\hat{f}$  can be solely targeted on the accuracy, and the coverage guarantee can be deferred and taken care of by the calibration steps such as these two algorithms.
Furthermore, we note the probability in (10) is taken with respect to only the validation data and the new sample. It means Proposition 1 holds even if the training data  $D_{l}$  (for obtaining  $\hat{f}$ ) is not sampled from the distribution  $\mathcal{P}$ . In other words, the result only requires the validation data  $\mathcal{D}_{\mathrm{val}} \sim \mathcal{P}$  and it is robust to a distribution shift of the training distribution when  $D_{l} \not\sim \mathcal{P}$ . In a broader context, when the training data, validation data, and test data come from different distributions, this situation is commonly referred to as the out-of-domain (OOD) problem. We refer to the survey papers Shen et al. [2021], Yang et al. [2021] for a more detailed discussion of this problem. Developing more coverage guarantees with OOD data will require future investigations, and we will list this problem as a future direction.
In addition, the coverage guarantee in Proposition 1 can be extended to a robustness guarantee for the contextual LP problem as in the following corollary.
Corollary 1. For a new sample  $(c, A, b, z)$  from distribution  $\mathcal{P}$ , denote the uncertainty set output from Algorithm 1 or Algorithm 2 by  $\mathcal{U}_{\alpha}(z)$ . Let  $x^{*}(z)$  and  $OPT(z)$  be the optimal solution and the optimal value of  $LP(\mathcal{U}_{\alpha}(z))$  (5). Then we have
$$
\mathbb {P} \left(c ^ {\top} x ^ {*} (z) \leq O P T (z)\right) \geq \alpha
$$
where the probability is taken with respect to the new sample  $(c, A, b, z)$  and  $\mathcal{D}_2$ .
To end this section, we remark that although the coverage guarantee (10) is in a population sense where the probability is taken with respect to both  $c$  and  $z$ , our framework of predict-then-calibrate also has the flexibility to achieve an individual coverage by changing the calibration algorithm under some assumptions. Here, the individual coverage means that the conditional probability  $\mathbb{P}\left(c\in \mathcal{U}_{\alpha}(z)|z\right) = \alpha$  holds for any covariate  $z$  as the constraint in (4), where  $\mathcal{U}_{\alpha}(z)$  denotes any given contextualized uncertainty set. However, the above-mentioned assumption to achieve individual coverage cannot be verified from the data prior, and achieving this coverage is, in general, a hard problem, although there have been many attempts in the literature on conformal prediction and model calibration. In addition, we want to emphasize that, compared to achieving individual coverage, the most significant focus of this paper is to introduce the idea of uncertainty quantification to solve robust optimization, which disentangles the prediction from the calibration. As a result, we delay the discussion of the individual coverage to Appendix B.1.
3.2 Value of Better Prediction and Contextual Uncertainty Set