A R ATE -D ISTORTION A PPROACH TO D OMAIN G ENERALIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁrst 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-speciﬁc image compared to its direct prediction from the domain-speciﬁc 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 veriﬁes the improvement over recent baselines.
I NTRODUCTION
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 signiﬁcant 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) domainspeciﬁc method: Domain2Vec (D2V) (Deshmukh et al., 2018) learns domain-speciﬁc embedding,
DMG (Chattopadhyay et al., 2020) aims to learn domain speciﬁc 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 deﬁne 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 classiﬁcation 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.
P RELIMINARY 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 ﬁnite 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 ∈ E all , denoted as ( X e , Y e ) ∼ P e
( X e , Y e ) .
We also assume that e ∈ E all is distributed as e ∼ P e .
Given samples from a ﬁnite subset E train (cid:40) E all 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 E train (cid:40)
E all be a ﬁnite subset of training domains.
We have access to the data for each training domain e train ∈ E train , but have no access to the data for each test domain
e test ∈ E all \E train .
Given a function class F and a loss function (cid:96) , 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 E all .
Deﬁne the risk of the predictor f on the domain e as R e ( f ) :
=
E P e
( X e ,Y e ) (cid:96) ( f ( X e ) , Y e ) .
We want to solve the following min-max optimization problem:
[Equation]
We establish the domain transformation model to characterize the relation between domain-aware instance X e and the domain-invariant latent instance X in the Assumption 1, which ﬁrst appears in Robey et al.
Assumption 1 (Domain transformation model) .
Let δ e denote a Dirac distribution for e ∈ E all .
We assume that there exists a measurable function G :
X
× E all → X , which we refer to as a domain transformation model , that parameterizes the inter-domain covariate shift via
[Equation]
where # denotes the push-forward measure and = d denotes equality in distribution.
The Assumption 1 can somewhat reﬂect the generation of domain speciﬁc 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 speciﬁc 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 classiﬁer and w ◦ Φ denote the full predictor.
The regret of the network on the domain e is denoted as R e ( w ◦ Φ) .
Next, we deﬁne standard properties related to the datasets used in the domain generalization literature (Ahuja et al., 2021).
For each e ∈ E all , the distribution ( X e , Y e ) ∼ P e satisﬁes the following properties: (1) ∃ a map Φ ∗ , which we call an invariant feature map , such that E (cid:2)
Y e (cid:12)(cid:12)
Φ ∗ (cid:0) X e (cid:1)(cid:3) is the same for all e ∈ E all and Y e (cid:54)⊥ Φ ∗ ( X e ) , where ⊥ means mutual independence.
(2) ∃ a map Ψ ∗ , which we call spurious feature map , such that E (cid:2)
Y e (cid:12)(cid:12)
Ψ ∗ (cid:0) X e (cid:1)(cid:3) is not the same for all e ∈ E all and
Y e (cid:54)⊥ Ψ ∗ ( X e ) for some domains.
Ψ ∗ often hinders learning predictors that only rely on Φ ∗ .
For example, in the CMNIST dataset, the Φ (cid:63) extracts the underlying digit and Ψ (cid:63) 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:
[Equation]
where |E train | denotes the number of training domains.
We say that a data representation Φ elicits an invariant predictor across the set of training domains E train if there is a predictor w that simultaneously achieves the minimum risk, i.e. w ∈ arg min w
[Equation]
R e ( w
[Equation]
◦ Φ) , ∀ e ∈ E train .
Using this notation, the main objective of Invariant Risk Minimization (IRM) is stated as:
[Equation]
Lastly, we rely on the notion of ‘informativeness’ about the datasets (Ahuja et al., 2021).
There are two such categories of informativeness.
In the ﬁrst case, the invariant features Φ ∗ ( X e ) are partially informative about the label, i.e. Y (cid:54)⊥ X e | Φ ∗ ( X e ) , 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 ⊥ X e | Φ ∗ ( X e ) , i.e., they contain all the information about the label that is contained in input X e .
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.
M ETHOD 3. R ATE D ISTORTION & I NFORMATION B OTTLENECK P RINCIPLE
Given a domain-free latent instance X ∈ R d X , its observation in a domain e is denoted as X e :=
G ( X, e ) .
We want to learn the feature Z e =
Φ( X e ) ∈ R
d
Z which is informative about the domainfree variable X , but invariant (i.e. uninformative) to the speciﬁc 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 X e .
We want the representation Z e =
Φ( X e ) to be domain-invariant, so we feed Z e into a decoder which outputs domain-invariant X .
We minimize the distortion between the original domain-aware instance
X e and the reconstructed domain-free instance X .
The distortion function measures how well X is predicted from a compressed representation Z e compared to its direct prediction from X e .
This trade-off is captured by the following loss function:
[Equation]
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 deﬁnition 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:
[Equation]
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
[Equation]
Continuous case : In terms of continuous variables, the differential entropy h ( · ) is not bounded below, which hinders our analysis.
To overcome this, we can deﬁne the lower bounded differential entropy (cid:98) h ( X ) :
= h ( X + ε ) , where ε is the independent bounded zero-entropy noise ε ∼ Uniform(0 , 1) .
Thus, (cid:98) h ( X ) ≥ h ( ε ) = 0 .
We can replace the Shannon entropy H ( · ) with lower bounded differential entropy (cid:98) h ( · ) in Equation (6):
[Equation]
For simplicity, we deﬁne H to be the Shannon entropy for the discrete variables, or lower bounded entropy (cid:98) h for continuous variables in the main text.
We deﬁne H e ( f ) :
=
E X e ∼ P e H
( f ( X e )) .
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:
[Equation]
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:
[Equation]
3. T HEORETICAL G UARANTEE
In this subsection [Equation] we establish the theoretical guarantee of our algorithm under the linear classiﬁcation case [Equation]
We consider the following standard 0-1 classiﬁcation model in literature (Ahuja et al., 2020; 2021):
Assumption 2.
Linear classiﬁcation structural equation model.
In each e ∈ E all ,
[Equation]
[Equation] where I( x ) is 1 if x is positive else 0, w (cid:63)
inv ∈ R m with (cid:107) w (cid:63) inv (cid:107) = 1 is the labelling hyperplane,
X e inv , X e inv ∈ R m , X e spu ∈ R
o , S ∈ R ( m + o ) × ( m + o ) and G is a continuous domain transformation.
Before presenting our main theorem, we ﬁrst add two assumptions on the support of invariant features.
Deﬁne the support of the invariant features X e inv in environment e as X e inv .
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. (cid:83) e ∈E all X e inv ⊆ [Equation]
e
[Equation] Assumption 4 (Strictly separable invariant features) .
The training support of invariant features (cid:83) e ∈E train X e inv is strictly separated by the labelling hyperplane w ∗ inv . In other words, min [Equation]
Assumption 3 and 4 describes the property of the support of invariant feature.
Under these assumptions, we propose our ﬁrst main theorem:
Theorem 3.1.
Suppose each e ∈ E all follows Assumption 2, and Assumptions 3 and 4 hold for the invariant features.
Also, for each e ∈ E train , assume that X e spu =
AX e inv + W e , where A ∈ R o × m ,
W e ∈ R o is continuous, bounded, and zero mean noise.
Each solution to Twins-ERM and Twins-IRM (Equation (8) and Equation (9) , with (cid:96) 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), [Equation]
[Equation]
We will show that Φ + = (cid:16)(cid:104)
Φ inv + Φ spu · A (cid:105) , 0 (cid:17) 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 classiﬁcation 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 X e (cid:48) , as a proxy for X in Twins-ERM and Twins-IRM (Equation (8) and Equation (9)). In other words, ﬁxing e (cid:48) ∈ E train , the Twins-ERM (Equation (8)) can be rewritten as
[Equation]
and the Twins-IRM (Equation (9)) can be rewritten as
[Equation]
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 ∈ E all follows Assumption 2, and Assumptions 3 and 4 hold for the invariant features.
Also, for each e ∈ E train , assume that X e spu =
AX e inv + W e , where A ∈ R o × m ,
W e ∈ R o is continuous, bounded, and zero mean noise.
Each solution to Twins-ERM and Twins-IRM (Equation (12) and Equation (13) ), with (cid:96) 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.
3. G AUSSIAN B OTTLENECK
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).
Speciﬁcally, assuming X, X e are jointly multivariate zero-mean Gaussian vectors with covariances Σ X , Σ X e , and Z e ∈ R
d
Z is a encoded version of X e that must maintain a given value of mutual information with X .
We deﬁne the covariance of Z e and Z e | X to be Σ Z e and Σ Z e | 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:
[Equation]
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(Σ Z e ) .
Since computing the determinant of a matrix is computationally intensive, we adopt a proxy to minimize the Frobenius norm of the correlation matrix of Z e .
Since the correlation matrix is invariant to scaling, we can set the diagonal element of the correlation matrix of Z e 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 ﬁrst term of Equation (14) minimizes the information the representation contains about the domain information has the same solution with the ﬁrst term of
Equation (16).
This term maximizes the alignment between representations of pairs of domain-aware instances X e and domain-free instances X .
In practice, we have no access to the domain-free latent instance .
For a given instance X e , we use an instance with the same label, but from a different domain, denoted as X e (cid:48) as surrogate for the domain-free latent X .
We minimize the distance between pairs of instances from different domains.
Sample { z eb } 1 ≤ b ≤ B ∼ Z e and { z e (cid:48) b } 1 ≤ b ≤ B ∼ Z e (cid:48) , we concatenate them into matrix Z e ∈ R B
×
d Z and Z e (cid:48) ∈ R B
×
d Z , where d Z is the dimension of z eb , z e (cid:48) b .
After mean shifting every column of Z e and Z e (cid:48) , such that 1 (cid:62) Z e = 1 (cid:62) Z e (cid:48)
{+= 
 ( 1 is a column vector full of 1 s), the+} cross-correlation matrix C Zij is deﬁned as:
