PaperlayerView exact revision on arXiv

Source-linked AI summary

Batch Normalization Embeddings for Deep Domain Generalization

Mattia Segu, Alessio Tonioni, Federico Tombari

arXiv:2011.12672v3cs.LGcs.CV

TL;DR

Models trained on one distribution can fail under domain shift, while collecting annotated samples for every possible test environment is often infeasible. The paper uses domain-specific batch-normalization statistics to embed domains in a shared latent space and train a lightweight ensemble. It reports better performance than alternatives across multiple benchmarks, including a +6.33% relative gain on PACS with AlexNet.

  • Problem

    Models trained on one distribution can fail under domain shift, while collecting annotated samples for every possible test environment is often infeasible.

    Source: §1 · p. 1

  • Method

    The paper collects domain-specific batch-normalization statistics, maps domains into a shared latent space, and uses them to train a lightweight ensemble.

    Source: §1 · p. 1

  • Results

    Across several domain-generalization benchmarks, the method outperforms many alternatives; on PACS with AlexNet, its relative gain is +6.33%, versus +4.88% for the second-best method.

    Source: §5 · p. 8, §6.3.1 · p. 11

  • Takeaways & Limitations

    The findings underline the advantage of maintaining domain-specific representations over forcing invariant ones.

    Source: §5 · p. 8

    The setup assumes target samples are unavailable during training and allows each target sample to come from a different unseen domain.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

Domain generalization aims at training machine learning models to perform robustly across different and unseen domains. Several recent methods use multiple datasets to train models to extract domain-invariant features, hoping to generalize to unseen domains. Instead, first we explicitly train domain-dependant representations by using ad-hoc batch normalization layers to collect independent domain's statistics. Then, we propose to use these statistics to map domains in a shared latent space, where membership to a domain can be measured by means of a distance function. At test time, we project samples from an unknown domain into the same space and infer properties of their domain as a linear combination of the known ones. We apply the same mapping strategy at training and test time, learning both a latent representation and a powerful but lightweight ensemble model. We show a significant increase in classification accuracy over current state-of-the-art techniques on popular domain generalization benchmarks: PACS, Office-31 and Office-Caltech.

1. Introduction

The paper addresses domain generalization by learning domain-specific representations rather than discarding domain information, using batch-normalization statistics to embed domains and guide predictions for unseen domains. Its lightweight ensemble shares parameters except for normalization statistics and applies to modern CNNs with batch normalization.

1. Introduction

  • Models trained on one data distribution often fail under domain shift, while collecting annotated samples from every possible test environment is often infeasible.Fine-tuning with explicit supervision or domain adaptation can help when a few annotated test-domain samples are available.Source: §1 · p. 1
  • Prior domain-generalization work learns domain-invariant features [14], adjusts model parameters through training policies [46] [27], or modifies architectures [25] [36].Source: §1 · p. 1
  • The proposed Batch Normalization Embedding maps each training domain to a latent-space point using domain-dependent batch-normalization statistics, while a lightweight ensemble shares parameters except for those statistics.The approach explicitly leverages domain-specific representations instead of aiming to discard domain-specific attributes.Source: §1 · pp. 1–2
  • At test time, instance-normalization statistics map an unseen sample into the domain space, and reciprocal distances to known-domain statistics weight the ensemble’s domain-specific predictions.The same model combination is used during training on known-domain samples so the ensemble learns a meaningful latent space and linearly combinable logits.Source: §1 · p. 2
  • The method applies to modern convolutional neural networks with batch-normalization layers and scales gracefully with the number of training domains.Source: §1 · p. 2

2. Related Work

The related work surveys domain-generalization strategies based on invariant feature learning, data augmentation, model design, and meta-learning, then positions domain-dependent batch normalization against earlier alignment methods that require target-domain information.

Feature-level

  • Feature-level methods learn invariant representations through alignment, shared/private features, embeddings, episodic training, or latent-domain discovery, but their invariance is limited by training-domain variety.Examples include multi-task autoencoders [14], adversarial feature alignment, Deep Separation Networks [4], Siamese embeddings, episodic training [27], and latent-domain clustering.Source: §2 · p. 2

Data-level

  • Data-level approaches augment training samples with domain-guided perturbations, adversarial examples, or randomized renderings, while the authors argue that discarding domain-specific information may hurt performance.The cited approaches include perturbations [46], adversarial examples, and domain randomization [33].Source: §2 · p. 2

