Domain Generalization under Conditional and Label Shifts via Variational Bayesian Inference

1 Introduction

Deep learning usually relies on the independent, identically distributed (i.i.d.) assumption of training and testing datasets, while target tasks are usually significantly heterogeneous and diverse [Che et al., 2021; Liu et al., 2021a]. This motivates many researchers to investigate domain adaptation (DA) [Liu et al., 2021b] and domain generalization (DG) [Matsuura and Harada, 2020a]. The source and target domain shifts are expected to be compensated for by a variety of adaptation steps. Despite the success of DA in several tasks, much of the prior work relies on utilizing the massive labeled/unlabeled target samples for its training.

The recently aroused DG task [Matsuura and Harada, 2020a] assumes that several labeled source domains are available in training without access to the target sample/label. Collectively exploiting these source domains can potentially lead to a trained system that can be generalized well on a target domain [Hu et al., 2019]. A predominant stream in DG is the domain invariant feature learning (IFL) [Ghifary et al., 2017], which attempts to enforce pi (f(x)) = pj (f(x)), where i and j index the two different source domains. The typical solution can be obtained via momentum or adversarial training [Ghifary et al., 2017].

However, the different classes can follow different domain shift protocols, e.g., the street lamp will be sparkly at night, while the pedestrian is shrouded in darkness. Therefore, we would expect the fine-grained class-wise alignment of the condition shift w.r.t. p(f(x)|y), where f(·) is a feature extractor [Zhang et al., 2013; Gong et al., 2016; Combes et al., 2020].

Assuming there is no concept shift (i.e., pi (y|f(x)) = pj (y|f(x))) and label shift (i.e., pi (y) = pj (y))), and given the Bayes' theorem, p(f(x)|y) = p(y|f(x))p(f(x)) p(y) , IFL is able to align p(f(x)|y) if pi (f(x)) = pj (f(x)).

However, the label shift pi (y) ̸= pj (y) [Liu et al., 2021b], i.e., class imbalance, is quite common in DG, as illustrated in Fig. 1. Since f(·) is a deterministic mapping function, IFL is able to encode the domain invariant representation under the covariate shift assumption (i.e., only pi (x) ̸= pj (x)) [Moreno-Torres et al., 2012]. Under the label shift, the covariate alignment cannot be used as an alternative of conditional alignment (i.e., pi (x|y) = pj (x|y)) [Li et al., 2018b]. Actually, both the conditional and label shifts are the realistic setting in most of DG tasks.

Recently, [Li et al., 2018b] proposes to align the conditional shift, assuming that there is no label shift. However, it is ill-posed to only consider one of conditional or label shift [Zhang et al., 2013; Kouw, 2018]. To mitigate this, both the conditional and label shifts are taken into account for DA from a causal interpretation view [Zhang et al., 2013; Gong et al., 2016]. However, its linearity assumption might be too restrictive for real-world challenges.

In this work, we first analyze the different shift conditions in real-world DG, and rethink the limitation of conventional IFL under different shift assumptions. Targeting the conditional and label shifts, we propose to explicitly align p(f(x)|y) and p(y) via variational Bayesian inference and posterior label alignment. Note that the fine-grained classwise p(f(x)|y) alignment can lead to p(f(x)) alignment following the law of total probability [Zhang et al., 2013].

Aligning the conditional distribution p(f(x)|y) across source domains under the label shift is usually intractable. Thus, we propose to infer the domain-specific variational approximations of these conditional distributions, and reduce the divergence among these approximations.

Specifically, we enforce the conditional domain invariance by optimizing two objectives. The first one enforces the approximate conditional distributions indistinguishable across domains by matching their reparameterized formulations (i.e., the mean and variance of Gaussian distribution). The second one maximizes the probability of observing the input x given the latent representation and domain label, which is achieved by a domain-wise likelihood learning network. Assuming that the conditional shift is aligned, we can then align the posterior classifier with the label distribution following a plug-and-play manner.

The main contributions are summarized as follows:

• We explore both the conditional and label shifts in various DG tasks, and investigate the limitation of conventional IFL methods under different shift assumptions.
• We propose a practical and scalable method to align the conditional shift via the variational Bayesian inference.
• The label shift can be explicitly aligned by the posterior alignment operation.

We empirically validate its effectiveness and generality of our framework on multiple challenging benchmarks with different backbone models and demonstrate superior performance over the comparison methods.

2 RelatedWork

DG assumes that we do not have prior knowledge about the target domain in that we do not have access to labeled/unlabeled samples of the target domain at the training stage [Matsuura and Harada, 2020a]. The conventional DG

