Source-linked AI summary
Representation via Representations: Domain Generalization via Adversarially Learned Invariant Representations
Zhun Deng, Frances Ding, Cynthia Dwork, Rachel Hong, Giovanni Parmigiani, Prasad Patil, Pragya Sur
TL;DR
The paper asks whether labeled data from multiple populations can reveal a shared predictor despite differences between studies. It learns representations with adversarial domain censoring, then analyzes the loss limit and conditions for prediction on unseen domains.
Problem
The paper asks whether labeled data from multiple populations can identify a shared high-quality predictor despite biological or data-acquisition differences between studies.
Source: §1 · pp. 1–2Method
It models domains as draws from a distribution and learns an encoder and classifier while adversarially reducing domain identification from representations.
Source: §1 · p. 2Results
The analysis characterizes the loss limit as the domain count grows, gives non-asymptotic bounds, and establishes worst-case unseen-domain guarantees under a bounded H-divergence condition.
Source: §1 · p. 3, §3.2 · p. 8Takeaways & Limitations
Good worst-case prediction on unseen domains is supported when each domain is close to one receiving relatively high probability under the domain distribution.
Source: §3.2 · p. 8The theoretical analysis does not yet fully capture the conditions under which the approach is expected to succeed.
Source: §5 · p. 11
Abstract
from arXiv · showhide
We investigate the power of censoring techniques, first developed for learning {\em fair representations}, to address domain generalization. We examine {\em adversarial} censoring techniques for learning invariant representations from multiple "studies" (or domains), where each study is drawn according to a distribution on domains. The mapping is used at test time to classify instances from a new domain. In many contexts, such as medical forecasting, domain generalization from studies in populous areas (where data are plentiful), to geographically remote populations (for which no training data exist) provides fairness of a different flavor, not anticipated in previous work on algorithmic fairness. We study an adversarial loss function for $k$ domains and precisely characterize its limiting behavior as $k$ grows, formalizing and proving the intuition, backed by experiments, that observing data from a larger number of domains helps. The limiting results are accompanied by non-asymptotic learning-theoretic bounds. Furthermore, we obtain sufficient conditions for good worst-case prediction performance of our algorithm on previously unseen domains. Finally, we decompose our mappings into two components and provide a complete characterization of invariance in terms of this decomposition. To our knowledge, our results provide the first formal guarantees of these kinds for adversarial invariant domain generalization.
1 Introduction
Studies can differ biologically or in data-acquisition methods, making predictors trained at one hospital unreliable elsewhere. The paper seeks shared predictive signals across populations using adversarially censored representations.
1 Introduction
- The central question is whether labeled data from several populations can reveal a common high-quality predictor, even when that predictor is not best for any one population.Source: §1 · p. 2
- Censoring domain information is intended to favor signals shared across domains over population-specific ones, extending fair-representation ideas to unseen populations.Source: §1 · p. 2
- Domains are sampled from a distribution, and an encoder and classifier are trained with an adversarial loss that penalizes prediction errors and correct domain guesses.Source: §1 · p. 2
- The paper characterizes the adversarial loss as k grows, gives non-asymptotic generalization bounds and optimizer consistency, and bounds unseen-domain performance under bounded H-divergence.Source: §1 · p. 3
- Prior work [2] uses probabilistic invariance with guarantees relying on linearity, while adversarial approaches lack the theoretical insights developed here.Source: §1 · p. 3
2 Formal setup
The setup treats observed domains as samples from a finite collection and represents each domain's data with corresponding samples. An encoder, discriminator, and predictor are trained through a prediction-plus-adversarial loss.
2 Formal setup
- The formal setup assumes finitely many possible domains, distinguishes their input marginals, and denotes the observed domains and their samples by D1:k and S1:k.Source: §2 · p. 3
- The encoder maps covariates into representations, the discriminator maps them to domain scores, and a predictor maps representations to outcome labels.Source: §2 · pp. 3–4
- The objective combines average prediction error with a λ-weighted adversarial loss, and optimization seeks parameters that reduce prediction error while limiting domain identification.Source: §2 · p. 4
- The adversarial loss generalizes H-divergence, and the paper assumes its function-class extrema are attained; experiments use neural-network classes.Source: §2 · pp. 4–5
3 Theoretical Results
The theoretical results characterize the adversarial loss as the number of sampled domains grows, then establish conditions for generalization to unseen domains and characterize invariant mappings. They show that more domains help generalization under stated assumptions, while worst-case guarantees require every domain to be well represented under the domain-sampling distribution.
3.1 Learning theoretic analysis
- As k grows, adversarial ID-guessing accuracy corresponds to assigning encoded domain marginals to N disjoint regions; the limiting loss is minimized when all encoded marginals are identical.The partition characterization holds over the finite set of true domains, including when sampled domains repeat.Source: §3.1 · pp. 5–6
- A geometric grid argument handles the discriminator’s changing structure: finer cells are associated with the underlying domain most prevalent in each cell.The final discriminator layer can assign the highest weight in each cell to its corresponding distribution.Source: §3.1 · p. 6
- Under continuity, boundedness, and bounded VC-dimension assumptions, Theorem 3.2 bounds empirical loss against population loss.The bound applies under Assumptions 3.1–3.3.Source: §3.1 · pp. 5, 7
- In the bound, term III vanishes when min_i n_i = Ω(k^α) for α ≥ 2, term I is small when k is much larger than N, and more domains help generalization.The authors also show that term II becomes small for sufficiently large k.Source: §3.1 · pp. 7–8
- Under the stated assumptions and a curvature condition on population loss, the empirical-loss estimates approximate the population-loss minimizers almost surely.The curvature condition is described as slightly weaker than strong convexity.Source: §3.1 · p. 8
3.2 Generalization to unseen domains
- For any unseen domain, Theorem 3.4 gives a worst-case guarantee when each domain is close to one receiving relatively high probability under μ; this also holds for countably infinite D.The well-represented condition requires a nearby domain with probability at least p_l, within distance δ in H = F × Φ.Source: §3.2 · p. 8
3.3 Characterization of invariant representation mappings
- Using a decomposition of mappings, Theorem 3.5 characterizes invariance when the discriminator’s weight regions cover R^s; that condition prevents ties and is not needed for earlier results.The decomposition is introduced as the basis for characterizing invariant mappings.Source: §3.3 · p. 9
4 Experiments
Experiments on synthetic data, colored MNIST, and PACS assess generalization to unseen domains. They report gains from observing more domains and favorable comparisons with baselines across several settings.
Experimental setup
- The experiments cover synthetic biomedical-style data, colored MNIST, and PACS, with comparisons to logistic regression and random forest on synthetic data and IRM and CIDDG on colored MNIST and PACS.The code drops the decoder to match the theoretical setting; extending the theory to include it is left for future work.Source: §4 · p. 9
Synthetic Data
- Synthetic-data test accuracy increases from k=4 to k=10, and the method uniformly outperforms both baselines by a notable margin.Tests use a new unseen domain of the same form as the training domains, with different parameters.Source: §4 · p. 10, Table 1 · p. 10
Colored MNIST
- On colored MNIST, the method beats IRM and CIDDG in multiple settings, performs comparably in others, and improves test accuracy with more seen domains.The comparison includes training with three versus six domains on the same test data.Source: §4 · p. 10, Table 2 · p. 11, Table 3 · p. 11
PACS
- On PACS, the method beats both baselines across all target-domain settings, with significant margins in most settings.Each of Photos, Art, Cartoon, and Sketch serves as the target domain in turn.Source: Table 4 · p. 11, §4 · p. 11
5 Discussion
The discussion frames unseen-population generalization as relevant to biomedical prediction and notes limits in the current theoretical analysis.
Discussion
- The theoretical analysis does not yet fully capture the conditions under which the approach is expected to succeed.Extending the theoretical results to cross-entropy or other loss notions is also identified as an open question.Source: §5 · p. 11
- The work targets biomedical prediction beyond studied populations, and the authors report initiating engagement with the applications that inspired it.They describe high-dimensional biomedical datasets and predictive algorithms used to inform personalized medical care.Source: §5 · p. 11