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: In this work, the authors study the problem of domain generalization. Recent works such as invariant risk minimization [Arjovsky et al.] have invigorated a lot of interest in the problem of domain generalization. The paper brings together ideas from rate distortion and information bottleneck principle (recently used in Ahuja et al. for domain generalization) and the work of model based domain generalization [Robey et al.].  In Ahuja et al., the authors combine invariance and information bottleneck to address domain generalization. The representation learned in Ahuja et al. is constrained by the label and information bottleneck constraints. In this work, the authors add some more structure to the generation of the observations X from the latent and say that there is a common instance X that is transformed by the domain transformation model to get the image in the current domain, i.e. X^e = G(X,e). The authors try to recover the common instance X that is shared across domains.