methods can be divided into two categories. The first strategy aims to extract the domain invariant features with IFL [Ghifary et al., 2017]. A typical solution is based on adversarial learning, which reduces the inter-domain divergence between the feature representation distributions [Li et al., 2018a; Liu et al., 2021a].

The other strategy focuses on the fusion of domain-specific feature representations (DFR-fusion). [Mancini et al., 2018] builds the domain-specific classifiers by multiple independent convolutional networks. Then, it uses a domain agnostic component to fuse the probabilities of a target sample belonging to different source domains. [Ding and Fu, 2017] infers the domain-invariant feature by matching its low-rank structure with domain-specific features. Typically, these DG methods assume that p(y) is invariant across domains. Therefore, aligning p(f(x)) can be a good alternative to align the conditional shift. However, this assumption is often violated due to the label shift in real-world applications. Therefore, independent conditional and label shift assumptions are more realistic in real-world applications.

Domain shifts in DA can be categorized into covariant, label, conditional, and concept shifts [Moreno-Torres et al., 2012; Zhang et al., 2015]. In this work, we examine these concepts and adapt their causal relationships to DG, as summarized in Fig. 2. Conventionally, each shift is studied independently, by assuming that the other shifts are invariant [Kouw, 2018]. For example, [Li et al., 2018b] aligns the conditional shift assuming that no label shift occurs. We note that the concept shift usually has not been considered in DG tasks, since most of the prior work assumes that an object has different labels in different domains. Some recent works [Zhang et al., 2013; Gong et al., 2016; Combes et al., 2020] assume that both conditional and label shifts exist in DA tasks and tackle the problem with a causal inference framework. However, its linearity assumption and location-scale transform are too restrictive to be applied in many real-world applications. It is worth noting that under the conditional and label shift assumption, y is the cause of x, and therefore it is natural to infer p(x|y) of different domains directly as in [Sch¨olkopf et al., 2012; Gong et al., 2018] as a likelihood maximization network.

In this work, we propose a novel inference graph as shown in Fig. 1 to explicitly incorporate the conditional dependence, which is trained via variational Bayesian inference.

3 Methodology

We denote the input sample, class label, and domain spaces as X , Y, and D, respectively. With random variables x ∈ X , y ∈ Y, and d ∈ D, we can define the probability distribution of each domain as p(x, y|d). For the sake of simplicity, we assume y and d are the discrete variables for which Y = {1, 2, . . . ,K} is the set of class labels. In DG, we are given N source domains {pi (x, y)}Ni=1 to train the latent representation encoder f(x) and the representation classifier p(y|f(x)) [Li et al., 2018a]. The trained and fixed encoder and classifier are used to predict the labels of samples drawn from the marginal distribution {pt (x)} of an "unseen" target domain {pt (x, y)}.

The conventional IFL assumes that p(y) and p(y|x) are invariant across domains. Since f(·) is a deterministic mapping function, p(y|f(x)) should also be invariant across domains. Therefore, if p(f(x)) of different domains are aligned, the conditional shift, p(f(x)|y) = p(y|f(x))p(f(x)) p(y) , is also aligned. However, with the conditional and label shift assumption, the alignment of p(f(x)|y) and p(y) is more challenging than the covariant shift which only requires to align the marginal distribution p(f(x)).

We note that with the law of total probability, p(f(x)) can be aligned if the fine-grained p(f(x)|y) is aligned [Zhang et al., 2013]. Moreover, p(y) for all source domains can be calculated by simply using the class label proportion in each domain. Besides, modeling p(f(x)|y) is natural, since it follows the inherent causal relation under the conditional and label shift (see Fig. 2).

For the simplicity and consistency with the notation of autoencoder, we denote f(x) as z1 , which is the latent variable encoded from x by f. We note that z = f(x) is dependent on its corresponding input sample x. The class conditional distribution p(f(x)|y) can be reformulated as p(z|x, y).

The corresponding inference graph and detailed framework are shown in Fig. 1 and Fig. 3, respectively. When inferring the latent representation z, we explicitly concatenate x and y as input. Moreover, the final class prediction is made by a posterior alignment of the label shift, which also depends on label distribution p(y).

3.1 Variational Bayesian Conditional Alignment

Although p(f(x)) and p(y) for all source domains can be modeled by IFL and class label proportion, respectively, p(z|x, y) is usually intractable for moderately complicated likelihood functions, e.g., neural networks with nonlinear hidden layers. While this could be solved by the Markov chain Monte Carlo simulation (MCMC), this requires expensive iterative inference schemes per data point and does not scalable to the large-scale high-dimensional data.

