Domain Adaptation with Conditional Distribution Matching and Generalized Label Shift

In spite of impressive successes, most deep learning models [24] rely on huge amounts of labelled data and their features have proven brittle to distribution shifts [59, 43]. Building more robust models, that learn from fewer samples and/or generalize better out-of-distribution is the focus of many recent works [5, 2, 57]. The research direction of interest to this paper is that of domain adaptation, which aims at learning features that transfer well between domains. We focus in particular on unsupervised domain adaptation (UDA), where the algorithm has access to labelled samples from a source domain and unlabelled data from a target domain. Its objective is to train a model that generalizes well to the target domain. Building on advances in adversarial learning [25], adversarial domain adaptation (ADA) leverages the use of a discriminator to learn an intermediate representation that is invariant between the source and target domains. Simultaneously, the representation is paired with a classifier, trained to perform well on the source domain [22, 53, 64, 36]. ADA is rather successful on a variety of tasks, however, recent work has proven an upper bound on the performance of existing algorithms when source and target domains have mismatched label distributions [66]. Label shift is a property of two domains for which the marginal label distributions differ, but the conditional distributions of input given label stay the same across domains [52, 62].

In this paper, we study domain adaptation under mismatched label distributions and design methods that are robust in that setting. Our contributions are the following. First, we extend the upper bound by Zhao et al. [66] to $k$-class classification and to conditional domain adversarial networks, a recently introduced domain adaptation algorithm [41]. Second, we introduce generalized label shift ($GLS$), a broader version of the standard label shift where conditional invariance between source and target domains is placed in representation rather than input space. Third, we derive performance guarantees for algorithms that seek to enforce $GLS$ via learnt feature transformations, in the form of upper bounds on the error gap and the joint error of the classifier on the source and target domains. Those guarantees suggest principled modifications to ADA to improve its robustness to mismatched label distributions. The modifications rely on estimating the class ratios between source and target domains and use those as importance weights in the adversarial and classification objectives. The importance weights estimation is performed using a method of moment by solving a quadratic program, inspired from Lipton et al. [35]. Following these theoretical insights, we devise three new algorithms based on learning importance-weighted representations, DANNs [22], JANs [40] and CDANs [41]. We apply our variants to artificial UDA tasks with large divergences between label distributions, and demonstrate significant performance gains compared to the algorithms’ base versions. Finally, we evaluate them on standard domain adaptation tasks and also show improved performance.

Notation  We focus on the general $k$-class classification problem. $\mathcal{X}$ and $\mathcal{Y}$ denote the input and output space, respectively. $\mathcal{Z}$ stands for the representation space induced from $\mathcal{X}$ by a feature transformation $g:\mathcal{X}\mapsto\mathcal{Z}$. Accordingly, we use $X,Y,Z$ to denote random variables which take values in $\mathcal{X},\mathcal{Y},\mathcal{Z}$. Domain corresponds to a joint distribution on the input space $\mathcal{X}$ and output space $\mathcal{Y}$, and we use $\mathcal{D}_{S}$ (resp. $\mathcal{D}_{T}$) to denote the source (resp. target) domain. Noticeably, this corresponds to a stochastic setting, which is stronger than the deterministic one previously studied [6, 7, 66]. A hypothesis is a function $h:\mathcal{X}\to[k]$. The error of a hypothesis $h$ under distribution $\mathcal{D}_{S}$ is defined as: $\varepsilon_{S}(h)\vcentcolon=\Pr_{\mathcal{D}_{S}}(h(X)\neq Y)$, i.e., the probability that $h$ disagrees with $Y$ under $\mathcal{D}_{S}$.

Domain Adaptation via Invariant Representations   For source ($\mathcal{D}_{S}$) and target ($\mathcal{D}_{T}$) domains, we use $\mathcal{D}_{S}^{X}$, $\mathcal{D}_{T}^{X}$, $\mathcal{D}^{Y}_{S}$ and $\mathcal{D}^{Y}_{T}$ to denote the marginal data and label distributions. In UDA, the algorithm has access to $n$ labeled points $\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}\in(\mathcal{X}\times\mathcal{Y})^{n}$ and $m$ unlabeled points $\{\mathbf{x}_{j}\}_{j=1}^{m}\in\mathcal{X}^{m}$ sampled i.i.d. from the source and target domains. Inspired by Ben-David et al. [7], a common approach is to learn representations invariant to the domain shift. With $g:\mathcal{X}\mapsto\mathcal{Z}$ a feature transformation and $h:\mathcal{Z}\mapsto\mathcal{Y}$ a hypothesis on the feature space, a domain invariant representation [22, 53, 63] is a function $g$ that induces similar distributions on $\mathcal{D}_{S}$ and $\mathcal{D}_{T}$. $g$ is also required to preserve rich information about the target task so that $\varepsilon_{S}(h\circ g)$ is small. The above process results in the following Markov chain (assumed to hold throughout the paper):

