Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARY ON DOMAIN GENERALIZATION
text: 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 cannot 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.

heading: 3 METHOD

heading: 3.1 RATE DISTORTION & INFORMATION BOTTLENECK PRINCIPLE
text: 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)

heading: 3.2 THEORETICAL GUARANTEE
text: In this subsection, we establish the theoretical guarantee of our algorithm under the linear classification case. We consider the following standard 0-1 classification model in literature (Ahuja et al., 2020; 2021):
Assumption 2. Linear classification structural equation model. In each e ∈ Eall,
Y e ← I(w?inv ·Xinv)⊕Ne, Ne ∼ Bernoulli(q), q ≤ 1
2 , Ne ⊥ (Xinv, Xespu)
Xe ← S(Xeinv, Xespu), Xeinv ← G(Xinv, e) (10)
where I(x) is 1 if x is positive else 0, w?inv ∈ Rm with ‖w?inv‖ = 1 is the labelling hyperplane, Xeinv, X e inv ∈ Rm, Xespu ∈ Ro, S ∈ R(m+o)×(m+o) and G is a continuous domain transformation.
Before presenting our main theorem, we first add two assumptions on the support of invariant features. Define the support of the invariant features Xeinv in environment e as X einv.
Assumption 3 (Invariant feature support overlap). The union of support of the invariant features of the training domains covers support of the invariant features of all the domains. i.e. ⋃ e∈Eall X e inv ⊆⋃
e∈Etrain X e inv.
Assumption 4 (Strictly separable invariant features). The training support of invariant features ⋃ e∈Etrain X e inv is strictly separated by the labelling hyperplane w ∗ inv. In other words,
minx∈ ⋃
e∈Etrain X einv sign(w
∗ inv · x) · (w∗inv · x) > 0.
Assumption 3 and 4 describes the property of the support of invariant feature. Under these assumptions, we propose our first main theorem: Theorem 3.1. Suppose each e ∈ Eall follows Assumption 2, and Assumptions 3 and 4 hold for the invariant features. Also, for each e ∈ Etrain, assume that Xespu = AXeinv +W e, where A ∈ Ro×m, W e ∈ Ro is continuous, bounded, and zero mean noise. Each solution to Twins-ERM and Twins-IRM (Equation (8) and Equation (9), with ` as 0-1 loss, and r = q) solves the domain generalization problem (Equation (DG)).
Sketch of Proof. The full proof is provided at Appendix A. We only present the main idea here. Denote Φ† as the solution to Equation (9),
Φ†Xe = Φ†S(Xeinv, X e spu) = Φ † invX e inv + Φ † spuX e spu = (Φ † inv + Φ † spu ·A)Xeinv + Φ†spuW e (11) We will show that Φ+ = ([ Φinv + Φspu · A ] , 0 ) S−1 can continue to achieve an error of q(= r)
across training domains, and have a lower information bottleneck penalty. Therefore, the optimal solution to Twins-ERM (Equation (8)) does not depend on the spurious noise W e, and hence solves the domain generalization problem (Equation (DG)).
It is known that ERM and IRM fails under the assumption of Theorem 3.1 (Theorem 3 in (Ahuja et al., 2021)). This theorem shows that our algorithm can provably solve the linear classification structural equation model.
In real world, however, we do not have direct access to the domain-free instance X . Hence, we practically adopt image from another domain, denoted as Xe ′ , as a proxy for X in Twins-ERM and Twins-IRM (Equation (8) and Equation (9)). In other words, fixing e′ ∈ Etrain, the Twins-ERM (Equation (8)) can be rewritten as
min w,Φ ∑ e∈Etrain He ( Φ ( Xe)|Xe ′ )− λHe ( Φ ) s.t. 1 |Etrain| ∑ e∈Etrain Re ( w ◦ Φ ) ≤ r, (12)
and the Twins-IRM (Equation (9)) can be rewritten as
min w,Φ ∑ e∈Etrain He ( Φ ( Xe)|Xe ′ )− λHe ( Φ )
s.t. 1 |Etrain| ∑ e∈Etrain Re ( w ◦ Φ ) ≤ r, , w ∈ arg min w̃ Re ( w̃ ◦ Φ ) ,∀e ∈ Etrain.
(13)
We will show adopting proxy from another domain will still be guaranteed to solve the domain generalization problem (Equation (DG)). Theorem 3.2. . Suppose each e ∈ Eall follows Assumption 2, and Assumptions 3 and 4 hold for the invariant features. Also, for each e ∈ Etrain, assume that Xespu = AXeinv +W e, where A ∈ Ro×m, W e ∈ Ro is continuous, bounded, and zero mean noise. Each solution to Twins-ERM and Twins-IRM (Equation (12) and Equation (13)), with ` as 0-1 loss, and r = q) solves the domain generalization problem (Equation (DG)).
The full proof of Theorem 3.1 and 3.2 can be found at Appendix A.

