TIDOT: A Teacher Imitation Learning Approach for Domain Adaptation with Optimal Transport

1 Introduction

Domain adaptation or covariate shift problem has emerged from the observation of significant degradation in predictive performance when there exists a shift between a source domain (over which a classifier is trained) and a target domain (over which the classifier does prediction). Many state-ofthe-art methods have been proposed for both shallow domain adaptation [Courty et al., 2017b] and deep domain adaptation [Ganin and Lempitsky, 2015; Long et al., 2015; French et al., 2018; Shu et al., 2018; Damodaran et al., 2018].

Imitation learning follows the principle of 'learning from demonstration'. In particular, there are two fundamental components: an expert teacher and a student. The former component knows how to do its job perfectly, whilst the latter one tries to imitate the teacher to solve its task. This learning principle has been applied successfully in reinforcement learning and sequence prediction [Abbeel and Ng, 2004; Ross et al., 2011; Ho and Ermon, 2016].

This work was supported by the US Air Force grant FA238619-1-4040.

In this work, using the principle of imitation learning and the theory of optimal transport we propose a novel model for unsupervised domain adaptation, named Teacher Imitation Domain Adaptation with Optimal Transport (TIDOT). The teacher in this scenario is apparently a classifier trained on the labeled source domain, but two questions naturally arise: i) which component is the student; and ii) what are the principle and mechanism to enable the student to mimic its teacher in this specific application? We address these two questions by developing a rigorous and intuitive theory based on the theory of optimal transport (OT) [Villani, 2008; Santambrogio, 2015; Peyr´e et al., 2019]. From an abstract interpretation, our theory and mechanism postulate that to predict an unlabeled target sample the student needs to match this target sample with a corresponding labeled source sample so as to conveniently imitate the prediction of the teacher on this source sample. We summarize our contributions in this work as follows:

• We propose a rigorous OT-based theory to leverage imitation learning and domain adaptation. This paradigm is sufficiently general to potentially and promisingly apply to many learning problems including reinforcement learning. In this paper, we demonstrate its application in the context of unsupervised domain adaptation.
• We conduct experiments to compare our TIDOT to the baselines, especially OT-based deep DA, e.g., DeepJDOT [Damodaran et al., 2018], SWD [Lee et al., 2019], DASPOT [Xie et al., 2019], ETD [Li et al., 2020] and RWOT [Xu et al., 2020]). The experimental results show that our proposed method outperforms existing methods on a variety of benchmark datasets including Digits, traffic sign, natural scenes, Office-31, OfficeHome, and ImageCLEF-DA.
• We empirically suggest a potential OT-inspired regularization technique for future work. In particular, as an intriguing side effect of our proposed model, by setting the target training set as the source validation set, we en-

force the teacher to not only predict well on the source training set, but also generalize to predict well on the unlabeled source validation set. We demonstrate that this strategy can yield a regularizer to mitigate the overfitting problem. Although this point is not the main claim of this work and we investigate it as an ablation study, its promising results reveal that this workaround is potentially a decent OT-inspired regularization technique.

2 RelatedWork

Deep domain adaptation (DA) has been intensively studied and shown appealing performance in various tasks and applications, notably [Ganin and Lempitsky, 2015; Long et al., 2015; French et al., 2018; Damodaran et al., 2018; Nguyen et al., 2019; 2020]. The core idea of deep DA is to bridge the gap between source and target distributions in a joint space by minimizing a divergence between distributions induced from the source and target domains. Popular choices of divergence include Jensen-Shannon divergence [Ganin and Lempitsky, 2015; Long et al., 2015; French et al., 2018]; maximum mean discrepancy distance [Long et al., 2015]; and WS distance [Courty et al., 2017b; Nguyen et al., 2021; Le et al., 2021].