Model-based

  • Model-based approaches use specialized architectures or domain-specific classifiers, but most require architectural changes that increase network parameters or complexity.Examples include low-rank CNNs [25], structured low-rank alignment, classifier fusion based on source-domain probabilities [36], and jigsaw-puzzle training [6].Source: §2 · p. 2

Meta-learning

  • Meta-learning methods simulate unseen domains by splitting available training samples into meta-training and meta-test sets, so their robustness is bounded by training-sample variety.The surveyed methods extend model-agnostic meta-learning [26] [13], introduce cross-domain regularization, discard dominant features [19], or use gradient-based meta-training [11].Source: §2 · pp. 2–3

Batch Normalization for distribution alignment

  • Earlier batch-normalization alignment methods use target-domain samples or metadata during training, whereas the proposed approach requires no external target-domain information.Separate batch-normalization statistics were first used for domain adaptation [7] [8], then extended to multi-domain settings [39] and graph-based alignment [38]; domain-specific normalization was also proposed for domain generalization [45].Source: §2 · p. 3

3. Method

The method uses domain-specific batch-normalization statistics to embed known domains and test samples in a shared latent space, then weights domain-specific predictions according to sample-to-domain similarity.

3.1. Problem Formulation

  • The model learns a lightweight ensemble for source domains and estimates an unavailable target distribution as a linear mixture of source distributions, weighted by sample-to-domain similarity.Target samples are unavailable during training, and each test sample may come from a different unseen domain.Source: §3.1 · p. 3, §3.3 · p. 4

3.2. Multi-Source Domain Alignment Layer

  • Domain-specific batch-normalization statistics normalize each domain’s activations to a common reference distribution, while training collects and applies domain-specific batch statistics to estimate population statistics.At inference, an unknown-domain sample can use instance statistics, the batch-statistic case with batch size one; instance and batch statistics estimate the same underlying distribution over different sample counts.Source: §3.2 · pp. 3–4

3.3. Domain Localization in the Batchnorm Latent Space

  • Each batch-normalization layer contributes population-statistic domain centroids and instance-statistic sample embeddings, which are stacked across layers into batch normalization embeddings (BNEs).Population statistics represent domains; a test sample is projected into the same space by computing its instance statistics during a forward pass.Source: §3.3 · pp. 4–5
  • The method uses Wasserstein distances between test-sample embeddings and domain centroids to derive reciprocal-distance weights and linearly combine domain-specific predictions.A sample from a source domain receives a high weight for its corresponding model; a sample from an unseen domain receives a weighted combination of the models.Source: §3.3 · p. 5

3.4. Training Policy

  • Training replicates distance-weighted prediction on known-domain samples using equal-sized domain batches, updating population statistics, computing instance-based similarities, and weighting predictions across domain branches.The procedure is intended to encourage a well-defined latent space; domain-specific normalization statistics are warmed up on the whole dataset before distance training to avoid convergence problems.Source: §3.4 · pp. 5–6

4. Experiments

The experiments evaluate BNE on PACS, Office-31, and Office-Caltech against DeepAll and domain-generalization methods, using leave-one-domain-out protocols and, for Office-Caltech, leave-two-domain-out testing.

4.1. Experimental Settings

  • Experiments use AlexNet and ResNet-18, compare against DeepAll and domain-generalization methods, and follow leave-one-domain-out testing, plus leave-two-domain-out testing for Office-Caltech.For published-result comparisons, Office-31 and Office-Caltech use ImageNet-initialized AlexNet; PACS uses ResNet-18. AlexNet is adapted with batch normalization on each convolutional layer.Source: §4.1 · p. 6

4.2. Domain Generalization for Classification

  • BNE reports 83.1% Avg. and +5.86 ∆% on PACS, best absolute accuracy in 2 of 3 Office-31 scenarios, and best average accuracy and baseline gain on Office-Caltech.On Office-31 Amazon, BNE improves absolute accuracy by +11.2% over baseline and +2.3% over the closest competitor. On PACS Sketch, it delivers a +9.6% absolute gain over its baseline.Source: §4.2 · p. 6, Table 1 · p. 7, Table 2 · p. 7

