One-Step Generalization Ratio Guided Optimization for Domain Generalization
2026-06-15 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address the problem of models overfitting to specific domains when trained to work on new, unseen domains. They introduce GENIE, an optimizer that measures how much each model parameter helps reduce errors and adjusts training so no small group of parameters dominates learning. This approach encourages the model to learn features that work well across different domains. Their method theoretically maintains good training speed and empirically improves performance compared to existing optimizers.
Domain GeneralizationOverfittingGradient AlignmentOptimizerOne-Step Generalization RatioPreconditioningDomain-Invariant FeaturesStochastic Gradient DescentRegularization
Authors
Sumin Cho, Dongwon Kim, Kwangsu Kim
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.