NormAUG: Normalization-guided Augmentation for Domain Generalization

TL;DR

NormAUG tackles domain generalization by employing normalization-guided augmentation that perturbs feature statistics through a BN bank and a classifier bank, creating diverse, domain-aware representations. The two-path training framework enables an auxiliary path to generate multiple BN sub-paths and corresponding classifiers, with a test-time ensemble that fuses auxiliary and main-path predictions for improved robustness. The approach is grounded in a domain generalization bound, showing that NormAUG can reduce the secondary term $(\gamma+\rho)/2$ and tighten the upper bound on target risk $\epsilon^{t}(h)$, supported by extensive experiments on PACS, Office-Home, mini-DomainNet, Digits-DG, and DomainNet. Ablation studies confirm the contributions of optimized normalization, augmentation, and ensemble components, and analyses reveal reduced domain gaps and more focused activations, highlighting practical impact for improving generalization in real-world DG settings.

Abstract

Deep learning has made significant advancements in supervised learning. However, models trained in this setting often face challenges due to domain shift between training and test sets, resulting in a significant drop in performance during testing. To address this issue, several domain generalization methods have been developed to learn robust and domain-invariant features from multiple training domains that can generalize well to unseen test domains. Data augmentation plays a crucial role in achieving this goal by enhancing the diversity of the training data. In this paper, inspired by the observation that normalizing an image with different statistics generated by different batches with various domains can perturb its feature, we propose a simple yet effective method called NormAUG (Normalization-guided Augmentation). Our method includes two paths: the main path and the auxiliary (augmented) path. During training, the auxiliary path includes multiple sub-paths, each corresponding to batch normalization for a single domain or a random combination of multiple domains. This introduces diverse information at the feature level and improves the generalization of the main path. Moreover, our NormAUG method effectively reduces the existing upper boundary for generalization based on theoretical perspectives. During the test stage, we leverage an ensemble strategy to combine the predictions from the auxiliary path of our model, further boosting performance. Extensive experiments are conducted on multiple benchmark datasets to validate the effectiveness of our proposed method.

Figures (9)

• Figure 1: Visualization of image's features by t-SNE van2008visualizing. In this figure, "A", "C", "S", and "P" represent images from the Art painting, Cartoon, Sketch, and Photo domains, respectively. For feature extraction, we utilize a ResNet-18 DBLP:conf/cvpr/HeZRS16 model pre-trained on the ImageNet dataset DBLP:conf/cvpr/DengDSLL009. It is important to note that we dynamically adjust the statistics ($\mu$ and $\sigma$) in all normalization layers during the feature extraction process. Each image is placed in different batches with different domains. For example, "A in ACS" indicates that images from the Art painting domain are combined with images from the Cartoon and Sketch domains for normalization during the test stage. As observed, this process perturbs the original feature representation.
• Figure 2: The pipeline of our method in the training stage. In our training setup, we assume that there are three domains in the training set. The network architecture consists of two paths: the main path and the auxiliary path. Specifically, the data augmentation is performed through the auxiliary path. During each iteration, we randomly select a set of batch normalization (BN) layers from the BN bank ( e.g., AUG-2 in the figure) and the corresponding classifier from the classifier bank to train the model. The main path and the auxiliary path share all parameters except for the BN layers. It is worth noting that BN layers with the same color in the BN bank share parameters, and all normalization layers in the auxiliary path are replaced by a BN bank. The corresponding algorithm is shown in Alg. \ref{['al01']}.
• Figure 3: The forward process of our method in the training stage. Here, we assume there are 3 source domains. This figure denotes that we randomly select the "AUG-2" in Fig. \ref{['fig05']} from the BN bank at an iteration.
• Figure 4: The pipeline of our method in the test stage. In the test stage, we fuse the results from sub-paths ( i.e., "AUG-1" in Fig. \ref{['fig05']}) to further improve the accuracy of the $\mathrm{C}_m$. In this figure, "AVG" denotes the average operation.
• Figure 5: Examples from PACS. We show some images from Sketch, Art painting, Cartoon, and Photo. As seen in this figure, the difference is obvious for these images with the same class from different domains.