PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning to Learn with Variational Information Bottleneck for Domain Generalization

Yingjun Du, Jun Xu, Huan Xiong, Qiang Qiu, Xiantong Zhen, Cees G. M. Snoek, Ling Shao

arXiv:2007.07645v1cs.CV

TL;DR

Domain generalization must handle both unseen-domain prediction uncertainty and distributional domain shift. The paper proposes probabilistic meta-learning with MetaVIB, which learns domain-invariant representations through a meta-learning variational bottleneck. Across three cross-domain visual-recognition benchmarks, the method consistently achieves the best performance against regular-learning and meta-learning approaches.

  • Problem

    Domain generalization lacks access to unseen-domain distributions, creating prediction uncertainty, while differing domain distributions produce domain shift.

  • Method

    The method models shared classifier parameters as distributions and uses MetaVIB, derived from variational mutual-information bounds, to learn domain-invariant representations.

  • Results

    The method consistently delivers the best performance on three cross-domain visual-recognition benchmarks, surpassing regular-learning and meta-learning methods.

  • Takeaways & Limitations

    Episodic MetaVIB training narrows domain gaps while the probabilistic model handles prediction uncertainty on previously unseen domains.

Abstract

from arXiv · show

Domain generalization models learn to generalize to previously unseen domains, but suffer from prediction uncertainty and domain shift. In this paper, we address both problems. We introduce a probabilistic meta-learning model for domain generalization, in which classifier parameters shared across domains are modeled as distributions. This enables better handling of prediction uncertainty on unseen domains. To deal with domain shift, we learn domain-invariant representations by the proposed principle of meta variational information bottleneck, we call MetaVIB. MetaVIB is derived from novel variational bounds of mutual information, by leveraging the meta-learning setting of domain generalization. Through episodic training, MetaVIB learns to gradually narrow domain gaps to establish domain-invariant representations, while simultaneously maximizing prediction accuracy. We conduct experiments on three benchmarks for cross-domain visual recognition. Comprehensive ablation studies validate the benefits of MetaVIB for domain generalization. The comparison results demonstrate our method outperforms previous approaches consistently.

1 Introduction

The paper targets domain generalization by jointly addressing prediction uncertainty on unseen domains and domain shift. It combines probabilistic meta-learning with MetaVIB and reports consistent gains across cross-domain recognition benchmarks.

  • Unseen-domain uncertainty arises because target-domain distributions are inaccessible during training.
  • Distinct distributions across domains create domain shift, while existing meta-learning methods do not address unseen-domain prediction uncertainty.
  • The proposed probabilistic meta-learning model represents shared classifier parameters as distributions inferred from seen domains.
  • MetaVIB uses variational mutual-information bounds within meta-learning to learn domain-invariant representations while maximizing target prediction.
  • The method consistently delivers the best performance across three cross-domain visual-recognition benchmarks, surpassing regular-learning and meta-learning approaches.

2 Related Work

Related work centers on domain-invariant representation learning, information bottleneck objectives, and meta-learning for generalization. The paper combines these directions with probabilistic modeling to address uncertainty and domain shift in domain generalization.

  • Domain generalization commonly learns representations invariant across source domains to generalize to unseen target domains.
  • Information bottleneck compresses inputs while maximizing information about target labels by minimizing I(Z; X) and maximizing I(Z; Y).
  • Variational approximations make information-bottleneck objectives parameterizable with neural networks, avoiding infeasible iterative optimization.
  • Meta-learning acquires common knowledge across related tasks to adapt models efficiently to new tasks and has shown effectiveness for domain generalization.
  • The paper introduces probabilistic meta-learning and derives MetaVIB for domain-invariant representation learning using episodic meta-train and meta-test splits.

3 Method