where $\widehat{Y}=h(g(X))$. We let $\mathcal{D}_{S}^{Z}$, $\mathcal{D}_{T}^{Z}$, $\mathcal{D}_{S}^{\widehat{Y}}$ and $\mathcal{D}_{T}^{\widehat{Y}}$ denote the pushforwards (induced distributions) of $\mathcal{D}_{S}^{X}$ and $\mathcal{D}_{T}^{X}$ by $g$ and $h\circ g$. Invariance in feature space is defined as minimizing a distance or divergence between $\mathcal{D}_{S}^{Z}$ and $\mathcal{D}_{T}^{Z}$.

Adversarial Domain Adaptation  Invariance is often attained by training a discriminator $d:\mathcal{Z}\mapsto[0,1]$ to predict if $z$ is from the source or target. $g$ is trained both to maximize the discriminator loss and minimize the classification loss of $h\circ g$ on the source domain ($h$ is also trained with the latter objective). This leads to domain-adversarial neural networks [22, DANN], where $g$, $h$ and $d$ are parameterized with neural networks: $g_{\theta}$, $h_{\phi}$ and $d_{\psi}$ (see Algo. 1 and App. B.5). Building on DANN, conditional domain adversarial networks [41, CDAN] use the same adversarial paradigm. However, the discriminator now takes as input the outer product, for a given $x$, between the predictions of the network $h(g(x))$ and its representation $g(x)$. In other words, $d$ acts on the outer product: $h\otimes g(x)\vcentcolon=(h_{1}(g(x))\cdot g(x),\dots,h_{k}(g(x))\cdot g(x))$ rather than on $g(x)$. $h_{i}$ denotes the $i$-th element of vector $h$. We now highlight a limitation of DANNs and CDANs.

Remark  The lower bound is algorithm-independent. It is also a population-level result and holds asymptotically with increasing data. Zhao et al. [66] prove the theorem for $k=2$ and $\widetilde{Z}=Z$. We extend it to CDAN and arbitrary $k$ (it actually holds for any $\widetilde{Z}$ s.t. $\widehat{Y}=\widetilde{h}(\widetilde{Z})$ for some $\widetilde{h}$, see App. A.3). Assuming that label distributions differ between source and target domains, the lower bound shows that: the better the alignment of feature distributions, the worse the joint error. For an invariant representation ($D_{\text{JS}}(\mathcal{D}_{S}^{\tilde{Z}},\mathcal{D}_{T}^{\tilde{Z}})=0$) with no source error, the target error will be larger than $D_{\text{JS}}(\mathcal{D}_{S}^{Y},\mathcal{D}_{T}^{Y})/2$. Hence algorithms learning invariant representations and minimizing the source error are fundamentally flawed when label distributions differ between source and target domains.

Common Assumptions to Tackle Domain Adaptation Two common assumptions about the data made in DA are covariate shift and label shift. They correspond to different ways of decomposing the joint distribution over $X\times Y$, as detailed in Table 1. From a representation learning perspective, covariate shift is not robust to feature transformation and can lead to an effect called negative transfer [66]. At the same time, label shift clearly fails in most practical applications, e.g. transferring knowledge from synthetic to real images [55]. In that case, the input distributions are actually disjoint.

In light of the limitations of existing assumptions, (e.g. covariate shift and label shift), we propose generalized label shift ($GLS$), a relaxation of label shift that substantially improves its applicability. We first discuss some of its properties and explain why the assumption is favorable in domain adaptation based on representation learning. Motivated by $GLS$, we then present a novel error decomposition theorem that directly suggests a bound minimization framework for domain adaptation. The framework is naturally compatible with $\mathcal{F}$-integral probability metrics [44, $\mathcal{F}$-IPM] and generates a family of domain adaptation algorithms by choosing various function classes $\mathcal{F}$. In a nutshell, the proposed framework applies a method of moments [35] to estimate the importance weight $\mathbf{w}$ of the marginal label distributions by solving a quadratic program (QP), and then uses $\mathbf{w}$ to align the weighted source feature distribution with the target feature distribution.

