Environment-Conditioned Tail Reweighting for Total Variation Invariant Risk Minimization

Out-of-distribution (OOD) generalization remains challenging when models simultaneously encounter correlation shifts across environments and diversity shifts driven by rare or hard samples. Existing invariant risk minimization (IRM) methods primarily address spurious correlations at the environment level, but often overlook sample-level heterogeneity within environments, which can critically impact OOD performance. In this work, we propose Environment-Conditioned Tail Reweighting for Total Variation Invariant Risk Minimization (ECTR), a unified framework that augments TV-based invariant learning with environment-conditioned tail reweighting to jointly address both types of distribution shift. By integrating environment-level invariance with within-environment robustness, the proposed approach makes these two mechanisms complementary under mixed distribution shifts. We further extend the framework to scenarios without explicit environment annotations by inferring latent environments through a minimax formulation. Experiments across regression, tabular, time-series, and image classification benchmarks under mixed distribution shifts demonstrate consistent improvements in both worst-environment and average OOD performance.

💡 Research Summary

The paper tackles the challenging problem of out‑of‑distribution (OOD) generalization when both correlation shifts across environments and diversity shifts driven by rare or hard samples occur simultaneously. Existing invariant risk minimization (IRM) approaches, especially recent total‑variation (TV) formulations, enforce invariance at the environment level but ignore heterogeneity within each environment, leading to poor performance under diversity shifts.

To address this gap, the authors propose Environment‑Conditioned Tail Reweighting for Total‑Variation Invariant Risk Minimization (ECTR). The method introduces a weight‑adversary θ that assigns a softmax probability πθ(i) to every training example based on a learned score. These global probabilities are then conditioned on environment membership m_{i,e} (hard one‑hot if environments are known, soft qη(e|i) if they are inferred) to produce environment‑specific distributions πθ(i|e) = πθ(i)·m_{i,e}/mass_e, where mass_e is the total probability mass allocated to environment e. This conditioning ensures that “tail” or hard examples receive higher weight within each environment without altering the relative contribution of different environments.

Using πθ(i|e), the environment‑conditioned tail risk R_{e,θ}(w∘Φ) = Σ_i πθ(i|e)·ℓ(w∘Φ(x_i), y_i) is defined, and the main objective is the average of these tail risks across environments. Crucially, the TV‑ℓ₁ invariance penalty is computed on the same tail risk, i.e., the L1 norm of the gradient with respect to a dummy scalar classifier w evaluated at w=1: P_TV(Φ,θ) = (1/E) Σ_e ‖∇w R{e,θ}(w∘Φ)‖₁|_{w=1}. An invariance adversary Ψ learns a non‑negative multiplier λ(Ψ,Φ) that scales only this TV term, preserving the original IRM spirit while focusing on the reweighted risks.

To prevent the adversarial reweighting from collapsing onto a few samples, an environment‑wise KL regularizer is added: KL_env(θ) = (1/E) Σ_e KL(πθ(·|e) ‖ Uniform_e). Uniform_e denotes the uniform distribution over the samples belonging to environment e. When environments are inferred, the KL term is detached from the soft assignments so that it regularizes only θ. This stabilizes training and maintains sufficient effective sample size.

If environment labels are unavailable, a second adversary η learns soft assignments qη(e|i). The same definitions of environment‑conditioned weights, tail risks, and TV penalties apply, with the KL regularizer still acting solely on θ. The overall min‑max objective becomes:

 min_Φ max_{θ,Ψ,η}  E_e