PaperlayerView exact revision on arXiv

Source-linked AI summary

Domain Generalization by Mutual-Information Regularization with Pre-trained Models

Junbum Cha, Kyungjae Lee, Sungrae Park, Sanghyuk Chun

arXiv:2203.10789v2cs.LGcs.CV

TL;DR

Domain generalization struggles to learn representations robust to large shifts from limited source domains. MIRO aligns target features with an oracle approximated by a pre-trained model through a tractable mutual-information regularizer, and experiments report stronger out-of-domain performance, especially with larger backbones.

  • Problem

    Domain generalization must handle unseen target domains despite significant distribution shifts that can make source-domain invariant representations insufficient.

  • Method

    MIRO maximizes mutual information between target representations and oracle representations, approximating the inaccessible oracle with a large pre-trained model and a variational lower bound.

  • Results

    MIRO outperforms evaluated domain-generalization methods across reported settings, improving over ERM by +1.7pp on average and by +6.1pp with a larger SWAG-pre-trained backbone.

  • Takeaways & Limitations

    The reported results support using oracle-guided mutual-information regularization to exploit pre-trained representations for domain generalization, with larger pre-trained models yielding larger gains.

  • Takeaways & Limitations

    MI alignment with the oracle has no theoretical guarantee of perfect alignment with domain-generalization performance, and the method assumes oracle knowledge transfers through mutual information.

Abstract

from arXiv · show

Domain generalization (DG) aims to learn a generalized model to an unseen target domain using only limited source domains. Previous attempts to DG fail to learn domain-invariant representations only from the source domains due to the significant domain shifts between training and test domains. Instead, we re-formulate the DG objective using mutual information with the oracle model, a model generalized to any possible domain. We derive a tractable variational lower bound via approximating the oracle model by a pre-trained model, called Mutual Information Regularization with Oracle (MIRO). Our extensive experiments show that MIRO significantly improves the out-of-distribution performance. Furthermore, our scaling experiments show that the larger the scale of the pre-trained model, the greater the performance improvement of MIRO. Source code is available at https://github.com/kakaobrain/miro.

1 Introduction

Domain generalization seeks representations that remain robust under unseen distribution shifts, but source-domain invariance can fail when training and test domains differ substantially. MIRO instead aligns target representations with an oracle approximation derived from a large pre-trained model.

  • Motivation: Domain generalization learns from multiple source domains but must remain robust on an unseen target domain under distribution shifts.Shifts can involve backgrounds, geography, demographics, textures, and day-to-night conditions.
  • Approach: MIRO reformulates the objective by maximizing mutual information between oracle and target representations while preserving source-domain training loss.The oracle is defined as an optimal model generalized to any possible domain.
  • Approach: MIRO approximates the inaccessible oracle with a large pre-trained model and derives a tractable variational lower bound.The resulting method combines an ERM objective with mutual-information regularization.
  • Results: MIRO improves robustness during fine-tuning, whereas naive fine-tuning of larger pre-trained models can fail to improve performance.ERM with ImageNet pre-trained ResNet achieves 64.2% averaged accuracy, compared with 61.1% for CLIP-pre-trained ViT.
  • Results: MIRO outperforms all evaluated domain-generalization methods in all settings and benefits from larger pre-trained models.The analysis includes varying optimizers and pre-trained models, including SWAG and CLIP.

2 Related works

Related work primarily learns domain-invariant features from source domains or exploits pre-trained models for transfer, distillation, continual learning, and out-of-distribution generalization.

  • Domain generalization: Traditional domain-generalization methods reduce source-domain discrepancies, simulate shifts, use robust optimization, or augment source examples.These approaches seek representations that preserve information shared across source domains.
  • Pre-trained models: Pre-trained models support transfer learning, knowledge distillation, continual learning, and emerging approaches to out-of-distribution generalization.Recent work also studies how naive fine-tuning affects pre-trained features.

3 Methods