To align the class-dependent p(z|x, y) across different domains, we first construct its approximate Gaussian distribution q(z|x, y), and resort to the variational Bayesian inference

1 We use f(x) and z interlaced to align with the conventional IFL and variational autoencoder literature, respectively.

[Kingma and Welling, 2013] to bridge it with a simple Gaussian family for which the inference is tractable. Specifically, we have the following proposition:

Proposition 1. The minimization of the inter-domain conditional shift p(z|x, y) is achieved when its approximate distribution q(z|x, y) is invariant across domains, and the KL divergence between p(z|x, y) and q(z|x, y) is minimized.

Proof 1. A simple justification can be:

if we have KL(q1(z|x, y)||p1(z|x, y)) = 0,

KL(q2(z|x, y)||p2(z|x, y)) = 0,

KL(q1(z|x, y)||q2(z|x, y)) = 0,

then, we have KL(p1(z|x, y)||p2(z|x, y)) = 0.

Following the variational bound [Kingma and Welling, 2013], minimizing the KL divergence between p(z|x, y) and q(z|x, y) is equivalent to maximizing the evidence lower bound (ELBO) [Domke and Sheldon, 2018] of the likelihood logp(x, y, z), denoted by L:

where the DKL term is the KL divergence of the approximate from the true posterior and the ELBO of the likelihood logp(x, y, z), i.e., L can be re-written as

which can be reformulated as

where E denotes the expectation. Therefore, approximating p(z|x, y) with q(z|x, y) requires two objectives, i.e., minimizing DKL(q(z|x, y)||p(z|y)), while maximizing the expectation of logp(x|y, z).

For N domains, pi (z|y), i ∈ {1, 2, . . . , N} is the prior distribution in the variational model, e.g., multivariate Gaussian distribution. When pi (z|y) is sampled from the same Gaussian distribution and is invariant across the source domains, the first objective in Eq. (3), i.e., DKL(q(z|x, y)||p(z|y)), can explicitly enforce qi (z|x, y) to be invariant across domains.

By further incorporating the second objective into Eq. (3), we attempt to minimize the KL divergence of p(z|x, y) and q(z|x, y) as in Eq. (1). Then, pi (z|x, y) should be invariant across the source domains, i.e.,

Actually, if we have p1(y|x) = p1(x|y)p1(y)/p1(x), p2(y|x) = p2(x|y)p2(y)/p2(x), and p1(x|y) = p2(x|y), p1(x) = p2(x), then p1(y|x)/p1(y) = p2(y|x)/p2(y).

The first optimization objective of Eq. (3) targets to align the conditional distribution q(z|x, y) across the source domains. Since the prior distribution p(z|y) is the multivariate Gaussian distribution, it is also natural to configure q(z|x, y) as multivariate Gaussian. Practically, we follow the reparametric trick of variational autoencoder [Kingma and Welling, 2013] in such a way that the inference model has two outputs, i.e., µij and σij which are the mean and variance of the Gaussian distribution q(zij |xij , yij ). Then, zij = µij + σij ⊙ ϵ, where ϵ ∈ N(0, I). Without loss generality, we have:

where Mi is the number of input in a batch from domain i. Usually, we set the prior pi (z|y) to be the standard multivariate Gaussian distribution, where the mean and variance are ¯µ = 0 and ¯σ = 1, respectively.

The second optimization objective of Eq. (3) aims to maximize the probability of observing the input x given y and z. We propose to configure a likelihood maximization network flm. It maximizes the likelihood that the latent feature representation of images in a specific domain can effectively represent the same-class images in this domain. Practically, our flm contains N sub-branches, each of which corresponds to a domain. At the training stage, we choose the corresponding branch according to the source domain label d. Its loss can be formulated as

which solves the maximum likelihood problem by minimizing the difference between the input data and the generated data in the corresponding domain. We note that flm is only used in training.

Posterior collapses is a long-lasting problem of VAEs using continuous Gaussian prior. The recent progress of discrete priors or flow model can be potentially added on our model. However, our training does not suffer from it significantly. Note that each of the N decoders is only trained with about 1/N data of encoder f and it is likely that the weak decoder can also be helpful. Using an approximation with multivariate Gaussian prior offers much better tractability as in variational autoencoder (VAE), which is good for posterior modeling.

3.2 Label-prior

Inferencing the latent representation z requires to know the label information y in advance, since we are modeling the

Algorithm 1 VB inference under conditional and label shifts

Initialize network parameters

repeat

