TRACE: Theoretical Risk Attribution under Covariate-shift Effects

2026-02-11

AI Summary

This work addresses the challenge of predicting performance changes when a source-domain model is replaced by a new one. To this end, the authors propose TRACE, a novel framework that, for the first time, decomposes the risk difference between two models under covariate shift into four interpretable components: two generalization gaps, a model change penalty, and a covariate shift penalty. The framework establishes a computable upper bound to diagnose the causes of performance degradation. TRACE estimates model sensitivity via high-quantile input gradients, quantifies data distribution shift using either optimal transport (OT) or maximum mean discrepancy (MMD), and measures model change through output distances on target samples. Experiments demonstrate that TRACE’s diagnostic scores exhibit strong monotonic correlation with actual performance degradation and achieve superior performance in deployment gating, as measured by AUROC and AUPRC, thereby enabling label-efficient and safe model replacement.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Adversarial Attacks & Robustness
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

When a source-trained model $Q$ is replaced by a model $\tilde{Q}$ trained on shifted data, its performance on the source domain can change unpredictably. To address this, we study the two-model risk change, $\Delta R := R_P(\tilde{Q}) - R_P(Q)$, under covariate shift. We introduce TRACE (Theoretical Risk Attribution under Covariate-shift Effects), a framework that decomposes $|\Delta R|$ into an interpretable upper bound. This decomposition disentangles the risk change into four actionable factors: two generalization gaps, a model change penalty, and a covariate shift penalty, transforming the bound into a powerful diagnostic tool for understanding why performance has changed. To make TRACE a fully computable diagnostic, we instantiate each term. The covariate shift penalty is estimated via a model sensitivity factor (from high-quantile input gradients) and a data-shift measure; we use feature-space Optimal Transport (OT) by default and provide a robust alternative using Maximum Mean Discrepancy (MMD). The model change penalty is controlled by the average output distance between the two models on the target sample. Generalization gaps are estimated on held-out data. We validate our framework in an idealized linear regression setting, showing the TRACE bound correctly captures the scaling of the true risk difference with the magnitude of the shift. Across synthetic and vision benchmarks, TRACE diagnostics are valid and maintain a strong monotonic relationship with the true performance degradation. Crucially, we derive a deployment gate score that correlates strongly with $|\Delta R|$ and achieves high AUROC/AUPRC for gating decisions, enabling safe, label-efficient model replacement.

Problem

• covariate shift
• model replacement
• risk change
• performance degradation
• distribution shift

Innovation

• covariate shift
• risk attribution
• model diagnostics
• optimal transport
• deployment gating