[Paper Review] Modeling Uncertain Feature Representation for Domain Generalization

Xiaotong Li, Zixuan Hu+4 more|arXiv (Cornell University)|Jan 16, 2023

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper proposes Domain Shifts with Uncertainty (DSU), a method that models feature statistics as multivariate Gaussian distributions to improve deep learning generalization under domain shifts. By sampling uncertain statistics during training and applying instance-wise adaptation at inference, DSU enhances robustness across diverse out-of-distribution domains without additional parameters, achieving state-of-the-art performance on image classification, segmentation, retrieval, and pose estimation tasks.

ABSTRACT

Though deep neural networks have achieved impressive success on various vision tasks, obvious performance degradation still exists when models are tested in out-of-distribution scenarios. In addressing this limitation, we ponder that the feature statistics (mean and standard deviation), which carry the domain characteristics of the training data, can be properly manipulated to improve the generalization ability of deep learning models. Existing methods commonly consider feature statistics as deterministic values measured from the learned features and do not explicitly model the uncertain statistics discrepancy caused by potential domain shifts during testing. In this paper, we improve the network generalization ability by modeling domain shifts with uncertainty (DSU), i.e., characterizing the feature statistics as uncertain distributions during training. Specifically, we hypothesize that the feature statistic, after considering the potential uncertainties, follows a multivariate Gaussian distribution. During inference, we propose an instance-wise adaptation strategy that can adaptively deal with the unforeseeable shift and further enhance the generalization ability of the trained model with negligible additional cost. We also conduct theoretical analysis on the aspects of generalization error bound and the implicit regularization effect, showing the efficacy of our method. Extensive experiments demonstrate that our method consistently improves the network generalization ability on multiple vision tasks, including image classification, semantic segmentation, instance retrieval, and pose estimation. Our methods are simple yet effective and can be readily integrated into networks without additional trainable parameters or loss constraints. Code will be released in https://github.com/lixiaotong97/DSU.

Motivation & Objective

• To address the performance degradation of deep networks in out-of-distribution testing scenarios, particularly due to domain shifts in feature statistics.
• To model the inherent uncertainty in domain shifts by treating feature statistics (mean and standard deviation) as probabilistic distributions rather than deterministic values.
• To improve model generalization by training on diverse, uncertainty-sampled feature statistics variants during training.
• To enhance inference-time robustness through an adaptive, instance-wise calibration of feature statistics without significant computational overhead.

Proposed method

• The method models each feature statistic (mean and standard deviation) as a multivariate Gaussian distribution centered on the original deterministic value, with variance estimated from mini-batch statistics in a non-parametric manner.
• During training, feature statistics are randomly sampled from the estimated Gaussian distributions to generate diverse perturbations, simulating potential domain shifts.
• The uncertainty estimation is channel-wise and adaptive, capturing varying shift potentials across different feature channels.
• At inference, an instance-wise adaptation strategy recalibrates the feature statistics based on the test sample’s uncertainty, improving robustness to unforeseen domain shifts.
• The approach is non-parametric and does not introduce additional trainable parameters or loss constraints, making it easily integrable into existing networks.
• Theoretical analysis supports the method’s generalization error bound and implicit regularization effect, justifying its robustness.

Experimental results

Research questions

• RQ1 Can modeling feature statistics as uncertain distributions improve model generalization under domain shifts?
• RQ2 How does uncertainty-aware feature statistics sampling during training affect robustness to unseen domain distributions?
• RQ3 Can instance-wise inference-time adaptation further enhance generalization with minimal computational cost?
• RQ4 How does the proposed uncertainty estimation compare to fixed or shared uncertainty distributions across channels?
• RQ5 What is the impact of different uncertainty distribution types (e.g., Gaussian vs. uniform) on model performance?

Key findings

• DSU achieves state-of-the-art performance on the PACS dataset, improving zero-shot generalization accuracy to 84.1%, surpassing the baseline of 79.0%.
• On the GTA5-to-Cityscapes semantic segmentation benchmark, DSU boosts performance from 37.0% to 43.1%, demonstrating strong robustness in domain shift scenarios.
• The method outperforms fixed Gaussian distributions with pre-defined standard deviations, indicating that adaptive uncertainty estimation is more effective than static perturbations.
• Sharing the same uncertainty distribution across channels reduces performance (39.3% vs. 43.1%), confirming that channel-wise uncertainty modeling is essential for optimal results.
• The method incurs only a 6% increase in training time and 1% in inference time, with less than 0.02% additional parameters, proving high efficiency.