A Rate-Distortion Approach to Domain Generalization

Domain generalization deals with the difference in the distribution between the training and testing datasets , i.e. , the domain shift problem . A principled approach to domain generalization is by extracting domain-invariant features . In this paper , we propose an information-theoretic approach for domain generalization . We first establish the domain transformation model , mapping a domain-free latent image into a domain . Then , we cast the domain generalization as a rate-distortion problem , and use the information bottleneck penalty to measure how well the domain-free latent image is reconstructed from a compressed representation of a domain-specific image compared to its direct prediction from the domain-specific image itself . We prove that the information bottleneck penalty guarantees that domain-invariant features can be learned . Lastly , we draw links of our proposed method with selfsupervised contrastive learning without negative data pairs . Our empirical study on two different tasks verifies the improvement over recent baselines . 1 INTRODUCTION . Deep neural networks ( DNNs ) are highly expressive models that reach state-of-the-art performance in challenging tasks , such as speech and visual recognition ( Devlin et al. , 2018 ; He et al. , 2016 ) , by capturing complex correlations among input elements , e.g. , pixels of an image . However , the correlations might also contain spurious features that hurt the generalization performance of DNNs on out-of-distribution samples ( Szegedy et al. , 2013 ; Beery et al. , 2018 ; Alcorn et al. , 2019 ) . Unfortunately , real-world applications often encounter such out-of-distribution samples , e.g. , when the training domain does not match the testing domain . A prominent example is deblurring , where models are trained on simulated blurs which differ substantially to real-world blurring ( Koh et al. , 2021 ) . In other words , generalization across domains is a critical task before deploying DNNs to real-world application . Learning features that are invariant across multiple training domains , and using those features for out-of-distribution generalization has emerged as a significant topic in domain generalization . In domain generalization , multiple source domains are accessible during training , but the target domains are not ( Blanchard et al. , 2011 ; Muandet et al. , 2013 ) . Invariant risk minimization ( IRM ) ( Arjovsky et al. , 2019 ) is a prominent approach for learning domain invariant features . However , IRM suffers from the case when the invariant features contains full information about the label ( Ahuja et al. , 2020 ) . To deal with this shortcoming , Ahuja et al . ( 2021 ) introduce the information bottleneck theory on neural networks ( Tishby & Zaslavsky , 2015 ) , and show that their method will be guaranteed to converge to the invariant features . On the empirical side , a series of works align source domain distributions for domain-invariant representation learning by either direct construct auxiliary penalty ( Duan et al. , 2012 ; Sun & Saenko , 2016 ; Li et al. , 2018b ; c ; 2017 ; Niu et al. , 2015 ) , or meta learning ( Li et al. , 2019 ; Balaji et al. , 2018 ; Li et al. , 2018a ) . There are also series of work do not rely on invariant features . They can be categorized as ( 1 ) domainspecific method : Domain2Vec ( D2V ) ( Deshmukh et al. , 2018 ) learns domain-specific embedding , DMG ( Chattopadhyay et al. , 2020 ) aims to learn domain specific masks ; and ( 2 ) augmentation method : ( Volpi et al. , 2018 ) augments the dataset adversarially , L2A-OT ( Zhou et al. , 2020 ) augments data with image information . Despite their success , there is no guarantee that empirical methods can solve the task across different environments . In this paper , we use an information-theoretic approach to tackle domain generalization . We assume there is a domain-free latent instance ( e.g. , an image ) that captures the invariant features we want to extract . We define a domain transformation model that maps the domain-free latent instance into a domain and then we apply the rate distortion theory to obtain a domain-invariant representation . The proposed method , called Twins , is guaranteed to converge to the invariant feature under the linear classification structural equation model ( Ahuja et al. , 2021 ) . We evaluate our method on linear unit tests ( Aubin et al. , 2021 ) and variants of MNIST dataset ( LeCun & Cortes , 2010 ; Xiao et al. , 2017 ; Clanuwat et al. , 2018 ) , which validates the theoretical analysis and demonstrates how the proposed method can outperform the previous ones . Our contributions can be summarized as follows : • We cast domain generalization as a rate distortion problem and prove how the proposed method can converge . • We illustrate how the proposed method extends previous results on domain generalization , and draw links to self-supervised contrastive learning . We demystify the success of contrastive learning by giving a contrastive learning based domain generalization algorithm with theoretical guarantee . • We evaluate our method on two datasets and observe consistent improvement over existing baselines . 2 PRELIMINARY ON DOMAIN GENERALIZATION . Assume that the instance-label pair ( X , Y ) is sampled from an unknown distribution P ( X , Y ) . The objective of standard supervised learning is to learn a predictor f that is able to predict the labels Y of corresponding instances X for each ( X , Y ) ∼ P ( X , Y ) , given the finite training samples drawn from the underlying distribution P ( X , Y ) . Unlike the standard supervised learning tasks , in domain generalization , we can not sample directly from the distribution P ( X , Y ) . Instead , we can only observe ( X , Y ) under different domains e ∈ Eall , denoted as ( Xe , Y e ) ∼ Pe ( Xe , Y e ) . We also assume that e ∈ Eall is distributed as e ∼ Pe . Given samples from a finite subset Etrain ( Eall of all the domains , the goal of the domain generalization problem is to learn a predictor f that generalizes across all possible domains . This can be summarized as follows : Problem 2.1 ( Domain generalization ) . Let Etrain ( Eall be a finite subset of training domains . We have access to the data for each training domain etrain ∈ Etrain , but have no access to the data for each test domain etest ∈ Eall\Etrain . Given a function classF and a loss function ` , our goal is to learn a predictor f ∈ F using the data from the training domain such that f minimizes the worst-case risk over Eall . Define the risk of the predictor f on the domain e as Re ( f ) : = EPe ( Xe , Y e ) ` ( f ( Xe ) , Y e ) . We want to solve the following min-max optimization problem : minimize f∈F max e∈Eall Re ( f ) . ( DG ) We establish the domain transformation model to characterize the relation between domain-aware instance Xe and the domain-invariant latent instance X in the Assumption 1 , which first appears in Robey et al . ( 2021 ) . Assumption 1 ( Domain transformation model ) . Let δe denote a Dirac distribution for e ∈ Eall . We assume that there exists a measurable function G : X × Eall → X , which we refer to as a domain transformation model , that parameterizes the inter-domain covariate shift via Pe ( X ) =d G # ( P ( X ) × δe ) ∀e ∈ Eall , ( 1 ) where # denotes the push-forward measure and =d denotes equality in distribution . The Assumption 1 can somewhat reflect the generation of domain specific instances . For example , the multiple different views of a 3D object ( Niu et al. , 2015 ) , different angles of the image ( Rotated MNIST ( Worrall et al. , 2017 ) ) . Besides , the MUNIT architecture ( Huang et al. , 2018 ) can effectively distangle the domain-free latent instance X and the specific environment e , and thus can be used as the domain transformation model G ( Robey et al. , 2021 ) . Let Φ denote the feature representation mapping , w denote the classifier and w ◦ Φ denote the full predictor . The regret of the network on the domain e is denoted as Re ( w ◦ Φ ) . Next , we define standard properties related to the datasets used in the domain generalization literature ( Ahuja et al. , 2021 ) . For each e ∈ Eall , the distribution ( Xe , Y e ) ∼ Pe satisfies the following properties : ( 1 ) ∃ a map Φ∗ , which we call an invariant feature map , such that E [ Y e ∣∣Φ∗ ( Xe ) ] is the same for all e ∈ Eall and Y e 6⊥ Φ∗ ( Xe ) , where ⊥ means mutual independence . ( 2 ) ∃ a map Ψ∗ , which we call spurious feature map , such that E [ Y e ∣∣Ψ∗ ( Xe ) ] is not the same for all e ∈ Eall and Y e 6⊥ Ψ∗ ( Xe ) for some domains . Ψ∗ often hinders learning predictors that only rely on Φ∗ . For example , in the CMNIST dataset , the Φ ? extracts the underlying digit and Ψ ? extracts background color . The baseline algorithm for domain generalization Equation ( DG ) is the Empirical Risk Minimization , i.e . directly minimizing the empirical risk on the training domains : min w , Φ 1 |Etrain| ∑ e∈Etrain Re ( w ◦ Φ ) , ( 2 ) where |Etrain| denotes the number of training domains . We say that a data representation Φ elicits an invariant predictor across the set of training domains Etrain if there is a predictor w that simultaneously achieves the minimum risk , i.e . w ∈ arg minw′ Re ( w′ ◦ Φ ) , ∀e ∈ Etrain . Using this notation , the main objective of Invariant Risk Minimization ( IRM ) is stated as : min w , Φ 1 |Etrain| ∑ e∈Etrain Re ( w ◦ Φ ) , s.t . w ∈ arg min w′ Re ( w′ ◦ Φ ) , ∀e ∈ Etrain . ( 3 ) Lastly , we rely on the notion of ‘ informativeness ’ about the datasets ( Ahuja et al. , 2021 ) . There are two such categories of informativeness . In the first case , the invariant features Φ∗ ( Xe ) are partially informative about the label , i.e . Y 6⊥ Xe|Φ∗ ( Xe ) , and color contains information about label not contained in the uncolored digit . In the second case , invariant features are fully informative about the label , i.e. , Y ⊥ Xe|Φ∗ ( Xe ) , i.e. , they contain all the information about the label that is contained in input Xe . Many real-world image datasets have fully informative invariant features , the labels are a deterministic function of the domain-invariant features and domain-aware spurious features do not affect the label . 3 METHOD . 3.1 RATE DISTORTION & INFORMATION BOTTLENECK PRINCIPLE . Given a domain-free latent instance X ∈ RdX , its observation in a domain e is denoted as Xe : = G ( X , e ) . We want to learn the feature Ze = Φ ( Xe ) ∈ RdZ which is informative about the domainfree variable X , but invariant ( i.e . uninformative ) to the specific domain e. We use rate–distortion theory ( Davisson , 1972 ; Blau & Michaeli , 2019 ) to formulate our domain generalization problem . Rate–distortion theory is a major branch of information theory which provides the theoretical foundations for lossy data compression . An encoder Φ encodes domain-aware instances Xe . We want the representation Ze = Φ ( Xe ) to be domain-invariant , so we feed Ze into a decoder which outputs domain-invariant X . We minimize the distortion between the original domain-aware instance Xe and the reconstructed domain-free instance X . The distortion function measures how well X is predicted from a compressed representation Ze compared to its direct prediction from Xe . This trade-off is captured by the following loss function : LIB ( θ , e ) = EX∼PX , e∼PeI ( Ze ; Xe ) − βI ( Ze ; X ) IB objective ( 4 ) where I denotes the mutual information , θ is the parameter of the representation function Φ , and β is a constant . In the following , we consider two cases : discrete and continuous variables , owing to their different definition of entropy . Discrete case : Since the representation function is deterministic with respect to θ , we can rewrite Equation ( 4 ) through a classical identity of mutual information : I ( X ; Y ) = H ( X ) − H ( X|Y ) , where H denotes the Shannon entropy for discrete variables , as follows : LIB ( θ , e ) = EX , eI ( Ze ; Xe ) − βI ( Ze ; X ) = EX , eH ( Ze ) −H ( Ze|Xe ) − β ( H ( Ze ) −H ( Ze|X ) ) = EX , eH ( Ze|X ) + 1− β β H ( Ze ) , ( 5 ) where in the last equality we omit the overall scaling factor of the loss function . If 0 ≤ β ≤ 1 , since H ( · ) is bounded below by 0 , setting Φ to be constant will clearly minimize the penalty , which is uninformative about the representations we want to learn . Hence , we set β > 1 , and replace 1−ββ with −λ , where 0 ≤ λ < 1 . The IB objective can be rewritten as LIB ( θ , e ) = EX , eH ( Ze|X ) − λH ( Ze ) . ( 6 ) Continuous case : In terms of continuous variables , the differential entropy h ( · ) is not bounded below , which hinders our analysis . To overcome this , we can define the lower bounded differential entropy ĥ ( X ) : = h ( X + ε ) , where ε is the independent bounded zero-entropy noise ε ∼ Uniform ( 0 , 1 ) . Thus , ĥ ( X ) ≥ h ( ε ) = 0 . We can replace the Shannon entropy H ( · ) with lower bounded differential entropy ĥ ( · ) in Equation ( 6 ) : LIB ( θ , e ) = EX , eĥ ( Ze|X ) − λĥ ( Ze ) . ( 7 ) For simplicity , we define H to be the Shannon entropy for the discrete variables , or lower bounded entropy ĥ for continuous variables in the main text . We define He ( f ) : = EXe∼PeH ( f ( Xe ) ) . We can extend the Empirical Risk Minimization ( ERM ) algorithm to include the IB Penalty , and the resulting algorithm , denoted as Twins-ERM method , is the following : min w , Φ ∑ e∈Etrain He ( Φ ∣∣X ) − λHe ( Φ ) s.t . 1 |Etrain| ∑ e∈Etrain Re ( w ◦ Φ ) ≤ r ( 8 ) where r is the threshold on the empirical risk on the training domains . In addition to ERM , another popular minimization framework is the invariant risk minimization ( Arjovsky et al. , 2019 ) . The proposed penalty can be readily incorporated into the IRM framework , we call the resulting algorithm Twins-IRM : min w , Φ ∑ e∈Etrain He ( Φ ∣∣X ) − λHe ( Φ ) , s.t . 1 |Etrain| ∑ e∈Etrain Re ( w ◦ Φ ) ≤ r , w ∈ arg min w̃ Re ( w̃ ◦ Φ ) , ∀e ∈ Etrain . ( 9 )

Summary: The paper proposes a method for domain generalization. They formulate the problem as a rate-distortion problem and introduce the information bottleneck penalty. Specifically, they realize the objective as a contrastive loss that encourges embeddings across domains to be similar if they share the same label. In the experiments, the method improves the corresponding baselines in several toy datasets.