[Paper Review] Iterative Feature Matching: Toward Provable Domain Generalization with Logarithmic Environments

Yining Chen, Elan Rosenfeld+3 more|arXiv (Cornell University)|Jun 18, 2021

Domain Adaptation and Few-Shot Learning35 references4 citations

TL;DR

This paper proposes Iterative Feature Matching (IFM), a novel algorithm for domain generalization that provably achieves generalization using only O(log dₛ) training environments—significantly fewer than the linear dₛ required by IRM or ERM—under a smoothed covariance model of spurious features. The method iteratively projects features to remove spurious variation, ensuring invariant prediction with high probability.

ABSTRACT

Domain generalization aims at performing well on unseen test environments with data from a limited number of training environments. Despite a proliferation of proposal algorithms for this task, assessing their performance both theoretically and empirically is still very challenging. Distributional matching algorithms such as (Conditional) Domain Adversarial Networks [Ganin et al., 2016, Long et al., 2018] are popular and enjoy empirical success, but they lack formal guarantees. Other approaches such as Invariant Risk Minimization (IRM) require a prohibitively large number of training environments -- linear in the dimension of the spurious feature space $d_s$ -- even on simple data models like the one proposed by [Rosenfeld et al., 2021]. Under a variant of this model, we show that both ERM and IRM cannot generalize with $o(d_s)$ environments. We then present an iterative feature matching algorithm that is guaranteed with high probability to yield a predictor that generalizes after seeing only $O(\log d_s)$ environments. Our results provide the first theoretical justification for a family of distribution-matching algorithms widely used in practice under a concrete nontrivial data model.

Motivation & Objective

• To address the lack of theoretical guarantees in distribution-matching algorithms widely used in practice for domain generalization.
• To overcome the prohibitive environment complexity of IRM, which requires Ω(dₛ) environments even in simple models.
• To provide the first provable guarantee for a distribution-matching approach under a nontrivial, realistic data model with adversarial spurious feature covariances.
• To demonstrate that iterative feature projection can achieve invariant prediction with sublinear, logarithmic environment complexity.

Proposed method

• Proposes an iterative feature matching (IFM) algorithm that projects features onto the orthogonal complement of the subspace spanned by spurious feature covariances across environments.
• Uses pairwise covariance differences between environments to estimate the subspace of spurious features, enabling their removal via projection.
• Applies singular value decomposition (SVD) to the covariance difference matrix to extract the projection matrix onto the spurious subspace.
• Projects input features using the orthogonal projection matrix to eliminate spurious components while preserving signal features.
• Employs a multi-layered iterative process that matches feature distributions across environments at multiple network layers to improve robustness.
• Uses a simple algorithm (Algorithm 2) based on covariance subtraction and SVD to recover the invariant predictor with only two environments under idealized assumptions.

Experimental results

Research questions

• RQ1 Can a distribution-matching algorithm achieve provable domain generalization with sublinear environment complexity?
• RQ2 Why do existing methods like IRM and ERM fail to generalize when the number of environments is o(dₛ)?
• RQ3 Can iterative feature projection outperform single-layer matching in removing spurious features across environments?
• RQ4 Is there a data model where IRM and ERM fail even with infinite samples, while a feature-matching approach succeeds with O(log dₛ) environments?
• RQ5 What is the minimal number of environments required to provably recover an invariant predictor under realistic feature covariance structures?

Key findings

• Under a smoothed covariance model with adversarial biases, both ERM and IRM fail to generalize when the number of environments is o(dₛ), even with infinite data per environment.
• The proposed IFM algorithm achieves generalization with high probability using only O(log dₛ) environments, providing the first provable sublinear environment complexity for a distribution-matching method.
• Empirical results show IFM outperforms ERM, IRM, and CORAL on both Gaussian and Noised MNIST datasets, with test accuracy improvements of up to 10% over baselines.
• Matching features at multiple layers (e.g., all or disjoint layers) yields better performance than matching only at the final layer, with CORAL (matching all layers) achieving 77.9% accuracy on a 4-layer MLP.
• A simple two-environment algorithm (Algorithm 2) based on covariance subtraction and SVD achieves exact recovery of the invariant predictor under Assumption 3.2, proving ŵ = w*.
• Experiments show that IFM generalizes well even with limited data, with all runs completing in under 10 minutes on a single GPU.