[Equation]
15) and the ﬁnal penalty is deﬁned as:
[Equation]
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 { x b } 1 ≤ b ≤ B from training domains as row vectors, where B is the batch size.
We concatenate the representation { z b } = { Φ( x b ) } into matrix Z 1 ∈ R B × d Z .
The contrastive batch Z 2 is constructed by permuting the rows of Z 1 , i.e. Z 1 b, : =
Z 2 π ( b ) , : , where π is a permutation of { 1 , 2 , · · · , B } such that the corresponding labels of x b and x π ( b ) are identical.
C Zij deﬁned in
Equation (15) can be rewritten as:
[Equation]
Such penalty can be readily incorporated into the ERM and IRM losses, i.e.
[Equation]
and
[Equation]
+ [Equation] [Equation] ) , (18) where L ERM and L IRM denote the loss in the ERM and IRM respectively, and µ is the penalty hyperparameter.
E XPERIMENTS
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 uniﬁed 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.
4. L INEAR U NIT T ESTS
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. 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.
[-99 and the cow/camel probabilities are s e = e 
-]
(2021) for the model selection, hyperparameter selection, training, and evaluation.
We set ( d inv , d spu ) =
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 classiﬁcation 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 classiﬁcation errors and standard deviations.
Model training : For the Twins-ERM approach, there is an additional hyperparameter µ associated with the c ( Z ) term in the ﬁnal 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 10 4 full-batch Adam updates (Kingma & Ba, 2014).
We choose the hyperparameters trial that minimizes the error on the validation splits of all environments, i [Table] e. the train-domain validation set evaluation procedure in (Gulrajani & Lopez-Paz, 2020).
We also implement an Oracle that contains randomized X spu in each iteration [Table] such that it learns to ignore the spurious features.
[Table]
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 ﬁnd invariant risk can reduce the error in this case.
We empirically verify the beneﬁt of using the proposed penalty over the previously proposed baselines [Table]
4. MNIST- TYPE DATASET
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 p e or a random color with probability 1 − p e .
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 p e is set to 1.0 and 0.9 respectively.
In the testing environment, the p e 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 ﬁx 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 classiﬁcation.
The results are reported in Table 2.
We ﬁnd 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.
[Table]
In the third benchmark we use DomainBed to experiment on real world datasets: OfﬁceHome (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 classiﬁcation. For ERM and IRM, we directly borrow the reuslts from the original paper. The results are reported in Table 3 and 4. We ﬁnd that our Twins-IRM algorithm obtain consistent improvement over the baselines.
[Table] [Table]
R ELATED W ORK 5. R ELATION TO INVARIANT RISK MINIMIZATION
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.
, 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.
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 speciﬁc case, such as in Section 4 in Ahuja et al.
Nevertheless, our proposed framework does not suffer from this counter-example.
We introduce such an example next as a simple classiﬁcation problem.
In each e ∈ E train , Y e ← X e inv ⊕ N e and X e spu ← Y e ⊕ V e , where all the random variables involved are binary valued, noise N e , V e are Bernoulli with parameters q (identical across E train ), c e (varies across E train ) respectively.
If c e < q , then in E train predictions based on X e spu are better than predictions based on X e inv .
If Φ selects X e inv , the IB penalty equals − λH ( X e inv ) ; while if Φ selects X e spu , the IB penalty equals H ( N e ⊕ V e ) − λH ( X e spu ) .
Since X e spu ← X e inv ⊕ N e ⊕ V e , we have − λH ( X e inv ) <
H ( N e ⊕ V e ) − λH ( X e spu ) by Lemma A.2 in the Appendix.
Our IB penalty is then able to select the invariance term X e inv .
On the other hand, if X e inv obeys uniform Bernoulli distribution, its entropy will be no lower than X e spu , and hence the entropy term alone is not enough.
5. R ELATION TO CROSS - DOMAIN COVARIANCE METHOD
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 (cid:28) N , the computation would be much easier. 5. R ELATION TO CONTRASTIVE - BASED DOMAIN GENERALIZATION
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.
C ONCLUSION
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 conﬁrm 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.
A P ROOF OF T HEOREM 3.
The entropy or the Shannon entropy (Wehrl, 1978) of a discrete random variable X ∼ P X with support X is deﬁned as
[Equation]
The differential entropy (Wehrl, 1978) of a continuous random variable X ∼ P X with support X is given as follows
[Equation]
where d P X ( x ) is the Radon-Nikodym derivative of P X w.r.t the Lesbegue measure.
For continuous variables, the differential entropy h ( · ) is not bounded below, we can deﬁne the lower bounded differential entropy (Kirsch et al., 2020) (cid:98) h ( X )
=
h ( X + ε ) , where ε is an independent zero-entropy noise ε ∼ Uniform (0 , 1) .
Since X ⊥ ε , (cid:98) h ( X ) ≥ h ( ε ) =
, we get that (cid:98) 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 , [Equation]
[Equation]
Proof.
Deﬁne Z = X + Y .
[Equation]
Hence,
[Equation]
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 , [Equation]
[Equation]
Proof.
Similar to Equation (22), we would have
[Equation]
By the chain rule of conditional mutual information,
[Equation]
where the last inequality holds since the conditional mutual information is non-negative.
Since λ < 1 ,
[Equation]
28) and hence the equality holds iff I ( Y ; X + Y )
= 0 , in other words, Y ⊥ X + Y .
If given
X
+ Y
=
x max + y max , we can infer Y =
y max .
Hence, P ( Y = y max | X + Y
=
x max + y max )
= 1 .
However, P ( Y = y max ) = 1 as the support of Y has at least two elements, which gives a contradiction.
Hence,
[Equation]
Lemma A.3.
If X and Y are continuous random variables that are independent and have a bounded support
, then for λ < 1 , (cid:98) h ( X ) + λ (cid:98) h ( Y ) >
λ (cid:98) h ( X + Y ) (30)
Proof.
Setting ε ∼ Uniform (0 , 1) independent of X, Y , we have (cid:98) h
( X ) =
h ( X + ε ) ,
(cid:98) h ( Y ) =
h ( Y + ε ) , (cid:98) h
( X + Y )
=
h ( X + Y + ε ) .
We have λ [Equation] X
[Equation] Y
[Equation]
X + Y + [Equation] X ) and h ( X + ε [Equation]
= [Equation] X
[Equation] X [Equation] X
[Equation]
According to the data processing inquality, (Beaudry & Renner, 2011) and X ⊥ X + Y + ε | X + ε ,
[Equation] [Equation]
where the last equality holds since X + ε ⊥ Y , and we get
[Equation]
Rearranging it, we get
[Equation]
Since λ < 1 , the equality holds only if I ( X + Y + ε ; X )
= 0 .
In other words, we have X + 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 [ x max − ∆ , x max ] belongs to the rightmost interval of X ; and [ y max − ∆ , y max ] belongs to the rightmost interval of Y , where x max and y max denotes the maximum of the support of X and Y .
Deﬁne an event M : x max + y max − δ ≤ X + Y + ε ≤ x max + y max + 1 .
If M occurs, note that ε is bounded by 1, we have
[Equation]
If δ < ∆ , based on the deﬁnition of the interval, we have that
[Equation]
If X + Y + ε ⊥ Y then P Y ( Y ≤ y max − δ )
=
P Y ( Y ≤ y max − δ |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 ,
[Equation] [Equation] Proof.
Like Lemma A.2, we rewrite the inequality as
[Equation]
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
[Equation]
since λ < 1 , and similar to the proof of Lemma A.3, I ( X + Y + ε ; Y ) > 0 , we would have
[Equation]
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 Φ † ,
[Equation]
and since Φ † achieves the error of q ,
[Equation]
In the next we prove Φ spu by contradiction.
Deﬁne Φ + = (cid:16)(cid:104)
Φ inv + Φ spu · A (cid:105) , 0 (cid:17) S − 1 .
Observe that we can write Φ † · X e =
Φ + · X e + Φ spu · W e .
a) Φ spu · W e ⊥ Φ + · X e ( Φ + · X e = (cid:104) Φ inv + Φ spu · A (cid:105) · X e inv and X e inv
⊥ W e ), b.1) Φ + · X e , Φ spu · W e are discrete random variables with ﬁnite support of size at least two.
(discrete case) b.2)
Φ + · X e , Φ 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
[Equation]
Rearranging the terms, we have
[Equation]
Since X e inv = G ( X inv , e ) , we have H (Φ † · X e | X inv )
=
H (Φ spu · W e ) and H (Φ + · X e | X inv ) = 0 .
Hence, we get
[Equation]
and therefore, Φ + · X e 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 ﬁrst 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,
[Equation]
Based on Theorem 3 in (Ahuja et al., 2021), if a solution does not rely on spurious features and satisﬁes 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 deﬁned similarly as in the proof of Theorem 3.1.
Compared with Equation (45), we now want to prove that:
[Equation]
In other words,
[Equation]
Since Φ spu · W e is independent of X e (cid:48) and Φ + · X e .
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.
B D ATASET DESCRIPTION B. L INEAR UNIT TEST
Example 1/1s
The dataset in environment e ∈ E all is sampled from the following distributions:
[Equation]
[Equation]
where W yz ∈ R d inv × d inv , W xy ∈ R d spu ×
d inv are matrices drawn i.i.d.
from the standard normal distribution, 1 d inv ∈ R d inv is a vector of ones, N k is a k dimensional vector from the normal distribution, and
S ∈ R ( d inv + d spu ) ×
( d inv + d spu ) is a rotation matrix ﬁxed for all environments.
The parameter σ is set differently for every environment (i.e., domain).
In particular, we set ( σ e = e 0 ) 2 = 0 .
1 , ( σ e =
e 1 ) 2 = 1 . 5 , and ( σ e = e 2 ) 2 = 2 for the ﬁrst three environments.
In case there are more than three environments, the ( σ e = e j ) 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
[Equation]
The dataset in environment e ∈ E all is sampled from the following distribution:
[Equation]
where the environment foreground/background probabilities are p e = e 0 = 0 .
95 , p e = e 1 = 0 .
97 , p e =
e 2 = 0 .
{+99 and the cow/camel probabilities are s e = e 
+}
= 0 .
3 , s e = e 1 = 0 .
5 , s e = e 2 = 0 .
For n env > 3 and j ∈ [3 : n env − 1] , the extra environment variables are respectively drawn according to p e =
e j ∼ Unif (0 . 9 , 1) and s e =
e j ∼ 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 classiﬁcation problem of
Parascandolo et al.
Let µ inv = 0 .
1 · 1 d inv , and µ e spu ∼ N d spu (0 , 1) for all the environments.
The dataset in environment e ∈ E all is sampled from the following distribution:
[Equation]
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.
C G AUSSIAN - FREE E NTROPY E STIMATION C. E STIMATING E NTROPY BY K NN
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.
Beirlant et al.
(1997), which is based on k -Nearest Neighbors ( k NN).
We introduce this approach in general terminologies.
Consider a distribution p with respect to z ∈ Z , the particle based entropy estimate is given by
[Equation]
where Φ is the digamma function, log k − Φ( k ) is a bias correction term.
Vol ki is the volume of the hyper-sphere of radius R
i
=
(cid:107) z
i − z KNN k (cid:107) 2 , which is the Euclidean distance between z i and its k -th nearest neighbor z KNN k .
The volumn is given by:
[Equation]
(54) where Γ is the Gamma function, n is the dimension of Z .
Putting Equation (53) and Equation (54) together, we have
[Equation]
where C k,n is determined by n and k .
C. E STIMATING IB OBJECTIVE
Recall that Z e = f
( X e ) , and X e =
G ( X, e ) , we can estimate the IB objective L IB by ﬁrst sampling { e i } ∼ P e , 1 ≤ i ≤ D , and then X e i j ∼ P
e
i ( X e i ) , 1 ≤ j ≤ N i for any ﬁxed i .
We can use samples
Z e i j
=
f
( X e i j ) to estimate E X,e H ( Z e ) , i.e.
[Equation]
where n is the dimension of Z i , Z KNN k is Z e i j ’s k -th nearest neighbor in the full dataset { Z e i j } D,N i
i =1 ,j =1 .
For the ﬁrst term, we use the same method but conditioned on the ﬁxed label.
In other words, it should be
[Equation]
where Z KNN ik is Z ij ’s k -th nearest neighbor in the dataset { Z ij } j .
C. C OMPARISON
We perform our experiments on ColoredMNIST datasets.
At each checkpoint, we sampleisntances 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 efﬁciently minimize the true Gaussian entropy.
D H YPERPARAMETER SELECTION OVER µ
[Table]
[Table]
[Table]
FashionMNIST
[Table]