heading: 3.3 GAUSSIAN BOTTLENECK
text: The IB obejctive (4) can be directly estimated by k-nearest-neighbor method, as described in Appendix C. However, such direct estimation requires large memory and is computationally intensive.
We denote the algorithm by kNN direct estimation as Twins-Direct. In order to facilitate the implementation of our penalty, we make the simplifying assumption that the datasets follow a Gaussian distribution (Chechik et al., 2005). Specifically, assuming X,Xe are jointly multivariate zero-mean Gaussian vectors with covariances ΣX ,ΣXe , and Ze ∈ RdZ is a encoded version of Xe that must maintain a given value of mutual information with X . We define the covariance of Ze and Ze|X to be ΣZe and ΣZe|X .
Next, we simplify the Equation (6). The entropy of a Gaussian distribution is simply given by the logarithm of the determinant of its covariance function (up to a constant that we ignore). The loss function becomes:
LIB(θ, e) = EX log det(ΣZe|X)− λ log det(ΣZe). (14)
Practical considerations: We reformulate our algorithm so that it resembles the contrastive learning method, i.e. Barlow Twins (Zbontar et al., 2021). The second term of the loss in Equation (14) maximizes det(ΣZe). Since computing the determinant of a matrix is computationally intensive, we adopt a proxy to minimize the Frobenius norm of the correlation matrix of Ze. Since the correlation matrix is invariant to scaling, we can set the diagonal element of the correlation matrix of Ze to be 1. Then, the second term of Equation (14) amounts to minimizing the off-diagonal term, i.e. the second term in Equation (16). This term essentially decorrelates the different dimensions of the representation and prevents these dimensions from encoding similar information.
Besides, it can also easily be shown that the first term of Equation (14) minimizes the information the representation contains about the domain information has the same solution with the first term of Equation (16). This term maximizes the alignment between representations of pairs of domain-aware instances Xe and domain-free instances X .
In practice, we have no access to the domain-free latent instance. For a given instance Xe, we use an instance with the same label, but from a different domain, denoted as Xe ′ as surrogate for the domain-free latent X . We minimize the distance between pairs of instances from different domains. Sample {zeb}1≤b≤B ∼ Ze and {ze ′ b }1≤b≤B ∼ Ze ′ , we concatenate them into matrix Ze ∈ RB×dZ and Ze ′ ∈ RB×dZ , where dZ is the dimension of zeb , ze ′ b . After mean shifting every column of Z e and Ze ′ , such that 1>Ze = 1>Ze ′ = 0 (1 is a column vector full of 1s), the cross-correlation matrix CZij is defined as:
CZij = 〈Ze:,i,Ze
′
:,j〉 ‖Ze:,i‖2‖Ze ′ :,j‖2 , 1 ≤ i, j ≤ dZ , (15)
and the final penalty is defined as: c(Z) = ∑ i (1− CZii )2 + λ ∑ i 6=j (CZij) 2. (16)
However, we do not have access to domain transformation model either. We construct the contrastive instances by permuting the instances that have the same label in each iteration. In particular, we sample B instances {xb}1≤b≤B from training domains as row vectors, where B is the batch size. We concatenate the representation {zb} = {Φ(xb)} into matrix Z1 ∈ RB×dZ . The contrastive batch Z2 is constructed by permuting the rows of Z1, i.e. Z1b,: = Z 2 π(b),:, where π is a permutation of {1, 2, · · · , B} such that the corresponding labels of xb and xπ(b) are identical. CZij defined in Equation (15) can be rewritten as:
CZij = 〈Z1:,i,Z2:,j〉 ‖Z1:,i‖2‖Z2:,j‖2 , 1 ≤ i, j ≤ dZ (17)
Such penalty can be readily incorporated into the ERM and IRM losses, i.e.
LTwins−ERM = LERM + µ · c(Z), and LTwins−IRM = LIRM + µ · c(Z), (18)
where LERM and LIRM denote the loss in the ERM and IRM respectively, and µ is the penalty hyperparameter.