MIRO formulates domain generalization as mutual-information alignment with an oracle representation while constraining source-domain performance. It approximates the oracle with a pre-trained model, derives a variational bound, and optimizes the resulting regularized objective.

  • 3.1 Mutual information regularization with oracle: The standard DG formulation minimizes empirical loss averaged over source domains, but can fail when target distributions differ substantially.CORAL reaches about 50% top-1 accuracy on several DomainNet domains but only 13.4% on QuickDraw and 19.7% on Infographics.
  • 3.1 Mutual information regularization with oracle: MIRO introduces an oracle feature extractor and maximizes mutual information between oracle and target features subject to a source-loss constraint.The oracle is assumed to generalize to any possible domain, while the target model learns from oracle representations of source samples.
  • 3.1 Mutual information regularization with oracle: The inaccessible oracle is approximated with a large pre-trained feature extractor, including ImageNet, CLIP, and SWAG models.The approximation assumes the pre-trained extractor lies near the oracle in feature-extractor hypothesis space.
  • 3.1 Mutual information regularization with oracle: The variational lower bound replaces oracle features with pre-trained features and includes an approximation penalty proportional to d2,∞(f ∗, f 0).A smaller distance between oracle and pre-trained extractors makes the approximated lower bound tighter.
  • 3.1 Mutual information regularization with oracle: The final objective combines empirical risk with a Gaussian variational mutual-information regularizer optimized by stochastic gradient methods.The Gaussian uses mean encoder µ(Zf) and covariance encoder Σ(Zf), with λ as the regularization coefficient.
  • 3.2 Mutual information analysis with the oracle model: MIRO achieves the highest empirical mutual information with the oracle across both ImageNet-ResNet and SWAG-RegNet settings.Larger pre-trained backbones show higher mutual information, while fine-tuning can substantially reduce it for powerful models.

4 Experiments

Across five domain generalization benchmarks, MIRO improves out-of-domain performance over ERM and prior methods, with larger gains for larger pre-trained models. Additional analyses examine pre-training choices, feature variances, and robustness under substantial distribution shifts.

  • Benchmark evaluation: MIRO improves performance over ERM on all five benchmarks, with a +1.7pp average gain and best results on all but PACS.It also surpasses prior methods by +1.3pp on OfficeHome and +1.8pp on TerraIncognita.
  • Benchmark evaluation: Combining MIRO with SWAD achieves the best performance on every dataset, improving over previous best results by +0.8pp on average.SWAD is treated as orthogonal to MIRO because it seeks flat minima.
  • Scaling pre-trained models: Using the SWAG-pre-trained RegNetY-16GF increases MIRO’s average improvement over ERM from +1.7pp to +6.1pp, reaching 77.3% with SWAD.The authors associate the larger gain with preserving mutual information during fine-tuning of large-scale models.
  • Pre-training robustness: MIRO improves ERM across supervised, self-supervised, image-language, and weakly supervised pre-training settings, with gains up to +16.3pp for CLIP.Reported gains include +1.7pp for ImageNet ERM, +3.5pp for Barlow Twins, +1.3pp for MoCo v3, +12.6pp for CLIP-ViT, and +6.1pp for SWAG.
  • Pre-training robustness: MIRO leverages large-scale pre-trained knowledge: CLIP-ViT improves from 61.1% with naive fine-tuning to 73.7%, while ImageNet fine-tuning reaches 64.2%.With MIRO, the ImageNet pre-trained model reaches 65.9%.
  • Feature analysis: Figure 2 shows layerwise variances increasing toward the output, indicating stronger preservation of pre-trained information in lower layers and more task-specific learning in higher layers.The variance term controls the distance loss between pre-trained and current features.
  • Distribution-shift analysis: On Camelyon17, MIRO consistently outperforms ERM despite the large visual and task gap between ImageNet pre-training and medical fine-tuning.Camelyon17 contains histological images from five hospitals and requires breast-cancer tumor classification.
  • Distribution-shift analysis: The effect of λ depends on the pre-trained model: ImageNet shows negative correlation with performance difference, whereas SWAG and CLIP-ViT improve with larger λ.Large-scale models also show less sensitivity to λ across target domains.

5 Conclusion

MIRO reframes domain generalization around an oracle model approximated by pre-trained backbones, achieving state-of-the-art DomainBed performance and larger gains with CLIP or SWAG.

  • MIRO approximates an oracle model with a pre-trained model to derive a tractable variational bound for robust domain generalization.
  • MIRO achieves state-of-the-art performance on DomainBed benchmarks.
  • Performance improvements increase when MIRO is combined with larger-scale pre-trained backbones such as CLIP or SWAG.
  • The authors position MIRO as a research direction for exploiting pre-trained backbones to learn robust domain-generalized representations.

A Derivation of Lower Bound

