Environment-Agnostic IRM via Unsupervised Clustering and Adaptive Penalty Scaling
摘要
Generalization under data shifts remains a critical challenge in machine learning. The invariant risk minimization (IRM) paradigm enhances robustness by searching invariant features across environments, but its practical application is constrained by the need to predefine environment. To overcome this limitation, we propose a clustering-based method that enables IRM training without prior environment knowledge by treating clusters as environments. Our experiments show that this approach improves model robustness to data shifts compared to empirical risk minimization (ERM). Specifically, on a weather prediction task, the mean squared error (MSE) was reduced by 10%, while in a language modeling task involving long texts, perplexity improved by up to 75%. Additionally, we introduce an adaptive hyperparameter tuning strategy for the IRM penalty term, which stabilizes training and further enhances robustness. This adaptive IRM achieves an additional 10% MSE improvement for weather prediction and a 460% perplexity gain on long textual inputs compared to classical IRM. An analysis of linear dependence between input variables and targets reveals that adaptive IRM encourages learning more complex, nonlinear invariant features, which underpins its superior generalization under distributional shifts. These results demonstrate that combining environment discovery via clustering with adaptive IRM substantially improves model generalization under distributional shifts.