heading: 4 EXPERIMENTS
text: In this section, we conduct experimentation on two benchmarks: Linear Unit Tests (in Section 4.1, Section 4.3) and DomainBed (in Section 4.2). Linear Unit Tests (Aubin et al., 2021) consists of several toy datasets to evaluate algorithms for domain generalization and invariance learning, while DomainBed (Gulrajani & Lopez-Paz, 2020) is a unified testbed for evaluating domain generalization algorithms. We use the following four baselines across our experiments: ERM, IB-ERM, IRM, IB-IRM (Ahuja et al., 2021) in the synthetic data, and use ERM, IRM as baselines in real-world datasets.

heading: 4.1 LINEAR UNIT TESTS
text: The dataset describes six linear low-dimensional problems, named Example 1/1s, Example 2/2s and Example 3/3s, where the ’s’ dictates a different rotation matrix. Each example, called unit test, is designed to test different types of out-of-distribution generalization. We describe in Appendix B.1 the precise distributions and the invariances captured by each example.
Benchmark details: We follow the same pipeline as those used in Aubin et al. (2021); Ahuja et al. (2021) for the model selection, hyperparameter selection, training, and evaluation. We set (dinv, dspu) = (5, 5). For all three examples, the models used are linear. The training loss is the square error for the regression setting (Example 1/1s), and binary cross-entropy for the classification setting (Example 2/2s, 3/3s). For the evaluation of performance on Example 1/1s, we report mean square errors and standard deviations. For the evaluation of performance on Example 2/2s, Example 3/3s, we report classification errors and standard deviations.
Model training: For the Twins-ERM approach, there is an additional hyperparameter µ associated with the c(Z) term in the final objective in Equation (18). We sample the µ from logµ ∼ Uniform(−3,−1). For each algorithm, we run a random hyperparameter search for 20 trials, and average the results over 50 data seeds. We train each algorithm and hyperparameter trial on the train splits of all environments, for 104 full-batch Adam updates (Kingma & Ba, 2014). We choose the hyperparameters trial that minimizes the error on the validation splits of all environments, i.e. the train-domain validation set evaluation procedure in (Gulrajani & Lopez-Paz, 2020).
We also implement an Oracle that contains randomized Xspu in each iteration, such that it learns to ignore the spurious features.
The experimental results are reported in Table 1. In the Example 2/2s, since the invariant feature contains full information about the label, we do observe that IB penalty in (Ahuja et al., 2021) and Twins penalty in our paper performs the best. In the Example 1/1s and 3/3s, the spurious feature contains partial information about the label, we generally find invariant risk can reduce the error in this case. We empirically verify the benefit of using the proposed penalty over the previously proposed baselines.

heading: 4.2 MNIST-TYPE DATASET
text: In the second benchmark we use DomainBed to experiment on MNIST-type datasets inspired by the construction of CS-CMNIST (Ahuja et al., 2021) to evaluate covariate shift. In addition to the
origin MNIST (LeCun & Cortes, 2010), we extend the benchmark to include FashionMNIST (Xiao et al., 2017) and KMNIST (Clanuwat et al., 2018), where the images of the latter two are used as drop-in replacements for MNIST images. The idea in CS-CMNIST is to associate each class with a color, and each image is assigned the color associated to its class with probability pe or a random color with probability 1 − pe. We construct three environments for this experiment: two training environments containing 20,000 data points each, one test containing 20,000 points. In the two training environments, the pe is set to 1.0 and 0.9 respectively. In the testing environment, the pe is set to 0, i.e., all the images are colored completely at random. A grid search in the range of {10−4, 10−3, 10−2, 10−1, 1} is used to determine the optimal penalty parameter µ. We fix the trade-off parameter λ = 5 × 10−3. We run the experiments using 5 different seeds and report the mean and the standard deviation of the classification.
The results are reported in Table 2. We find that generally setting µ = 10−2 would be the best choice. See Appendix D Notice that both the IRM and the ERM versions in each case perform similarly. The results reveal that ERM and IRM have the weakest performance. IB-ERM and IB-IRM increase the accuracy of ERM and IRM respectively, with Twins-ERM and Twins-IRM outperforming all the compared methods.

heading: 4.3 REAL WORLD DATASETS
text: In the third benchmark we use DomainBed to experiment on real world datasets: OfficeHome (Venkateswara et al., 2017), PACS (Li et al., 2017). For our Twins algorithm, we ran a hyperparameter search in the range of {10−4, 10−3, 10−2, 10−1} for µ, and 20 hyperparameter seeds for remaining hyperparameters in the DomainBed suite (Gulrajani & Lopez-Paz, 2020). We run the experiments using 3 different seeds and report the mean and the standard deviation of the classification. For ERM and IRM, we directly borrow the reuslts from the original paper.
The results are reported in Table 3 and 4. We find that our Twins-IRM algorithm obtain consistent improvement over the baselines.

heading: 5 RELATED WORK

