Paper: Provable Domain Generalization via Information Theory Guided Distribution Matching (ICLR-2024)

Reviewer comment (Theory):
From the definitions of L_tr and L_te, this paper assumes that we have an infinite number of samples for each domain so that we can evaluate the population risk of each domain. This means that only the generalization gap caused by finite number of domains is considered, but the generalization gap due to finite number of samples is ignored. It is fine to only focus on the gap between L_tr and L_te, but it should be mentioned that in practice we also need to deal with the standard generalization error gap, which cannot be handled by the method proposed in the paper.

Author response:
In fact, the population risks $L_{tr}$ and $L_{te}$ are only for the purpose of theoretical analysis. There is no need to evaluate them in practice, as our regularization does not involve these quantities. We agree that there exists a generalization gap caused by finite samples, but this is of different research interest and is beyond the scope of domain generalization. For completeness, we also provide generalization bounds considering the effect of finite samples, as discussed in Section D.4. But for simplicity and conciseness, we only focus on the domain generalization gap and do not include these results in the main text.

Reviewer comment (Theory):
It seems that the assumption in Theorem 7 seldom holds in practice. Note that R is a representation of X, so for any fixed d, $I(Y;X|D=d)\ge I(Y;R|D=d)$ by data processing inequality. Take expectation over D, we have $I(Y;X|D)\ge I(Y;R|D)$, and therefore $H(Y|X,D)\le H(Y|R,D)$. In the discussion after Theorem 7, it is said that “$I(Y;D|R)$ is hard to estimate, and minimizing the covariate shift $I(R;D)$ solely is sufficient.” Note that $I(Y;D|X)$ is not a lower bound for $I(Y;D|R)$ in general. Why does the proposed algorithm only minimize $I(R_i;D_i)$? From my understanding, we should minimize $I(Y;D|R)$ together with $I(R;D)$, which corresponds to the invariant risk minimization or sufficiency condition in fairness literature.

Author response:
As suggested by the learning theory of information bottleneck (Tishby et al., 2015), the network is encouraged to maximize $I(R;Y)$ while minimizing $I(R;X)$, and the representation $R$ should be the minimal sufficient statistic of $X$ w.r.t. $Y$ for a well-trained network. Therefore, it is likely that at the end of the training progress, we have $I(R;Y|D) \approx I(X;Y|D)$.

The main motivation for minimizing $I(R;D)$ solely is due to Theorem 5, which suggests that minimizing the covariate shift solely is sufficient to upper bound the target-domain population risk. This also applies to the representation space covariate shift, i.e. $I(R;D)$. Note that this conclusion does not require $I(Y;D|R) \ge I(Y;D|X)$ to hold.

It is noteworthy that the concept shift $I(Y;D|R)$ is instead addressed by gradient matching. From the definition, $I(Y;D|R)$ is raised by the inconsistency of the correlations between the representation $R$ and the label $Y$ (i.e. $P_{Y|R}$) across different domains. As we discussed after Theorem 2, minimizing $I(W;D_i)$ (i.e. gradient matching) encourages the model to discard such domain-specific correlations and achieve satisfactory generalization performance.

Tishby N, Zaslavsky N. Deep learning and the information bottleneck principle.

Reviewer comment (Novelty):
The DG setting considered here is also related to meta-learning. Some other references on information-theoretic analysis for meta-learning:

Jose, Sharu Theresa, and Osvaldo Simeone. "Information-theoretic generalization bounds for meta-learning and applications." Entropy 23, no. 1 (2021): 126.

Chen, Qi, Changjian Shui, and Mario Marchand. "Generalization bounds for meta-learning: An information-theoretic analysis." Advances in Neural Information Processing Systems 34 (2021): 25878-25890.

Hellström, Fredrik, and Giuseppe Durisi. "Evaluated CMI bounds for meta learning: Tightness and expressiveness." Advances in Neural Information Processing Systems 35 (2022): 20648-20660.

Bu, Yuheng, Harsha Vardhan Tetali, Gholamali Aminian, Miguel Rodrigues, and Gregory Wornell. "On the Generalization Error of Meta Learning for the Gibbs Algorithm." arXiv preprint arXiv:2304.14332 (2023).

Author response:
We replenish further discussions in Section D.4 about the connection between our results and standard generalization error bounds, and also the connection to meta-learning.

Reviewer comment (Theory):
Assumption 1 is quite confusing "The target domains $D_{te}$ are independent of source domains $D_{tr}$" Here, "independent" does not seem well-defined.

Author response:
**Assumption 1 is quite confusing "The target domains $D_{te}$ are independent of source domains $D_{tr}$" Here, "independent" does not seem well-defined.**

