One-Step Generalization Ratio Guided Optimization for Domain Generalization
Sumin Cho, Dongwon Kim, Kwangsu Kim
cs.LG, stat.ML
Submitted: 2026-06-15
Comments: 29 pages, accepted at the 42nd International Conference on Machine Learning (ICML 2025)
License: http://creativecommons.org/licenses/by/4.0/
The gist: Domain Generalization (DG) aims to train models that generalize to unseen target domains but often overfit to domain-specific features, known as undesired correlations.
Terminology
Abstract
Domain Generalization (DG) aims to train models that generalize to unseen target domains but often overfit to domain-specific features, known as undesired correlations. Gradient-based DG methods typically guide gradients in a dominant direction but often inadvertently reinforce spurious correlations. Recent work has employed dropout to regularize overconfident parameters, but has not explicitly adjusted gradient alignment or ensured balanced parameter updates. We propose GENIE (Generalization-ENhancing Iterative Equalizer), a novel optimizer that leverages the One-Step Generalization Ratio (OSGR) to quantify each parameter's contribution to loss reduction and assess gradient alignment. By dynamically equalizing OSGR via a preconditioning factor, GENIE prevents a small subset of parameters from dominating optimization, thereby promoting domain-invariant feature learning. Theoretically, GENIE balances convergence contribution and gradient alignment among parameters, achieving higher OSGR while retaining SGD's convergence rate. Empirically, it outperforms existing optimizers and enhances performance when integrated with various DG and single-DG methods.
Sources
- Invariant Risk Minimization
- Adam: A Method for Stochastic Optimization
- When is invariance useful in an Out-of-Distribution Generalization problem ?
- UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction
- Preconditioning for Accelerated Gradient Descent Optimization and Regularization
- SAND-mask: An Enhanced Gradient Masking Strategy for the Discovery of Invariances in Domain Generalization
- The Clever Hans Mirage: A Comprehensive Survey on Spurious Correlations in Machine Learning
- Improve Unsupervised Domain Adaptation with Mixup Training
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks