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

Domain generalization seeks robustness to unseen domains despite the difficulty of obtaining target-domain samples. The paper uses domain-specific batch-normalization statistics to embed known and unknown domains, then combines domain-specific classifiers according to latent-space distances. It reports improved performance across domain-generalization benchmarks, with a reported +6.33% relative gain on PACS using AlexNet.

  • Problem

    Models trained on one data distribution often fail on different distributions, while collecting target-domain samples for adaptation is not always feasible.

  • Method

    The method collects domain-specific batch-normalization statistics, maps domains into a shared latent space, and combines domain-specific models using distances to known domains.

  • Results

    The method outperforms many alternatives on domain-generalization benchmarks, including a +6.33% relative gain on PACS with AlexNet.

  • Takeaways & Limitations

    Maintaining domain-specific representations can outperform forcing invariant representations while enabling unknown samples to be localized relative to known domains.

  • Takeaways & Limitations

    The domain-generalization setting assumes target samples are unavailable during training, and each target sample may belong to a different unseen domain.

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

Domain generalization addresses domain shift by learning models that remain robust on unseen domains, without requiring target-domain samples. This work instead exploits domain-specific batch-normalization statistics to embed known and unknown domains and combine domain-specific classifiers.

  • Target-domain fine-tuning or adaptation can mitigate domain shift, but collecting samples for every possible environment is often infeasible.
  • Domain generalization seeks models that remain robust when test samples come from unseen domains.
  • Unknown-domain properties guide classification through a linear combination of domain-specific classifiers weighted by inverse distances to known domains.
  • The method reverses domain-invariance strategies by collecting domain-dependent batch-normalization statistics for each training domain.
  • Batch Normalization Embeddings map known domains into a latent domain space and project unknown test samples into it using instance-normalization statistics.
  • The lightweight ensemble shares model parameters while retaining separate normalization statistics, and the approach applies to CNNs that use batch normalization.

2. Related Work

Related domain-generalization methods address domain shift through feature-level, data-level, model-based, and meta-learning strategies. These approaches reduce discrepancy, diversify training data, modify architectures, or alter training policies to improve robustness across domains.

  • Across these approaches, domain generalization methods seek robustness to domain shift using feature objectives, data variation, architectural changes, or training procedures.
  • Feature-level: Feature-level methods learn domain-invariant representations by minimizing discrepancies between multiple training domains.
  • Data-level: Data-level methods reduce training-set domain bias by increasing the cardinality and variety of samples through guided perturbations, adversarial examples, or domain randomization.
  • Model-based: Model-based methods use specialized architectures, including low-rank parameterized networks and aligned domain-specific and domain-invariant models.
  • Meta-learning: Meta-learning methods use specialized training policies that expose optimization to domain shift, often by splitting training data into meta-training and meta-test sets.

3. Method

The method uses domain-specific batch-normalization statistics to embed source domains and unseen samples in a shared latent space. Distances between these embeddings weight a lightweight ensemble that estimates unknown-domain predictions as mixtures of source-domain models.

  • Domain alignment: Batch-normalization layers collect domain-specific batch statistics during training and maintain moving averages that approximate domain population statistics for inference.At inference, unknown-domain samples use instance statistics, equivalent to batch statistics with batch size one.
  • Domain-specific ensemble: Unlike domain-invariant approaches, it retains domain-specific batch-normalization statistics and shares all other parameters across a lightweight model ensemble.Each ensemble member corresponds to a source domain through its normalization statistics.
  • Domain embeddings: The method maps known domains and unseen samples into a shared latent space using domain population statistics and sample instance statistics.Population statistics represent source-domain centroids, while instance statistics project target samples into the same space.
  • Domain localization: The method measures sample-to-domain similarity as the reciprocal of the distance between their activation-statistics embeddings.Distances are computed across batch-normalization layers using a Wasserstein distance between Gaussian representations of means and variances.
  • Prediction: Similarity weights combine learned source distributions and domain-specific predictions to estimate the unknown target distribution and final output.The final lightweight-ensemble prediction is a linear combination of domain-dependent models weighted by sample-specific domain similarities.

4. Experiments

Experiments evaluate BNE for domain generalization across PACS, Office-31, and Office-Caltech, using leave-one-domain-out classification and comparisons with prior methods and variants. Results also examine model variants and distance functions, with Wasserstein distance supporting the strongest performance.

  • Experimental Settings: BNE is evaluated on PACS, Office-31, and Office-Caltech using standard leave-one-domain-out domain generalization protocols with AlexNet and ResNet-18.PACS contains four domains, while Office-31 and Office-Caltech provide additional cross-domain classification benchmarks.
  • Domain Generalization for Classification: 83.1% average accuracy and a +5.86% relative gain make BNE the second most effective algorithm on PACS with ResNet-18.BNE achieves the best absolute accuracy on one of four target domains and its largest gain on Sketch is +9.6%.
  • Ablation Studies: Ablation experiments compare BNE with DeepAll and DNet to measure the effects of its components and domain-membership strategy.DNet learns domain membership with a domain-classification network, whereas BNE uses distances between latent batch-normalization embeddings.
  • Choosing a Distance Metric: Wasserstein distance consistently delivers the best average and left-out-domain performance among the evaluated distance choices.Distance-based sample-wise domain attribution is more effective than uniformly averaging predictions, while requiring fewer parameters than DNet.

5. Conclusions

The paper concludes that domain-specific normalization statistics can represent training domains in a latent space and localize unseen samples as combinations of known domains. Across several benchmarks, this approach outperforms many alternatives while retaining a lightweight ensemble structure.

  • Conclusions: BNE maps unknown domains as combinations of known domains in a latent space of batch-normalization statistics.The method uses domain-specific normalization layers to disentangle training-domain representations and localize unseen samples.
  • Conclusions: BNE outperforms many alternatives on several domain generalization benchmarks, supporting domain-specific representations instead of forcing invariant representations.The authors identify further exploration of batch-normalization properties and extension to domain adaptation as future directions.

6. Supplementary Material

The supplementary material adds algorithmic and training details, expands evaluations across architectures and datasets, and provides qualitative analyses of batch sizes, distances, and normalization strategies.

  • Supplementary Material: Supplementary sections provide an algorithmic overview, additional training details, and experiments on Office-31, Office-Caltech, and PACS.The additional experiments use ResNet-18 and AlexNet configurations.
  • Supplementary Material: The supplementary material qualitatively analyzes batch sizes and distance measures and validates BNE against other normalization strategies.

6.1. Training Policy

Training alternates between updating domain embeddings from domain batches and computing instance-based similarities for individual samples. Domain-specific predictions are then combined into a final prediction whose cross-entropy loss updates the shared model parameters.

  • Training Policy: Each training batch contains K domain batches, with n samples from each corresponding domain.
  • Training Policy: Domain batches first update their corresponding embeddings, after which individual samples provide instance statistics for target embeddings and domain similarities.
  • Training Policy: Each sample is propagated through K domain-specific branches, and their predictions are weighted according to the computed domain similarities.
  • Training Policy: The weighted final prediction is trained with cross-entropy loss, which is back-propagated to update the shared model weights.

6.2. Training Settings

The experiments use ImageNet-pretrained AlexNet and ResNet-18 models, fine-tuning their final fully connected layers before domain-generalization training. Optimization, regularization, normalization, and augmentation settings differ by architecture.

  • Both AlexNet and ResNet-18 are initialized with ImageNet-pretrained weights, and their last fully connected layers are fine-tuned for 20 epochs.
  • AlexNet uses SGD with momentum 0.95, weight decay 5 × 10−5, and an initial learning rate of 10−3 with exponential decay.
  • ResNet-18 uses Adam with weight decay 10−6 and an initial learning rate of 10−4.
  • Gradients are computed through batch-normalization mean and standard-deviation operations, and inputs are normalized using ImageNet statistics.
  • Training augmentation resizes images to 256 pixels before random cropping to architecture-specific input sizes.
  • The models are implemented in TensorFlow 2.0 and initialized from converted publicly available Caffe ImageNet weights.