//Construct the mini-batch for the training {x, y, d} ← random sampling from the dataset //Forward pass to compute the inference output ˆ y ← flp(x); z ← f(x, ˆ y); ˜x ← flm(z, ˆ y); y ← fcls(z, y) //Calculate the loss functions: L1, L2, LCE1, LCE2 //Update parameters according to gradients f ← L1 + αL2 + βLCE2 + θLˆ y; flp ← LCE1; flm ← L2; fcls ← LCE2

until deadline

approximate conditional distribution q(z|x, y). Although y is always available in training, the ground truth y is not available in testing. We note that flm is only used in training which always has y.

To alleviate this limitation, we infer the label from the input image x as a prior to control the behavior of the conditional distribution matching module. Specifically, we configure a label-prior network flp : x → ˆ y to infer the pseudo-label, and use ˆ y as input to the posterior label alignment classifier fcls in both training and testing. Our label prior network flp is trained by the cross-entropy loss LCE1 with the ground-truth label y.

Moreover, at the training stage, we can further utilize the ground-truth y and minimizing DKL(q(z|x, y)||q(z|x, ˆ y)) to update the encoder f. We denote the to be minimized KL divergence as Lˆ y. Minimizing Lˆ y is not mandatory, while Lˆ y can encourage the encoder to be familiar with the noisy ˆ y [Lian et al., 2019] and learn to compensate for the noisy prediction. We note that assigning ˆ y to an uniform histogram as the dialog system [Lian et al., 2019] to fill the missing variate y can degenerate the modeling of p(f(x)|y) to p(f(x)) in our DG task. Therefore, the pseudo-label will be post-processed by both encoder and classifier, which may adjust the unreliable ˆ y.

3.3 Posterior Alignment with Label Shift

Finally, we align p(y) to obtain the final classifier fcls. Since the classifier is deployed on all of the source domains, we can regard all of the source domains as a single domain, and denote the classifier as

where ¯ p(f(x)|y), ¯ p(f(x)), and ¯ p(y) are its class-conditional, latent representation, and label distribution, respectively.

Suppose that all the conditional distribution pi (f(x)|y) and the latent distribution pi (f (x)) are aligned to each other using variational Bayesian inference, they should also be aligned with ¯ p(f(x)|y) and ¯ p(f(x)). Therefore, the posterior alignment and the final prediction of the sample f(x) from domain i can be formulated as

where ¯ p(y = k|f(x)) is the k-th element value of the classifier's softmax prediction. Here, we also calculate the cross entropy loss LCE2 between pi (y = k|f(x)) and the ground truth label y.

As detailed in Algorithm 1, we update f with L1 + αL2 + βLCE2 + θLˆ y, update flp with LCE1, update flm with L2, and update fcls with LCE2, respectively.

4 Experiments

This section demonstrates the effectiveness of our variational Bayesian inference framework under conditional and label shift (VBCLS) on the classic VLCS DG benchmark for image classification and the PACS benchmark for object recognition with domain shift.

4.1 Implementation Details

The domain invariant encoder f and posterior alignment classifier fpa use the encoder and classifier structure as our compared models (e.g., AlexNet and ResNet18), and the label prior network is simply the concatenation of f and fpa, and the likelihood maximization network flm uses the reversed CNN decoder structure of f. We implement our methods using PyTorch, we empirically set α = 0.5, β = 1, θ = 0.1 and t = 100 via grid searching. In our experiments, the performance does not sensitive to these hyperparameters for a relatively large range.

Following previous work on domain generalization [Li et al., 2017; Li et al., 2019], we use models pre-trained on the ILSVRC dataset and label smoothing on the task classifier in order to prevent overfitting. The model is trained for 30 epochs via the mini-batch stochastic gradient descent (SGD) with a batch size of 128, a momentum of 0.9, and a weight decay of 5e - 4. The initial learning rate is set to 1e - 3, which is scaled by a factor of 0.1 after 80% of the epochs. For the VLCS dataset, the initial learning rate is set to 1e-4, since a high learning rate leads to early convergence due to the overfitting in the source domain. In addition, the learning rate of the domain discriminator and the classifier is set to be larger (i.e., 10 times) than that of the feature extractor.

For the ablation study, VBCLS-fpa denotes without posterior alignment and VBCLS-Lˆ y denotes without minimizing DKL(q(z|x, y)||q(z|x, ˆ y)).

4.2 VLCS Dataset

VLCS [Ghifary et al., 2017] contains images from four different datasets including PASCAL VOC2007 (V), LabelMe (L), Caltech-101 (C), and SUN09 (S). Different from PACS, VLCS offers photo images taken under different camera types or composition bias. The domain V, L, C, and S have 3,376, 2,656, 1,415, and 3,282 instances, respectively. Five shared classes are collected to form the label space, including bird, car, chair, dog, and person. We follow the previous works to exploit the publicly available pre-extracted DeCAF6 features (4,096-dim vector) [Donahue et al., 2014] for leaveone-domain-out validation by randomly splitting each domain into 70% training and 30% testing. We report the mean over five independent runs for our results.