heading: 5.1 RELATION TO INVARIANT RISK MINIMIZATION
text: Relation to IB-ERM/IB-IRM (Ahuja et al., 2021). Ahuja et al. (2021) introduce the information bottleneck method. However, similar to Tishby et al. (2000), Ahuja et al. (2021) utilize the information
bottleneck to learn a representation that compresses the input as much as possible while preserving all the relevant information about the target label. Ours instead is label-free and tries to preserve the relevant information about the domain-free latent image during compression, and uses another image with the same label but different domain as surrogate for the unknown domain-free latent image. Essentially, setting β = 0 will reduce our penalty into the one in Ahuja et al. (2021), but β > 1 can help eliminate the trivial solution that the representation mapping is constant, and is shown to achieve better performance on various datasets.
Previous work has demonstrated that the entropy penalty alone (i.e. β = 0) might fail in specific case, such as in Section 4 in Ahuja et al. (2021). Nevertheless, our proposed framework does not suffer from this counter-example. We introduce such an example next as a simple classification problem. In each e ∈ Etrain, Y e ← Xeinv⊕Ne and Xespu ← Y e⊕V e, where all the random variables involved are binary valued, noise Ne, V e are Bernoulli with parameters q (identical across Etrain), ce (varies across Etrain) respectively. If ce < q, then in Etrain predictions based on Xespu are better than predictions based on Xeinv. If Φ selects X e inv, the IB penalty equals −λH(Xeinv); while if Φ selects Xespu, the IB penalty equals H(N e ⊕ V e) − λH(Xespu). Since Xespu ← Xeinv ⊕Ne ⊕ V e, we have −λH(Xeinv) < H(Ne ⊕ V e)− λH(Xespu) by Lemma A.2 in the Appendix. Our IB penalty is then able to select the invariance term Xeinv. On the other hand, if X e inv obeys uniform Bernoulli distribution, its entropy will be no lower than Xespu, and hence the entropy term alone is not enough.

heading: 5.2 RELATION TO CROSS-DOMAIN COVARIANCE METHOD
text: Aligning the cross-domain distribution has been studied extensively in the domain generalization community both empirically and theoretically (Sun & Saenko, 2016; Li et al., 2018b; Rahman et al., 2020; Kpotufe & Martinet, 2018). Despite the fact that we use covariance as well, the motivation and implication of the proposed regularization scheme are different. Domain aligning method, such as CORAL (Sun & Saenko, 2016), aligns the correlation matrix from different domains. However, our method (Equation (16)) tries to decorrelate each dimension of the representation by minimizing the off-diagonal term of cross-correlation matrix. In our penalty, we do not want to align the covariance between different domains. For example, given a batch of data of size B ×N , where B is the batch size and N is the feature dimension. Cross-domain covariance tries to deal with the row vector and the correlation matrix is of size N ×N . Our covariance penalty deals with column vector, and the correlation matrix is of size B ×B. Since usually B N , the computation would be much easier.

heading: 5.3 RELATION TO CONTRASTIVE-BASED DOMAIN GENERALIZATION
text: Our method can also be regarded as a contrastive-based domain generalization problem. Contrastive learning (Chopra et al., 2005; Caron et al., 2020; Grill et al., 2020; Chen et al., 2020; Zbontar et al., 2021) has been a successful paradigm in self-supervised leaning. Contrastive learning aims at bringing positive pair samples closer together, while moving negative samples further away in a learned embedding space. Essentially, the aim of domain generalization is to extract domaininvariant features, similarly aiming to minimize the distance of features within the same class in the embedding space, while maximizing the distance of features from different classes. Such aim is closely related to the domain generalization. SelfReg (Kim et al., 2021) uses only positive data pairs and introduces inter-domain curriculum learning to prevent representation collapse (Grill et al., 2020). (Jeon et al., 2021) uses domain-aware supervised contrastive to ensure domain invariance while increasing class discriminability, Compared to previous works, our method instead introduces a much simpler framework to ensure convergence to domain invariant features with theoretical guarantee.

heading: 6 CONCLUSION
text: In this work, we introduce an information-theoretical approach for domain generalization. We cast the task of domain generalization as a rate distortion problem and then use information bottleneck penalty to obtain guarantees on the existence of features we want to learn. We link our method, called Twins, with self-supervised learning, which can provide a theoretical perspective in the success behind selfsupervised learning. We conduct an empirical study on Twins-ERM and Twins-IRM under various datasets and confirm the consistent improvement of the proposed method over existing baselines. In the future, we intend to further analyze domain generalization in the rate distortion framework and conduct large scale experiments to verify our IB formulation in real-world applications.

