Source-linked AI summary
Learning to Balance Specificity and Invariance for In and Out of Domain Generalization
Prithvijit Chattopadhyay, Yogesh Balaji, Judy Hoffman
TL;DR
Domain generalization often favors features shared across source domains, although source-specific characteristics can improve recognition when test data resembles particular sources. DMG learns domain-specific masks over a shared network to balance specificity and invariance, achieving competitive in- and out-of-domain performance on PACS and DomainNet.
Problem
Domain generalization needs a single model for unseen domains, but relying only on shared features can overlook source-specific characteristics useful for in-domain recognition.
Method
DMG learns distribution-specific masks over neurons in a shared CNN, optimizing classification and reduced cross-domain mask overlap to balance invariant and specific features.
Results
DMG achieves competitive in- and out-of-domain performance on PACS and DomainNet, with DMG-KD improving over baselines when test-domain labels are known.
Takeaways & Limitations
Balancing specialized features with components shared across source domains supports one model’s performance across multiple distinct domains.
Takeaways & Limitations
On DomainNet’s hardest shifts, Quickdraw performance is below 25%, and DMG performs comparably to naive baselines on some such shifts.
Abstract
from arXiv · showhide
We introduce Domain-specific Masks for Generalization, a model for improving both in-domain and out-of-domain generalization performance. For domain generalization, the goal is to learn from a set of source domains to produce a single model that will best generalize to an unseen target domain. As such, many prior approaches focus on learning representations which persist across all source domains with the assumption that these domain agnostic representations will generalize well. However, often individual domains contain characteristics which are unique and when leveraged can significantly aid in-domain recognition performance. To produce a model which best generalizes to both seen and unseen domains, we propose learning domain specific masks. The masks are encouraged to learn a balance of domain-invariant and domain-specific features, thus enabling a model which can benefit from the predictive power of specialized features while retaining the universal applicability of domain-invariant features. We demonstrate competitive performance compared to naive baselines and state-of-the-art methods on both PACS and DomainNet.
1 Introduction
Domain generalization must balance domain-invariant features for unseen distributions with domain-specific features that can improve recognition when test data resembles source domains. DMG learns this balance through domain-specific masks and achieves competitive in- and out-of-domain performance on PACS and DomainNet.
- Domain generalization seeks models that perform beyond held-out in-domain data across diverse test-time distributions.
- Invariant-feature methods may overlook domain-specific characteristics that help predict instances resembling one or more source domains.The usefulness of source-specific features depends on the test instance’s relative similarity to the available source domains.
- DMG automatically balances domain-invariant and domain-specific features in a single model across multiple distinct domains.
- DMG casts balanced feature selection as learning distribution-specific binary masks over neurons in a shared CNN, optimized end-to-end with network parameters.The masks determine whether individual task-network neurons are active during a forward pass.
- DMG targets both unseen-domain performance and source-domain performance because realistic deployment can include data from both distributions.
- DMG achieves competitive out-of-domain performance on PACS and DomainNet, while also supporting in-domain use and conversion to a domain-specific model with a known test-time label.The model can serve as a drop-in replacement for an aggregate model on in-domain samples.
2 Related Work
Related work largely promotes domain-invariant representations for generalization or disentangles domain-specific and domain-independent factors in adaptation settings. DMG instead performs an implicit disentanglement through masking while drawing connections to pruning, sparsification, dropout, and attention.
- Domain adaptation methods align source and target distributions in feature space using adversarial or MMD-based losses, alongside pixel-level alignment approaches.
- Prior adaptation work disentangles domain-specific and domain-independent factors, whereas DMG performs a similar implicit disentanglement through masking.
- DMG’s domain-specific masks are related to techniques from network pruning, sparsification, dropout, and attention.
3 Approach
DMG learns domain-specific masks over a shared task network to balance domain-specific and domain-invariant features. It trains these masks with predictive loss and minimal overlap, then combines source-domain masks for test-time prediction.
- 3 Approach: DMG learns distribution-specific binary masks over features of a shared CNN to balance domain-specific and domain-invariant components.Each source domain has mask parameters that determine whether task-network neurons are active during a forward pass.
- 3 Approach: For each source-domain input, mask probabilities are converted into sampled binary neuron masks, whose elementwise products with activations determine the next-layer inputs.The sampled masks are applied throughout the task network, producing domain-specific subnetworks during training.
- 3 Approach: The mask parameters and network parameters are learned end-to-end with classification loss and a soft-overlap loss that encourages source-domain masks to overlap minimally.The straight-through estimator approximates gradients through the discrete mask samples, while sIoU penalizes overlap between mask distributions.
- 3 Approach: The soft-overlap objective encourages predictions from different source domains to use different mask-defined subnetworks.The overlap approximation uses the inner product of mask distributions for each pair of source-domain masks at a layer.
- 3 Approach: At test time, DMG scales neurons by each source domain’s soft mask and averages the resulting predictions instead of enumerating all sampled subnetworks.This scheme is used because the number of possible domain-specific subnetworks is exponential.
4 Experiments
Experiments evaluate DMG on PACS and DomainNet using out-of-domain and, for DomainNet, in-domain performance. DMG is compared with aggregate, multi-headed, and prior domain-generalization methods, while additional analyses examine mask specialization, robustness to λO, and challenging shifts.
- Datasets and Metrics: The proposed PACS splits lack an in-domain test split, preventing in-domain performance measurement there.This limitation applies to PACS under the specified experimental splits.
- Datasets and Metrics: DomainNet contains approximately 0.6 million images across 345 classes and six domains, with out-of-domain accuracy measured on the unseen-domain test split and in-domain accuracy averaged over source-domain test splits.Experiments use leave-one-out shifts and evaluate ImageNet-pretrained AlexNet, ResNet-18, and ResNet-50 backbones.
- Out-of-Domain Generalization: DMG improves overall DomainNet out-of-domain performance by 2.7% over MetaReg and 2.6% over Aggregate with AlexNet.With ResNet-18 and ResNet-50, DMG is reported as competitive with MetaReg and Aggregate-SGD, respectively.
- Out-of-Domain Generalization: On PACS, DMG with AlexNet outperforms the listed baselines and prior approaches, including a 4.09% improvement over MASF on the A,C,P→S shift.ResNet-18 and ResNet-50 produce comparable or improved overall performance, with reported margins of 0.04% and 0.7%, respectively.
- Out-of-Domain Generalization: DomainNet remains difficult on shifts targeting Quickdraw and Infograph, with Quickdraw performance below 25% and DMG often comparable to naive baselines.The authors identify substantial room for improvement on these hardest shifts.
- In-Domain Generalization: For DomainNet in-domain evaluation, DMG-KD achieves significant improvements over baselines when the test domain is known, consistently improving over DMG by approximately 1–2%.The result supports specialization of the learned masks for individual source domains.
5 Analysis
DMG’s source-domain masks become specialized: matching a test instance with its corresponding source mask improves in-domain performance, while some specialized masks remain competitive for unseen domains.
- Domain Specialization: Matching each source domain with its corresponding mask improves in-domain accuracy over mismatched mask-domain pairs.Differences from the second-best source mask are approximately 2–4% on I,P,Q,R,S→C and 3–6% on C,I,P,R,S→Q.
- Domain Specialization: Averaging predictions from all source masks approaches DMG-KD but remains approximately 2–3% lower than it.
- Domain Specialization: Some individual source masks achieve out-of-domain accuracies within 1% of the combined setting.mQuickdraw does so for I,P,Q,R,S→C, while mClipart, mInfograph, and mSketch do so for C,I,P,R,S→Q.
- Domain Specialization: These results support using domain-specific characteristics alongside invariant features for generalization.
6 Conclusion
The paper proposes DMG, which balances domain-specific and domain-invariant representations through neuron masks. It reports competitive out-of-domain performance on PACS and competitive in- and out-of-domain performance on DomainNet.
- DMG balances domain-specific and domain-invariant feature representations to produce one model for multi-source domain generalization.
- DMG learns this balance with domain-specific neuron masks optimized to minimize cross-domain feature overlap.
- DMG achieves competitive out-of-domain performance on PACS and competitive in- and out-of-domain performance on DomainNet.
- The paper notes that blending specificity and invariance may also support model interpretability, but leaves this beyond its scope.
7 Appendix
The appendix analyzes mask specificity, incentive choices, test-time ensembling, benchmark comparisons, and implementation details. It finds that sparsity is sensitive and can produce overly sparse, highly overlapping masks, whereas alternative ensembling remains close in performance.
- 7.1 Domain Specificity: Discrete source-domain masks are obtained by thresholding soft-mask values at 0.5, and pairwise IoU measures their overlap.The appendix uses IoU to quantify mask specificity after convergence.
- 7.1 Domain Specificity: For the I,P,Q,R,S→C and C,I,P,R,S→Q shifts, Fig. 4 reports pairwise mask IoU and average activated-neuron fractions on DomainNet.
- 7.1 Domain Specificity: Pairwise IoU categorizes neurons as off across domains, shared across domains, or specific to selected source domains.
- 7.1 Domain Specificity: Average IoU decreases as λO increases across fc6, fc7, and fc8, indicating greater domain-specificity in individual layers.
- 7.2 Choice of Incentive: sIoU vs Sparsity: Replacing sIoU with an L1 sparsity penalty makes predictive performance and specificity highly sensitive to λS.
- 7.2 Choice of Incentive: sIoU vs Sparsity: Performance is comparable to DMG only at λS = 10^-5, while higher λS values significantly degrade in- and out-of-domain accuracy.
- 7.2 Choice of Incentive: sIoU vs Sparsity: At high λS, masks become extremely sparse and highly overlapping, effectively encouraging a bottleneck layer; sIoU better preserves the specificity–performance balance.
- 7.3 Ensembling Choices: Mask-Ens performs within approximately 1% of Pred-Ens for both in- and out-of-domain performance on DomainNet.