Optimal transport theory has been applied in domain adaptation in [Courty et al., 2017b; 2017a; Damodaran et al., 2018; Redko et al., 2019; Lee et al., 2019; Xie et al., 2019; Li et al., 2020; Xu et al., 2020]. Particularly, [Lee et al., 2019] proposes using sliced Wasserstein distance for domain adaption, whereas [Xie et al., 2019] introduces SPOT in which the optimal transport plan is approximated by a pushforward of a reference distribution, and cast the optimal transport problem into a minimax problem. Recently, ETD [Li et al., 2020] measures the domain discrepancy by minimizing attention-aware transport distance while RWOT [Xu et al., 2020] exploits spatial prototypical information and intra-domain structure to reduce the negative transfer brought by the target samples near decision boundaries. Moreover, [Courty et al., 2017b] proposes a brilliant idea to connect the theory of optimal transport and domain adaptation [Courty et al., 2017a], which later inspires an OT-based deep DA method (DeepJDOT) [Damodaran et al., 2018] and learning from multiple data sources [Redko et al., 2019]. Different from [Courty et al., 2017b], our theory originates from OTbased imitation learning for which we develop a rigorous theory to explain the intuition of OT-based imitation DA and also theoretically analyze the general loss of OT-based imitation deep DA, wherein we employ deep neural networks for transfer learning. We note that this makes our theory significantly distinguish from [Courty et al., 2017b] which only limits in the standard setting of transfer learning. In terms of modeling, we propose TIDOT which encourages OTbased imitation learning via a teacher and student in which the teacher guides and offers pseudo labels to the student, whereas the student tries to imitate the teacher. Furthermore, we invoke the clustering view of OT as an intuitive tool to explain why TIDOT can mitigate the label shift problem. Last but not least, our OT-based imitation learning viewpoint together with its developed theory is potential to apply to a broader context such as adversarial machine learning (AML), generative models, and imitation learning in reinforcement learning.

3 Related Background

3.1 OT with Entropic Regularized Duality

Consider two distributions P and Q which operate on the domain Ω ⊆ Rd , let d (x, y) be a non-negative and continuous cost function or metric. Wasserstein distance [Santambrogio, 2015; Villani, 2008] between P and Q w.r.t the metric d is defined as

where γ is a coupling that admits Q, P as its marginals.

To enable the application of optimal transport in machine learning and deep learning, [Genevay et al., 2016] developed an entropic regularized dual form. First, they proposed to add an entropic regularization term to primal form (1)

where ϵ is the regularization rate, DKL (·‖·) is the KullbackLeibler (KL) divergence, and Q ⊗ P represents the specific coupling in which Q and P are independent.

Using Fenchel-Rockafellar theorem, they obtained the following entropic regularized dual form of (2)

3.2 Clustering View of Optimal Transport

This view of optimal transport has been utilized to study a rich class of hierarchical and multilevel clustering problems [Ho et al., 2019; 2017]. We now present the clustering view of optimal transport which assists us to interpret our method developed in the sequel. Let P and Q be two discrete distributions defined as

where δx indicates a Dirac measure centered at x.

The clustering view reveals that if we learn the atoms of Q to minimize Wd (P,Q) w.r.t the metric d, the optimal atoms of Q become the centroids of the clusters formed by the atoms of P or the atoms of Q are moving to find the groups of atoms of P with the aim to minimize the distortion w.r.t the metric d (see our supplementary material for more detail).

4 Our ProposedMethod

4.1 Optimal Transport Based Imitation Learning

In what follows, we present OT-based imitation learning which lays foundation for our proposed TIDOT. Consider two data domains XA and XB with two data distributions PA and PB respectively, we assume that hA : XA → Y△ (where Y△ := {π ∈ RM : ‖π‖1 = 1 and π ≥ 0} and M is the number of classes), which is a well-qualified classifier that gives accurate prediction for data instances on XA sampled from PA. We wish to learn a classifier hB to predict accurately data instances sampled from PB by imitating what is done by hA on (XA, PA).

To serve the development of OT-based imitation learning, given two pairs z1 = (x1, y△ 1 ) ∈ XS × Y△ and z2 = (x2, y△ 2 ) ∈ XT × Y△, we define the cost (distance) function between them as:

where dx is ground metric (cost) defined on XS ×XT and dy is a divergence defined on Y△.

Based on the data distribution PA and classifier hA, we define a distribution PA,hA over XA × Y△ including sample pair (x, hA (x)) by first sampling x ∼ PA and then computing hA (x). Similarly, we can define another distribution PB,hB over XB ×Y△ using the data distribution PB and the classifier hB. To allow hB to imitate the behavior of hA, we propose to inspect Wasserstein distance between PA,hA and PB,hB w.r.t the cost (metric) function d defined as above. The following proposition is crucial for us to speculate what it really means by OT-based imitation learning.

