Source-linked AI summary
Deep Stable Learning for Out-Of-Distribution Generalization
Xingxuan Zhang, Peng Cui, Renzhe Xu, Linjun Zhou, Yue He, Zheyan Shen
TL;DR
Deep models can fail when testing distributions shift because they exploit spurious feature-label correlations, especially when domains are unknown or unbalanced. StableNet learns sample weights to decorrelate features, combining Random Fourier Features with global feature-and-weight saving and reloading. Across distribution-generalization benchmarks and varied settings, the paper reports effective generalization compared with state-of-the-art methods.
Problem
Deep models can fail under distribution shifts, while conventional approaches assume known training-domain heterogeneity or approximately balanced domains.
Method
StableNet learns sample weights to decorrelate features, using Random Fourier Features and a global mechanism that saves and reloads features and weights during training.
Results
StableNet’s extensive experiments across distribution-generalization settings demonstrate its effectiveness compared with other methods, including strong results on unbalanced PACS and VLCS settings.
Takeaways & Limitations
Removing statistical correlations between relevant and irrelevant features can support more stable deep-model generalization under distribution shifts.
Abstract
from arXiv · showhide
Approaches based on deep neural networks have achieved striking performance when testing data and training data share similar distribution, but can significantly fail otherwise. Therefore, eliminating the impact of distribution shifts between training and testing data is crucial for building performance-promising deep models. Conventional methods assume either the known heterogeneity of training data (e.g. domain labels) or the approximately equal capacities of different domains. In this paper, we consider a more challenging case where neither of the above assumptions holds. We propose to address this problem by removing the dependencies between features via learning weights for training samples, which helps deep models get rid of spurious correlations and, in turn, concentrate more on the true connection between discriminative features and labels. Extensive experiments clearly demonstrate the effectiveness of our method on multiple distribution generalization benchmarks compared with state-of-the-art counterparts. Through extensive experiments on distribution generalization benchmarks including PACS, VLCS, MNIST-M, and NICO, we show the effectiveness of our method compared with state-of-the-art counterparts.
1. Introduction
Distribution shifts expose models’ reliance on spurious correlations between irrelevant features and labels. StableNet targets this problem without requiring known domains or balanced latent domains by decorrelating features through sample weighting.
- I.I.D.-trained models can fail under unknown distribution shifts caused by data-selection biases, confounding factors, and other real-world peculiarities.
- Saliency maps show vanilla ResNet-18 attends to both dogs and water, whereas StableNet focuses mostly on dogs.
- Spurious correlations between irrelevant and relevant features can make models associate context, such as water, with category labels.
- Domain generalization seeks invariant representations by using multiple domains where irrelevant features vary while relevant features remain stable.
- Stable learning addresses unknown domains and unbalanced latent domains by decorrelating relevant and irrelevant features, extending prior linear approaches toward deep models.
- StableNet combines nonlinear feature decorrelation with Random Fourier Features and global iterative saving and reloading of features and weights.
2. Related Works
Domain generalization studies how models trained on multiple source domains can generalize to unseen domains. Existing approaches emphasize invariant features, domain-specific modules, augmentation, and regularization.
- Domain generalization evaluates deep models’ ability to generalize to unseen domains after training on multiple source domains.
- Common approaches extract domain-invariant features, aggregate domain-specific modules, enlarge source-domain data through augmentation, or apply meta-learning and IRM regularization.
3. Sample Weighting for Distribution Generalization
StableNet learns sample weights that decorrelate representation features, using Random Fourier Features to capture nonlinear dependence. A global saving-and-reloading mechanism makes this feasible during deep-model training.
- StableNet weights training samples globally to decorrelate all representation features and remove dependencies between relevant and irrelevant features.
- Random Fourier Features enable StableNet to measure and eliminate both linear and nonlinear feature dependencies.
- The independence statistic is a nonnegative Frobenius norm that approaches zero as two variables become independent.
- StableNet optimizes sample weights by minimizing dependence statistics across feature pairs under a normalized weight constraint.
- The global mechanism merges saved features and weights with current-batch information, keeping global weights fixed while learning local weights.
- Preserving k feature-weight groups reduces storage and computation from O(N) to O(kB), while smoothing parameters combine long- and short-term information.
- During inference, StableNet predicts directly without calculating sample weights.
4. Experiments
Experiments evaluate StableNet across increasingly challenging distribution-shift settings, including unbalanced, flexible, adversarial, and classic domain configurations. Results show consistent gains on several benchmarks, while ablations support nonlinear decorrelation and global feature preservation as useful components.
- Unbalanced: StableNet outperforms other methods in all target domains on PACS and VLCS when source-domain capacity follows a 5:1:1 ratio.The result is reported for the unbalanced setting and remains consistent across other ratios in the appendix.
- Unbalanced + flexible: StableNet outperforms state-of-the-art methods on PACS and VLCS under unbalanced, flexible domains, especially when category-specific source domains differ.M-ADA and DG-MMLD fail to outperform ResNet-18 on NICO in this setting.
- Unbalanced + flexible + adversarial: As the dominant context ratio increases on MNIST-M, predictive performance drops; StableNet provides noticeable improvement under stronger spurious correlations.RSC can perform worse than vanilla CNNs because regularization weakens supervision from minor domains.
- Classic: In the classic setting, StableNet achieves the highest average accuracy on VLCS and the highest accuracy on the PACS photo target domain.Its PACS average accuracy is 0.46% lower than RSC while remaining comparable overall.
- Ablation study: Increasing Random Fourier Feature dimensions improves nonlinear dependence removal, while removing them leaves a regularizer that models only linear correlation.The ablation identifies nonlinear dependence as common in vision features and relevant to learning true input-label dependence.
- Ablation study: Increasing the presaved feature size slightly raises accuracy and significantly reduces variance by enabling more global sample-weight learning.With no presaved features, weights are learned within each batch and show noticeable variance.
5. Conclusion
StableNet is proposed to improve deep-model generalization under distribution shifts by eliminating statistical correlations between relevant and irrelevant features through sample weighting. Extensive experiments across varied settings support its effectiveness.
- StableNet eliminates statistical correlation between relevant and irrelevant features via sample weighting to improve generalization under distribution shifts.
- Extensive experiments across a wide range of settings demonstrated the effectiveness of StableNet.
A.1. Detailed Training Procedure of StableNet
StableNet trains with globally informed sample weights and iteratively optimizes model parameters and weights. Its procedure combines saved representations, local and global weights, weighted losses, and direct inference without sample weighting.
- StableNet learns batch sample weights using global feature correlations saved from earlier training.
- For each batch, current representations are concatenated with pre-saved global representations to form the combined representation input.
- The method learns local weights jointly with pre-saved global weights by optimizing the specified objective.
- The representation extractor and classifier are optimized using sample weights and penalties in a weighted prediction loss.
- Present features and weights are integrated with previously saved global features and weights during training.
- During inference, backpropagation is disabled, so StableNet bypasses sample weighting and predicts directly.
A.4. Correlation between our method and feature disentanglement
StableNet differs from feature-disentanglement methods by adjusting sample weights rather than forcing features to become disentangled, preserving their semantic meaning.
- Feature-disentanglement methods such as VAE impose strong constraints that alter the semantic implication of features.
- StableNet learns sample weights to adjust data structure while preserving feature semantics, using disentanglement only as a route toward true correlations.
- StableNet’s training procedure repeatedly reloads global features, optimizes sample weights, backpropagates weighted loss, and saves updated features and weights.
B.1. Datasets
The experiments use four datasets spanning digit images, object categories across domains, varied image styles, and context-rich non-I.I.D. classification.
- The experiments use four datasets: PACS, VLCS, MNIST-M, and NICO.
- MNIST-M blends MNIST digits over image patches extracted from BSDS500 images.
- VLCS contains five shared object categories from PASCAL VOC 2007, LabelMe, Caltech, and Sun.
- PACS contains seven object categories across photo, art-painting, cartoon, and sketch image styles.
- NICO is designed for non-I.I.D. image classification, with category images varying across ten contexts.
B.2. Training Details
The experiments use dataset-specific ResNet-18 backbones and training schedules, with ImageNet pretraining for PACS and VLCS but random initialization for NICO.
- PACS and VLCS use ResNet-18 pretrained on ImageNet because their images are insufficient for training a randomly initialized model.
- NICO uses a randomly initialized ResNet-18 backbone.
- All methods train for 30 epochs on PACS and VLCS, versus 60 epochs on NICO.
- The implementation uses authors’ published code and original hyperparameters for JiGen, DG-MMLD, and RSC, while M-ADA uses 5,000 or 8,000 iterations depending on the dataset.M-ADA uses 5,000 iterations on MNIST-M, PACS, and VLCS, and 8,000 on NICO.
B.3. More Results and Data Split Details of Unbalanced Setting
The unbalanced setting varies the relative capacities of source domains, and StableNet remains strongest across the evaluated PACS and VLCS conditions.
- Source-domain capacity ratios of 5:1:1, 3:1:1, and 2:1:1 simulate increasingly varied imbalance across PACS and VLCS.The 5:1:1 condition is used in the main setting, while 3:1:1 and 2:1:1 provide additional evaluations.
- StableNet outperforms all other state-of-the-art methods on every VLCS target domain and three of four PACS target domains at ratio 3:1:1.
- StableNet achieves the highest average accuracy across four domains on both PACS and VLCS at ratio 3:1:1.
- StableNet shows strong generalization consistently across the 5:1:1, 3:1:1, and 2:1:1 unbalanced settings.
B.4. More Results and Data Split Details of Unbalanced + Flexible Setting
The unbalanced + flexible setting tests models across all domains and varied domain ratios, while also documenting dataset construction, implementation details, and qualitative saliency findings.
- Setting: The unbalanced + flexible setting evaluates domain ratios of 5:1:1, 3:1:1, 2:1:1, and 1:1:1, with 1:1:1 reducing to the flexible setting.
- Setting: Unlike classic domain generalization, this setting trains once and tests on all domains for given source domains.
- Results: StableNet outperforms other state-of-the-art methods on almost all target domains and average accuracy across PACS and VLCS, including the 1:1:1 condition.
- Datasets: NICO contains 19 categories and 10 contexts per category, with contexts defined by backgrounds or object postures depending on the category.
- Limitations: CNBB results are omitted because the method is designed for AlexNet and did not achieve reasonable results in the authors’ ResNet framework.
- Analysis: Saliency maps indicate that StableNet focuses on objects rather than context, whereas JiGen emphasizes possible jigsaw-puzzle margins.