6.3. Additional Results

Additional experiments extend evaluation to PACS, Office-31, and Office-Caltech with AlexNet and ResNet-18. Across these comparisons, BNE improves over DeepAll and prior methods, including a +6.33% relative gain on PACS with AlexNet and a +5.5% gain on Office-Caltech with ResNet-18.

  • Office-Caltech: Additional Office-Caltech experiments evaluate both ResNet-18 and AlexNet architectures.
  • PACS: +6.33% relative performance gain makes the proposed method outperform prior solutions on PACS with AlexNet.The second-best method achieves +4.88%.
  • PACS: +13% absolute accuracy gain is obtained over the baseline when Sketch is the unseen PACS domain.
  • Office-31: The method improves performance over DeepAll across all three Office-31 tests with ResNet-18.
  • Office-Caltech: +5.5% gain over DeepAll is reported for Office-Caltech with ResNet-18.

6.4. Ablation Study

Ablation studies on PACS examine batch size, method components, and normalization strategies. Larger batches improve domain-statistics estimates, while BNE achieves the strongest normalization results.

  • Batch size: Larger training batch sizes improve generalization because they better approximate true population statistics and domain embeddings.
  • Method components: The ablations compare separate domain batch-normalization statistics and the projection-weighting strategy against the DeepAll baseline using AlexNet on PACS.
  • Normalization strategies: The normalization comparison evaluates InstanceNorm, BatchNorm, Freeze BatchNorm, and BNE on PACS with ResNet-18.
  • Normalization strategies: +3.5% overall performance over Freeze BatchNorm makes BNE the best-performing normalization strategy.BNE also outperforms the other normalization strategies on every specific domain.
  • Normalization strategies: Fine-tuning population statistics on training datasets can overfit, whereas freezing ImageNet statistics generalizes better than the DeepAll-equivalent setting.

6.5. Latent Space Validation

The latent-space validation tests whether instance statistics capture domain-specific attributes and whether distances identify useful domain branches. Samples assigned to closer domain embeddings yield higher prediction accuracy.

  • Latent-space projection: The validation projects training samples into the batch-normalization latent space using instance statistics after training ResNet-18 on PACS.
  • Experimental setting: The experiments use PACS with ResNet-18 without distance training and warm-up, considering Photo or Sketch as unseen domains.
  • Distance-based evaluation: Prediction accuracy is compared when using the closest, second-closest, or farthest domain branch under a leave-one-domain-out protocol.
  • Distance-based evaluation: The closest domain branch produces higher accuracy than the second-closest and farthest branches, showing a clear correlation between distance and accuracy.

6.6. Domain Discovery Net

The section compares BNE with a Domain Discovery Network (DNet) for domain and image classification across unseen-domain tests on PACS. BNE uses a parameter-free representation based on activations throughout the network, whereas DNet uses a lateral branch and first-block activations.

  • DNet takes first-convolution-block activations and predicts the probability that each sample belongs to a training domain, which weights domain-specific ensemble predictions.
  • The comparison reports domain classification and image classification accuracy across four PACS tests with different unseen domains.
  • Without cross-entropy on domain logits, BNE largely outperforms DNet, whose 33% average domain accuracy indicates it predicts the same domain class.
  • Adding cross-entropy on DNet domain logits improves domain prediction but causes a remarkable drop in image classification accuracy, partly attributed to overfitting.
  • BNE uses activations throughout the network and provides a parameter-free latent-space representation, while DNet relies on first-layer activations and a lateral branch.
Loading 2011.12672v3…