Adversarial Teacher-Student Representation Learning for Domain Generalization

1 Introduction . Deep neural networks have achieved promising performance on a wide variety of tasks . However , these networks assume the training and testing samples fall in the same data distribution . Such a strong assumption would limit the applicability of the learned models in real-world scenarios ( e.g. , cross-city autonomous driving or multi-cite medical imaging task ) , in which training and testing data are typically observed under different conditions . In other words , the generalizability of the model at unseen target domains might be poor due to unexpected domain shifts . To tackle the domain discrepancy , domain generalization ( DG ) has been proposed and drawn increasing attention recently . The aim of DG is to train models using data observed from single or multiple source domains , while expecting that the model would be generalized to unseen target domains . Most existing DG approaches focus on deriving domain-invariant features [ 1 ] among multiple source domains or adopting meta-learning techniques [ 2 , 3 , 4 , 5 ] , which would simulate domain shifts during the meta-training stage . However , the features derived by the above methods are generally guaranteed to be invariant to the seen source domains , not the generalizability of the learned representation to describe unseen domain data . To overcome the limitation , [ 6 , 7 , 8 , 9 ] turn to leverage data generation techniques for diversifying the source distributions , and thus avoid overfitting on source domains yet improve the generalization ability of models . Specifically , several works [ 6 , 7 , 8 ] choose to generate novel-domain images by either perturbing the style of source data to confuse the domain classifier [ 6 , 7 ] , or transporting the source data to novel styles via optimal transport based objective [ 8 ] . [ 9 ] adopts Mixup [ 10 ] to interpolate the feature statistics between samples from different domains . While the above methods perform well , designing an objective for generating samples with DG guarantees remains a challenging and open problem . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Recently , self-supervised pre-training manifests the potential to derive generalizable representation , which serves as a promising start point for downstream tasks ( e.g. , image segmentation or object detection ) . In domain generalization , a number of self-supervision techniques have been introduced [ 11 , 12 , 13 ] to improve network transferability by discovering the intrinsic properties within images . For instance , [ 11 , 12 , 13 ] adopt jigsaw puzzles as the pretext task , which predicts the relative positions of image patches to constrain the semantic feature learning in a multi-task training fashion . Recently , contrastive learning approaches [ 14 , 15 , 16 , 17 , 18 ] have been proposed and widely applied , which establish the representation learning from multiple views of an image to extract the task-relevant information and discard task-irrelevant noise . However , the concept of such multiview learning [ 19 , 18 ] is simply realized by hand-crafted image transformations ( e.g. , RandomResizedCrop , Color Jittering , or Gaussian Blur ) . The effectiveness of these hand-crafted image transformations for benefiting the generalization to unseen distributions is still not guaranteed . In this paper , we propose a unique Adversarial Teacher-Student Representation Learning framework for tackling domain generalized visual classification . Based on the recent success of contrastive learning , we advance the concept of multi-view learning into DG regime for augmenting source instances to out-of-source styles and diversifying training distributions . To be more precise , with the goal of learning representations which are robust to unseen domain shift , we propose to jointly perform Domain Generalized Representation Learning and Novel Domain Augmentation in an adversarial learning manner . Based on Teacher-Student learning schemes [ 20 , 16 , 21 ] , our framework utilizes original images as inputs to the teacher network and takes stylized augmentations as input to the student network . To ensure both learning stages produce domain generalized representation , we adopt the Teacher-Student co-training scheme , which progressively refines Teacher by the distilled knowledge learned from Student by observing augmented novel-domain data , enabling Teacher to be generalizable to data with out-of-source distributions . On the other hand , Adversarial Novel Domain Augmentation aims at augmenting unseen domain data using source-domain training instances . The objective is to maximize the discrepancy between the input and augmented data , derived from the teacher and student modules , respectively . In order to have such augmented data exhibit sufficient domain differences , the above discrepancy will be calculated using features derived from data across different source domains . By iteratively training the above two stages in an adversarial learning fashion , the resulting model ( Teacher ) would be able to derive domain generalizable representations . The contributions of this paper are highlighted as below : • Different from existing meta-learning based approaches , we propose a teacher-student adversarial learning scheme for addressing domain generalization classification problems . • In the stage of Domain Generalized Representation Learning , the student network observes augmented novel-domain data and distills the information to update the teacher network , allowing derivation of domain generalizable representation . • In the stage of Novel Domain Augmentation , the generator aims at producing unseen yet plausible domain data , which maximizes the discrepancy between augmented and existing domains while the semantic information is preserved . • Evaluations on several benchmark datasets in multiple and single source domain settings verify that our method performs favorably against existing DG approaches and exhibits sufficient domain generalization capability . 2 Related Works . Domain Generalization ( DG ) Different from domain adaptation ( DA ) , which observes both source and target-domain training data for performing learning tasks across domains [ 22 , 23 , 24 , 25 , 26 ] , DG deals with a more realistic yet challenging setting . More precisely , DG aims at generalizing the model trained only on single or multiple source domains to recognize the test instance in unseen but similar target domain . With only source-domain data observed during training , a number of works [ 3 , 2 , 27 ] apply meta-learning for learning domain-invariant features . These methods typically partition source domains into meta-train and meta-test splits to simulate the domain shifts during training . Feature-Critic [ 27 ] meta-learns a critic network to evaluate the generalized degree of extracted features for encouraging robust feature derivation . [ 4 ] introduces an episodic training that cross-trains domain-specific feature extractors and classifiers to let the learned model invariant to the domain shift . MLDG [ 2 ] and MASF [ 5 ] both adopt gradient based meta-learning to simulate the domain shift , while [ 5 ] additionally enforces local and global constraints in meta-training . In addition to meta-learning approaches , [ 11 ] jointly solves jigsaw puzzle as an auxiliary task with standard classification in a multi-task fashion . RSC [ 28 ] iteratively discards the dominant features on the training data to improve generalization . Nevertheless , these approaches employ solely limited source domains to derive generalizable features , which still draws a concern about over-fitting on source domains [ 8 , 29 ] and restricts the generalization ability to unseen domains . Recent research works consider data generation as an alternative technique for domain generalization , which increase the diversity of training data distribution . To achieve this goal , [ 6 , 7 ] are inspired by adversarial attack [ 30 ] . CrossGrad [ 6 ] perturbs source data by adding adversarial gradients ; DDAIG [ 7 ] learns a transformation network that outputs perturbations to confuse the domain classifier . However , such perturbed images do not necessarily exhibit sufficient data domain diversity . In contrast of adding perturbation to images , L2A-OT [ 8 ] learns a conditional generator that transforms images from a source distribution to a pseudo-novel distribution by an optimal transport based objective . MixStyle [ 9 ] produces image features with mixed feature statistics across source domains . Very recently , PDEN [ 31 ] utilizes a progressive learning strategy for single-source domain generalization , which iteratively expands the training data set by adding augmented data . Note that although they adopt contrastive and adversarial learning objectives which are similar to ours , our proposed approach is able to tackle both multi-source and single-source DG problems , and also comes with superior memory efficient performance and comparably stable training process . Self-Supervised Learning ( SSL ) Self-supervision is a recent paradigm for unsupervised learning . The idea is to design pretext tasks for feature learning to facilitate the downstream task learning . Such auxiliary pretext tasks can be predictions of the image colors [ 32 ] , relative locations of patches from the same image [ 11 , 12 ] , and image rotation [ 33 ] . Very recently , contrastive learning [ 14 , 15 , 16 , 17 , 18 ] has achieved promising results on network pre-training to learn generalized image features . [ 19 , 18 ] reveal that the success of contrastive learning is typically built on the multi-view perspective , and prove theoretically and experimentally that the compact and robust representations can be learned by deriving the invariance among multiple views of an image . We adapt the above concept of multi-view learning into DG regime . We focus on learning novel-domain data augmentations across source-domain instances in an adversarial training fashion . As detailed and verified later , our proposed learning scheme would produce domain generalizable representation for unseen target-domain data , and performs favorably against state-of-the-art DG approaches . 3 Proposed Method . 3.1 Problem Formulation and Model Overview . For the sake of completeness , we first define the problem setting and notations used in this paper . We assume that training data are observed from N source domains Dtr = { D1 , D2 , ... , DN } , each of which contains a set of image and label pairsDi = { Xi , Yi } . Our goal is to learn a model which would exhibit sufficient generalization capability , so that classification of test data in unseen target domains can be performed . In order to derive domain-generalized feature representations , we present a novel Adversarial Teacher-Student Representation Learning framework , which is a min-max deep learning framework alternating between the following two stages : domain generalized representation learning ( Sec . 3.2 ) and novel domain augmentation ( Sec . 3.3 ) , as depicted in Fig . 1 . For domain generalized representation learning , we learn a domain-generalized teacher network ( Teacher ) FT with the help from a student network ( Student ) FS , which observes synthesized novel-domain augmentation and distills knowledge to Teacher . As for novel-domain augmentation , the novel-domain augmenter G is learned to observe the discrepancy of Teacher-Student encoders , which progressively generates strong novel stylized augmentations to diversify training distributions . Once the learning of the above framework is complete , the teacher network would extract domain-generalized features for the task network ( e.g. , classifier ) , and thus classification of unseen target-domain data can be performed accordingly . We now detail our proposed learning schemes in the following subsections . 3.2 Teacher-Student Domain Generalized Representation Learning . While techniques based on learning across multiple source domains for DG exist ( e.g. , using metalearning techniques like [ 3 , 4 , 5 ] ) , it is not clear how the learned model and feature representations would be generalized to unseen target domains . Instead of directly fitting models across source domains , we propose Domain Generalized Representation Learning based on the Teacher-Student learning scheme , with the goal of extracting domain generalizable feature representations . To ensure our teacher encoder to gain generalizability by observing out-of-source domain information , we deploy a Student FS for exploring novel-domain augmentation synthesized from the novel-domain augmenter G , while distilling the associated knowledge to update FT . To address this representation learning task , we first train the teacher module together with a singlelayer classifier C using multiple source-domain data . The standard cross-entropy loss Lce is utilized to initialize FT as warm-up . As illustrated in Fig . 1 , we then input training images x sampled from the source domains into the novel-domain augmenterG ( detailed in the following sub-section ) , producing the style ( or domain ) perturbed augmentation x̃ yet preserving its semantic information . While such a domain augmented x̃ would be fed into the student module resulting in feature z̃ = FS ( x̃ ) , we also feed the original input x into Teacher to derive z = FT ( x ) . To ensure that z̃ would contain the same semantic information as z does , we particularly propose an objective to minimize the discrepancy between z̃ and z . To be more specific , we define the discrepancy loss LFdis to minimize the distance between the normalized features z̃ and z : min FS LFdis ( z , z̃ ) = ∥∥∥∥ z‖z‖2 − z̃‖z̃‖2 ∥∥∥∥2 2 = ∥∥∥∥ FT ( x ) ‖FT ( x ) ) ‖2 − FS ( x̃ ) ‖FS ( x̃ ) ‖2 ∥∥∥∥2 2 . ( 1 ) In addition , we calculate the cross-entropy loss on the domain-augmented feature z̃ , i.e. , Lce ( C ( z̃ ) , y ) , which further enforces the classification capability of the student module ( note that C indicates the single-layer classifier , and y denotes the corresponding class label ) . We note that , in this representation learning stage , only the student network FS is updated by the above two objectives LFdis ( z , z̃ ) and Lce ( C ( z̃ , y ) , and we apply a stop-gradient strategy to forbid FT and G from being updated by gradients . Thus , optimization of FS with learning rate γ can be expressed as follows : θS ← θS − γ ∂ ( LFdis ( z , z̃ ) + Lce ( C ( z̃ ) , y ) ) ∂θS . ( 2 ) As for the teacher network FT , we adopt exponential moving average ( EMA ) [ 20 , 16 , 21 ] to progressively refine the associated model parameter θT . That is , the learned knowledge from Student ’ s parameter θS is distilled to update θT as follows , θT ← τθT + ( 1− τ ) θS , where τ ∈ [ 0 , 1 ) , ( 3 ) Note that τ controls the updates on the teacher network . Finally , it is also worth pointing out that , such a refinement strategy would avoid the teacher module from directly observing unrealistic domain augmentations , which might degrade its domain generalization capability .

Summary: This paper proposes a new data augmentation method for domain generalization. The main idea is to train a neural network-based generator to generate novel images that can be correctly classified by a student model and in the meantime, make the feature discrepancy larger between the student model and the EMA teacher model. The student model, which is fed with the generated images, is trained to minimize both the classification error and the feature discrepancy with its teacher. Only the EMA teacher is deployed at test time. This method is evaluated on both multi- and single-source domain generalization settings, outperforming previous SOTAs by a clear margin.