4.3. Ablation Study

  • The PACS ablation tests per-domain batchnorm statistics with inference-time projection and weighting, adds training-time distance weighting, then adds warm-up; BNE outperforms DNet with fewer parameters.DNet learns domain membership with a domain-classification network, unlike BNE’s distance between latent batch-normalization embeddings.Source: §4.3 · p. 7

4.4. Choosing a Distance Metric

  • Wasserstein distance performs best on average and on each left-out domain; sample-wise distance attribution beats uniform prediction averaging, so effectiveness derives from attribution rather than normalization choice.The evaluated alternatives are uniform distance, equivalent to averaging domain-specific branch predictions, and Bhattacharya distance; uniform is similar to [45], which uses learned instance/batch normalization mixtures.Source: §4.4 · pp. 7–8, Table 5 · p. 8

5. Conclusions

The method uses domain-specific batch-normalization statistics to represent unknown domains as combinations of known ones. It outperforms alternatives on domain-generalization benchmarks and could be extended to domain adaptation.

5. Conclusions

  • Domain-specific normalization layers disentangle training-domain representations, while batch-normalization statistics locate unseen samples as combinations of known domains.The method navigates a latent space of batch-normalization statistics to describe unknown domains.Source: §5 · p. 8
  • The method outperforms alternatives on several domain-generalization benchmarks ([25]), supporting specific domain representations over forced invariant representations.Source: §5 · p. 8
  • The formulation could extend to domain adaptation by injecting unlabeled target samples during training, and multiple available target samples could improve unseen-domain statistics estimates.The authors identify these extensions as directions for future work.Source: §5 · p. 8
  • The work highlights potentially interesting properties of batch-normalization layers that have not yet been extensively explored.Source: §5 · p. 8

6. Supplementary Material

The supplementary material extends the main paper’s experiments and documents BNE’s training policy, additional evaluations, design analyses, and comparisons. It also evaluates the latent-space proposal and compares BNE with DNet.

Supplementary Material

  • The supplement compares BNE’s performance with DNet, a variant that uses a domain discovery network.Source: §6 · p. 11
  • Additional results evaluate ResNet-18 on Office-31 and Office-Caltech, and AlexNet on PACS.Source: §6 · p. 11
  • Qualitative analyses examine the choices of batch size and distance measure, and compare BNE with other normalization strategies.Source: §6 · p. 11
  • The supplement provides an algorithmic overview of BNE’s training policy and additional training details.Source: §6 · p. 11
  • A quantitative analysis evaluates the proposed latent space.Source: §6 · p. 11

6.1. Training Policy

The training policy combines predictions from domain-specific branches using sample-derived domain similarities, then optimizes the final prediction to encourage a batch-normalization latent space.

6.1. Training Policy

  • A training batch T contains K domain batches, each with n samples from its corresponding domain.Source: §6.1 · p. 11
  • Each training step first updates domain embeddings from domain batches, then uses instance-normalized sample statistics to compute target embeddings and domain similarities.Source: §6.1 · p. 11
  • Each sample is run through K domain-specific branches, whose predictions are weighted by domain similarities; cross-entropy on the final prediction updates model weights and encourages a batch-normalization latent space.Source: §6.1 · p. 11

6.2. Training Settings

Both networks are initialized with ImageNet-pretrained weights and trained with architecture-specific optimization settings, batch sizes, and input augmentation.

6.2. Training Settings

  • AlexNet and ResNet-18 [17] are initialized with ImageNet-pretrained weights, and their final fully connected layers are fine-tuned for 20 epochs.The weights are publicly available Caffe weights, carefully converted for use in the implementation.Source: §6.2 · p. 11
  • AlexNet uses SGD with momentum 0.95, weight decay 5 × 10−5, and an initial learning rate of 10−3 exponentially decayed at rate 0.95.Source: §6.2 · p. 11
  • ResNet-18 uses Adam [22] with weight decay 10−6 and an initial learning rate of 10−4.Source: §6.2 · p. 11
  • Training backpropagates through batch-normalization mean and standard-deviation computation, and uses ImageNet-normalized images with random crops and horizontal flips.Images are resized to 256 before cropping to 224×224 for ResNet-18 or 227×227 for AlexNet; batches contain 16 samples per domain for ResNet-18 and 6 for AlexNet.Source: §6.2 · p. 11
  • All models are implemented in TensorFlow 2.0 ([1]) using carefully converted, publicly available Caffe weights pretrained on ImageNet.Source: §6.2 · p. 11

