Mapping conditional distributions for domain adaptation under generalized target shift

1 INTRODUCTION . Unsupervised Domain Adaptation ( UDA ) methods ( Pan & Yang , 2010 ) train a classifier with labelled samples from a source domain S such that its risk on an unlabelled target domain T is low . This problem is ill-posed and simplifying assumptions have been considered . Initial contributions focused on two settings which decompose differently the joint distribution over X × Y : covariate shift ( CoS ) , when input marginals differ across domains while class-posteriors are unchanged i.e . pS ( Y |X ) = pT ( Y |X ) , pS ( X ) 6= pT ( X ) and target shift ( TarS ) when label distributions differ while conditionals are unchanged i.e . pS ( Y ) 6= pT ( Y ) , pS ( X|Y ) = pT ( X|Y ) . These restrictive assumptions were recently extended into model shift i.e . pS ( Y |X ) 6= pT ( Y |X ) , pS ( X ) 6= pT ( X ) ( Wang & Schneider , 2014 ; 2015 ) and generalized target shift ( GeTarS ) i.e . pS ( X|Y ) 6= pT ( X|Y ) , pS ( Y ) 6= pT ( Y ) , which are more realistic in real-world applications . Our work addresses GeTarS for which a key challenge is to learn how to map the source domain onto the target one to minimize both conditional and label shifts , without using target labels . The current state-of-the-art in Gong et al . ( 2016 ) ; Combes et al . ( 2020 ) ; Rakotomamonjy et al . ( 2020 ) ; Shui et al . ( 2021 ) learns domain-invariant representations and uses , as importance weights in the training objectives , the estimated class-ratios between domains . However , this approach has several limitations . First , to transfer representations , the domain-invariance constraint breaks the original problem structure and it was shown that this may degrade the discriminativity of target representations ( Liu et al. , 2019 ) . Some approaches proposed to better enforce target discriminativity when learning invariant representations e.g . Xiao et al . ( 2019 ) ; Li et al . ( 2020 ) ; Chen et al . ( 2019 ) , yet they were not applied to the GeTarS setting . Second , generalization guarantees are derived under strong assumptions , detailed in Section 2.3 , which may not hold in practice . In this paper , we address these limitations with a new general approach , named Optimal Sample Transformation and Reweight ( OSTAR ) , which maps pretrained representations using Optimal Transport ( OT ) . OSTAR proposes an alternative to constraining representation invariance and performs jointly three operations : given a pretrained encoder , ( i ) it learns an OT map , implemented as a neural network ( NN ) , between encoded source and target conditionals , ( ii ) it estimates target proportions for sample reweighting and ( iii ) it learns a classifier for the target domain using source labels . OSTAR has several benefits : ( i ) it is flexible , scalable and preserves target discriminativity and ( ii ) it provides strong theoretical guarantees under mild assumptions . In summary , our contributions are : • We propose an approach , OSTAR , to align pretrained representations under GeTarS . Without constraining representation-invariance , OSTAR jointly learns a classifier for inference on the target domain and an OT map , which maps representations of source conditionals to those of target ones under class-reweighting . OSTAR preserves target discriminativity and experimentally challenges the state-of-the-art for GeTarS . • OSTAR implements its OT map as a NN shared across classes . Our approach is thus flexible and has native regularization biases for stability . Moreover it is scalable and generalizes beyond training samples unlike standard linear programming based OT approaches . • OSTAR has strong theoretical guarantees under mild assumptions : its solution is unique , recovers target proportions and correctly matches source and target conditionals at the optimum . It also explicitly controls the target risk with a new Wasserstein-based bound . Our paper is organized as follows . In Section 2 , we define our problem , approach and assumptions . In Section 3 , we derive theoretical results . In Section 4 , we describe our implementation . We report in Section 5 experimental results and ablation studies . In Section 6 , we present related work . 2 PROPOSED APPROACH . In this section , we successively define our problem , present our method , OSTAR and its main ideas and introduce our assumptions , used to provide theoretical guarantees for our method . 2.1 PROBLEM DEFINITION . Denoting X the input space and Y = { 1 , . . . , K } the label space , we consider UDA between a source S = ( XS , YS , pS ( X , Y ) ) with labelled samples Ŝ = { ( x ( i ) S , y ( i ) S ) } ni=1 ∈ ( XS × YS ) n and a target T = ( XT , YT , pT ( X , Y ) ) with unlabelled samples T̂ = { x ( i ) T } mi=1 ∈ XmT . We denote Z ⊂ Rd a latent space and g : X → Z an encoder from X to Z . ZS and ZT are the encoded source and target input domains , ZŜ and ZT̂ the corresponding training sets and Z a random variable in this space . The latent marginal probability induced by g on D ∈ { S , T } is defined as ∀A ⊂ Z , pgD ( A ) , g # ( pD ( A ) ) 1 . For convenience , pYD ∈ RK denotes the label marginal pD ( Y ) and pD ( Z ) , p g D ( Z ) . In all generality , conditional distributions in this latent space and label marginals differ across domains ; this is the GeTarS assumption ( Definition 1 ) illustrated in Figure 1a . This assumption , made in feature space Z rather than input space X , states that latent representations for both domains from a given class are different with different label proportions . Operating in the latent space has several practical advantages such as improved discriminativity and dimension reduction . Definition 1 ( GeTarS ) . GeTarS is characterized by conditional mismatch across domains i.e . ∃j ∈ { 1 , . . . , K } , pS ( Z|Y = j ) 6= pT ( Z|Y = j ) and label shift i.e . pYS 6= pYT . Our goal is to learn a classifier in Z with low target risk , using source labels . This is challenging as ( i ) target labels are unknown and ( ii ) there are two shifts to handle . We will show that this can be achieved with pretrained representations if we recover two key properties : ( i ) a map which matches source and target conditional distributions and ( ii ) target proportions to reweight samples by classratios and thus account for label shift . Our approach , OSTAR , achieves this objective . 2.2 MAPPING CONDITIONAL DISTRIBUTIONS UNDER LABEL SHIFT . We now present the components in OSTAR , their objective and the various training steps . Components The main components of OSTAR , illustrated in Figure 1b , are detailed below . These components are learned and estimated using the algorithm detailed in Section 4 . They include : • a fixed encoder g : X → Z , defined in Section 2.1 . • a mapping φ : Z → Z , acting on source samples encoded by g. • a label proportion vector pYN on the simplex ∆K . • a classifier fN : Z → { 1 , . . . , K } for the target domain in a hypothesis classH over Z . 1f # ρ is the push-forward measure f # ρ ( B ) = ρ ( f−1 ( B ) ) , for all measurable set B . Objective g encodes source and target samples in a latent space such that it preserves rich information about the target task and such that the risk on the source domain is small . g is fixed throughout training to preserve target discriminativity . φ should map encoded source conditionals in ZS onto corresponding encoded target ones in ZT to account for conditional shift ; ZN denotes the mapped space . pYN should estimate the target proportions p Y T to account for label shift . Components ( φ , p Y N ) define a new labelled domain in latent space N = ( ZN , YN , pN ( Z , Y ) ) through a Sample Transformation And Reweight operation of the encoded S domain , as illustrated in Figure 1b . Indeed , the pushforward by φ of encoded source conditionals defines conditionals in domain N , pφN ( Z|Y ) : ∀k , pφN ( Z|Y = k ) , φ # ( pS ( Z|Y = k ) ) 1 Then , pYN weights each conditional in N . This yields a marginal distribution in N , p φ N ( Z ) : pφN ( Z ) , K∑ k=1 pY=kN p φ N ( Z|Y = k ) ( 1 ) Finally , classifier fN is trained on labelled samples from domain N . This is possible as each sample in N is a projection of a labelled sample from S. fN can then be used for inference on T . We will show that it has low target risk when components φ and pYN achieve their objectives detailed above . Training We train OSTAR ’ s components in two stages . First , we train g along a source classifier fS from scratch by minimizing source classification loss ; alternatively , g can be tailored to specific problems with pretraining . Second , we jointly learn ( fN , φ , pYN ) to minimize a classification loss in domain N and to match target conditionals and proportions with those in domain N . As target conditionals and proportions are unknown , we propose a proxy problem for ( φ , pYN ) to match instead latent marginals pT ( Z ) and p φ N ( Z ) ( 1 ) . We solve this proxy problem under least action principle measured by a Monge transport cost ( Santambrogio , 2015 ) , denoted C ( φ ) , as in problem ( OT ) : min φ , pYN∈∆K C ( φ ) , K∑ k=1 ∫ z∈Z ||φ ( z ) − z||22 pS ( z|Y = k ) dz subject to pφN ( Z ) = pT ( Z ) ( OT ) For any function φ e.g . parametrized by a NN , C ( φ ) is the transport cost of encoded source conditionals by φ . It uses a cost function , c ( x , y ) = ‖x − y‖p2 , where without loss of generality p = 2 . The optimal C ( φ ) is the sum of Wasserstein-2 distances between source conditionals and their mappings . Problem ( OT ) seeks to minimize C ( φ ) under marginal matching . We provide some further background on optimal transport in Appendix C and discuss there the differences between our OT problem ( OT ) and the standard Monge OT problem between pS ( Z ) and pT ( Z ) . Our OT formulation is key to the approach . First , it is at the basis of our theoretical analysis . Under Assumption 2 later defined , the optimal transport cost is the sum of the Wasserstein-2 distance between source and matched target conditionals . We can then provide conditions on these distances in Assumption 3 to formally define when the solution to problem ( OT ) correctly matches conditionals . Second , it allows us to learn the OT map with a NN . This NN approach : ( i ) generalizes beyond training samples and scales up with the number of samples unlike linear programming based OT approaches ( Courty et al. , 2017b ) and ( ii ) introduces useful stability biases which make learning less prone to numerical instabilities as highlighted in de Bézenac et al . ( 2021 ) ; Karkar et al . ( 2020 ) .

Summary: This paper proposes a method for unsupervised domain adaptation under conditional and label shift, one of the most challening versions of DA. At the core of this method is an optimal transport problem that seeks a mapping between source and target domains *and* a class proportion vector that minimize a total transportation cost between the domains. Based on this formulation, and a set of 4 sensible assumptions, the paper provides theoretical results in the form of unicity of the solution and classic generalization bound on the target domain risk. The implemented algorithm is inspired (though not identical) to this objective, which is extended with a few other terms: a constraint on the the label proportions based on a confusion matrix, a classification loss for the mapped source samples, and a relaxed version of the OT's pushforward constraint, and finally, two objectives that seek to improve the discriminativity of target representations. The paper then proceeds to validate the propose method in a wide array of benchmark DA datasets, with various levels of class unbalance, whereby it shows that the method repeatedly and significantly outperforms SOTA UDA alternatives.