# TaskMet: Task-Driven Metric Learning for Model Learning

arXiv: https://arxiv.org/abs/2312.05250 · HTML: https://mlsanity.com/paper/2312.05250/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 56 references that are in the map.
- First posted on arXiv 8 Dec 2023.

## Builds on (most cited first, 30 of 32)

- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Optimization Methods for Large-Scale Machine Learning](https://mlsanity.com/paper/1606.04838/) (2016)
- [Differentiable Convex Optimization Layers](https://mlsanity.com/paper/1910.12430/) (2019)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Smart "Predict, then Optimize"](https://mlsanity.com/paper/1710.08005/) (2017)
- [Optimizing Millions of Hyperparameters by Implicit Differentiation](https://mlsanity.com/paper/1911.02590/) (2019)
- [Few-shot Object Detection via Feature Reweighting](https://mlsanity.com/paper/1812.01866/) (2018)
- [Melding the Data-Decisions Pipeline: Decision-Focused Learning for Combinatorial Optimization](https://mlsanity.com/paper/1809.05504/) (2018)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [Smart Predict-and-Optimize for Hard Combinatorial Optimization Problems](https://mlsanity.com/paper/1911.10092/) (2019)
- [A Survey of Contextual Optimization Methods for Decision Making under Uncertainty](https://mlsanity.com/paper/2306.10374/) (2023)
- [Generalization Bounds in the Predict-then-Optimize Framework](https://mlsanity.com/paper/1905.11488/) (2019)
- [Tuning computer vision models with task rewards](https://mlsanity.com/paper/2302.08242/) (2023)
- [POPCORN: Partially Observed Prediction COnstrained ReiNforcement Learning](https://mlsanity.com/paper/2001.04032/) (2020)
- [Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation](https://mlsanity.com/paper/2106.03273/) (2021)
- [Value Gradient weighted Model-Based Reinforcement Learning](https://mlsanity.com/paper/2204.01464/) (2022)
- [Differentiable Gaussian Process Motion Planning](https://mlsanity.com/paper/1907.09591/) (2019)
- [Resolving Class Imbalance in Object Detection with Weighted Cross Entropy Losses](https://mlsanity.com/paper/2006.01413/) (2020)
- [SurCo: Learning Linear Surrogates For Combinatorial Nonlinear Optimization Problems](https://mlsanity.com/paper/2210.12547/) (2022)
- [Goal-Driven Dynamics Learning via Bayesian Optimization](https://mlsanity.com/paper/1703.09260/) (2017)
- [Leaving the Nest: Going Beyond Local Loss Functions for Predict-Then-Optimize](https://mlsanity.com/paper/2305.16830/) (2023)
- [Landscape Surrogate: Learning Decision Losses for Mathematical Optimization Under Partial Information](https://mlsanity.com/paper/2307.08964/) (2023)
- [Decision-Aware Learning for Optimizing Health Supply Chains](https://mlsanity.com/paper/2211.08507/) (2022)
- [Spectral, Probabilistic, and Deep Metric Learning: Tutorial and Survey](https://mlsanity.com/paper/2201.09267/) (2022)
- [Active Learning For Contextual Linear Optimization: A Margin-Based Approach](https://mlsanity.com/paper/2305.06584/) (2023)
- [Fixes That Fail: Self-Defeating Improvements in Machine-Learning Systems](https://mlsanity.com/paper/2103.11766/) (2021)
- [End-to-End Learning with Multiple Modalities for System-Optimised Renewables Nowcasting](https://mlsanity.com/paper/2304.07151/) (2023)

## Built on by (most cited first, 11 of 11)

- [Functional Bilevel Optimization for Machine Learning](https://mlsanity.com/paper/2403.20233/) (2024)
- [Decision-Focused Fine-Tuning of Time Series Foundation Models for Dispatchable Feeder Optimization](https://mlsanity.com/paper/2503.01936/) (2025)
- [Locally Convex Global Loss Network for Decision-Focused Learning](https://mlsanity.com/paper/2403.01875/) (2024)
- [What is the Right Notion of Distance between Predict-then-Optimize Tasks?](https://mlsanity.com/paper/2409.06997/) (2024)
- [Towards certifiable AI in aviation: landscape, challenges, and opportunities](https://mlsanity.com/paper/2409.08666/) (2024)
- [Prediction Loss Guided Decision-Focused Learning](https://mlsanity.com/paper/2509.08359/) (2025)
- [Bi-Level Decision-Focused Causal Learning for Large-Scale Marketing Optimization: Bridging Observational and Experimental Data](https://mlsanity.com/paper/2510.19517/) (2025)
- [DFF: Decision-Focused Fine-tuning for Smarter Predict-then-Optimize with Limited Data](https://mlsanity.com/paper/2501.01874/) (2025)
- [Task-Induced Riemannian Metrics for Vision Transformer Feature Spaces](https://mlsanity.com/paper/2609.27988/) (2026)
- [Estimate, Don't Imitate: Reusing Differentiable State-Based Policies for Visuomotor Control](https://mlsanity.com/paper/2609.34018/) (2026)
- [The Curvature of Regret in Contextual Linear Optimization](https://mlsanity.com/paper/2610.01980/) (2026)

## Abstract (arXiv)

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