SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dishank Bansal* Ricky T. Q. Chen Mustafa Mukadam Brandon Amos Meta
Deep learning models are often used with some downstream task. Models solely trained to achieve accurate predictions may struggle to perform well on the desired downstream tasks. We propose using the task loss to learn a metric which parameterizes a loss to train the model. This approach does not alter the optimal prediction model itself, but rather changes the model learning to emphasize the information important for the downstream task. This enables us to achieve the best of both worlds: a prediction model trained in the original prediction space while also being valuable for the desired downstream task. We validate our approach through experiments conducted in two main settings: 1) decision-focused model learning scenarios involving portfolio optimization and budget allocation, and 2) reinforcement learning in noisy environments with distracting states. The source code to reproduce our experiments is available here.
Machine learning models for prediction are typically trained to maximize the likelihood on a training dataset. While the models are capable of universally approximating the underlying data generating process to predict the output, they are prone to approximation errors due to limited training data and model capacity. These errors lead to suboptimal performance in downstream tasks where the models are used. Furthermore, even though a model may appear to have reasonable predictive performance on the metric and training data it was trained on, such as the mean squared error, employing the model for a downstream task may require the model to focus on different parts of the data that were not emphasized in the training for predictive performance. Overcoming the discrepancy between the model's prediction task and performance on a downstream task is the focus of our paper.
Figure 1: The MSE results in a model close to the true model in the prediction space, but may give poor task performance. Decision-focused learning (DFL) methods optimize the task loss, but may deviate from the prediction space. TaskMet optimizes the task loss while retaining the prediction task.
Examples of settings where the model's prediction loss  $\mathcal{L}_{\mathrm{pred}}$  is mis-matched from the downstream task  $\mathcal{L}_{\mathrm{task}}$  include the following, which table 1 also summarizes:
1. the portfolio optimization setting from Wilder et al. [2019], which predicts the expected returns from stocks for a financial portfolio. Here, the  $\mathcal{L}_{\mathrm{pred}}$  is the MSE and  $\mathcal{L}_{\mathrm{task}}$  is from the regret of running a portfolio optimization problem on the output;
2. the allocation setting from Wilder et al. [2019], which predicts the value of items that are being allocated, e.g. click-through-rates for recommender systems. Here,  $\mathcal{L}_{\mathrm{pred}}$  is the MSE and  $\mathcal{L}_{\mathrm{task}}$  measures the result of allocating the highest-value items.
*Work done as part of the Meta AI residency program.
Table 1: Settings we focus on where there is a discrepancy between the prediction task of a model and the downstream task where the model is deployed, i.e.,  $\mathcal{L}_{\mathrm{pred}} \neq \mathcal{L}_{\mathrm{task}}$ .
Setting	(x)
Features	(y)
Predictions	(Lpred)
Prediction Loss	(Ltask)
Task Loss
Portfolio optimization	Stock information	Expected return of a stock	MSE	Portfolio's performance
Budget allocation	Item information	Value of item	MSE	Allocation's performance
Model-based RL	Current state and action	Next state	MSE	Value estimation given the model
3. the model-based reinforcement learning setting of learning the system dynamics from Nikishin et al. [2022]. Here,  $\mathcal{L}_{\mathrm{pred}}$  is the MSE of dynamics model and the  $\mathcal{L}_{\mathrm{task}}$  measures how well the agent performs for downstream value predictions.
Motivated by examples such as in table 1, the research topics of end-to-end task-based model learning [Bengio, 1997, Donti et al., 2017], decision-focused learning [Wilder et al., 2019], and Smart "Predict, then Optimize" [Elmachtoub and Grigas, 2022] study how to use information from the downstream task to improve the model's performance on that particular task. Task-based learning has applications in financial price predictions [Bengio, 1997, Elmachtoub and Grigas, 2022], inventory stock, demand, and price forecasting [Donti et al., 2017, Elmachtoub and Grigas, 2022, El Balghiti et al., 2019, Mandi et al., 2020, Liu et al., 2023], dynamics modeling for model-based reinforcement learning [Farahmand et al., 2017, Amos et al., 2018, Farahmand, 2018, Bhardwaj et al., 2020, Voelcker et al., 2022, Nikishin et al., 2022], renewable nowcasting [Vohra et al., 2023], vehicular routing [Shi and Tokekar, 2023], restless multi-armed bandits for maternal and child care [Wang et al., 2022], medical resource allocation [Chung et al., 2022], and budget allocation, matching, and recommendation problems [Kang et al., 2019, Wilder et al., 2019, Shah et al., 2022].
Limitations of task-based learning. Task-based model learning comes with the goal of being able to discover task-relevant features and data-samples on its own without the need of explicit inductive biases. The current trend for end-to-end model learning uses task loss along with the prediction loss to train the prediction models. Though easy to use, these methods may be limited by 1) the prediction overfitting to the particular task, rendering it unable to generalize; 2) the need to tuning the weight combining the task and prediction losses as in eq. (1).
Our contributions. We propose one way of overcoming these limitations: use the task-based learning signal not to directly optimize the weights of the model, but to shape a prediction loss that is constructed in a way so that the model will always stay in the original prediction space. We do this in section 3 via metric learning in the prediction space and use the task signal to learn a parameterized Mahalanobis loss. This enables more interpretable learning of the model using the metric compared to learning with a combination of task loss and prediction loss. The learned metric can uncover underlying properties of the task that are useful for training the model, e.g. as in figs. 4 and 7. Section 4 shows the empirical success of metric learning on decision focused model learning and model-based reinforcement learning. Figure 1 illustrates the differences to prior methods.
Task-based model learning. We will mostly focus on solving regression problems where the dataset  $\mathcal{D} \coloneqq \{(x_i, y_i)\}_{i=1}^N$  consists of  $N$  input-output pairs, which we will assume to be in Euclidean space. The model makes a prediction  $\hat{y} \coloneqq f_\theta(x)$  and is parameterized by  $\theta$ . The model has an associated prediction loss,  $\mathcal{L}_{\mathrm{pred}}$ , and is used in conjunction with some downstream task that provides a task loss,  $\mathcal{L}_{\mathrm{task}}$ , which characterizes how well the model performs on the task. The most relevant related work to ours includes the approaches of Bengio [1997], Donti et al. [2017], Farahmand et al. [2017], Kang et al. [2019], Wilder et al. [2019], Nikishin et al. [2022], Shah et al. [2022], Voelcker et al. [2022], Nikishin et al. [2022], Anonymous [2023], Shah et al. [2023], which learn the optimal prediction model parameter  $\theta$  to minimize the task loss  $\mathcal{L}_{\mathrm{task}}$ :
$$
\theta^ {\star} := \underset {\theta} {\arg \min } \mathcal {L} _ {\text {t a s k}} (\theta) + \alpha \mathcal {L} _ {\text {p r e d}} (\theta), \tag {1}
$$
where  $\alpha$  is a regularization parameter to weigh the prediction loss which is MSE error (eq. (2)) in general. Alternatives to eq. (1) include 1) Smart, "Predict, then Optimize" (SPO) methods [Elmachtoub and Grigas, 2022, El Balghiti et al., 2019, Mandi et al., 2020, Liu et al., 2023], which
consider surrogates for when the derivative is undefined or uninformative, or 2) changing the prediction space from the original domain into a latent domain with task information, e.g. task-specific latent dynamics for RL [Hafner et al., 2019b,a, Hansen et al., 2022]. Extensions such as Gupta and Zhang [2023], Zharmagambetov et al. [2023], Ferber et al. [2023] learn surrogates to overcome computationally expensive losses in eq. (1). Sadana et al. [2023] provide a further survey of this research area.
Separate from above line of work, the computer vision and NLP communities have also considered task-based losses for models: [Pinto et al., 2023] tune vision models with task rewards, e.g. for detection, segmentation, colorization, and captioning; Wu et al. [2021] consider representation learning for multiple tasks, Fernando and Tsokos [2021], Phan and Yamamoto [2020] consider weighted loss for class imbalance problems in classification, object detection.
Works such as Farahmand et al. [2017], Voelcker et al. [2022] use task loss in a different way compared to the above methods. They use task loss as a weighting term in the MSE loss itself. So the models are trained to focus more on samples with higher task loss. In their work, the task is the estimation of the value function in model-based RL. This can be seen as the instantiation of our work where the task loss is directly used as a metric instead of learning a metric.
Other related work on metric learning such as Hastie and Tibshirani [1995], Yang and Jin [2006], Weinberger and Tesauro [2007], Kulis et al. [2013], Hauberg et al. [2012], Kaya and Bilge [2019] often learns a non-Euclidean metric or distance that captures the geometry of the data and then solves a prediction task such as regression, clustering, or classification in that geometry. Other methods such as Voelcker et al. [2022] can handcraft metrics based on domain knowledge. In contrast to these, in the task-based model learning, we propose that the downstream task (instead of the data alone) gives the relevant metric for the prediction, and that it is possible to use end-to-end learning as in eq. (4) to obtain the task-based metric.
How our contribution fits in. The mentioned methods mainly deal with using task-based losses to condition the model learning either by differentiation through task loss to update the model or using it directly as weighing for MSE prediction loss. Whereas our work focuses on using task loss to learn a parameterized prediction loss which is then used to train the model. The task loss is not directly used for model training.
3 Task-driven metric learning for model learning
We first present why it's useful to see the prediction space as a non-Euclidean metric space with an unknown metric, then show how task-based learning methods can be used to learn that metric.
3.1 Metrics in the prediction space — Mahalanobis losses
When defining a loss on the model, we are forced to make a choice about the geometry to quantify how good a prediction is. This geometric information is often implicitly set in standard learning settings and there are often no other reasonable choices without more information. For example, a supervised model  $f_{\theta}$  parameterized by  $\theta$  is often trained with the mean squared error (MSE)
$$
\theta_ {\text {M S E}} ^ {*} := \underset {\theta} {\arg \min } \mathbb {E} _ {(x, y) \sim \mathcal {D}} \left[ \left(f _ {\theta} (x) - y\right) ^ {2} \right]. \tag {2}
$$
The MSE makes the assumption that the geometry of the prediction space is Euclidean. While it is a natural choice, it may not be ideal when the model needs to focus on important parts of the data that are under-emphasized under the Euclidean metric. This could come up by needing to emphasize some samples over others, or some dimensions of the prediction space over others.
While there are many possible geometries and metric spaces that could be defined over prediction spaces, they are difficult to specify without more information. We focus on the metric space defined by the Mahalanobis norm  $\|z\|_M \coloneqq (z^\top M z)^{1/2}$ , where  $M$  is a positive semi-definite matrix. The Mahalanobis norm results in the prediction loss
$$
\mathcal {L} _ {\text {p r e d}} (\theta , \phi) := \mathbb {E} _ {(x, y) \sim \mathcal {D}} \left[ \| f _ {\theta} (x) - y \| _ {\Lambda_ {\phi} (x)} ^ {2} \right], \tag {3}
$$
where  $\Lambda_{\phi}$  is a metric parameterized by  $\phi$  and this is also conditional on the feature  $x$  so it can learn the importance of the regression space from each part of the feature space.
y0
Figure 2: Examples of the Mahalanobis loss from eq. (3) in a 2-dimensional prediction task. The model's loss is zero only when  $\hat{y} = y^{\star}$ . Here, the metric  $\Lambda_{\phi}(x)$  increases the weighting on the  $y_0$  component of the loss and thus emphasizes the predictions along this dimension.
Many methods can be seen as hand-crafted ways of setting a Mahalanobis metric, including: 1) normalizing the input data by making the metric appropriately scale the dimensions of the prediction, 2) re-weighting the samples as in Donti et al. [2017], Lambert et al. [2020] by making the metric scale each sample based on some importance factor, or 3) using other performance measures, such as the value gradient in Voelcker et al. [2022].
More generally beyond these, the Mahalanobis metrics help emphasize the:
1. relative importance of dimensions. the metric allows for down- or up-weighting different dimensions of the prediction space by changing the diagonal entries of the metric. Figure 2 illustrates this.
2. correlations in the prediction space. the quadratic nature of the loss with the metric allows the model to be aware of correlations between dimensions in the prediction space.
3. relative importance of samples. heteroscedastic metrics  $\Lambda (x)$  enable different samples to be weighted differently for the final expected cost over the dataset.
Without more information, parameterizing and specifying the best metric for learning the model is challenging as it involves the subproblem of understanding the relative importance between predictions. We suggest that when it is available, the downstream task information characterizing the overall model's performance can be used to learn a metric in the prediction space. Hence, learning model parameters with a metricized loss can be seen as conditioning the learning problem. The ability to learn the metric end-to-end enables the task to condition the learning of the model in any or all of the three ways described above. This approach offers an interpretable method for the task to guide the model learning, in contrast to relying solely on task gradients for learning model parameters, which may or may not align effectively with the given prediction task.
3.2 End-to-end metric learning for model learning
The key idea of the method is to learn a metric end-to-end with a given task, which is then used to train the prediction model as shown in eq. (3). The learning problem of the metric and model parameters are formulated as the bilevel optimization problem
$$
\phi^ {\star} := \underset {\phi} {\arg \min } \mathcal {L} _ {\text {t a s k}} \left(\theta^ {\star} (\phi)\right), \tag {4}
$$
$$
\text {s u b j e c t} \theta^ {\star} (\phi) = \underset {\theta} {\arg \min } \mathcal {L} _ {\operatorname {p r e d}} (\theta , \phi) \tag {5}
$$
where  $\phi$  and  $\theta$  are (respectively) the metric and model parameters,  $\mathcal{L}_{\mathrm{pred}}$  is the metricized prediction loss (eq. (3)) to train the prediction model, and  $\mathcal{L}_{\mathrm{task}}$  is the task loss defined by the task at hand (which could be another optimization problem, e.g. eq. (8), or another learning task, e.g. eq. (10)).
Gradient-based learning. We learn the optimal metric  $\Lambda_{\phi^{\star}}$  with the gradient of the task loss, i.e.  $\nabla_{\phi}\mathcal{L}_{\mathrm{task}}(\theta^{\star}(\phi))$ . Using the chain rule and assuming we have the optimal  $\theta^{\star}(\phi)$  for some metric
Figure 3: TaskMet learns a metric for predictions with the gradient from a downstream task loss.
Algorithm 1 TaskMet: Task-Driven Metric Learning for Model Learning
Models: predictor  $f_{\theta}$  and metric  $\Lambda_{\phi}$  with initial parameterizations  $\theta$  and  $\phi$   
while unconverged do  
// approximate  $\theta^{\star}(\phi)$  given the current metric  $\Lambda_{\phi}$   
for  $i$  in 1...K do  
 $\theta \gets \text{update}(\theta, \nabla_{\theta} \mathcal{L}_{\text{pred}}(\theta, \phi)) // \text{fit the predictor } f_{\theta}$  to the current metric loss (eq. (3))  
end for  
 $\phi \gets \text{update}(\phi, \nabla_{\phi} \mathcal{L}_{\text{task}}) // \text{update the metric } \Lambda_{\phi}$  with the task loss (eq. (6))  
end while  
return optimal predictor  $f_{\theta}$  and metric  $\Lambda_{\phi}$  solving the bi-level problem in eq. (4)
parameterization  $\phi$ , this derivative is
$$
\nabla_ {\phi} \mathcal {L} _ {\text {t a s k}} \left(\theta^ {\star} (\phi)\right) = \left. \nabla_ {\theta} \mathcal {L} _ {\text {t a s k}} (\theta) \right| _ {\theta = \theta^ {\star} (\phi)} \cdot \frac {\partial \theta^ {\star} (\phi)}{\partial \phi} \tag {6}
$$
To calculate the term  $\nabla_{\phi}\mathcal{L}_{\mathrm{task}}(\theta^{\star}(\phi))$ , we need to compute two gradient terms:  $\left.\nabla_{\theta}\mathcal{L}_{\mathrm{task}}(\theta)\right|_{\theta = \theta^{\star}(\phi)}$  and  $\partial \theta^{\star}(\phi) / \partial \phi$ . The former can be estimated in standard way since  $\mathcal{L}_{\mathrm{task}}(\theta)$  is an explicit function of  $\theta$ . However, the latter cannot be directly calculated because  $\theta^{\star}$  is a function of optimization problem which is multiple iterations of gradient descent, as shown in eq. (5). Backpropping through multiple iterations of gradient descent can be computationally expensive, so we use the implicit function theorem (appendix A) on the first-order optimality condition of eq. (5), i.e.  $\frac{\partial\mathcal{L}_{\mathrm{pred}}(\theta,\phi)}{\partial\theta} = 0$ . Combining these,  $\nabla_{\phi}\mathcal{L}_{\mathrm{task}}(\theta^{\star}(\phi))$  can be computed with
$$
\nabla_ {\phi} \mathcal {L} _ {\text {t a s k}} \left(\theta^ {\star} (\phi)\right) = \nabla_ {\theta} \mathcal {L} _ {\text {t a s k}} (\theta) \cdot \underbrace {- \left(\frac {\partial^ {2} \mathcal {L} _ {\text {p r e d}} (\theta , \phi)}{\partial \theta^ {2}}\right) ^ {- 1} \left. \frac {\partial^ {2} \mathcal {L} _ {\text {p r e d}} (\theta , \phi)}{\partial \phi \partial \theta} \right| _ {\theta = \theta^ {\star} (\phi)}} _ {\partial \theta^ {\star} / \partial \phi} \tag {7}
$$
The implicit derivatives in eq. (7) may be challenging to compute or store in memory because the Hessian term  $\partial^2\mathcal{L}_{\mathrm{pred}}(\theta ,\phi) / \partial \theta^2$  is the Hessian of the prediction loss with respect to the model's parameters. Approaches such as Lorraine et al. [2020] are able to scale related implicit differentiation problems to models with millions of hyper-parameters. The main insight is that the Hessian does not need to be explicitly formed or inverted and the entire implicit derivative term needed for backpropagation can be obtained with an implicit solver. We follow Blondel et al. [2022] and compute the implicit derivative by using conjugate gradient on the normal equations.
We evaluate our method in two distinct settings: 1) when the downstream task involves an optimization problem parameterized by the prediction model output, and 2) when the downstream task is another learning task. For the first setting, we establish our baselines by replicating experiments from previous works such as Shah et al. [2022] and Wilder et al. [2019]. These baselines encompass tasks like portfolio optimization and budget allocation. In the second setting, we focus on model-based
reinforcement learning. Specifically, we concentrate on learning a dynamics model (prediction model) and aim to optimize the Q-value network using the learned dynamics model for the Cartpole task [Nikishin et al., 2022]. Appendix B provides further experimental details and hyper-parameter information.
4.1 Metric parameterization
We parameterize the metric using a neural network with parameters  $\phi$ , denoted as  $\Lambda_{\phi} \coloneqq L_{\phi}^{\top}L_{\phi}$  where  $L_{\phi}$  is an  $n \times n$  matrix, where  $n$  is the dimension of the prediction space. This particular factorization constraint ensures that the matrix is positive semi-definite, which is crucial for it to be considered a valid metric. The neural network parameters are initialized to make  $\Lambda_{\phi}$  closer to the identity matrix  $\mathbb{I}$ , representing the Euclidean metric. The learned metric can be conditional on the input, denoted as  $\Lambda_{\phi}(x)$ , or unconditional, represented as  $\Lambda_{\phi}$ , depending on the problem's structure.
4.2 Decision-Focused Learning
4.2.1 Background and experimental setup
We use three standard resource allocation tasks for comparing task-based learning methods [Shah et al., 2022, Wilder et al., 2019, Donti et al., 2017, Futoma et al., 2020]. In this setting, resource utility prediction based on some input features constitute a prediction model, resource allocation constitutes the downstream task which is characterized by  $\mathcal{L}_{\mathrm{task}}$  The prediction model's output parameterized the downstream resource optimization. The settings are implemented exactly as in Shah et al. [2022] and have task losses defined by
$$
\mathcal {L} _ {\text {t a s k}} (\theta) := \mathbb {E} _ {(x, y) \sim \mathcal {D}} [ g (z ^ {\star} (\hat {y}), y) ] \tag {8}
$$
where  $z^{\star}(\hat{y}) \coloneqq \arg \min_{z} g(z, \hat{y})$  and  $g(z, y')$  is some combinatorial optimization objective over variable  $z$  parameterized by  $y'$ . The task loss  $\mathcal{L}_{\mathrm{task}}$  is the expected value of objective function with decision variable  $z^{\star}(\hat{y})$  induced by the prediction model  $\hat{y} = f_{\theta}(x)$  under the ground truth parameters  $y$ . We use corresponding surrogate losses to replicate the  $z^{\star}(\hat{y})$  optimization problem as in Shah et al. [2022], Wilder et al. [2019], Xie et al. [2020] and differentiate through the surrogate using cvxpylayers [Agrawal et al., 2019].
These settings evaluate the ability of TaskMet to capture the correlation between model predictions and differentiate between different data-points in accordance to their importance for the optimization problem. Hence, we consider a heteroscedastic metric, i.e.,  $\Lambda_{\phi}(x)$ .
Baselines. We compare with standard baseline losses for learning models:
1. The standard MSE loss  $\theta^{\star} = \arg \min_{\theta}\mathbb{E}_{(x,y)\sim \mathcal{D}}[(f_{\theta}(x) - y)^2 ]$  . This method doesn't use any task information.
2. DFL [Wilder et al., 2019], which trains the prediction model with a weighted combination of  $\mathcal{L}_{\mathrm{task}}$  and  $\mathcal{L}_{\mathrm{pred}}$  as in eq. (1).
3. LODL Shah et al. [2022], which learns a parametric loss for each point in the training data to approximate the  $\mathcal{L}_{\mathrm{task}}$  around that point. That is,  $LODL_{\psi_n}(\hat{y}_n) \approx \mathcal{L}_{\mathrm{task}}(\hat{y}_n)$  for all  $n$ . They create a dataset of  $\{(\hat{y}_n, \mathcal{L}_{\mathrm{task}}(\hat{y}_n))\}$  for  $\hat{y}_n$  sampled around the  $y_n$ . After this they learn the LODL loss for each point as  $\psi_n^{\star} = \arg \min_{\psi_n} \frac{1}{K} \sum_{k=1}^{K} (LODL_{\psi_n}(y_n^k) - \mathcal{L}_{\mathrm{task}}(y_n^k))^2$ . We chose the "Quadratic" variant of their method which is the closest to ours, where  $LODL_{\psi_n}(\hat{y}) = (\hat{y} - y)^{\top} \psi_n(\hat{y} - y)$  where  $\psi_n$  is a learned symmetric Positive semidefinite (PSD) matrix. LODL also uses eq. (1) to learn the model parameters, but using  $LODL_{\psi_n}(\hat{y}_n) \approx \mathcal{L}_{\mathrm{task}}(\hat{y}_n)$
Experimental settings. We use the following experimental settings from [Shah et al., 2022]: