TRACE: Theoretical Risk Attribution under Covariate-shift Effects
Hosein Anjidani, S. Yahya S. R. Tehrani, Mohammad Mahdi Mojahedian, Mohammad Hossein Yassaee
TL;DR
TRACE introduces a principled, computable attribution framework for the risk change |ΔR| that occurs when replacing a source-trained model with one trained on covariate-shifted data. It decomposes |ΔR| into four interpretable components—source/generalization gap, target/generalization gap, model-change penalty, and covariate-shift penalty—and instantiates each with data-driven proxies based on Optimal Transport or Maximum Mean Discrepancy, feature-space transport, and high-quantile gradient norms. The bound is integrated into a practical diagnostic and deployment gate, validated on synthetic and DomainNet vision benchmarks with strong monotonic relationships to true risk degradation and near-perfect gating performance. The approach supports safe, label-efficient model replacement by offering actionable diagnostics that separate geometric data shift from algorithmic retraining effects. The framework is demonstrated theoretically in ridge regression and empirically across synthetic and real-world shifts, underscoring its relevance for anchor-aware deployment decisions and broader risk-sensitive ML workflows.
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, $Δ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 $|Δ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 $|ΔR|$ and achieves high AUROC/AUPRC for gating decisions, enabling safe, label-efficient model replacement.