Geometric Rate–Distortion Invariance for Domain Generalization
Abstract
Domain generalization (DG) aims to learn representations that remain predictive under distribution shifts. A key challenge is that the target domain is unobserved during training, which complicates the search for invariant representations: alignment objectives that ignore discriminative structure can become ill-conditioned under finite samples. This calls for shaping the geometry of class-conditional representations across domains, not merely matching their distributions. We propose Geometric Rate–Distortion Invariance (RDI), a DG framework that realizes this principle by generalizing classical rate–distortion theory to Grassmann manifolds. RDI models class-conditional representations as low-dimensional subspaces and formulates DG as a joint optimization of (i) cross-domain subspace alignment (geometric distortion) and (ii) spectral–volumetric complexity (a capacity-regularized rate term), promoting stable alignment while preventing the collapse of discriminative geometry. We provide finite-sample stability guarantees under bounded shifts and show on DomainBed that RDI is competitive with strong DG baselines, with ablations confirming that both alignment and complexity control are necessary for reliable generalization.
Lay Summary
Imagine training a system to recognize dogs from professional photographs, then deploying it on grainy phone snapshots: accuracy often collapses. This problem—called domain generalization—asks how to build AI that stays reliable when the environment shifts. Existing methods mostly try to make models behave similarly across training environments by matching statistical summaries, but this can blur the very features that distinguish one class from another. We propose RDI, which represents each class not as a cloud of points but as a small plane floating in the model's internal feature space, capturing the direction along which examples of that class spread. When the same class appears under different environments, these planes can tilt in different directions—our method rotates them into alignment while preventing them from flattening into a single line, which would erase the variations that distinguish one class from another. We adapt classical ideas from signal compression theory to balance these two goals automatically. On standard benchmarks for cross-environment generalization, RDI performs competitively with leading methods. More broadly, the work suggests that the geometry of how a model represents each class—not just its statistics—is a useful design lever for building AI that remains reliable when conditions change.