heading: A PROOF OF THEOREM 3.1
text: The entropy or the Shannon entropy (Wehrl, 1978) of a discrete random variable X ∼ PX with support X is defined as
H(X) = − ∑ x∈X PX(X = x) log ( PX(X = x) ) . (19)
The differential entropy (Wehrl, 1978) of a continuous random variable X ∼ PX with support X is given as follows h(X) = − ∫ x∈X log ( PX(x) ) dPX(x), (20)
where dPX(x) is the Radon-Nikodym derivative of PX w.r.t the Lesbegue measure.
For continuous variables, the differential entropy h(·) is not bounded below, we can define the lower bounded differential entropy (Kirsch et al., 2020) ĥ(X) = h(X + ε), where ε is an independent zero-entropy noise ε ∼ Uniform(0, 1). Since X ⊥ ε, ĥ(X) ≥ h(ε) = 0, we get that ĥ(·) is bounded below. Lemma A.1. If X and Y are discrete random variables that are independent with the supports satisfying 2 ≤ |X | <∞, 2 ≤ |Y| <∞, where | · | denotes the number of element in a set, then for λ < 1,
λH(X) +H(Y ) > λH(X + Y ) (21)
Proof. Define Z = X + Y . H(Z|X) = − ∑ x∈X PX(x) ∑ z∈Z PZ|X(Z = z|X = x) log ( PZ|X(Z = z|X = x) ) = −
∑ x∈X PX(x) ∑ z∈Z PY |X(Y = z − x|X = x) log ( PY |X(Y = z − x|X = x) ) = −
∑ x∈X PX(x) ∑ z∈Z PY |X(Y = z − x|X = x) log ( PY |X(Y = z − x|X = x) ) (use X ⊥ Y )
= − ∑ x∈X PX(x) ∑ z∈Z PY (Y = z − x) log ( PY (Y = z − x) ) = H(Y )
(22)
Hence,
H(X+Y )−H(X) = H(X+Y )−H(X+Y |Y ) = I(X+Y ;Y ) = H(Y )−H(Y |X+Y ) ≤ H(Y ) (23)
λ(H(X + Y )−H(X)) ≤ λH(Y ) ≤ H(Y ) (24) when λ < 1.
The equality holds if and only if H(Y ) = 0, which is impossible since 2 ≤ |Y| <∞.
Lemma A.2. If X , Y and Z are discrete random variables with the supports satisfying 2 ≤ |X | < ∞, 2 ≤ |Y| <∞ and 2 ≤ |Y| <∞, where | · | denotes the number of element in a set. Besides, Y is independent of X and Z. Then for λ < 1,
H(X + Y |Z)− λH(X + Y ) > H(X|Z)− λH(X) (25)
Proof. Similar to Equation (22), we would have
H(X + Y )−H(X) = I(X + Y ;Y ), H(X + Y |Z)−H(X|Z) = I(X + Y ;Y |Z), (26)
By the chain rule of conditional mutual information,
I(X + Y ;Y |Z) = I(Y ;X + Y, Z)− I(Y ;Z) = I(Y ;X + Y, Z)− 0 ( since Y ⊥ Z) = I(Y ;X + Y ) + I(Y ;Z|X + Y ) ≥ I(Y ;X + Y )
(27)
where the last inequality holds since the conditional mutual information is non-negative. Since λ < 1,
I(X + Y ;Y |Z) ≥ I(Y ;X + Y ) ≥ λI(Y ;X + Y ) (28) and hence the equality holds iff I(Y ;X + Y ) = 0, in other words, Y ⊥ X + Y . If given X +Y = xmax + ymax, we can infer Y = ymax. Hence, P(Y = ymax|X +Y = xmax + ymax) = 1. However, P(Y = ymax) = 1 as the support of Y has at least two elements, which gives a contradiction. Hence,
I(X + Y ;Y |Z) > λI(Y ;X + Y ) (29)
Lemma A.3. If X and Y are continuous random variables that are independent and have a bounded support, then for λ < 1,
ĥ(X) + λĥ(Y ) > λĥ(X + Y ) (30)
Proof. Setting ε ∼ Uniform(0, 1) independent of X,Y , we have ĥ(X) = h(X + ε), ĥ(Y ) = h(Y + ε), ĥ(X + Y ) = h(X + Y + ε). We have
λ(h(X + Y + ε)− h(Y + ε)) = λI(X + Y + ε;X) and h(X + ε) = h(X + ε)− h(ε) = I(X + ε;X) (31)
According to the data processing inquality, (Beaudry & Renner, 2011) and X ⊥ X + Y + ε|X + ε,
I(X + Y + ε;X) ≤ I(X + ε, Y ;X) = I(X + ε;X) (32)
where the last equality holds since X + ε ⊥ Y , and we get
λ(h(X + Y + ε)− h(Y + ε)) ≤ I(X + Y + ε;X) ≤ I(X + ε;X) = h(X + ε) (33)
Rearranging it, we get ĥ(X) + λĥ(Y ) ≥ λĥ(X + Y ) (34)
Since λ < 1, the equality holds only if I(X+Y +ε;X) = 0. In other words, we haveX+Y +ε ⊥ X . In the next, we show that it is not possible.
The support of X can be divided into the union of intervals. We assume ∆ > 0 such that [xmax − ∆, xmax] belongs to the rightmost interval of X; and [ymax − ∆, ymax] belongs to the rightmost interval of Y , where xmax and ymax denotes the maximum of the support of X and Y . Define an eventM : xmax + ymax− δ ≤ X + Y + ε ≤ xmax + ymax + 1. IfM occurs, note that ε is bounded by 1, we have PX(X ≤ xmax − δ|M) = 0, PY (Y ≤ ymax − δ|M) = 0 (35) If δ < ∆, based on the definition of the interval, we have that
PX(X ≤ xmax − δ) > 0, PY (Y ≤ ymax − δ) > 0 (36)
If X + Y + ε ⊥ Y then PY (Y ≤ ymax − δ) = PY (Y ≤ ymax − δ|M), which is not the case from the above equations (35) and (36).
Lemma A.4. If X ,Y and Z are continuous random variables that have a bounded support, and Y is independent of X and Z, then for λ < 1,
ĥ(X + Y |Z)− λĥ(X + Y ) ≥ ĥ(X|Z)− λĥ(X) (37)
Proof. Like Lemma A.2, we rewrite the inequality as
I(X + Y + ε;Y |Z) > λI(X + Y + ε;Y ) (38)
Similar to Equation (27), Y is independent of Z we could write I(X + Y + ε;Y |Z) = I(Y ;Z,X + Y + ε). We then use the data processing inequality, we would get
I(X + Y + ε;Y |Z) = I(Y ;Z,X + Y + ε) ≥ I(X + Y + ε;Y ) (39)
since λ < 1, and similar to the proof of Lemma A.3, I(X + Y + ε;Y ) > 0, we would have
I(X + Y + ε;Y |Z) > λI(X + Y + ε;Y ) (40)
Proof of Theorem 3.1. The proof of the theorem resembles the proof of Theorem 4 in (Ahuja et al., 2021). Consider a solution to equation Φ†,
Φ† ·Xe = Φ† · S(Xeinv, Xespu) = Φinv ·Xeinv + Φspu ·Xespu = [ Φinv + Φspu ·A ] ·Xeinv + Φspu ·W e.
(41)
and since Φ† achieves the error of q,
I(w+inv ·X e inv) = I(Φinv ·Xeinv + Φspu ·Xespu) (42)
In the next we prove Φspu by contradiction. Define Φ+ = ([ Φinv + Φspu · A ] , 0 ) S−1. Observe that we can write Φ† · Xe = Φ+ · Xe + Φspu · W e. a) Φspu · W e ⊥ Φ+ · Xe (Φ+ · Xe =[ Φinv + Φspu ·A ] ·Xeinv and Xeinv ⊥W e),
b.1) Φ+ ·Xe, Φspu ·W e are discrete random variables with finite support of size at least two. (discrete case)
b.2) Φ+ ·Xe, Φspu ·W e are continuous bounded random variables. (continuous case) In the discrete case, from a), b.1), and Lemma A.1 it follows that
λH(Φ+ ·Xe) +H(Φspu ·W e) > λH(Φ† ·Xe) (43)
Rearranging the terms, we have
H(Φspu ·W e)− λH(Φ† ·Xe) > −λH(Φ+ ·Xe) (44)
Since Xeinv = G(Xinv, e), we have H(Φ † ·Xe|Xinv) = H(Φspu ·W e) and H(Φ+ ·Xe|Xinv) = 0. Hence, we get
H(Φ† ·Xe|Xinv)− λH(Φspu ·W e) > H(Φ+ ·Xe|Xinv)−H(Φ+ ·Xe). (45)
and therefore, Φ+ ·Xe would have a lower penalty. In the continuous case, the argument is similar by invoking a), b.2) and Lemma A.3. Φ+ can achieve strictly lower penalty than Φ†.
Following the proof of the first part of Theorem 4 in (Ahuja et al., 2021), we can show that Φ+ achieves the same error of q in all the training environments. Thus Φ+ is a strictly better solution Φ†, which contradicts the optimality of Φ†. Therefore, it follows that Φspu = 0. And hence,
I(w+inv ·X e inv) = I(Φinv ·Xeinv) (46)
Based on Theorem 3 in (Ahuja et al., 2021), if a solution does not rely on spurious features and satisfies equation (46) for all the points in the support, then under the Assumption 3 such a solution solves the domain generalization problem (DG).
Proof of Theorem 3.2. The major difference here is that we do not condition on the unknown invariant feature. Notations are defined similarly as in the proof of Theorem 3.1. Compared with Equation (45), we now want to prove that:
H(Φ† ·Xe|Xe ′ )− λH(Φspu ·W e) > H(Φ+ ·Xe|Xe ′ )−H(Φ+ ·Xe). (47)
In other words,
H(Φ+ ·Xe + Φspu ·W e|Xe ′ )− λH(Φspu ·W e) > H(Φ+ ·Xe|Xe ′ )−H(Φ+ ·Xe). (48)
Since Φspu ·W e is independent of Xe ′
and Φ+ ·Xe. We use Lemma A.2 for the discrete case and Lemma A.4 for the continuous case to prove Equation (48). The rest of the proof is the same as the proof in Theorem 3.1.

