# Environment-Conditioned Tail Reweighting for Invariant Learning under Mixed Shifts

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 30 references that are in the map.
- First posted on arXiv 30 Jan 2026.

## Builds on (most cited first, 24 of 24)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [Deeper, Broader and Artier Domain Generalization](https://mlsanity.com/paper/1710.03077/) (2017)
- [Out-of-Distribution Generalization via Risk Extrapolation (REx)](https://mlsanity.com/paper/2003.00688/) (2020)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Deep Stable Learning for Out-Of-Distribution Generalization](https://mlsanity.com/paper/2104.07876/) (2021)
- [Heterogeneous Risk Minimization](https://mlsanity.com/paper/2105.03818/) (2021)
- [OoD-Bench: Quantifying and Understanding Two Dimensions of Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.03721/) (2021)
- [NICO++: Towards Better Benchmarking for Domain Generalization](https://mlsanity.com/paper/2204.08040/) (2022)
- [Stable Learning via Sample Reweighting](https://mlsanity.com/paper/1911.12580/) (2019)
- [ZIN: When and How to Learn Invariance Without Environment Partition?](https://mlsanity.com/paper/2203.05818/) (2022)
- [In-N-Out: Pre-Training and Self-Training using Auxiliary Information for Out-of-Distribution Robustness](https://mlsanity.com/paper/2012.04550/) (2020)
- [Model Agnostic Sample Reweighting for Out-of-Distribution Learning](https://mlsanity.com/paper/2301.09819/) (2023)
- [Discovering environments with XRM](https://mlsanity.com/paper/2309.16748/) (2023)
- [Domain Generalisation via Risk Distribution Matching](https://mlsanity.com/paper/2310.18598/) (2023)
- [Invariant Risk Minimization Is A Total Variation Model](https://mlsanity.com/paper/2405.01389/) (2024)
- [Out-of-distribution Generalization for Total Variation based Invariant Risk Minimization](https://mlsanity.com/paper/2502.19665/) (2025)
- [Weighted Risk Invariance: Domain Generalization under Invariant Feature Shift](https://mlsanity.com/paper/2407.18428/) (2024)

## Abstract (arXiv)

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.