Proposition 1. The WS distance of interest Wd (PA,hA , PB,hB ) can be expressed as:

According to Proposition 1, when computing Wd (PA,hA , PB,hB ), we need to find an optimal transport K∗ : K∗#PB = PA that moves PB to PA so as to minimize the difference in predictions of hB and hA. This further implies that given x ∼ PB, the prediction behavior of hB(x) imitates that of hA (K∗ (x)) for K∗ (x) ∼ PA. Figure 1 provides an intuitive explanation for our proposed imitation learning viewpoint based on optimal transport.

In addition, from Proposition 1, it is obvious that

In the context of unsupervised domain adaptation, PB defined on a latent space via a feature extractor (generator) is hence an unfixed distribution, whereas hB is a trainable classifier. Therefore, from Inequalities (5) and (6), when minimizing Wd (PA,hA , PB,hB ), we minimize WdX (PA, PB) to reduce the data shift between two data distributions and simultaneously find the transport map K to allow hB (x) (target example x ∼ PB) imitating hA (K (x)) (K (x) ∼ PA) for mitigating the label shift between two domains.

4.2 OT-based Imitation Learning on Domain Adaptation

In what follows, we present how to apply our OT-based imitation learning mechanism to unsupervised domain adaptation (UDA). The UDA setting include two datasets, a labeled dataset DS = {(xSi , yS i )}NS i=1 from a source domain with xSi ∈ Rd and yS i ∈ {1, 2, ...,M} and an unlabeled dataset DT = {xTi }NT i=1 from a target domain. We denote PS and PT as the empirical data distributions for the source and target domains, i.e., PS = 1 NS ∑NS i=1 δxSi and PT = 1 NT ∑NT i=1 δxTi , where δx indicates a Dirac measure centered at x.

Following [Ganin and Lempitsky, 2015; Long et al., 2013; Long et al., 2015], we employ a generator (feature extractor) G to map both source and target examples into a latent space. On the latent space, we train a teacher classifier hS using the labeled source dataset DS and a student classifier hT using our proposed OT-based imitation mechanism to reduce the data and label shifts. More specifically, we propose minimizing the objective function consisting of the empirical loss of the teacher hS and an OT-based imitation learning term involving both the teacher hS and the student hT :

where α > 0 is a trade-off parameter and we have defined

for which ℓ is a loss function (e.g., the cross-entropy loss) and

for which PS,hS is the joint distribution constituted by the pairs (G(x) , hS (G(x))) where x ∼ PS and PT,hT is the joint distribution constituted by the pairs (G (x) , hT (G (x))) where x ∼ PT . Note that the ground metric (cost) d now involves the latent space and is defined as:

where z1 = (G(x1) , hS (G(x1))) and z2 = (G(x2) , hT (G(x2))) and dX (·, ·) is a distance between two data examples on the latent space.

Generally, we train a teacher classifier hS to predict well on the source domain, while using the OT-based imitation mechanism to train a student classifier hT to move target representations to source representations on the latent space (i.e., reducing data shift), whereas encouraging hT to mimic the predictions of hS (i.e., reducing label shift). In our experiments, the student classifier hT consistently outperforms the teacher classifiers. The reason is possibly that though target representations tend to move to source representations, there always exists a gap, hence the teacher classifier hS trained on source examples is hard to predict perfectly target examples, whereas the student classifier performs better because it is trained on target examples via imitating the predictions of hS on relevant source examples.

Clustering View Explanation. We now explain why minimizing RWS = Wd (PT,hT , PS,hS ) can help to mitigate the data shift and label shift, two thorny issues existing in UDA. This is intuitively explainable from the clustering view of the optimal transport (see Section 3.2). More specifically, let us denote zSi = (G (xSi ) , hS (G (xSi ))) , ∀i = 1, ...,NS and zTi = (G (xTi ) , hT (G (xTi ))) , ∀i = 1, ...,NT . It appears

Referred to the clustering view of OT, when minimizing RWS = Wd (PT,hT , PS,hS ), we encourage each zTi in the target domain to find an appropriate group (cluster) of zSj with j ∈ J for some set of indices J ⊂ {1, ..., NS} so that the total distortion ∑j∈J d (zTi , zSj ) defined as

is minimized. This further implies that (i) G (xTi ) moves toward the group or cluster of G (xSj ) , j ∈ J and (ii) the predictions of hS for the source data examples in that group or cluster (i.e., hS (G (xSj )) , j ∈ J) need to be a consensus. The first conclusion (i) is straight-forward from minimizing dx (G (xTi ) ,G (xSj )), whilst the second conclusion (ii) comes from the fact that the prediction hT (G (xTi )) mimics those of hS (G (xSj )) for all j ∈ J, hence hS (G (xSj )) , ∀j ∈ J should reach a consensus on their predictions. Eventually, each G (xTi ) is encouraged to move to a group or cluster of G (xSj ) , j ∈ J in the source domain which shares the same prediction label to imitate their common prediction. This would help to mitigate the label shift issue (see Figure 2).

4.3 Entropic Regularized Solution

To tackle the OT-based regularization term RWS , we use entropic regularized duality form (see Eq. (3)) of optimal transport. Specifically, we approximate RWS = Wϵd (PT,hT , PS,hS ) which has the following form:

where d (zTi , zSj ) = λdx (G (xTi ) ,G (xSj )) + dy (hT (G (xTi )) , hS (G (xSj ))) is the transportation cost, φ is a neural net named Kantorovich potential network (see Eq. (3)).

4.4 Ensuring Clustering Assumption for TIDOT

Clustering assumption [Chapelle and Zien, 2005] is a technique that encourages the classifier to preserve its predictions for data examples in a cluster. Basically, the clustering assumption enforces the decision boundary of a given classifier to lie in the gap among the data clusters and never crosses over any clusters. We observe that coupling the clustering assumption for the classifiers hS, hT with TIDOT helps to boost its performance.

The reason for this complementary collaboration is that although minimizing the OT-based regularization term RWS helps to move the target example G (xTi ) to a group or cluster of the source examples G (xSj ) with the same label, the prediction of hT (G (xTi )) as in Eq. (9) is encouraged to mimic the predictions of hS (G (xSj )) including diverge data examples of different classes. Therefore, the prediction of hT (G (xTi )) for those lying on the cluster boundary tends to be possibly less confident and misleading.

With the assistance of the clustering assumption, the classifier hT is strengthened to predict well the target examples lying on the clustering boundary. Specifically, this encourages the classifier hT to predict those target examples using the same label as others in the cluster, hence correcting the predictions for those examples. To enforce the clustering assumption, we employ Virtual Adversarial Training (VAT) [Miyato et al., 2019] in conjunction with minimizing the entropy of prediction [Grandvalet and Bengio, 2005] as in [Shu et al., 2018; Kumar et al., 2018]

where with H to be the entropy, we have defined

with which DKL represents a Kullback-Leibler divergence and θ is a very small positive number.

4.5 Teacher Imitation Domain Adaptation Based on Optimal Transport

The final optimization problem of our TIDOT is as follows:

where β > 0 is a trade-off parameter. Under the clustering view of OT, it is worth noting that when minimizing RWS = Wϵd (PT,hT , PS,hS ) in Eq. (10), we aim to push the representations of source and target data to be intermingled in the joint space and encourage each G (xTi ) to find its corresponding G (xSj ) to mimic the prediction of hS. Finally, the pseudocode for the training process of TIDOT is presented in Algorithm 1 which is placed in our supplementary material due to the space limitation.

5 Experiments

5.1 Model Evaluation

In this section, we conduct experiments on four main datasets to evaluate our TIDOT with state-of-the-art domain adaptation methods: (1) ResNet-50 [He et al., 2016]; (2) DANN [Ganin and Lempitsky, 2015]; (3) P-model [French et al., 2018]; (4) iCAN [Zhang et al., 2018]; (5) CDAN [Long et al., 2018]; (6) SWD [Lee et al., 2019]; (7) DeepJDOT [Damodaran et al., 2018]; (8) DASPOT [Xie et al., 2019]; (9) ETD [Li et al., 2020]; (10) RWOT [Xu et al., 2020].

Digits, Traffic Sign, and Natural Scenes Datasets include MNIST (MN), USPS (US), MNIST-M (MM), Synthetic Digits (SN), Street View House Numbers (SV), Synthetic Traffic Signs (SS), German Traffic Signs Recognition Benchmark (GT), CIFAR-10 (CI), and STL-10 (ST).

