TaskMet: Task-Driven Metric Learning for Model Learning

Deep learning models are often deployed in downstream tasks that the training procedure may not be aware of. For example, models solely trained to achieve accurate predictions may struggle to perform well on downstream tasks because seemingly small prediction errors may incur drastic task errors. The standard end-to-end learning approach is to make the task loss differentiable or to introduce a differentiable surrogate that the model can be trained on. In these settings, the task loss needs to be carefully balanced with the prediction loss because they may have conflicting objectives. We propose take the task loss signal one level deeper than the parameters of the model and use it to learn the parameters of the loss function the model is trained on, which can be done by learning a metric in the prediction space. 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 at https://github.com/facebookresearch/taskmet

Paper

References (56)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer wsqn7/10 · confidence 4/52023-07-06

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.

Strengths

- The manuscript is well-structured, and the subject of research is relevant - The authors provide an introduction referencing a comprehensive set of related work

Weaknesses

We have not identified strong weaknesses in this paper.

Questions

GENERAL COMMENTS: - (1) Normalized Test Decision Quality: please provide an equation and reference this metric. - (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. TABLES: - (3) Table 1: align reported results to the right. Ensure they all have the same number of decimals. - (4) Table 2: same comment as for Table 1. Add arrows (up/down) beside the metric name to indicate whether higher/lower is better. FIGURES: - (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. - (6) Figure 6: add a brief interpretation of the results. SPELLING/WORDING: - (7) "parametrize generalized mahalanobis distance based loss for training the prediction models" -> "parametrize generalized Mahalanobis distance based loss for training the prediction models" - (8) "value network learning insteead of minimizing" -> value network learning instead of minimizing - (9) "Given the correlatoin matrix Q" -> "Given the correlation matrix Q" - (10) "We use three resource allocation tasks that has been used" -> "We use three resource allocation tasks that have been used" - (11) "Budget Allcoation can be found" -> "Budget Allocation can be found" - (12) "value network learning insteead of minimizing" -> "value network learning instead of minimizing" - (13) "fig. X" -> "Fig. X" - (14) "eq. X" -> "Eq. X" - (15) "table X" -> "Table X" - (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" - (17) "underlying properties of the task for useful for training the model." -> "underlying properties of the task useful for training the model."

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors have identified and acknowledged limitations of the proposed methods.

Reviewer jBQF5/10 · confidence 3/52023-07-06

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.

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?). The 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). In 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. Some related work could be added, particularly work in meta-learning loss functions, but this is not critical. The 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.

Questions

Does TaskMet scale to more real-world settings, for example more difficult RL problems? Does the learning metric produce useful information for real prediction problems, other than the artificial cartpole with noise setting? Is TaskMet limited to learning a Mahalanobis norm? Why not a more complex/expressive learned loss function? Does TaskMet require significantly more compute than MSE?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Yes, although addressing the differences in computational requirements (if any) would be helpful.

Authorsrebuttal2023-08-18

Dear reviewer, Thank you again for taking the time to go through our paper. We've taken your comments into account in the latest (and hopefully clearer) version of our paper. When you get a chance, can you please read through our response and let us know if you are curious in discussing anything else? We still agree that the experimental settings generally studied by the DFL/OMD literature are relatively small-scale relative to other parts of the ML field, but scoped to this sub-field, our experimental settings are standard. We think our results sufficiently demonstrate and validate the idea of learning a metric rather than the model using the end-to-end loss. Can you please take this into consideration for your final evaluation of our paper? Best regards, The authors.

Reviewer WKyg6/10 · confidence 3/52023-07-07

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.

Strengths

+ The paper is well-written and easy to follow. + Motivations are clear and reasonable. + The proposed method is novel and has the potential to be widely applicable. + 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.

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? - Are the baselines results re-implemented by the authors? how did you pick hyper parameters for them?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Authors adequately addressed the limitations.

Reviewer uXLc7/10 · confidence 2/52023-07-26

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.

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. - The results are promising. The proposed method has outperformed the other baseline approaches.

Weaknesses

I'm not an expert in this field. - The writing of the methodology is a bit difficult to understand. It would be great if the writing can be polished with more illustrations.

Questions

I'm not an expert in this field. 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The paper has provided the limitations of the proposed method.

Reviewer uXLc2023-08-17

Post rebuttal

Thanks for the rebuttal! The rebuttal has addressed my concerns -- I will keep my previous score.

Reviewer WKyg2023-08-17

I'd like to thanks the authors for providing clarifications on the raised points. Having reviewed both the authors' response and the feedback provided by other reviewers, I find that all of my concerns have been adequately addressed. As a result, I am maintaining my original recommended score.

Reviewer wsqn2023-08-17

We thank the authors for the detailed response. We consider they have addressed all of our comments.

Reviewer jBQF2023-08-19

I appreciate the authors' clarifications about the experimental settings, usage of Mahalanobis norm, and tuning of the baselines. These have addressed some of my primary concerns. My primary remaining concern is the experimental settings. Although they may have been used in prior work, at some point, experimental settings must adapt to the complexity and capabilities of our methods, and the settings considered in this paper feel very simple and synthetic to me. Nonetheless, while I think the paper has shortcomings in evaluations, I would not be upset if it were accepted. Therefore I will raise my score to 5.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC