PGrad: Learning Principal Gradients For Domain Generalization

TL;DR

PGrad introduces a novel training strategy that learns principal gradient directions to enhance neural network robustness and generalization across unseen domains, effectively handling out-of-distribution challenges.

Contribution

The paper proposes PGrad, a new gradient aggregation method that improves domain generalization by focusing on main parameter dynamics and reducing domain-specific noise.

Findings

PGrad achieves competitive results on seven datasets.

It effectively reduces loss curves during training.

Demonstrates robustness across synthetic and real-world shifts.

Abstract

Machine learning models fail to perform when facing out-of-distribution (OOD) domains, a challenging task known as domain generalization (DG). In this work, we develop a novel DG training strategy, we call PGrad, to learn a robust gradient direction, improving models' generalization ability on unseen domains. The proposed gradient aggregates the principal directions of a sampled roll-out optimization trajectory that measures the training dynamics across all training domains. PGrad's gradient design forces the DG training to ignore domain-dependent noise signals and updates all training domains with a robust direction covering main components of parameter dynamics. We further improve PGrad via bijection-based computational refinement and directional plus length-based calibrations. Our theoretical proof connects PGrad to the spectral analysis of Hessian in training neural networks. Experiments on DomainBed and WILDS benchmarks demonstrate that our approach effectively enables robust DG optimization and leads to smoothly decreased loss curves. Empirically, PGrad achieves competitive results across seven datasets, demonstrating its efficacy across both synthetic and real-world distributional shifts. Code is available at https://github.com/QData/PGrad.