Covariate shift has been studied and used in many adaptation algorithms [30, 26, 3, 1, 53, 65, 48]. While less known, label shift has also been tackled from various angles over the years: applying EM to learn $\mathcal{D}^{Y}_{T}$ [14], placing a prior on the label distribution [52], using kernel mean matching [61, 20, 46], etc. Schölkopf et al. [50] cast the problem in a causal/anti-causal perspective corresponding to covariate/label shift. That perspective was then further developed [61, 23, 35, 4]. Numerous domain adaptation methods rely on learning invariant representations, and minimize various metrics on the marginal feature distributions: total variation or equivalently $D_{\text{JS}}$ [22, 53, 64, 36], maximum mean discrepancy [27, 37, 38, 39, 40], Wasserstein distance [17, 16, 51, 33, 15], etc. Other noteworthy DA methods use reconstruction losses and cycle-consistency to learn transferable classifiers [67, 29, 56]. Recently, Liu et al. [36] have introduced Transferable Adversarial Training (TAT), where transferable examples are generated to fill the gap in feature space between source and target domains, the datasets is then augmented with those samples. Applying our method to TAT is a future research direction.

Other relevant settings include partial ADA, i.e. UDA when target labels are a strict subset of the source labels / some components of $\mathbf{w}$ are $0$ [11, 12, 13]. Multi-domain adaptation, where multiple source or target domains are given, is also very studied [42, 18, 45, 63, 28, 47]. Recently, Binkowski et al. [8] study sample reweighting in the domain transfer to handle mass shifts between distributions.

Prior work on combining importance weight in domain-invariant representation learning also exists in the setting of partial DA [60]. However, the importance ratio in these works is defined over the features $Z$, rather than the class label $Y$. Compared to our method, this is both statistically inefficient and computationally expensive, since the feature space $\mathcal{Z}$ is often a high-dimensional continuous space, whereas the label space $\mathcal{Y}$ only contains a finite number ($k$) of distinct labels. In a separate work, Yan et al. [58] proposed a weighted MMD distance to handle target shift in UDA. However, their weights are estimated based on pseudo-labels obtained from the learned classifier, hence it is not clear whether the pseudo-labels provide accurate estimation of the importance weights even in simple settings. As a comparison, under $GLS$, we show that our weight estimation by solving a quadratic program converges asymptotically.

We have introduced the generalized label shift assumption, $GLS$, and theoretically-grounded variations of existing algorithms to handle mismatched label distributions. On tasks from classic benchmarks as well as artificial ones, our algorithms consistently outperform their base versions. The gains, as expected theoretically, correlate well with the JSD between label distributions across domains. In real-world applications, the JSD is unknown, and might be larger than in ML datasets where classes are often purposely balanced. Being simple to implement and adding barely any computational cost, the robustness of our method to mismatched label distributions makes it very relevant to such applications.

Extensions  The framework we define in this paper relies on appropriately reweighting the domain adversarial losses. It can be straightforwardly applied to settings where multiple source and/or target domains are used, by simply maintaining one importance weights vector $\mathbf{w}$ for each source/target pair [63, 47]. In particular, label shift could explain the observation from Zhao et al. [63] that too many source domains hurt performance, and our framework might alleviate the issue. One can also think of settings (e.g. semi-supervised domain adaptation) where estimations of $\mathcal{D}_{T}^{Y}$ can be obtained via other means. A more challenging but also more interesting future direction is to extend our framework to domain generalization, where the learner has access to multiple labeled source domains but no access to (even unlabelled) data from the target domain.

The authors thank Romain Laroche and Alessandro Sordoni for useful feedback and helpful discussions. HZ and GG would like to acknowledge support from the DARPA XAI project, contract #FA87501720152 and a Nvidia GPU grant. YW would like acknowledge partial support from NSF Award #2029626, a start-up grant from UCSB Department of Computer Science, as well as generous gifts from Amazon, Adobe, Google and NEC Labs.

Our work focuses on domain adaptation and attempts to properly handle mismatches in the label distributions between the source and target domains. Domain Adaptation as a whole aims at transferring knowledge gained from a certain domain (or data distribution) to another one. It can potentially be used in a variety of decision making systems, such as spam filters, machine translation, etc.. One can also potentially think of much more sensitive applications such as recidivism prediction, or loan approvals.

While it is unclear to us to what extent DA is currently applied, or how it will be applied in the future, the bias formalized in Th. 2.1 and verified in Table 17 demonstrates that imbalances between classes will result in poor transfer performance of standard ADA methods on a subset of them, which is without a doubt a source of potential inequalities. Our method is actually aimed at counter-balancing the effect of such imbalances. As shown in our empirical results (for instance Table 18) it is rather successful at it, especially on significant shifts. This makes us rather confident in the algorithm’s ability to mitigate potential effects of biases in the datasets. On the downside, failure in the weight estimation of some classes might result in poor performance on those. However, we have not observed, in any of our experiments, our method performing significantly worse than its base version. Finally, our method is a variation over existing deep learning algorithms. As such, it carries with it the uncertainties associated to deep learning models, in particular a lack of interpretability and of formal convergence guarantees.