heading: B DATASET DESCRIPTION
text: B.1 LINEAR UNIT TEST
Example 1/1s The dataset in environment e ∈ Eall is sampled from the following distributions: Xeinv ∼ Ndinv(0, (σe)2), Ỹ e ∼ Ndinv(WyxXeinv, (σe)2), Xespu ∼ Ndspu(WxyỸ e, 1), Xe ← S · (Xeinv, Xespu),
Y e ← 2 (dinv + dspu) 1Tdinv Ỹ e,
(49)
where Wyz ∈ Rdinv×dinv , Wxy ∈ Rdspu×dinv are matrices drawn i.i.d. from the standard normal distribution, 1dinv ∈ Rdinv is a vector of ones, Nk is a k dimensional vector from the normal distribution, and S ∈ R(dinv+dspu)×(dinv+dspu) is a rotation matrix fixed for all environments. The parameter σ is set differently for every environment (i.e., domain). In particular, we set (σe=e0)2 = 0.1, (σe=e1)2 = 1.5, and (σe=e2)2 = 2 for the first three environments. In case there are more than three environments, the (σe=ej ) for j > 3 is uniformly from Unif(10−2, 10). The rotation matrix S is set to the identity matrix in Example 1 and a random unitary matrix in Example 1s.
Example 2/2s Following the notation of the original paper (Aubin et al., 2021), let µcow = 1dinv , µcamel = −µcow, νanimal = 10−2, µgrass = 1dspu , µsand = −µgrass, νbackground = 1.
(50)
The dataset in environment e ∈ Eall is sampled from the following distribution: je ∼ Categorical ( pese, (1− pe)se, pe(1− se), (1− pe)(1− se) ) ,
Xeinv ∼ {
(Ndinv(0, 0.1) + µcow) · νanimal if je ∈ {1, 2}, (Ndinv(0, 0.1) + µcamel) · νanimal if je ∈ {3, 4},
Xespu ∼ {
(Ndspu(0, 0.1) + µgrass) · νbackground if je ∈ {1, 4}, (Ndspu(0, 0.1) + µsand) · νbackground if je ∈ {2, 3},
Xe ← S · (Xeinv, Xespu), Y e ← I(1TdinvX e inv),
(51)
where the environment foreground/background probabilities are pe=e0 = 0.95, pe=e1 = 0.97, pe=e2 = 0.99 and the cow/camel probabilities are se=e0 = 0.3, se=e1 = 0.5, se=e2 = 0.7. For nenv > 3 and j ∈ [3 : nenv − 1], the extra environment variables are respectively drawn according to pe=ej ∼ Unif(0.9, 1) and se=ej ∼ Unif(0.3, 0.7). The rotation matrix S is set to the identity matrix in Example 2 and a random unitary matrix in Example 2s.
Example 3/3s The example is meant to present a linear version of the spiral classification problem of Parascandolo et al. (2020). Let µinv = 0.1 · 1dinv , and µespu ∼ Ndspu(0, 1) for all the environments. The dataset in environment e ∈ Eall is sampled from the following distribution:
Y e ∼ Bernoulli (1
2
) , Xe ← S · (Xeinv, Xespu)
Xeinv ∼ { Ndinv(+µinv, 0.1) if Y e = 0, Ndinv(−µinv, 0.1) if Y e = 1,
Xespu ∼ { Ndspu(+µespu, 0.1) if Y e = 0, Ndspu(−µespu, 0.1) if Y e = 1,
,
(52)
The rotation matrix S is set to the identity matrix in Example 3 and a random unitary matrix in Example 3s. In the above dataset, the invariant features are anti-causally related to the label Y e.
Remark on Linear unit test: In the Example 1/1s and Example 3/3s, the invariant features are causal and partially informative about the label. The spurious features carry extra information about the label not contained in the invariant features. In the Example 2/2s, the invariant features are causal and carry full information about the label.

heading: C GAUSSIAN-FREE ENTROPY ESTIMATION
text: C.1 ESTIMATING ENTROPY BY KNN
Since the feature is in high-dimensional spaces it is challenging to estimate the density of Z, preventing us from directly computing the exact entropy. To remedy this issue, we resort to the particle-based entropy estimator from Singh et al. (2003); Beirlant et al. (1997), which is based on k-Nearest Neighbors (kNN). We introduce this approach in general terminologies. Consider a distribution p with respect to z ∈ Z , the particle based entropy estimate is given by
Ĥk(p) = − 1
N N∑ i=1 log k NVolki + log k − Φ(k) ∝ N∑ i=1 log Volki (53)
where Φ is the digamma function, log k − Φ(k) is a bias correction term. Volki is the volume of the hyper-sphere of radius Ri = ‖zi − zKNNk ‖2, which is the Euclidean distance between zi and its k-th nearest neighbor zKNNk . The volumn is given by:
Volki = ‖zi − zKNNk ‖n2 · πn/2
Γ(n2 + 1) (54)
where Γ is the Gamma function, n is the dimension of Z . Putting Equation (53) and Equation (54) together, we have
Ĥk(p) = n
N N∑ i=1 log ‖zi − zKNNk ‖2 + logN + C (55)
where Ck,n is determined by n and k.
C.2 ESTIMATING IB OBJECTIVE 4
Recall that Ze = f(Xe), and Xe = G(X, e), we can estimate the IB objective LIB by first sampling {ei} ∼ Pe, 1 ≤ i ≤ D, and then Xeij ∼ Pei(Xei), 1 ≤ j ≤ Ni for any fixed i. We can use samples Zeij = f(X ei j ) to estimate EX,eH(Ze), i.e.
EX,eH(Ze) ≈ n∑D i=1Ni D∑ i=1 Ni∑ j=1 log ‖Ze i j − ZKNNk ‖2 + log( D∑ i=1 Ni) + Ck,n (56)
where n is the dimension of Zi, ZKNNk is Z ei j ’s k-th nearest neighbor in the full dataset {Z ei j } D,Ni i=1,j=1.
For the first term, we use the same method but conditioned on the fixed label. In other words, it should be
EX,eH(Ze|X) ≈ 1
D D∑ i=1  n Ni Ni∑ j log ‖Zij − ZKNNik ‖2 + logNi + Ck,n (57) where ZKNNik is Zij’s k-th nearest neighbor in the dataset {Zij}j .
C.3 COMPARISON
We perform our experiments on ColoredMNIST datasets. At each checkpoint, we sample 1024 isntances and set kNN parameter to be 5 to estimate the IB penalty (4) by Equation (56) and (57) at every checkpoints. We plot the trajectory of kNN based penalty in Figure (1). Clearly, our Twins method is able to efficiently minimize the true Gaussian entropy.
D HYPERPARAMETER SELECTION OVER µ
µ = 10−4 10−3 10−2 10−1 1
MNIST 79.24 ± 0.69 80.44 ± 2.70 82.83 ± 2.73 77.19 ± 6.49 72.11 ± 4.08 KMNIST 49.48 ± 4.27 52.24 ± 3.94 52.29 ± 3.26 43.02 ± 0.68 36.31 ± 3.20 FashionMNIST 54.88 ± 1.57 53.87 ± 2.41 56.04 ± 1.79 52.25 ± 2.70 50.96 ± 1.61

Review summaries:
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.

This paper studies the problem of domain generalization via invariant feature learning studied from the lens of information theory. The motivation is to use rate distortion theory to obtain a domain-invariant representation. The authors propose an algorithm titled Twins that uses the aforementioned rate distortion theory coupled with a multivariate Gaussian assumption to provide an additional penalty based on cross-correlation, i.e., minimize the distances between pairs of points belonging to different domains but the same label. The authors demonstrate, that under suitable assumptions, this approach will recover the invariant feature representation. The authors construct several hand-crafted problems where indeed this assumption (and the proposed algorithm) work well, and also provide some experiments on MNIST-type data.

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.