What is the Right Notion of Distance between Predict-then-Optimize Tasks? (2409.06997v2)

Abstract: Comparing datasets is a fundamental task in machine learning, essential for various learning paradigms-from evaluating train and test datasets for model generalization to using dataset similarity for detecting data drift. While traditional notions of dataset distances offer principled measures of similarity, their utility has largely been assessed through prediction error minimization. However, in Predict-then-Optimize (PtO) frameworks, where predictions serve as inputs for downstream optimization tasks, model performance is measured through decision regret rather than prediction error. In this work, we propose OTD$^3$ (Optimal Transport Decision-aware Dataset Distance), a novel dataset distance that incorporates downstream decisions in addition to features and labels. We show that traditional feature-label distances lack informativeness in PtO settings, while OTD$^3$ more effectively captures adaptation success. We also derive a PtO-specific adaptation bound based on this distance. Empirically, we show that our proposed distance accurately predicts model transferability across three different PtO tasks from the literature. The code is available at https://github.com/paularodr/OTD3.

• The paper proposes a decision-aware distance measure that integrates features, labels, and decisions to predict model transferability in PtO tasks.
• It leverages Optimal Transport theory with a custom cost function to compare joint feature-label-decision triplets across datasets.
• Empirical evaluations in settings like Linear Model Top-K, Warcraft, and Inventory Stock show superior performance over traditional metrics.

Notion of Distance between Predict-then-Optimize Tasks

The paper presents a new dataset distance measure specifically tailored to the context of Predict-then-Optimize (PtO) tasks. Traditional measures of dataset similarity, which typically focus on the feature and label space, fail to capture the distinct requirements of PtO frameworks where the key performance indicator is decision regret minimization. This work (i) demonstrates the inadequacy of classical dataset distances in PtO contexts and (ii) introduces a decision-aware dataset distance that integrates the effects of downstream decisions.

The PtO framework involves making predictions that are subsequently used as inputs to an optimization model, shifting the evaluation criterion from prediction error to decision quality regret. Traditional distances, such as Total Variation distance, Wasserstein distance, and Integral Probability Metrics, do not incorporate the decision-making component and thus fall short in capturing the nuances of PtO tasks. In response, the authors propose a novel dataset distance grounded in Optimal Transport (OT) theory, which factors in the feature, label, and decision dimensions.

Decision-Aware Dataset Distance

The novel distance measure is constructed by first defining a cost function that compares feature-label-decision triplets in the joint space $X \times Y \times \Omega$. This cost function is a convex combination of the individual costs in each dimension: $$c_{PtO}((x,y,z),(x',y',z')) = \alpha_X d_X(x,x') + \alpha_Y d_Y(y,y') + \alpha_W g(z,z';y,y')$$ where $d_X$ and $d_Y$ are metrics on the feature and label spaces, respectively, and $g(z,z';y,y')$ is a decision quality disparity function. The decision quality disparity captures the impact of different decisions under varying labels, making this distance particularly suited for PtO contexts.

The decision-aware distance is then obtained by solving the OT problem with the above cost function, yielding a metric that effectively measures similarities between datasets in PtO settings. This metric demonstrates superior performance in predicting the success of transferring models between different PtO tasks.

Empirical Evaluation

Empirical validation of the proposed distance involves three types of PtO tasks: Linear Model Top-K, Warcraft Shortest Path, and Inventory Stock Problem. In each of these settings, the decision-aware distance surpasses traditional feature-label distances in predicting transferability.

For instance, in the Linear Model Top-K setting, models trained on different source distributions exhibited diverse performance on a target distribution. Traditional distances failed to distinguish between source datasets with drastically different transferabilities, whereas the decision-aware distance could accurately predict which source model would perform better on the target task. Similarly, in the Warcraft setting, incorporating the decision component significantly enhanced the correlation between dataset distance and transfer performance compared to using only feature-label distances.

Implications and Future Directions

This work highlights the inadequacies of traditional dataset distances in PtO contexts and provides a principled approach for developing more suitable distance measures. The decision-aware distance offers a key tool for tasks that require fine-grained understanding of decision impacts, extending the range of applications beyond classic machine learning problems to more complex, real-world decision-making scenarios.

Future research could expand this framework to more intricate PtO structures, explore methods for tuning the component weights ($\alpha_X, \alpha_Y, \alpha_W$) independently of transferability measures, and contend with decision spaces of varying structures and dimensions. Additionally, hierarchical OT frameworks could be leveraged to improve the robustness and versatility of this distance measure for various PtO scenarios.

Conclusion

The paper makes a compelling case for the necessity of incorporating decision components into dataset distance measures for PtO tasks. By integrating features, labels, and decisions into a unified distance metric, the proposed decision-aware distance provides a more accurate and predictive measure of dataset similarity in PtO contexts, addressing a crucial gap in current methodology. This contribution lays the groundwork for enhanced model transferability and robustness in optimized decision-making environments.