The mean accuracy on the test partition using leave-onedomain-out validation on VLCS is reported in Table 1. We also compare our method with the DG models trained using the same architecture and source domains.

Our results indicate that our VBCLS outperforms the covariant shift setting methods (e.g., CCSA [Motiian et al., 2017], MMD-AAE [Li et al., 2018a]) by a large margin. The improvement over the conditional shift only method CIDDG [Li et al., 2018b] is significant, demonstrating the necessity of incorporating both conditional and label shift. When compared with the recent SOTA methods, e.g., selfchallenging based RSC [Huang et al., 2020], information bottleneck based MetaVIB [Du et al., 2020], and self-training based EISNet [Wang et al., 2020], we can observe that our VBCLS yields better performance in almost all cases. We note that JiGen [Carlucci et al., 2019] uses the Jigsaw Puzzles solving, which are essentially different from the IFL.

The good performance of our strategy indicates the invariant feature learning works well, and the conditional and shift assumption can be an inherent property that needs to be addressed for real-world DG challenges. The discrepancy between marginal distributions ¯ p(y) and pt (y) is measured via the KL-divergence as the semi-supervised learning with the selective bias problem [Zadrozny, 2004]. The impact of the label shift is empirically illustrated in Fig. 4 left. In Fig. 4 right, we show the label refinement can effectively estimate the testing label distribution without tuning the network. The label alignment can be relatively accurate after 3 epochs and is almost stable after 5 epochs.

In the ablation study, we can see that posterior alignment is necessary if there is label shift. Besides, the performance of the label prior network can be improved by minimizing DKL(q(z|x, y)||q(z|x, ˆ y)) on encoder f.

Another evaluation protocol on VLCS is to examine whether eliminating examples from one source domain impacts the performance on the target domain. This protocol is designed to evaluate the impact of the diversity of the source domains on the target accuracy. In this experiment, each target domain on models by training all combinations of the remaining domains as the source domains is evaluated, as shown in Table 2. In addition, the CIDDG baseline is included for reference. Results in Table 2 demonstrate that for all target domains, reducing the number of source domains from 3 (see Table 1) to 2 degrades the performance for all combinations of the source domains. We can see that, in some cases, excluding a particular source from the training substantially degrades the target loss. However, we can see that our VBCLS can still be more robust under these cases.

4.3 PACS Dataset

The object recognition benchmark PACS [Li et al., 2017] consists of images divided into 7 classes from four different datasets, including Photo (P), Art painting (A), Cartoon (C), and Sketch (S). In Tab. 3, we provide the detailed statistics of the number of samples in each domain. It is clear that the class ratio is different across domains, which indicates the significant label shift.

As shown in the results (c.f. Table 4), the Sketch domain produces the lowest accuracy when used as the target domain, and therefore it is deemed the most challenging one. In light of this, we follow the previous work to tune the model using the S domain as the target domain and reuse the same hyperparameters for the experiments with the remaining domains. In Table 4, we show the result using our method, which was averaged over 5 different initializations alongside all the other comparison methods. Overall, we can see that our method yields better average performance over all source domains compared to previous SOTA methods. Among them, CrossGrad [Shankar et al., 2018] synthesizes data for a new domain, MMLD [Matsuura and Harada, 2020b] using a mixture of multiple latent domains. It is also comparable to the recent self-challenging based RSC [Huang et al., 2020], information bottleneck based MetaVIB [Du et al., 2020], and self-training based EISNet [Wang et al., 2020]. More importantly, we can observe that our method outperforms CIDDG [Li et al., 2018b], an adversarial conditional IFL strategy, by a large margin. The ablation studies are also consistent with the results in VLCS and PACS. We also provide the results using the ResNet18 backbone.

5 Conclusion

In this work, we target to establish a more realistic assumption in DG that both conditional and label shifts arise independently and contemporarily. We theoretically analyze the inequilibrium of conventional IFL under the different shift assumptions. Motivated by that, a concise yet effective VBCLS framework based on variational Bayesian inference with the posterior alignment is proposed to reduce both the conditional shift and label shifts. Extensive evaluations verify our analysis and demonstrate the effectiveness of our approach, compared with IFL approaches.

Acknowledgments

This work was supported in part by the Hong Kong GRF 152202/14E, PolyU Central Research Grant G-YBJW, and Jiangsu NSF (BK20200238).