The method combines episodic probabilistic meta-learning with a meta variational information bottleneck to model classifier uncertainty and learn domain-invariant representations. Its objective balances meta-test prediction accuracy with a KL-based alignment between target-domain encodings and meta-train class distributions.

  • Probabilistic Modeling: Source domains are episodically split into meta-train Ds and meta-test Dt sets, with the model optimized on prediction errors from Dt.The learned model is evaluated on held-out target domains during testing.
  • Probabilistic Modeling: Classifier parameters ψ are treated as stochastic variables inferred from meta-train data Ds and applied to the meta-test domain.The posterior q(ψ|Ds) replaces a classifier distribution conditioned directly on meta-test inputs.
  • Meta Variational Information Bottleneck: MetaVIB extends the information bottleneck principle to meta-learning by deriving new variational mutual-information bounds for domain-generalization representations.The method uses a meta prior conditioned on meta-train data rather than an uninformative prior.
  • Meta Variational Information Bottleneck: The objective combines meta-test negative log predictive likelihood with a KL term aligning target-domain encodings to distributions estimated from same-class meta-train samples.Minimizing the first term promotes prediction accuracy, while minimizing the KL term supports domain-invariant representations.
  • Learning with Stochastic Neural Networks: Training uses stochastic neural networks, class-wise pooled meta-train features, Monte Carlo sampling, and re-parameterization for end-to-end optimization.The feature extractor and inference networks parameterize the stochastic representations and classifier distributions.

4 Experiments

Experiments on VLCS, PACS, and Rotated MNIST evaluate MetaVIB through dataset benchmarks, ablations, representation visualizations, prediction cases, and state-of-the-art comparisons. The results support probabilistic modeling, MetaVIB regularization, domain-invariant representations, and consistent performance advantages across benchmarks.

  • Experimental Setup: Experiments use three domain-generalization benchmarks with leave-one-domain-out evaluation, AlexNet features, and average classification accuracy.The benchmarks are VLCS, PACS, and Rotated MNIST; source domains train the model and one target domain is held out for testing.
  • Ablation Study: Probabilistic modeling improves over pre-trained AlexNet by 1.35% on VLCS and 3.35% on PACS.The probabilistic baseline models classifier parameters as distributions and generalizes better to target domains than the pre-trained AlexNet baseline.
  • Ablation Study: 74.01% on VLCS and 73.37% on PACS show that regular VIB improves over the probabilistic baseline, while MetaVIB outperforms VIB by up to 2.37% on PACS.MetaVIB uses episodic training and a KL term to align meta-target representations with same-class meta-source representations across domains.
  • Ablation Study: β = 0.01 yields the best MetaVIB performance, combining minimal I(Z; X) with high I(Z; Y) and, in Fig. 2, the highest I(T; Y).The authors select β using validation sets and report that this value performs best across datasets.
  • Ablation Study: MetaVIB largely reduces feature discrepancies across domains while separating samples from different classes, unlike pre-trained AlexNet and regular VIB.The visualization compares feature representations across four PACS domains and seven classes, including unseen domains.
  • State-of-the-Art Comparison: MetaVIB handles uncertainty in difficult cases and consistently outperforms previous regular-learning and meta-learning methods across VLCS, PACS, and Rotated MNIST.It surpasses MASF on VLCS by 0.43% and can assign reasonable probabilities to competing labels even when its final prediction is incorrect.

5 Conclusion

The paper proposes MetaVIB, a probabilistic meta-learning model that addresses prediction uncertainty and domain shift in domain generalization. It models classifier parameters probabilistically and learns domain-invariant representations, achieving high performance across three benchmarks.

  • MetaVIB models shared classifier parameters as probabilistic distributions inferred from source domains to address prediction uncertainty on target domains.
  • MetaVIB learns domain-invariant representations through a Meta Variational Information Bottleneck principle derived from a variational mutual-information bound.
  • MetaVIB integrates meta-learning, variational inference, and probabilistic modeling within one domain-generalization framework.
  • Extensive experiments on three cross-domain visual-recognition benchmarks and ablations validate the proposed contributions.
  • MetaVIB consistently achieves high performance and advances the state of the art on all three benchmarks.

A Algorithms of MetaVIB for Training