Office-31 contains 3 domains Amazon (A), Webcam (W), and DSLR (D). There are 31 common classes for all domains and the total number of images is 4,110.

Office-Home consists of roughly 15,500 images in a total of 65 object classes and belonging to 4 different domains: Artistic (Ar), Clip Art (Cl), Product (Pr) and Real-world (Rw). ImageCLEF-DA contains three domains: Caltech-256 (C), ImageNet ILSVRC 2012 (I) and Pascal VOC 2012 (P). There are total 600 images in each domain and 12 common classes.

5.2 Implementation Detail

Architecture. We employ small, medium and large network architectures whose detail in the supplementary material. To compare with baselines on Office-Home and Office-31, all transfer tasks use the pre-trained ResNet-50 [He et al., 2016] features which have 2,048 dimensions.

Hyperparameters. The specifications of hyperparameters are described in the supplementary material.

5.3 Result and Discussion

We first evaluate TIDOT on Digits, traffic sign, and natural scene datasets and report the results in Table 1. The experimental results show that whilst TIDOT teacher outperforms almost all state-of-the-art baselines, TIDOT student further enhances the accuracy of TIDOT teacher with various degrees of improvement (i.e., from as marginal as 0.1% to as significant as 4.6%). It is noticeable that although the transfer task MN→SV is extremely challenging in which the source dataset includes grayscale handwritten digits whereas the target dataset is created by real-world digits, our TIDOT is still capable of mitigating the shift of data and label between domains and outperforms the second-best method by a sizeable margin (15.4%).

We further testify TIDOT's performance on Office-31 and report the classification results in Table 2. In general, our model achieves 94.1% on average and significantly outperforms on challenging adaptation where the source and target images are dissimilar in the background, i.e., D→A, W→A.

The results on Office-Home are reported in Table 3. TIDOT student exceeds almost comparison methods and achieves state-of-the-art performance, experiencing a go up by 3.3% on average compared with RWOT. More specifically, our model sees a remarkable improvement on challenging adaptation tasks, namely Ar→Pr, Cl→Pr.

The full comparison table for Office-Home and results on ImageCLEF-DA are shown in the supplementary material due to the limit of space.

5.4 Ablation Study

An Intuitive OT-based Loss. We further plot the values of RWS = Wϵd (PT,hT , PS,hS ) when simultaneously training teacher and student on the pair MN→SV. As shown in Figure 3, the cost of the optimal transport plan for moving from PT,hS to PS,hS is minimized, which means the data shift and label shift tends to be mitigated during the training process. Another remarkable advantage of OT-based methods as our TIDOT is that due to the effect of the envelope theorem, RWS smoothly decreases while losses developed based on generative adversarial network (GAN) [Goodfellow et al., 2014] always sees largely fluctuates.

Effect of Clustering Assumption. We investigate the effectiveness of the VAT loss (Lvat ) w.r.t. source and target distribution, and conditional entropy loss w.r.t. target distribution (Lent ) on ImageCLEF-DA. The results in Table 4 show that by adding VAT loss (fourth row), the model sees a rise averagely by 1.4% compared to the basic setting (second row). Moreover, the figures are better when Lvat is combined with Lent (fifth row). Via this ablation study, we find that VAT in conjunction with minimizing entropy supports our TIDOT to predict well on target samples lying on the decision boundary, and hence boots model performance further.

Parameter Sensitivity. We further evaluate the effects of the trade-off parameters α, β in Figure 4. We search α, β in the grid of {0.005, 0.01, 0.05, 0.1, 0.5, 1.0} and report the test accuracy on transfer tasks A→W and I→P. The results show that the model yields high performances when α from 0.005 to 0.1 and β from 0.05 to 1.0. However, with the other values of α and β, our model still achieves significant performances, which demonstrates the robustness and flexibility of TIDOT.

6 Conclusion

In this paper, we leverage the perspective of imitation learning and the theory of optimal transport to propose Teacher Imitation Domain Adaptation with Optimal Transport (TIDOT). Via two fundamental components of TIDOT, a teacher and a student, we apply our proposed method to unsupervised domain adaptation and conduct comprehensive experiments to compare TIDOT against the baselines. The experimental results show that our TIDOT outperforms the existing state-ofthe-art OT-based method. Additionally, as a side effect of our developed theory, we interestingly discover a novel regularization technique for deep networks based on optimal transport, which is potential for future work.