Environment-Conditioned Tail Reweighting for Invariant Learning under Mixed Shifts

Abstract:Out-of-distribution generalization becomes challenging when spurious correlations vary across environments while difficult or underrepresented samples remain insufficiently emphasized within them. Invariant learning and reweighting address complementary aspects of this mixed-shift problem, yet applying reweighting only to prediction leaves the invariance constraint evaluated on a different training risk. We introduce Environment-Conditioned Tail Reweighting (ECTR), which couples cross-environment TV invariance with within-environment adversarial tail weighting through a shared weighted risk. Sample weights are normalized within each environment, and the resulting risk is used consistently for both predictive learning and the total-variation (TV) stationarity penalty, while an environment-wise KL term controls adversarial concentration. Our analysis characterizes the stationarity-dependent weighting signal and its KL-regularized distributional interpretation, alongside conditional-risk and optimization properties under the stated assumptions. Controlled mixed-shift ablations consistently support the shared-risk coupling across three difficulty settings. Across synthetic and real-world benchmarks, ECTR achieves favorable results under both observed- and inferred-environment settings, including improvements over TV-based parent methods and strong performance against additional invariant-learning baselines. Training-corruption experiments further show that ECTR assigns less excess weight to corrupted samples than the tested fixed-tail rules. Together, these results provide a unified theoretical and empirical account of environment-conditioned tail reweighting for invariant learning under mixed shifts.

Machine Learning (cs.LG)

arXiv:2601.22944 [cs.LG]

(or arXiv:2601.22944v3 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2601.22944 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.