6.3. Additional Results

Additional experiments extend comparisons to AlexNet on PACS and ResNet-18 on Office-31, and report Office-Caltech results with both architectures.

6.3.1 PACS

  • +6.33% relative performance gain exceeds the second-best +4.88% on PACS with AlexNet, while the unseen Sketch domain gains +13% absolute accuracy over the baseline.The method achieves absolute performance comparable to the state of the art despite starting from a weaker baseline.Source: §6.3.1 · p. 11, Table 6 · p. 12

6.3.2 Office-31

6.3.3 Office-Caltech

  • +5.5% gain over DeepAll is reported for Office-Caltech with ResNet-18, confirming the good property also observed with AlexNet.Source: §6.3.3 · p. 12, Table 8 · p. 13

6.4. Ablation Study

The PACS ablations examine batch size, method components, and normalization strategies. They show that larger batches improve performance and that BNE leads the compared normalization strategies.

6.4.1 Batch Size

  • +2.9 Average accuracy separates batch size 64 from 16; larger training batches improve domain-embedding statistics and final performance.The study attributes the improvement to better approximation of true population statistics.Source: §6.4.1 · p. 13

6.4.2 Method Components

  • On PACS with AlexNet, each added component raises average accuracy over DeepAll across the four target sets.The ablation adds distance weighting at training time and then a warm-up phase to stabilize population statistics, building on separate domain batch-normalization statistics and projection-weighting.Source: §6.4.2 · p. 13

6.4.3 Normalization Strategies

  • +3.5% overall: BNE outperforms Freeze BatchNorm and also achieves the best results on every specific domain among the compared normalization strategies.Freeze BatchNorm generalizes better than fine-tuning population statistics, while InstanceNorm performs worst; BNE combines instance- and batch-normalization properties.Source: §6.4.3 · p. 13

6.5. Latent Space Validation

The batchnorm latent space encodes domain membership, and distance to a domain embedding correlates with the accuracy of that domain’s prediction branch.

6.5. Latent Space Validation

  • The validation trained ResNet-18 on PACS with Photo or Sketch treated as unseen domains, then projected training samples’ instance statistics into the batchnorm latent space.Domain membership was assessed by whether the nearest embedding matched the sample’s real domain.Source: §6.5 · p. 13
  • Higher diagonal reciprocal-distance similarities in Table 12 show that the batchnorm latent space encodes domain membership for PACS training samples.The analysis used ResNet-18 without distance training or warm-up, projecting instance statistics and measuring reciprocal distance to the centroids of the three training domains.Source: §6.5 · p. 13, Table 12 · p. 14
  • Across all four leave-one-domain-out test settings, the closest domain branch achieved higher average accuracy than the second-closest or farthest branch.This establishes a clear correlation between distance and prediction accuracy.Source: §6.5 · pp. 13–14

6.6. Domain Discovery Net

The section compares BNE with a domain discovery network (DNet), finding that domain-logit supervision improves DNet’s domain classification but reduces image classification accuracy, while BNE provides a domain representation without that loss.

6.6. Domain Discovery Net

  • Following, DNet uses first-convolution-block activations in a lateral pooling–ReLU–fully connected–softmax branch to predict training-domain probabilities that weight ensemble predictions.It is trained end-to-end with the main classifier, either using image-classification cross-entropy alone or additionally supervising domain labels.Source: §6.6 · pp. 14–15
  • 33% Avg. Domain for DNet without domain-logit cross-entropy reflects one-domain predictions; adding that loss improves domain accuracy but markedly lowers image accuracy, whereas BNE achieves higher image accuracy without it.The authors suggest DNet may overfit to training data and generalize less well to test data; BNE’s nonparametric representation shows no visible Avg. Domain increase from applying cross-entropy to its assigned domain memberships.Source: Table 14 · p. 15, §6.6 · p. 15
  • BNE estimates domain membership from activations throughout the network, whereas DNet relies on first-layer activations and a lateral branch; BNE also maps samples into a latent space for distances from domain embeddings.DNet can only output the distance of an input sample from the training domains, while BNE’s domain representation is parameter-free.Source: §6.6 · p. 15

Related papers

Loading 2011.12672v3…