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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

Environment-Agnostic IRM via Unsupervised Clustering and Adaptive Penalty Scaling

  • Bratenkov Miron,
  • Ivan Bondarenko

摘要

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.