TaskMet: Task-driven Metric Learning for Model Learning
Authors: Dishank Bansal
Ricky T. Q. Chen
Mustafa Mukadam
Brandon Amos
Venue: NeurIPS 2023 poster (2023)
Keywords: task-based learning
decision-focused learning
TL;DR: Using task loss to learn parameterized prediction loss, which are then used to train prediction models

Abstract: 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.

MeskMet:
Task-Driven Metric Learning for Model Learning
Dishank Bansal"
Ricky T. Q. Chen
Mustafa Mukadam
Brandon Amos
Meta
Abstract
Deep learning models often used with some downstream task. Models solely trained to achieve accurate predictions may struggle to perform well on the de- sired downstream tasks. We propose using the task loss learn 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: 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 cerisioonffousedd 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.
Introduction
task loss
Machine learning models for prediction are typically trained to maximize the likelihood training dataset. While the models are capable of universally approximating the under- lying data generating process to predict the output, they are prone to approximation errors due to limited training data and model capacity. These errors lead suboptimal performance in downstream tasks where the models are used. Furthermore, even though model may appear have reasonable predictive performance on the metric and training data it was trained on, squared error, employing the model for
model space
DEL
Oe TaskMet
true model MSE
Figure The MSE results model close to true model in the predic- tion space, but may give poor task per- formance. 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.
such the 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.
Examples of settings where the model's prediction loss Cpeed is mische from the downstream task Lrask include the following, which table also summarizes:
1. the portfolio optimization setting from Wilder et al. [2019], which predicts the expected returns from stocks for financial portfolio. Here, the Epred is the MSE and Ciss is from the regret of running a portfolio optimization problem on the output;
2. the allocation setting from Wilder al. [2019], which predicts the value of items that being allocated, e.g. tlick-through-ratee for recommender systems. Here, Spod is the MSE and Link measures the result of allocating the highest-value items.
"Work done as part of Meta Al residency program.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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., Pped # L
(2)
(3)
(Equz)
(uaa) Task Loss
Setting
Features
Predictions
Prediction Loss
Portfolio optimization
Stock information
Expected return a stock
MSE
Portfolio's performance
Budget allocation
Item information
Value of item Next state
MSE
Allocation's iss performance
Model-based RL
Current state and action
MSE
Value estimation given the model
3. the model-based reinforcement learning setting of learning the system dynamics from Nikishin et al. [2022]. Here, Speed is the MSE of dynamics model and the Luss 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 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 [Vooha 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 unable to generalize; 2) the need 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 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. and 7. Section 4 shows the empirical success of metric learning on decision focused model learning and model-based reinforcement learning. Figure illustrates the differences to prior methods.
2 Background and related work
Task- model learning. We will mostly focus solving regression problems where the dataset D: {(z,, } consists of N input-output pairs, which we will assume to be in Euclidean space. The model makes a prediction := fe(z) and is parameterized by 0. The model has associated prediction loss, fpeed and is used in conjunction with some downstream task that provides a task loss, Lusk, which characterizes how well the model performs the task. The most relevant related work to ours includes the approaches of Bengio [1997], Donti et al. [2017], Farahmand et al. [2017], Kang al. [2019], Wilder et al. [2019], Nikishin al. [2022], Shah al. [2022], Voelcker et al. [2022], Nikishin et al. [2022], Anonymous [2023], Shah et al. [2023], which learn the optimal prediction model parameter 0 to minimize the task loss Lrask:
0* == arg min (o) + (P).
where is 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 latent domain with task information, e-g. task- specific latent dynamics for RI. [Hafner et al., 2019b,a, Hansen et al., 2022]. Extensions such as Gupta and Zhang [2023], Zharmagambetov et al. [2023], Ferber et [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 Farahmand al. [2017], Voelcker al. [2022] task loss in 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 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 al. [2013]. Hauberg et al. [2012], Kaya and Bilge [2019] often learns a non- non-Euclidean metric or distance that captures the geometry of the data and then solves prediction task such as regression, clustering, classification in that geometry. Other methods such as Voeleker 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 using it directly as weighing for MSE prediction loss. Whereas our work focuses on using task loss to learn 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 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 choice about the geometry to quantify how good a prediction is. This geometric information is often implicitly set standard learning settings and there are often no other reasonable choices without more information. For example, a supervised model fo parameterized by is often trained with the mean squared error (MSE)
BNSE := arg min ((ry))> [(fo(x) y)
The MSE makes the assumption that the geometry of the prediction space is Euclidean. While natural choice, it may not be ideal when the model needs to focus on important parts of the data that are eeppr-eeppaaizzed under the Euclidean metric. This could come up by needing to emphasize some samples over others, 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 || M := (z M=) where M is a positive semi-definite matrix. The Mahalanobis norm results in the prediction loss
Cpea (0, 6) E [lf((())
- v v
where Ag is metric parameterized by & 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.
3

---

1 0
2
0
10
0
(a)
0
0
1
1
y1
V1
1/1
yo
yo
yo
Figure 2: Examples of the Mahalanobis loss from eq. (3) in 2-dimensional prediction task. The model's loss is zero only when ò - y"" Here, the metric Ag (x) increases the weighting on the yo 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 some importance factor, 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- op-weighting different dimensions of the prediction space by changing the diagonal entries of the metric. Figure illustrates this.
2. correlations 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 A(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 challenging as 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 used to learn a metric in the prediction space. Hence, learning model parameters with 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 interpretable method for the task to guide the model learning, 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 learn a metric end-to-end with given task, which 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
arg min Lusk (aax(0(0)) (Ø)).
(4)
subject 8" (*()) arg min (0, (8,
(5)
where á and 0 are (respectively) the metric and model parameters, Lored is the metricized prediction loss (eq. (3)) to train the prediction model, and Cisc 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 Ag. with the gradient of the task loss, i.e. V Lusk (0' (o)). Using the chain rule and assuming have the optimal 0* (Ø) for some metric

---

arg min prccd (0,0)
û (x)
æ: features
Ao
0x (Φ)
L Lask
y: targets
Ag: metric
C0 () == F: (((((())) c
00 (Ø)
via IFT
V Lask
OΦ
Figure 3: TaskMet learns a metric for predictions with the gradient from a downstream task loss.
Algorithm 1 TaskMet: Task-Driver Driven Metric Learning for Model Learning
Models:
predictor So and metric As with initial parameterizations and Ø
while unconverged do // approximate 8* (Ø) given the current metric As for in 1. .... K do
0 t- update(8; (0, Vo o)) // fit the predictor fo to the current metric loss (eq. (3))
end for Φ < (p, V &uLus) // update the metric As with the task loss (eq. (6))
end while
return optimal predictor fo and metric A₂ solving the bi-level problem in eq. (4)
parameterization o, this derivative is
86 (Φ)
(f VgLusk(()) (0)
(6)
dΦ
To calculate the term (f* need compute two gradient terms: V VoLusk(0) 0=0-
and 00 (Ø) The former can be estimated in standard way since Lus(() ()) is an explicit function of 9. However, the latter cannot be directly calculated because 8" is function of optimization problem which multiple iterations of gradient descent, as shown in eq. (5). Backpropping through multiple iterations of gradient descent
be computationally expensive, we use the implicit
(0
function theorem (appendix A) on the first-order optimality condition of eq. (5), i.e. Combining these, &u (f*()) (Ø)) can be computed with
0.
89
(
-(82 Cpeed 002
Φ)
-1
02 Ereed Ø)
V VoLina (0) V Crask( (0)
0o00
6= (6)
80 a
The implicit derivatives in eq. (7) may be challenging to compute store in memory because the Hessian term 8 (0, (0000.))/000 /0/2 the Hessian the prediction loss with respect 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 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. follow Blondel et al. [2022] and compute the implicit derivative by using conjugate gradient the normal equations.
4 Experiments
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 [2022] and Wilder like portfolio optimization and budget allocation. In the second setting, focus on model-based
[2019]. These baselines encompass tasks
5

---

reinforcement learning. Specifically, we concentrate on learning a dynamics model (prediction model) and aim 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 ø, denoted as Ag. := L Le. where L an n x n matrix, where n is the dimension of the prediction space. This particular factorization constraint ensures that the matrix is positíve semi-definite, which is crucial for it to be considered a valid metric. The neural network parameters are initialized to make Ag closer to the identity matrix 1, representing the Euclidean metric. The learned metric can be conditional on the input, denoted as (z), or unconditional, represented as Ag, 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 al., 2022, Wilder al., 2019, Donti al., 2017, Futoma al., 2020]. In this setting, resource utility prediction based on some input features constitute prediction model, resource allocation constitutes the downstream task which characterized by Lusk 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
(8) where 2* =9) arg min, g(2, 9) and g(z. ) is some combinatorial optimization objective over variable parameterized by y/. The task loss Lans is the expected value of objective function with decision variable 2" (f) induced by the prediction model أ fo(x)) under the ground truth parameters We use corresponding surrogate losses to replicate the = (0) optimization problem as in Shah et al. [2022], Wilder et al. [2019], Xie et al. [2020] and diffcrentiate through the surrogate using cvxpylayers [Agrawal et al., 2019].
Lusk == (E(e,))~ [g(== (i),v)
These settings evaluate the ability of TaskMet capture the correlation between model predictions and differentiate between different data- in accordance to their importance for the optimization problem. Hence, we consider a heteroscedastic metric, i.e., (z).
Baselines. We compare with standard baseline losses for learning models:
1. The standard MSE loss 0* = arg ming E(e.g)=D((o(()) - y) This method doesn't 'tt use any task information.
2. DFL [Wilder et al., 2019], which trains the prediction model with a weighted combination Lrask and Cpred as in eq. (1).
3. LODL Shah et al. [2022], which learns parametric loss for each point in the training data to approximate the Lusk around that point. That is, LODLS, (fn) s Lask ((n) for all They create dataset of {(in, Lusk ( (ūn))) for Ån sampled around the yn. After this they learn the LODL loss for each point oă arg min & " (LODL (3h) Cas ?. We chose the "Quadratic" variant of their method which is the closest to ours, where LODLY (v) = (i y) "h - y) where v is learned symmetric Positive semidefinite (PSD) matrix. LODL also uses eq. (1) to learn the model parameters, but using LODL (fn) Re Lrosk(0
Experimental settings. We use the following experimental settings from [Shah et al., 2022]:
1. Cubic:
This setting evaluates methods under model mismatch scenario where the model being learned suffers with severe approximation error. this task, it important for methods to allocate model capacity to the points more critical for the downstream tasks. Prediction Model: A linear prediction model fx( := Or is learned for the problem where the ground truth data is generated by cubic function, i.e., ys 10x² - 6.52z; E 1 1
Downstream task: Choose top B 1 out of M 50 resources ý [91,- 9M
z" (8) : arg max; y

---

Table 2: Normalized test. decision quality (DQ) on the decision oriented learning problems.
Table 3: Test prediction errors (MSE) on the deci- sion oriented learning problems.
Problems
Problems
Method
Cubic
Budget
Portfolio
Method
Cubic
Budget Oxte 5 Portfolio oxte "
MSE
0.96=0002
0.54+0.17
0.33++0.0
MSE
2. 2.30=0.03
4,32+2.33
4.03=0,21
DFL
0.61=0.74
0.91
0.25:40.02
DFT.
0
2.89 2..32
71.7+-5.33
4-0
DFL
10
0.62+00.74
0.81+00111
0.344-0.33
DFL
10
2.41 10.05
8.09
5.18-.0.6
LODL
0.96+0.000
0.84+0. 105
0.17+0.05
LODL
2.88 10.03
35.9112.9
55.6-99.5
LODL
0.955 005
0.58:0011
0.30:5003
LODI.
10
2.29-0011
5.05++1.85
4.31-0.31
Met
0.96+0.008
0.83+0.12 12
0.33-+0.03
TaskMet a is the prediction loss weight in eq. (1)
2.89=0.03
9.74+13.79
4.69=0.56
O=random model ==oracle model
2. Budget Allocation: Choose top = 2 websites to advertise based on Click-through--ates
(CTRs) predictions of K users on M websites.
Prediction Model: ym = fo(x [Tm) where I'm is given features of m website and y. 1m........ Dm,, x is the predicted CTRs for m website for all K users.
Downstream task: Determine B = websites such that the expected number of users that click on the ad at least once is maximized, i.e., z" (fnm) ) arg max; 2j= > II 24- 0is)
where 2, E 10.
3. Portfolio Optimization:
The task is to choose a distribution over M stocks in Markowitz portfolio optimization [Markowitz and Todd, 2000, Michaud, 1989] that maximizes the expected return under the risk penalty.
Prediction Model: Given the historical data Im about a stock m, predict the future stock price dm. Combining prediction over M stocks to get y = ........ Mas]
Downstream Task: Given the correlation matrix Q of the stocks, choose distribution over Az T Qz L. Σ ž
stocks z (5) = aarg max,
and 0 S = S 1, Vi
We run our own experiments for LODL [Shah et al., 2022] using their public code.
4.2.2 Experimental results
Table 2 presents a summary of the performance of differ- ent methods on all the tasks. Each problem poses unique challenges for the methods. The cubic setting suffers from severe approximation errors, hence the learning method needs to allocate limited model capacity more towards higher utility points compared lower utility points. The MSE method performs the worst it lacks task information and only care about prediction error. DFL with 0 performs better than MSE, but it can get trapped in local optima, leading to higher variance in the problem [Shah et al., 2022]. LODL (a 0) per- forms among the highest in this problem since it uses learned loss for each point. TaskMet performs as good as LODL as it can capture the relative importance of higher utility points versus lower utility points using the learned metric, resulting in more accurate predictions for those points (see fig. 4). In budget allocation, DFL (with = 0) performs the best, since it is solely optimizer over Luask but on the other hand it has 10 orders of larger prediction error as shown in table 3 indicating that the model is overfit the task, LODL (a 0) suffers from the same problem. TaskMet has the 2nd best Decision Quality without overfitting on the task, i.e., low prediction error. In Porifolio Optimization, the decision quality correlates highly with the model accuracy/prediction emror as in this setting the optimization problem mostly depends upon the accurate prediction of the stocks. This is the reason that MSE, DFL (a 10) performs the best, but DFL (a 0) performs the worst, since it has solely
TaskMet
MSE
3
v(x)
A(x)
1
0
x
Figure 4: (Cubic problem) TaskMet learns a metric that prioritizes points that are the most important the downstream task. The euclidean metric (MSE) puts equal weight all points and leads to bad model with respect to the downstream task.
7

---

TaskMet
Learning
Planning
Acting
Metric
Model
Q
Loss
IFT
IFT
Backprop
OMD
Figure 5: OMD [Nikishin et al., 2022] uses planning task loss to learn the model parameters using implicit gradients. TaskMet add one more optimization step over OMD and instead of learning the model parameters using task loss, we learn the metric which then is used to learn model parameters.
real states
state dimension
2
3 4
5
7
9
10 11
Eploode Return
16
metric values (darker=higher)
300
distractors
32
64
128
* 256
512
# Distractors
TaskMet
MLE
OMD
Figure 7: Our learned metric successfully distinguishes the real states from the distracting states, i.c. the real states take higher metric value.
Figure 6: Results the cartpole with distracting states [Nikishin et al., 2022].
heing trained Lusk without any Fped- As shown in table and table 3, TaskMet is the only method that consistently performs well considering both task loss and prediction loss, across all the problem settings, this is due to the ability of the metric to infer problem-specine features without manual tuning, unlike other methods.
4.3
Model Based Reinforcement Learning
4.3.1 Background and experimental setup
Model-based RL suffers from sejective-mismattee [Bansal al., 2017, Lambert et al., 2020]. This is because dynamics models trained for data likelihood maximization do not translate to optimal policy. To reduce bojective-mismatee
different losses [Farahmand et al., 2017, Voelcker al., 2022] have been proposed to learn the model which is better suited to learning optimal policies. TaskMet provides alternative approach towards reducing mbjistive-mismatch, using task loss. We set the MBRL problem as follows. Given the current state st and control at timestep of a discrete-time MDP, the dynamics model predicts the next state transition, i.e. $141 fe(st, at). The prediction loss is commonly the squared error loss 2. || and the downstream task is to find the optimal Q-value/policy. Nikishin al. [2022] introduces idea of optimal model design (OMD) to learn the dynamics model end-to-end with the policy objective via implicit differentiation. Let Qu (s, a) be the aection-conitionn value function parameterized by w. The Q network is trained to minimize the Bellman error induced by the model fø:
the prediction loss is directly learnt
fe a
Lo(w, 0) == Ex.alQu a) - B² (s,
where w is moving average of and Bd is the model-induced Bellman operator B² (s re +) B [log E., exp Q(s', a') a')]. Q-network optimality defines w as an implicit function of the model parameters 0 as (0) arg min Co (w(w,) 0) (Cq(u.0) 0. Now have task dw
loss which is optimized to find optimal Q-values:
Luck (w (0)) : E., (o) (s) BQa a)l²
(10)
where the Bellman operator induced by ground-truth trajectory and reward is BQa r(s,
Ey,,, log E., exp (s('',''))

---

500
Episode Return
400
Dynamics (MSE)
300
10
200
10-1
100
50K
100K
150K
200K
50K
100K
150K
200K
Environment Steps
Environment Steps
TaskMet
MLE
OMD
Figure 8: Results on cartpole with a reduced model capacity from Nikishin et al. [2022].
OMD setup. OMD en-to-end optimizes the model for the task loss, i.e. 6* = arg ming Lusk (w= (f)).
TaskMet setup. For metric learning, extend OMD to learn metric using task gradients, train the model parameters, see fig. 5. Metric learning just adds one more level of optimization to OMD) and results in the tri-level problem
arg min (usk("")
subject to w" (0*) = arg min (o(w, 8*)
(11)
0" (6) arg min (p.
where (wr)("") and (w, 8*) are defined in eq. (10) and eq. (9), respectively, and Coee 0) Eecae || - fo( az) (B22(())))) the metricized prediction loss in eq. (3).
To learn of using gradient descent, we estimate V Cask as
Ow a0*
V V Lusk = whrnsslu"
00 0Φ
(8 (FL(W,00) Ow?
0 BLw, 0*)
(FFCpca (0,
8 Cprcd (0, 6)
-VuLma(w"))
00Ow
002
8000
(w-(0-)
(0)
(12)
4.3.2 Experimental results
We replicated experiments from Nikishin et al. [2022] on the Cartpole environment. The first experiment involved state distractions, where the state of the agent was augmented with noisy and uninformative values. In this setting, we considered an unconditional diagonal metric of dimension n, which is the dimension of the state space, i.e. Ag := diag(ø), where & E R As shown in fig. the MLE method performed the worst across different numbers of distracting states, as it allocated its capacity to learn distracting states as well. TaskMet outperformed the other methods in all scenarios. The superior performance of TaskMet with distracting states can be attributed to the metric's ability to explicitly distinguish informative states from noise states using the task loss and then train the model using the given metric, as shown in fig. 7. The learned metric in fig. 7 assigned the highest weight to the third dimension the state space, which corresponds to the pole angle the most indicative dimension for the reward. This shows that the metric can differentiate state dimensions based on their importance to the task.
We also consider a setting with reduced model capacity, where the network is under- parametrized, foreing the model to prioritize how it allocates its capacity. In this scenario, we employ a full conditional metric, denoted as A = (x), which enables the metric to weigh dimensions and state-action pairs differently. We conducted the experiment using a model size hidden units in the layer. As depicted in fig. 8, TaskMet achieves a better return on evaluation compared to MLE and OMD. Additionally, it evident that TaskMet achieves a lower MSE on the model predictions compared to OMD, indicating that learning with the metric contributes to better dynamics model.

---

Conclusion and discussion
In conclusion, this paper addresses the challenge of combining task and prediction losses in task- based model learning. While task-based learning methods offer the advantage of discovering task- relevant features and data samples without explicit inductive biases, the current trend of using task loss alongside prediction loss has potential limitations. These limitations include overfitting of the prediction model to a specific task, rendering it ineffective for other tasks, and the lack of interpretability in the task-relevant features learned by the prediction model.
To overcome these limitations, the paper introduces the concept of task-driven metric learning, which integrates the task loss into a parameterized prediction loss. This approach enables end-to-end learning of metrics train prediction models, allowing the models focus on task-relevant features and dimensions in the prediction space. Moreover, the resulting prediction models become more interpretable, as metric learning serves as a preconditioning step for gradient-based model training. The effectiveness of the method is shown using different scales of experimental setting decision oriented tasks as well as downstream learning tasks.
One of the limitations of the method is stability of learning the metric. Bad gradients can lead collapsed metric which can lead unrecoverable bad predictions. Hence, hyper-parameter tuning of learning rate for metric learning and parameterization choices of the metric are crucial. Possible extensions to this work includes end-to-end metric learning with multiple task losses, learning metric for training dynamics models to be used for long-horizon planning tasks, etc.
Acknowledgments
We would like to thank Brian Karrer, Claas Voelcker, Karen Ullrich, Leon Bottou, Maximilian Nickel, and Mike Rabbat insightful comments and discussions.
References
Akshay Agrawal, Brandon Amos, Shane Barratt, Stephen Boyd, Steven Diamond, and Zico Kolter. Differentiable convex optimization layers. Advances in neural information processing systems, 32, 2019. Cited on page 6.
Brandon Amos, Ivan Jimenez, Jacob Sacks, Byron Boots, and Zico Kolter. Differentiable mpc for end-to-end planning and control. Advances neural information processing systems, 31, 2018. Cited on page 2.
Anonymous. Predict- peedict-teen-optiiiz via learning to optimize from features. In Submitted to The Twelfth International Conference on Learning Representations, 2023. URL httphtt//opeererrrtit net/ forum?id--jDvv3333je under review. Cited on page 2.
Somil Bansal, Roberto Calandra, Ted Xiao, Sergey Levine, and Claire J Tomlin. Goal-driven dynamics learning via bayesian optimization. In 2017 IEEE 56th Annual Conference on Decision and Control (CDC), pages 5168-5173. IEBE, 2017. Cited on page 8.
Yoshua Bengio. Using financial training criterion rather than a prediction criterion. International journal of neural systems, 8(04):433-443 1997. Cited page 2.
Mohak Bhardwaj, Byron Boots, and Mustafa Mukadam. Differentiable gaussian process motion planning. In 2020 IEEE international conference on robotics and automation (ICRA). pages 10598-10604, IEEE, 2020. Cited on page 2.
Mathieu Blondel, Quentin Berthet, Marco Cuturi. Roy Frostig, Stephan Hoyer, Felipe Llinares-óppzz Fabian Pedregosa, and Jean-Philippe Vert. Efficient and modular implicit differentiation. Advances in neural information processing systems, 35:5230-5244 5242, 2022. Cited on page
"sai-Hsuan Chung, Vahid Rostami, Hamsa Bastani, and Osbert Bastani. Decision-aware learning for optimizing health supply chains. arXiv preprint arXiv:221 108507, 2022. Cited on page
Ulisse Dini. Analisi infinitesimale. Lithografia Gorani, 1878. Cited on page 14.
10

---

Asen L Dontchev and R Tyrrell Rockafellar. Implicit functions and solution mappings, volume 543. Springer, 2009, Cited on page 14.
Priya Donti, Brandon Amos, and Zico Kolter. Task-based end-to-end model learning in stochastic optimization. Advances in neural information processing systems, 30, 2017. Cited on pages 2, 4, and 6.
Othman EI Balghiti, Adam N Elmachtoub, Paul Grigas, and Ambuj Tewari. Generalization bounds in the 2019. Cited on page 2.
framework. Advances in neural information processing systems, 32,
Adam N Elmachtoub and Paul Grigas. Smart "predict, then optimize". Management Science, 68(1): 9-26, 2022. Cited page 2.
Amir-massoud Farahmand. Iterative value-aware model learning. Advances in Neural Information Processing Systems, 31, 2018. Cited on page
Amir-massouee Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model- based reinforcement learning. In Artificial Intelligence and Statistics, pages 1486-1494.
PMLR, 2017 Cited on pages 2, 3, and
Aaron M Ferber, Taoan Huang. Daochen Zha, Martin Schubert, Benoit Steiner, Bistra Dilkina, and Yuandong Tian. Surco: Learning linear surrogates for combinatorial nonlinear optimization problems. In International Conference on Machine Learning, pages 10034-10052 PMLR, 2023.
Cited on page 3.
K Ruwani M Fernando and Chris P Tsokos. Dynamically weighted balanced loss: class imbalanced learning and confidence calibration of deep neural networks. IEEE Transactions on Neural Networks and Learning Systems, 33(7):2940-2951, 2021. Cited page 3.
Joseph Futoma, Michael C Hughes, and Finale Doshi-Vekke Popcorn: Partially observed prediction constrained reinforcement learning. arXiv preprint arXiv:2001.04N0 044332, 2020. Cited on page
Rishabh Gupta and Qi Zhang. Data-driven decision-focused surrogate modeling. arXiv preprint arXiv:2308. 12161, 2023, Cited on page 3.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1rXiv:1912.01603, 2019a. Cited on page 3.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555-2565, PMLR, 2019b. Cited on page 3.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprini arXiv:2203.0499 2022. Cited on page 3.
Trevor Hastie and Robert Tibshirani. Discriminant adaptive nearest neighbor classification and regression. Advances in neural information processing systems, 8, 1995. Cited on page 3.
Søren Hauberg, Oren Freifeld, and Michael Black. geometric take on metric learning. Advances in Neural Information Processing Systems, 25, 2012. Cited on page 3.
Bingyi Kang, Zhuang Liu, Xin Wang, Fisher Yu, Jiashi Feng, and Trevor Darrell. Few-shot object detection via feature reweighting. Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 8420-8429, 2019. Cited on page 2.
Mahmut Kaya and Hasan Şakir Bilge. Deep metric learning: A survey, Symmetry, 11(9):1066 2019.
Cited on page
Brian Kulis al. Metric learning: A survey. Foundations and Trends@ in Machine Learning, 5(4): 287-364, 2013. Cited page 3.
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523, 2020. Cited on pages and 8.

---

Mo Liu, Paul Grigas, Heyuan Liu, and Zuo-Jun Max Shen. Active learning in the predict-then- optimize framework: A margin-based based approach. arXiv preprint arXiv:3305.06584, 2023. Cited on page 2.
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, pages 1540-1552. PMLR, 2020. Cited on page 5.
Jayanta Mandi, Peter J Stuckey. Tias Guns, et al. Smart predict-and for hard combina- torial optimization problems. Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 1603-1610, 2020. Cited on page 2.
Harry Markowitz and Peter Todd. Mean-variance analysis in portfolio choice and capital markets, volume 66. John Wiley & Sons, 2000. Cited on page 7.
Richard Michaud. The markowitz optimization enigma: Is 'optimized' optimal? Financial analysts journal, 45(1):31-45 1989. Cited on page 7.
Evgenii Nikishin, Romina Abachi, Rishabh Agarwal, and Pierre-Lue Bacon. Control model- based reinforcement learning with implicit differentiation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7886-7894, 2022. Cited on pages 6, 8, 9, 14, and 15.
Trong Huy Phan and Kazuma Yamamoto. Resolving class imbalance in object detection with weighted cross entropy losses. arXiv preprint arXiv:2006.014/3 2020. Cited on page
André Susano Pinto, Alexander Kolesnikov, Yuge Shi, Lucas Beyer, and Xiaohua Zhai. Tuning computer vision models with task rewards. arXiv preprint arXiv:2302.0822 082422 2023. Cited on page 3.
Utsav Sadana, Abhilash Chenreddy, Erick Delage, Alexandre Forel, Emma Frejinger, and Thibaut Vidal. A survey of contextual optimization methods for decision making under uncertainty, 2023.
Cited on page
Sanket Shah, Kai Wang, Bryan Wilder, Andrew Perrault, and Milind Tambe. Decision-focuseee learning without decision-making: Learning locally optimized decision losses. In Alice Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https::/ http:/ppaooeeeeeeeeeeooorraaaaaa pages 2, 5, 6, 7, and 14.
Cited
Sanket Shah, Andrew Perrault, Bryan Wilder, and Milind Tambe. Leaving the nest: Going beyond local loss functions for predict-then-optimiz. arXiv preprint arXiv:2305. 2023. Cited on page 2.
Guangyao Shi and Pratap Tokekar. Decision-oriented learning with differentiable submodular maximization for vehicle routing problem. arXiv preprint wrXiv:2303.01543 2023. Cited on page 2.
Claas Voelcker, Victor Liao, Animesh Garg, and Amir-massoud Farahmand. Value gradient weighted model-based reinforcement learning. arXiv preprint 200vv2204.001664, 2022. Cited on pages 2, 3, 4, and
Rushil Vohra, Ali Rajaei, and Jochen L Cremer. End-to-end learning with multiple modalities for system-optimised renewables noweasting. arXiv preprint wrXiv:2304.07151, 2023. Cited on page 2.
Kai Wang, Shresth Verma, Aditya Mate, Sanket Shah, Aparna Taneja, Neha Madhiwalla, Aparna Hegde, and Milind Tambe. Decision-focuse learning in restless multi-armed bandits with ap- plication to maternal and child care domain. arXiv preprint arXiv:2202.00916 2022. Cited on page 2.
Kilian Q Weinberger and Gerald Tesauro. Metric learning for kernel regression. In Artificial intelligence and statistics, pages 612-619. PMLR, 2007. Cited on page
12

---

Bryan Wilder, Bistra Dilkina, and Milind Tambe. Melding the data-decisions pipeline: Decision- focused learning for combinatorial optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 1658-1665, 2019. Cited on pages 1, 2, 5, and 6.
Ruihan Wu, Chuan Guo, Awni Hannun, and Laurens van der Maaten. Fixes that fail: Self-defeating improvements in machine-learning systems. Advances in Neural Information Processing Systems, 34:11745-11756 2021. Cited on page 3.
Yujia Xie, Hanjun Dai, Minshuo Chen, Bo Dai, Tuo Zhao, Hongyuan Zha, Wei Wei, and Tomas Pfister. Differentiable top-k with optimal transport. Advances in Neural Information Processing Systems, 33:20520-20531, 2020. Cited page 6.
Liu Yang and Rong Jin. Distance metric learning: A comprehensive survey. Michigan State Universiy, 2(2):4, 2006. Cited on page 3.
Arman Zharmagambetov, Brandon Amos, Aaron Ferber, Taoan Huang, Bistra Dilkina, and Yuandong Tian. Landscape surrogate: Learning decision losses for mathematical optimization under partial information. arXiv preprint
2023. Cited on page 3.
13

---

A The implicit function theorem
We used the implicit function theorem to compute the derivative of the prediction model with respect to the metric's parameters in eq. (7). For completeness, this section briefly presents the standard implicit function theorem, cf. Dini [1878] and Dontchev and Rockafellar [2009, Theorem IB.1]:
Theorem (Implicit Function Theorem) Implicit Function Theorem: Let f : x R - IR be --"= continuous differentiable function, and let x", ,y" be point satisfying , O. If the Jacobian 8(f(=" .²)) + O, then there exists an open set around (x",y") and a unique continuously differentiable by function such that y² g(x*) ) and ((x)) = 0. Additionally, the following relation holds:
0g(z)
Of (ff(x,))) (f
-===((
(13)
Or
Oy
dr
B Implementation Details
B.1 Decision Oriented Model Learning
We replicated our experiments using the codebase provided by Shah et al. [2022], which can be found on github. To ensure consistency, used the same hyperparameters code or article for the baselines. Our metric learning pipeline was added on top of their code, and thus focused on tuning hyperparameters related metric learning. The metric is parameterized as (x(x) = (x(x)) L (x) + fgllnxn, where Es is a learnable parameter that explicitly controls the amount of Euclidean metric in the predicted metric. This helps ensure the stability of metric learning. We initialize the parameters in such way that the predicted metric is close to the Euclidean metric. For each outer loop of metric update, we perform K inner updates to train the predictor. Following the methodology of Shah et al. [2022], conducted 50 runs with different seeds for each of the experiments, where each method was evaluated on 10 different datasets, with 5 different seeds used for each dataset.
mentioned in the
Table 4:
for Decision Oriented Learning Experiments
Hyper-Partee Parameter
Values
Ag. learning rate
10-3
As hidden layer sizes
[200]
Warmup steps
500
Inner Iterations (K)
100
Implicit derivative batchsize
10
Implicit derivative solver
Conjugate gradient on the normal equations (5 iterations)
B.2 Model-Based Reinforcement Learning
We consider the work of Nikishin et al. [2022] as the baseline for replicating the experiments, and build upon their source code. Our metrie learning is just one additional step to their method. We adopt exact same hyperparameters as their for dynamics learning and Action-Value function learning. We focus on exploring and tuning the hyper-parameters specific to the metric learning component of the method.
Table Hyper-paran for the CartPole experiments
Hyper-Parameter
Values
Ag. learning rate
10-3
Ao hidden layer sizes
[32, 32]
Warmup steps
5000
Inner iterations (K)
1
Implicit derivative batchsize
256
Implicit derivative solver
Conjugate gradient on the normal equations (10 iterations)
14

---

For the state distractor experiments, we parameterize the metric an unconditional diagonal matrix, denoted as Ag diag(() where E In and n is the dimension of the state space. In addition. we consider hyper-parameter of metric parumeterization, for which we either take normalize or unnormalized metric. When we refer to normalizing the metric, we mean constraining the norm of the Φ vector to be equal to the 1.2 norm of an euclidean metric which is used by MSE method. This constrains the family of learnable metrics. To achieve this, we set : TØl2 v ensuring
||ol|2 = ||_nxn ll2 = vī. We also used LI regularization on the metric output, to induce sparsity in the metric. We sweep over three values of the regularization coefficient [0.00 0.001, 0.1]. We ran sweep over the 6 combinations of hyperparameters junnormalized, normalized] x [0.0, 0.001, 0.1] and choose the best hyper- parameter combination for each of the experiment. All the number reported in the experiments are calculated over 10 random seeds.
Our metric learning approach uses two implicit gradient steps. Firstly, we take the implicit derivative through action-value network parameters, approximating the inverse hessian the identity, similar to Nikishin et al. [2022]. Secondly, for the step through dynamics network parameters, calculate the exact implicit derivative.
15


Reviews:
[{"review_id": "nIQip46A1J", "reviewer": "Reviewer_uXLc", "summary": "This paper presents a method for end-to-end learning of metrics to train prediction models. It proposes a concept of task-driven metric learning. The main idea is to let the model focus on task-relevant features and dimensions in the prediction space. In addition, the resulting prediction model can be more interpretable, as the proposed method uses metric learning as a preconditioning step for gradient-based model training. Experiments have shown the effectiveness of the proposed method.", "questions": "I'm not an expert in this field.\nThe proposed method proposes to parametrize generalized Mahalanobis distance-based loss for training the prediction model.\nWhy do have to use generalized Mahalanobis distance-based loss?\nAre there any other options for the metric? If yes, it would be better to provide some discussion and comparisons.", "limitations": "The paper has provided the limitations of the proposed method.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed idea of only modifying the metric that the prediction model uses for training is quite interesting. It enables the model to focus more on task-relevant features and dimensions without explicitly altering the optimal prediction model itself.  Hence, it allows the prediction model trained on its original prediction problem while being valuable to the downstream tasks.\n\n- The results are promising. The proposed method has outperformed the other baseline approaches.", "weaknesses": "I'm not an expert in this field.\n- The writing of the methodology is a bit difficult to understand. It would be great if the writing can be polished with more illustrations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for the review! We are delighted you liked the idea and found the results promising. Here are a few responses to your comments, please let us know if there is anything else:\n\n> The writing of the methodology is a bit difficult to understand. It would be great if the writing can be polished with more illustrations.\n\nWe agree the methodology may be difficult to read without broader knowledge about other end-to-end DFL and implicit differentiation methods. We’ll make a pass through and try to explain as much as possible so it’s easier to understand, and we’ll think about places we can add illustrations. Are there any particular places in the methodology where you would want polishing? \n\n> The proposed method proposes to parametrize generalized Mahalanobis distance-based loss for training the prediction model. Why do have to use generalized Mahalanobis distance-based loss? Are there any other options for the metric? If yes, it would be better to provide some discussion and comparisons.\n\nThis is a very good question that we will clarify in the paper. For us, the generalized Mahalanobis distance was a natural extension of the MSE loss: any parameterization of the Mahalanobis loss retains the optimal solution of the MSE loss and the metric learning only changes how modeling errors are weighted, which is learned using task gradients.\nBeyond the parameterized Mahalanobis distance, many other parameterized losses could be learned. [This paper](https://arxiv.org/abs/1906.05374) on meta-learning losses simply uses a neural network as the loss and uses the meta-learning signal to update it. While we didn’t go to this level of generality because the neural network loss loses some interpretability, exploring other loss models are certainly an interesting direction for future work, and perhaps there are other intermediates between the Mahalanobis and neural network losses.\n"}, {"review_id": "p17Xa8ybQg", "reviewer": "Reviewer_WKyg", "summary": "The paper proposes a novel method, called TaskMet, for learning a metric that can be used to improve the performance of prediction models on downstream tasks. The key idea is to use the task loss to guide the learning of the metric. This is done by using a gradient-based approach to optimize the metric parameters. The authors show that TaskMet can improve the performance of prediction models on a variety of tasks, including portfolio optimization, budget allocation, and model-based reinforcement learning.", "questions": "- Authors mentioned that hyper parameter tuning is important for metric learning. I was wondering what is the significance of the hyperparameters that were used in the experiments? Is this the reason for quite high standard deviations shown in Table 2 for TaskMet on Budget Allocation?\n- Are the baselines results re-implemented by the authors? how did you pick hyper parameters for them?\n", "limitations": "Authors adequately addressed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "+ The paper is well-written and easy to follow.\n+ Motivations are clear and reasonable.\n+ The proposed method is novel and has the potential to be widely applicable.\n+ The authors provide extensive experimental results to support their claims.", "weaknesses": "It is not clear how sensitive the results are to the choice of hyperparameters. Please provide a more detailed analysis of the hyperparameter sensitivity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the review of our paper. We’re happy you found the paper well-written, well-motivated, and easy to follow. Here are some clarifications about the hyperparameters and baselines. To further help with reproducibility in the field we will also publish the entire source code necessary to reproduce every part of our paper.\n\n> It is not clear how sensitive the results are to the choice of hyperparameters. Please provide a more detailed analysis of the hyperparameter sensitivity.\n\nThe mentioned sensitivity of hyperparameters was due to their usage in reinforcement learning experiments. As we know that reinforcement learning methods need careful hyperparameter tuning to work, hence it was the combination of TaskMet with RL which caused hyperparameter sensitivity. When using TaskMet for Decision focused model learning experiments, we didn’t find any particular sensitivity to hyperparameters. We didn’t perform any meaningful hyperparameter sensitivity analysis as it was not something core to the introduced method but rather something application/task oriented.\n\n> Authors mentioned that hyper parameter tuning is important for metric learning. I was wondering what is the significance of the hyperparameters that were used in the experiments? Is this the reason for quite high standard deviations shown in Table 2 for TaskMet on Budget Allocation?\n\nThe significance of hyperparameters tuning in TaskMet is similar to what of other bilevel optimization problems. For example, number of inner iterations when solving for model parameters is critical, as to employ implicit gradient theorem we have to make sure that the inner optimization has converged, i.e,$ \\nabla_\\theta L(\\theta, \\phi)|_{\\theta = \\theta^\\star(\\phi)} = 0$. Other details about hyper-parameter can be found in supplementary material section A. The high standard deviation in Table 2 on Budget Allocation task is not due to hyper-parameters tuning of TaskMet. All other methods also have similar standard deviation, the higher standard deviation on this task is due the particular nature of task and dataset being used to train and test.   \n\n> Are the baselines results re-implemented by the authors? how did you pick hyper parameters for them?\n\nBaselines of decision focused model learning were not re-implemented but were reproduced using the code and hyper-parameters provided by the original paper. We were in consultation with the author of the baseline to ensure fair reproduction of results. For model based RL experiments, we directly used the numbers reported in the original paper. \n"}, {"review_id": "qWX2lUG8SB", "reviewer": "Reviewer_jBQF", "summary": "The authors propose an approach to loss re-shaping for a supervised prediction model with the goal of improving the performance of a larger system that performs a downstream task using the predictions of the prediction model. Such shaping is necessary because reducing loss on the prediction task may only roughly correspond to improved performance on the downstream task. The authors show that their approach mostly outperforms standard maximum likelihood estimation and achieves similar performance to existing baselines.", "questions": "Does TaskMet scale to more real-world settings, for example more difficult RL problems?\n\nDoes the learning metric produce useful information for real prediction problems, other than the artificial cartpole with noise setting?\n\nIs TaskMet limited to learning a Mahalanobis norm? Why not a more complex/expressive learned loss function?\n\nDoes TaskMet require significantly more compute than MSE?", "limitations": "Yes, although addressing the differences in computational requirements (if any) would be helpful.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "TaskMet is an interesting method that enables injection of knowledge about downstream tasks without ignoring the potentially useful training data for the base prediction task.", "weaknesses": "The clarity of the presentation could be significantly improved. For example, the distinction between \"task learning\" and \"prediction\" is not clear to readers who are not very familiar with the terminology; these are interchangeable synonyms in many fields. The descriptions of the benchmark tasks are also not very accessible to those not already familiar with them (what are \"resources\" in the Cubic problem?).\n\nThe experimental results in section 5.2 are relatively weak; TaskMet does not show improved performance in any of the three benchmark tasks. Although the average performance across tasks is slightly better than existing approaches, it's not clear if TaskMet outperforms the baselines if they are well-tuned (and the authors note that hyperparameter tuning for TaskMet is key).\n\nIn section 5.3 (the model-based RL experiments), the results are stronger. However, the setting is somewhat contrived, in that the authors study either a) a large number of artificial noise dimensions added to the state, which is likely to improve TaskMet's performance because it is specifically designed to rescale the loss of each dimension and b) when the network is unrealistically small (3 hidden units). Although the results in these settings are positive, they do not decisively show the benefit of the proposed method in real-world settings. Further, the baselines considered are either only using the prediction task loss (MSE) or only using the downstream loss (OMD). It's not clear if TaskMet outperforms other methods that utilize both objectives.\n\nSome related work could be added, particularly work in meta-learning loss functions, but this is not critical.\n\nThe advantages of TaskMet re: interpretability are not actually validated experimentally, other than on a synthetic problem showing that TaskMet learns to ignore dimensions of the targets that are artificial pure noise.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and comments on our paper. We were surprised by your decision to reject the paper based on the clarity and experimental weaknesses, as we have evaluated the idea of learning a metric on standard experimental settings in the most recent DFL and end-to-end model learning papers. We do agree that these experimental settings are relatively small-scale, but they are standard in the DFL/task-based learning communities and we chose them so that we could focus our contribution only on adding metric learning to existing experimental settings. Even though small-scale, these tasks/experiments are representative of challenges present in task-based model learning paradigm. In light of this information, we hope you will agree that the work is at a stage ready to disseminate to the community and that you will consider raising your score to help with this. We address your specific questions below. \n\n> The clarity of the presentation could be significantly improved [...]\n\nWe will distinguish between prediction and task loss for the camera-ready version. We have cited the papers from which we introduced these benchmarks, we have tried our best to mention all details about the benchmark tasks within the confines of the page constraints. Additionally, further details can be checked in prior works, e.g. we used the exact same benchmark as used in Shah et al [2022] and Nikishin et al. [2022]. Cubic problem is a toy experiment, we refer model outputs as resources in the example, as it’s consistent with the portfolio optimization and budget allocation terminology and enables us to ground the cubic example in a practical sense. \n\n\n> The experimental results in section 5.2 are relatively weak; [...]\n\nThe goal in these demonstrations was to show that metric learning works well for standard DFL tasks while retaining better prediction accuracy on the tasks. While Table 1 shows that TaskMet’s end-to-end loss is similar to the DFL’s loss, Table 2 shows that TaskMet retains far better prediction accuracy than the model learned with the DFL, where the predictive model loses the ability to make predictions in the original space. We hope this way of thinking about using the end-to-end signal for the loss instead of the model will be a helpful shift for the larger settings DFL is used where it is important to keep the predictions in the original space, and we think that the idea of a metricized loss is far more reliable than just optimizing the DFL+MSE loss as in eq (1).\n\nFurthermore, the baselines are well-tuned in this experiment setting. In fact we reran the hyperparam sweep over the baselines to get best performance as the default parameters in their official codebase were not giving optimal performance for the baselines. Also, in this experiments, hyperparameter tuning is not very critical for TaskMet. That was a remark for model based RL setting. \n \n> In section 5.3 (the model-based RL experiments), the results are stronger. However, [...]\n\nWe replicated the benchmark which has been used in previous works. We did artificial noise dimension according to the experiments done in the previous work [Nikishin et al. [2022] for fair comparison. Even though artificially injected noise, the experiments show that metric learns to distinguish between task relevant and irrelevant features. We agree that an MLP with 3 hidden units is small, but Nikishin et al. [2022] uses a hidden size 3; We chose this setting so we could directly compare our experimental results to OMD. We understand desire to see TaskMet's benefits in more complex settings. We share this enthusiasm and are actively exploring scaling TaskMet to challenging MDPs, and we will ensure that future work addresses this.\n\n> Some related work could be added, particularly work in meta-learning loss functions [...]\n\nWe agree and will add references to similar meta learning loss functions, such as [this paper](https://arxiv.org/abs/1906.05374) on meta-learning a neural network as the loss. Please let us know if there are any others you think are relevant.\n\n> The advantages of TaskMet re: interpretability are not actually validated [...]\n\nDue to character constraints, we ask to please refer to **$2^{nd}$ point of our reply to Reviewer wsqn.**\n\n> Does TaskMet scale to more real-world settings, for example more difficult RL [...]\n\nEven though our experimental settings are relatively small-scale, they are standard settings in the field, e.g. as in Shah et al. [2022], that gave us a testbed for quantifying the impact of learning a metric. The datasets and tasks are reflective for what scaled versions of the tasks would look like: for budget allocation the dataset was real-world Yahoo! Website dataset, as well as the portfolio optimization is also on real stock data from 2004 to 2017 available with QuandlWIKI dataset. Training predictive models for stock prices from QuandlWIKI or budget allocations from the Yahoo! Website are already challenging tasks and there are a lot of tradeoffs that a model can make related to the types of errors to emphasize or parts of the prediction space to be more accurate. Our perspective is that the learned metric is capable of encoding some of this information.\n\nFurthermore, the cartpole setting is taken from OMD and a first step towards scaling task-based dynamics learning to larger MDPs. Our learned metric was not only able to distinguish between artificial noise but also produce useful information such as finding the important features for the task which in this case was pole angle. \n\n> Is TaskMet limited to learning a Mahalanobis norm? Why not a more complex/expressive learned loss function?\n\nPlease refer to **$2^{nd}$ point of our reply to reviewer uXLc.**\n\n> Does TaskMet require significantly more compute than MSE?\n\nTaskMet requires minor extra compute compared to MSE. We have benchmarked this in terms of compute time per gradient step of TaskMet vs MSE. Details can be found in supplementary section B. \n"}, {"review_id": "18tACZeCfL", "reviewer": "Reviewer_wsqn", "summary": "The authors propose using task information to modify a metric that the prediction model uses for training and include task information. This way, they encourage the model to generate more accurate predictions specifically for inputs relevant to the downstream task.", "questions": "GENERAL COMMENTS:\n - (1) Normalized Test Decision Quality: please provide an equation and reference this metric.\n - (2) The authors claim that one of the key contributions is a \"more interpretable learning of the model using the metric compared to learning with a combination of task loss and prediction loss.\" Nevertheless, there is no comparison or assessment on how/how much more is the proposed approach more interpretable.\n \nTABLES:\n - (3) Table 1: align reported results to the right. Ensure they all have the same number of decimals.\n - (4) Table 2: same comment as for Table 1. Add arrows (up/down) beside the metric name to indicate whether higher/lower is better.\n\nFIGURES: \n - (5) Figure 3: provide more spacing between both Figures to clarify that they refer to different analyses. Furthermore, both can be referenced as (a) and (b). Moreover, there is no reference for the black curve at the top. Provide some short explanation of both Figures in the caption. \n - (6) Figure 6: add a brief interpretation of the results.\n\nSPELLING/WORDING:\n - (7) \"parametrize generalized mahalanobis distance based loss for training the prediction models\" -> \"parametrize generalized Mahalanobis distance based loss for training the prediction models\"\n - (8) \"value network learning insteead of minimizing\" -> value network learning instead of minimizing\n - (9) \"Given the correlatoin matrix Q\" ->  \"Given the correlation matrix Q\"\n - (10) \"We use three resource allocation tasks that has been used\" -> \"We use three resource allocation tasks that have been used\"\n - (11) \"Budget Allcoation can be found\" -> \"Budget Allocation can be found\"\n - (12) \"value network learning insteead of minimizing\" -> \"value network learning instead of minimizing\"\n - (13) \"fig. X\" -> \"Fig. X\"\n - (14) \"eq. X\" -> \"Eq. X\"\n - (15) \"table X\" -> \"Table X\"\n - (16) \"the ability of metric to infer problem specific features without manual tuning unlike other methods\" -> \"the ability of the metric to infer problem-specific features without manual tuning unlike other methods\"\n - (17) \"underlying properties of the task for useful for training the model.\" -> \"underlying properties of the task useful for training the model.\"", "limitations": "The authors have identified and acknowledged limitations of the proposed methods. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": " - The manuscript is well-structured, and the subject of research is relevant\n - The authors provide an introduction referencing a comprehensive set of related work", "weaknesses": "We have not identified strong weaknesses in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review of our paper and detailed comments. We are happy to hear you enjoyed the paper. Thanks also for  pointing out the other minor problems about the typos and figures. We will correct them. Here are some clarifications to your comments, please let us know if there is anything else. \n\n> (1) Normalized Test Decision Quality: please provide an equation and reference this metric.\n\nThis is the same metric which is used by the baseline method in previous work [Shah et al., 2022]. Test Decision Quality is the negative of test task loss. Normalized Test Decision Quality basically means normalizing the decision quality such that decision quality of random prediction is mapped to 0 whereas decision quality of oracle predictions is mapped to 1. Hence the normalized DQ is given by $DQ_{normalized}=(DQ-DQ_{random})/(DQ_{oracle}-DQ_{random}) \\in (-\\infty, 1]$ \n\n> (2) The authors claim that one of the key contributions is a \"more interpretable learning of the model using the metric compared to learning with a combination of task loss and prediction loss.\" Nevertheless, there is no comparison or assessment on how/how much more is the proposed approach more interpretable.\n\nWhen we mentioned the interpretability of TaskMet, we were thinking about (1) the prediction model making predictions in the original space rather than a latent space as in DFL, and (2) the meaning of the metric’s values across parts of the prediction space. We realize now that our use of the word “interpretability” could have other meanings, such as somehow having better interpretations of *why* the model is making certain predictions, but we will clarify that this is not what we mean.\n\nFor #2, the learned metric reveals inherent properties of the task. For example, in Figure 5, the metric puts the largest weight on dimension 3 in all the experiments, indicating that dimension 3 is the most important feature for the task. This is aligning very well with the fact that dimension 3 is actually pole angle, which is the most indicative feature for the cumulative reward since the task is to keep the pole straight. Similarly, in Figure 3, the learned metric, reveals the important prediction space for the task, allowing for the higher metric value at extreme high utility  ends, indicating that these samples should be priorities for accuracy as they are more critical for the downstream task, which makes sense as we look at the optimization problem we will see that the higher utility samples are the one which downstream task is more sensitive to. \n"}]

Author rebuttal:
We extend our gratitude to the reviewers for investing their time and expertise into our paper. The feedback provided has been immensely valuable, and we greatly appreciate the positive remarks regarding the novelty and impact of our proposed method. While we include more detailed responses inline below, we wish to highlight and emphasize the following points:

1. We are encouraged to see that the *majority of the reviewers resonated with us on the idea* – learning a metric end-to-end with a task to train models and also found the idea well corroborated with the experiments. We are delighted that our contributions and the soundness of our method have been recognized, as indicated by the favorable ratings of 3 (good) in both the areas.
2. *On the evaluation of our method in more challenging and real-world scenarios:* We would like to highlight that the decision focused model learning tasks are considered difficult and are widely used as benchmarks in task-based model learning literature. Since they incorporate solving combinatorial optimization problems, they represent a wide variety of real-world downstream optimization problems. Furthermore, while our model-based learning experiments may appear modest in scale, they effectively demonstrated the method's ability to distinguish between relevant and irrelevant features, as well as rank the task relevant features in accordance to their task relevance. At the same time, we would agree that it will be interesting to showcase these capabilities on more complex RL systems, which we assure to work towards. 
3. Lastly, we wish to emphasize the distinct capability of TaskMet in case any reviewer overlooked – *achieving task performance on par with or exceeding that of DFL methods, while having minimal prediction error and making predictions in original prediction space unlike other DFL methods.* The idea – learning metric for the model loss function instead of learning model parameters using the task loss – allows TaskMet to achieve optimal task performance while being relevant for the original prediction problem. 

In conclusion, we are genuinely appreciative of the reviewers' insights, which will be instrumental in refining our work. We look forward to more fruitful discussion in the coming week.


Meta-review: This paper adapts the end-to-end task-based model learning paradigm by defining a metric loss that measures distance between model predictions and ground truth (similar to MSE minimization) but via a metric that is parameterized by task-based information. (This is in contrast to existing task-based learning approaches that weight between an MSE loss and a task-based loss.) Experimental results show that the approach achieves task performance competitive to task-based models, while still achieving prediction error (MSE error?) competitive to MSE-optimizing methods (i.e., without overfitting to the task). The reviewers agreed the work was interesting and well-motivated and the results were promising, and were unanimous in recommending acceptance. That said, they asked that the presentation be made more accessible (e.g., by polishing the writing and adding illustrations/flow charts/algorithm boxes) and recommended the implementation of larger-scale experiments to further validate the efficacy of the method.
Decision: Accept (poster)