The training algorithm initializes feature-extraction and inference networks, repeatedly samples source domains into meta-train and meta-test sets, and optimizes the resulting model episodically.

  • Algorithm 1 takes training data from K source domains, a learning rate λ, and an iteration count Niter as inputs.
  • The model initializes parameters Θ = {θ, φ1, φ2} for the feature-extraction network hθ(·) and inference networks gφ1(·), gφ2(·).
  • Each iteration randomly selects a meta-test domain Dt and uses the remaining source domains as meta-train domains Ds.
  • The algorithm samples training data and constructs classifier-parameter components ψc before assembling ψ = [ψ1, · · ·, ψc, · · ·, ψC].

B Learning Architecture

The architecture uses a feature-extraction network and two inference networks to produce classifier-parameter and latent-variable distributions, with cross-entropy loss guiding prediction.

  • The feature-extraction network hθ(·) extracts representations from the input data.
  • The architecture uses cross-entropy loss Lcls, with solid and dashed lines distinguishing meta-test and meta-train data flow.

C Training Details

Training uses episodic batches formed from randomly selected source domains, Adam optimization, validation-based model selection, and dataset-dependent bottleneck settings.

  • Each training batch randomly selects three domains: two meta-train domains and one meta-test domain.
  • Adam optimization uses a learning rate of 10−4, 256 samples per domain, batch size 256 × 3, and 25,000 iterations.
  • The model with the highest validation accuracy is used to evaluate the meta-test domain.
  • The optimal bottleneck size β is dataset-dependent: β = 0.01 for VLCS and β = 0.001 for Rotated MNIST.

E Influence of the number of Monte Carlo Influence of the number of Monte Carlo samples

MetaVIB’s PACS performance depends on the number of Monte Carlo samples used to draw z. Performance is best and converges at Lz = 10, while larger sample counts degrade results.

  • Performance is inferior with Lz = 1, improves consistently at Lz = 5 and 10, then worsens at Lz = 50 and 100.The experiments set Lz = 10 and average over 20 test-domain runs; variance reflects Monte Carlo sampling error.
  • Lz = 10 yields MetaVIB’s best PACS results.Table E.3 identifies Lz = 10 as the best setting.

F Network Architectures

The appendix specifies separate feature-extraction and inference-network architectures for the evaluated datasets. PACS and VLCS share architecture references, while Rotated MNIST uses separate configurations.

  • Inference networks: Inference network gφ1(·) has dedicated architecture specifications for PACS/VLCS and Rotated MNIST.These configurations are listed in Tables F.6 and F.7.
  • Inference networks: Inference network gφ2(·) likewise has separate architecture specifications for PACS/VLCS and Rotated MNIST.These configurations are listed in Tables F.8 and F.9.
  • Feature extraction networks: Feature extraction networks are specified separately for PACS/VLCS and Rotated MNIST.They are denoted hθ(·) in Tables F.4 and F.5, respectively.

G Prediction Uncertainty Analysis

The prediction-uncertainty analysis examines MetaVIB’s success and failure cases using multiple sampled classifiers. These classifiers produce different category probabilities, and some samples can remain correct even when the final prediction is wrong.

  • Sampled classifier behavior: Different sampled classifiers ψ produce different prediction probabilities for each image category.ψ µ denotes the classifier’s mean value, while Figures G.3–G.10 show probabilities across sampled classifiers.
  • Success cases: For the fourth success case, classifiers ψ 1 and ψ 2 assign higher probability to horse than giraffe although the final classification is giraffe.The relevant sampled-classifier probabilities appear in Figure G.6.
  • Failure cases: For the fourth failure case, classifier ψ 4 assigns higher probability to elephant than dog, while some sampled classifiers still predict correctly.The final model classification is dog, and the sampled probabilities are shown in Figure G.10.
  • Case analysis: MetaVIB’s success and failure cases report the top two prediction probabilities alongside red ground-truth labels.Figure G.2 associates each image with its top two predicted probabilities.
Loading 2007.07645v1…