The section derives a lower bound on the conditional log-density term under a regularity assumption, using Taylor expansion and a distance between functions. It also verifies that the Gaussian distribution satisfies the required regularity condition.

  • The derivation assumes that the variational distribution q(·|z) satisfies a regularity condition.
  • For Gaussian q, finite μ(z) and Σ(z), together with bounded E_{X|z}[|X|^2], establish the regularity condition.
  • Taylor’s theorem expands log q(x|y) around x_0 and introduces an intermediate point c between x and x_0.
  • The gradient bound converts the expansion into a Euclidean-norm product involving ∥∇_x log q(x|y)|_{x=c}∥_2 and ∥x−x_0∥_2.
  • The resulting lower bound subtracts C d_{2,∞}(f*,f^0) from E_{Z_{f^0},Z_f}[log q(Z_{f^0}|Z_f)].

B Additional Implementation Details

The implementation separates hyperparameters into algorithm-specific and algorithm-agnostic groups, tuning them independently to reduce search cost. Despite this protocol, some experiments remain computationally intensive, and hardware limits affect one baseline’s batch size.

  • MIRO has one algorithm-specific hyperparameter, λ, while batch size, learning rate, dropout, and weight decay are algorithm-agnostic.
  • The search first tunes algorithm-specific hyperparameters with default algorithm-agnostic settings, then tunes the algorithm-agnostic hyperparameters.
  • λ is tuned only for non-main experiments because the search remains computationally expensive despite the efficient protocol.
  • SWAG uses batch size 16 because of GPU memory limitations, while more intensive tuning and GPU memory could further improve performance.

B.2 Implementation details

The implementation specifies initialization and optimization choices for the variance encoder and estimates mutual information with MINE. Features are collected across domains, and fine-tuned models receive averaged mutual-information estimates.

  • The variance encoder starts with an estimated variance of 0.1, uses Softplus for non-negativity, and receives a 10× larger learning rate than other model components.
  • MINE estimates mutual information between the oracle-model features Z_{f*} and model features Z_f using a statistics network.
  • The statistics network has two 512-dimensional hidden linear layers with ELU activations, following the cited MINE setup.
  • After global average pooling, Z_{f*} and Z_f features are uniformly collected across domains.
  • For fine-tuned ERM−, ERM+, and MIRO models, mutual information is estimated separately per target-domain model and averaged.

C Additional Analysis and Discussion

MIRO can be adapted to alternative domain-generalization assumptions, including class-conditional distribution shifts, by selecting an appropriate mutual-information strategy.

  • The broader DG literature varies assumptions about source domains, domain labels, and label sets.
  • C-MIRO uses class-conditional mutual information, I(Zf ∗; Zf|Y ), instead of unconditional mutual information.
  • C-MIRO achieves comparable scores with MIRO and outperforms ERM even when the problem setting is not class-conditional.

C.2 The relationship between mutual information and domain generalization performance

The analysis compares mutual-information rankings with domain-generalization accuracies for ERM variants and MIRO, finding matching orderings in reported settings.

  • Table 6 reports PACS accuracies of 51.6% for ERM−, 84.2% for ERM+, and 85.4% for MIRO with an ImageNet-pre-trained ResNet.
  • Table 6 reports PACS accuracies of 51.5% for ERM−, 89.6% for ERM+, and 97.4% for MIRO with an Instagram-3.6B-pre-trained RegNet.
  • The MI rankings of ERM−, ERM+, and MIRO match their accuracy rankings in the reported experiments.
  • The method assumes that high mutual information with the oracle model transfers beneficial knowledge to the target model, but provides no theoretical guarantee of perfect alignment with DG performance.

D Additional Results

The additional results examine MIRO across a substantial pre-training-to-fine-tuning shift between ImageNet and Camelyon17, whose visual domains and tasks differ markedly.

  • ImageNet and Camelyon17 exhibit a huge visual gap between pre-training and fine-tuning datasets.
  • MIRO shows consistent performance improvement despite the large gap between the pre-training and fine-tuning distributions.
  • ImageNet is a multiclass object-recognition task, whereas Camelyon17 is binary classification for detecting tumor tissue.

D.2 Relationship between the pre-training scale and the intensity of the mutual information regularization

The analysis studies how pre-trained backbone scale interacts with the strength of MI regularization across additional domain-generalization benchmarks.

  • The comparisons across OfficeHome, TerraIncognita, and DomainNet show trends similar to those observed in PACS.
  • ImageNet-pre-trained backbones show a negative correlation between MIRO's performance difference and λ in some target domains.
  • Figure 5 compares ResNet-50, SWAG-pre-trained RegNetY-16GF, and CLIP-pre-trained ViT-B according to λ.
  • The figure evaluates how mutual-information regularization intensity relates to MIRO's performance difference over ERM.
Loading 2203.10789v2…