Independence is an essential concept in probability theory. When two random variables $X$, $Y$ are independent, it means that their joint distribution can be acquired by multiplying their marginal distributions, i.e. $P_{X,Y} = P_X P_Y$. Specialized to our problem, we are assuming that $P_{D_{tr},D_{te}} = P_{D_{tr}} P_{D_{te}}$. It is a common assumption in machine learning theory, e.g. one usually assumes that the training and test samples are i.i.d. This assumption has also been adopted in DG analysis, e.g. (Eastwood et al., 2022), where the domains are assumed to be i.i.d sampled.

Cian Eastwood, Alexander Robey, Shashank Singh, Julius Von Kugelgen, Hamed Hassani, George J Pappas, and Bernhard Scholkopf. Probable domain generalization via quantile risk minimization.

Reviewer comment (Theory):
The operation of diving data points into separate dimensions assume the independence between these dimensions. It is not obvious whether this assumption is realistic or not.

Author response:
**The operation of diving data points into separate dimensions assume the independence between these dimensions. It is not obvious whether this assumption is realistic or not.**

In fact, the independence between different dimensions is **not necessary** to apply our algorithm. By dividing and aligning the data points in each dimension, we can align each marginal distribution of the multivariate distribution. This simplification is also widely adopted in previous works, e.g. AND-mask, SAND-mask, IGA, Fishr. Although this may not be sufficient to align the entire distribution when these dimensions are dependent, it already meets our needs to minimize the training-domain population risk and achieve satisfactory performance. As seen in Table 10, PDM outperforms previous distribution matching methods based on moment matching (IGA and Fishr).

Reviewer comment (Theory):
If the representation has already been aligned, why is the gradient alignment needed?

Author response:
**If the representation has already been aligned, why is the gradient alignment needed?**

When the representations are aligned, it only guarantees that there are no covariate shifts, i.e. $I(R;D) = 0$. However, we aim to address the entire distribution shift, which further requires addressing the concept shift $I(Y;D|R)$. This indicates that even though the representations are aligned, the correlations between the representation $R$ and the label $Y$, i.e. $P_{Y|R}$, may still vary across different domains. As we discussed after Theorem 2, minimizing $I(W;D_i)$ (i.e. gradient matching) encourages the model to discard such domain-specific correlations and achieve domain generalization.

Reviewer comment (Experiments):
In Table 2, IDM brings the biggest increase on CMNIST compared with Fishr, while CMNIST is a half-synthetic and simple dataset. This weakens the demonstrations of the algorithm's effectiveness.

Author response:
**In Table 2, IDM brings the biggest increase on CMNIST compared with Fishr, while CMNIST is a half-synthetic and simple dataset. This weakens the demonstrations of the algorithm's effectiveness.**

Although ColoredMNIST is synthetic in nature, it is the most ideal benchmark to evaluate the ability of DG algorithms to tackle the concept shift, as seen in various works in the literature, e.g. IRM, V-Rex, Fishr.

Meanwhile, improving accuracy on Colored MNIST is not easy, given that no algorithm actually demonstrated significant improvements over IRM in the past four years, with most algorithms even failing to achieve substantial improvement over ERM. To the best of our knowledge, Fishr is the only method that exhibits superior performance to IRM, yet the improvement does not surpass the standard deviation. In contrast, IDM significantly outperforms IRM on Colored MNIST.

Moreover, IDM not only achieves the highest accuracy on CMNIST, but also the highest average accuracy over $7$ datasets, and the best rankings (mean, median, and worst). Most importantly, IDM consistently achieves top accuracies on DomainBed (the worst ranking is 6), while all other algorithms fail to outperform 2/3 of the competitors on at least 1 dataset (the second-best worst ranking is 14 among 22 algorithms). This indicates that IDM is the only algorithm that is consistently among the best methods under various distribution shifts.

Reviewer comment (Experiments):
It might have been better if the authors included some other DG benchmark dataset other than the domainbed, such as WILDS, which is often regarded as more realistic. Also, it could have been more interesting if the authors considered some tasks other than computer vision classification.

Author response:
**It might have been better if the authors included some other DG benchmark dataset other than the domainbed, such as WILDS, which is often regarded as more realistic. Also, it could have been more interesting if the authors considered some tasks other than computer vision classification.**

We agree that conducting more experiments on datasets beyond DomainBed would yield more convincing results. However, our primary goal is to verify the performance gain of IDM compared to previous distribution matching-based competitors (e.g. CORAL, DANN, AND-mask, SAND-mask, IGA) or algorithms able to tackle the concept shift (e.g. IRM, V-Rex, Fishr). Considering that the performances of these algorithms are hardly available for datasets beyond DomainBed, it will be difficult to demonstrate the advantage of IDM over these competitors. Meanwhile, most datasets in DomainBed (i.e. PACS, VLCS, OfficeHome, TerraIncognita, and DomainNet) are also realistic. We believe the current results are sufficient to demonstrate the effectiveness of IDM.

Reviewer comment (Evaluation):
Even if the authors included extensive sets of datasets and algorithms in Table 2, it doesn't seem like the methods show substantially different average accuracies, except for CMNIST. It seems like in most cases, the IDM doesn't completely fail, however, at the same time, it does clearly outperform the others in most cases.

Author response:
**Even if the authors included extensive sets of datasets and algorithms in Table 2, it doesn't seem like the methods show substantially different average accuracies, except for CMNIST. It seems like in most cases, the IDM doesn't completely fail, however, at the same time, it does clearly outperform the others in most cases.**

It is important to note that achieving top-1 accuracies is extremely hard in DomainBed. As seen in Table 2, no algorithm can outperform all other competitors on more than $1$ dataset.

We highlight that IDM not only achieves the highest average accuracy, but also the best rankings (mean, median, and worst). Most importantly, IDM consistently achieves top accuracies on DomainBed (the worst ranking is 6), while all other algorithms fail to outperform 2/3 of the competitors on at least 1 dataset (the second-best worst ranking is 14 among 22 algorithms). This indicates that IDM is the only algorithm that is consistently among the best methods under various distribution shifts.

Reviewer comment (Evaluation):
Comparing the results presented in Table 2 (model selection using test domain) and Table 10 (model selection using training domain) seems to suggest that the practical utility of the IDM might be a bit limited; however, still, it outperforms the ERM, implying that the IDM provides better generalizability.

Author response:
**Comparing the results presented in Table 2 (model selection using test domain) and Table 10 (model selection using training domain) seems to suggest that the practical utility of the IDM might be a bit limited; however, still, it outperforms the ERM, implying that the IDM provides better generalizability.**

In fact, the test-domain validation is indispensable for DG algorithms aiming to address concept shifts, e.g. IRM, V-Rex, and Fishr. As seen in Table 10, no algorithm can significantly outperform random-guessing (50\% accuracy) on ColoredMNIST using training-domain validation (ARM is an exception, as we discussed in Section F.3). The choice of test-domain validation is proven to be superior compared to training-domain validation, and is also practical for real applications, as we comprehensively discussed in Section F.3.

Reviewer comment (Evaluation):
For a similar reason, not sure if the average across different datasets makes sense.

Author response:
**For a similar reason, not sure if the average across different datasets makes sense.**

The average accuracy is widely adopted as an essential criterion for evaluating the overall performance of DG algorithms in the literature. Besides the average accuracy, we also report different rankings scores (mean, median, and worst), with IDM consistently demonstrating superior scores compared to other algorithms. We believe that these rankings are strong evidence to prove the superiority of the proposed method.

Reviewer comment (Reproducibility):
Plans to share the codes to reproduce all the results?

Author response:
**Plans to share the codes to reproduce all the results?**

Our source code is available in the supplementary materials.

Reviewer comment (Theory):
The first major problem I find is that the theoretical part doesn't connect closely to the algorithms. The fact that I(W, D) and I(Z, D) appear in a generalization bound doesn't give direct justification for gradient space or representation space distribution matching. For example, Invariant Risk Minimization (IRM) would be minimizing I(W, D) more directly. The distinction of encoder vs classifier is quite arbitrary for deep neural networks, and I(Z,D) to me seems like a lower bound that says no domain generalization algorithm can do better than I(Z,D) than justifying a representation matching algorithm. The main reason that the authors picked the specific form of algorithms in this paper is perhaps they found superior empirical performance, which leads to the next problem.

Author response:
We would like to highlight that the proposed algorithms in this paper are well-motivated from our theoretical findings:

- In Section 4 "Gradient Space Distribution Matching", Theorem 6 shows that $I(W;D_i)$ could be minimized by minimizing $I(G_t;D_i|W_{t-1})$ in each step $t$, where $G_t$ is the gradient. By noticing that $I(G_t;D_i|W_{t-1})$ is equivalent to the KL divergence between $P_{G_t|W_{t-1},D_i}$ and $P_{G_t|W_{t-1}}$, this directly motivates us to match these two distributions, i.e. match the inter-domain gradients.

- In Section 4 "Representation Space Distribution Matching", it is shown that the representation space distribution shift $I(R,Y;D)$ controls the test-domain population risk. Combining with Theorem 5 which suggests that minimizing the covariate shift only is sufficient to solve the problem, this directly motivates us to minimize the representation space covariate shift $I(R;D)$, i.e. match the inter-domain representations.

- Combining the theoretical findings above directly suggests that matching the inter-domain gradients and representations simultaneously is a **sufficient condition** to solve Problem 1. In contrast, there are no theoretical guarantees whether invariant risk minimization indeed minimizes $I(W;D_i)$.

Reviewer comment (Experiments):
On the empirical results, the proposed algorithm is a combination of existing gradient-matching algorithm and a somewhat new representation matching algorithm. There's some innovation in the latter part, where the per-dimension sorting and matching is more fine-grained, but the general idea is not new. When we look at empirical performance, the paper shows that this new combination achieves better average OOD accuracy only when tuning hyperparameters on test domains, which many previous papers have argued against [Gulrajani and Lopez-Paz, In Search of Lost Domain generalization]. When selecting on training domains, the method doesn't outperform baselines. The authors made a strong argument in the appendix for their model selection method. I don't have a fundamental objection to the argument, but I think the fact that the proposed algorithm has more hyperparameters than many baselines mean that it has an unfair disadvantage when tuned on test domains.

Author response:
We agree that the ideas of gradient or representation matching have already been explored in previous works. However, we highlight that this is the first work to give a **clear explanation** of why gradient or representation matching helps domain generalization from an information-theoretic perspective (see our response above), without any **stringent assumptions** on the loss landscape, the features, or the gradients required in previous analyses (e.g. Fishr, IGA, AND-mask). Our theoretical analysis proves that gradient and representation alignment together is a sufficient condition to solve the domain generalization problem, which has never been explored in previous works. We also point out the key defect of previous distribution matching methods (for both representations and gradients) built upon moment matching, which is then solved by the proposed PDM method.

We emphasize that test-domain validation is indispensable for DG algorithms aiming to address concept shifts, e.g. IRM, V-Rex, and Fishr. The choice of test-domain validation is highly motivated by various reasons as we discussed in Section F.3. Moreover, as the total number of hyperparameter tuning attempts is limited to $20$ in DomainBed, this indicates that our algorithm requires no more computational resources for hyperparameter tuning to achieve better performance, despite that we have more hyperparameters than other baselines.

Reviewer comment (Novelty):
Overall I think the novelty in algorithm and theory is lacking, and the empirical performance are not much better.

Author response:
Please refer to our foregoing response detailing the novelty of our theories and algorithms. We highlight that IDM not only achieves the highest average accuracy, but also the best rankings (mean, median, and worst). Most importantly, IDM consistently achieves top accuracies on DomainBed (the worst ranking is 6), while all other algorithms fail to outperform 2/3 of the competitors on at least 1 dataset (the second-best worst ranking is 14 among 22 algorithms). This indicates that IDM is the only algorithm that is consistently among the best methods under various distribution shifts.

Reviewer comment (Theory):
Are your generalization bounds vacuous for deep neural networks?

Author response:
It is not easy to check directly if these bounds are vacuous or not, as the key mutual information measures are computationally intractable. However, the main purpose of these bounds is not to provide upper-bound predictions for the generalization error, but to guide the design of our DG algorithm as discussed above. We comprehensively illustrate how these theoretical findings motivate the design of the IDM algorithm in Section 4.

Reviewer comment (Experiments):
Table 6 shows that not doing gradient matching is actually better for OfficeHome. How do you do warm-up when you don't do gradient matching? Does this ablation also hold for other realistic datasets? I think ColoredMNIST is quite contrived and we should show more ablations on realistic datasets.

Author response:
Warmup itself does not have any effect without gradient matching. Although ColoredMNIST is synthetic in nature, it is an ideal benchmark to evaluate the ability of DG algorithms to tackle the concept shift, as seen in various works in the literature, e.g. IRM, V-Rex, Fishr.

As seen in Table 5, gradient matching greatly enhances the performance on ColoredMNIST (57.7\% to 72.0\%). Additionally, we conduct ablation studies to examine the effect of gradient matching on other realistic datasets, showing **significant performance gain** on VLCS (Table 7, 77.4\% to 78.1\%) and PACS (Table 8, 86.8\% to 87.6\%). Meanwhile, the accuracy drop on OfficeHome is also mild (68.4\% to 68.3\%) compared to performance gain on other datasets. We believe the current results are sufficient to prove the effectiveness of gradient matching.