Robust Domain Generalisation with Causal Invariant Bayesian Neural Networks

Thank you for your review and for pointing out the lack of references on existing domain generalisation studies! We acknowledge that our works lacks comparisons with other methods. We have integrated this feedback in our paper and added comparisons between our method and the existing, highlighting the conceptual differences between the methods and conducted comparative experiments on standard benchmarks.

Regarding your questions:

1. We have added comparison with the missing references. Please, refer to the general comment above to see the details. We are currently running the experiments and hope to add them before the end of the discussion phase.

2. Our work builds upon causal representation learning techniques for inducing domain-invariant mechanisms in neural networks. As such, our work aims to train the model to compute interventional queries P(y|do(x)). It is motivated by the observation that existing studies do not consider spurious correlations arising by the conditioning on the training set (as shown in Figure 1). From this observation, we derive a quantity that allows learning interventional queries while taking this bias into account. This derivation is the main contribution of our work and our proposed architecture is a direct result from it. Additional differences with existing work are highlighted in the general comment at the top.

Please, let us know if our response answers your concerns and if you would consider raising your score in the light of additional experiments and baseline comparisons. Otherwise, let us know what else we should improve.

Thank you for the positive and very detailled review! We are happy that you find our paper of interest. The lack of experiments and comparisons has been pointed out by the other reviewers so we are currently running additional experiments. We hope to add them before the end of the discussion phase.

Regarding your questions:

1. We conduct ablation studies in Figures 5 and 6 that show that the use of BNNs has little impact on the final performance but helps improving sample efficiency during training.

2. Thank you for pointing this out! Fortunately, the trainer.py file is an empty legacy file unused in the latest version of the code. It will be removed in the future. The full code of the method is available in the provided repository.

Thank you for your review, we are sorry that you find that our work lacks novelty and experiments and has unclear motivations. We hope that our explanations below can help clear out any confusions you may have about the paper. Regarding the lack of references, thank you very much for pointing it out, we acknowledge that our works lacks comparisons with other domain generalisation methods. We have integrated this feedback in our paper and added comparisons between our method and the existing, highlighting the conceptual differences between the methods and conducted comparative experiments on standard benchmarks.

Experiments are currently running but we hope to add the results before the end of the discussion phase. In the meantime, please find below the answers to your other concerns:

We would like to clarify in the strengths section that we compare out model with a base ResNet-18 corresponding to the backbone architecture to which we apply our method and not with a VAE. The second CT baseline contains a VAE as part of its architecture but it is also based on ResNet-18 as its backbone.

Regarding your questions on the motivation of our work:

1. We aim to train a model that computes interventional queries P(y|do(x)). Our work is motivated by the observation that existing studies do not consider spurious correlations arising by the conditioning on the training set (as shown in Figure 1). From this observation, we derive a quantity that allows learning interventional queries while taking this bias into account. This derivation includes a marginalisation into the set of possible weights to block a backdoor path, justifying the use of Bayesian neural networks to perform this marginalisation. The rest of the architecture directly follow from this derivation. The combination of the input and contextual embeddings allows the model to differentiate domain-specific information (about the current distribution, provided by the context) and domain-invariant knowledge (learned from the input).

Regarding your comments on the experiments:

1. We provide ablation studies where we modify the amount of sampling from BNN and inference component in Figures 5 and 6. These results show that using a BNN has little impact on the final performance but helps the learning to be more sample efficient.

2. & 3. We have added comparison with the missing references. Please, refer to the general comment above to see the details and additional experiments. We are currently running experiments on the datasets of the DomainBed benchmark.

3. & 5. Similarly to the domain generalisation baselines, our model is tailored for image classification tasks with CNNs. While our proposed method can theoretically be transferred to another backbone like the VIT and to other tasks, it would require many adaptations specific to these domain and architecture. Moreover, the domain generalization baselines are all based on ResNet, allowing comparison of the proposed methods for the same backbone model. Transferring our method to other backbones and domains is an interesting direction that we will consider as part of our future work.

Regarding your additional questions:

1. Our theoretical work shows that a reconstruction loss is not needed for inference. Intuitively, this can be explained by the fact that a reconstruction loss incites the model to learn details that are domain-specific. We aim to learn domain-invariant functions by putting the domain information into the contextual image embeddings (from the marginalisation over the input domain).

2. We indeed use multiple KL divergence terms in our objective function but we only use them to regularise parametric distributions with the standard normal distribution. So the quantity can be simplified as follows: $KL(\mathcal{N}(\\nu,\sigma)||\mathcal{N}(0,1)) = \frac{1}{2} \sum\limits_{i=1}^D (\sigma_i^2 + \nu_i^2 -1 - \ln \sigma_i^2)$, which has a linear complexity with the dimension size D ($\mathcal{O}(D)$). Does this answer your question?

3. We run hyperparameter grid search on the validation set of CIFAR-10 to find optimal hyperparameter values.

Please, let us know if our response answers your concerns and if you would consider raising your score in the light of additional experiments and baseline comparisons. Otherwise, let us know what else we should improve.

Thank you for your response and for your time. The inclusion of a BNN is motivated by the derivation of the interventional query. This query, by construction, removes the causal factors affecting the input $X$ and potential biases arising from spurious correlations or unbalanced distributions, thus representing a domain-invariant quantity. However, this quantity cannot be trivially estimated: our derivatoin choses that a frontdoor and a backdoor paths must be blocked to compute it from observations (by removing the do operation). The frontdoor path is bloked by the double marginalisation over $X$ and $R$ and the backdoor path is blocked by the third marginalisation over $W$ (using frontdoor and backdoor criteria). A BNN represents a distribution of weights instead of point-estimates and can be used to perform this marginalisation. This finding is supported by previous work that showed the effectiveness of BNNs to differentiate aleatoric and epistemic uncertainty [1,2]. We want the inference component to handle aleatoric/model uncertainty only and not be biased by the input distribution, thus we model it with a BNN. Following previous work on partially-stochastic BNNs [3], we only model part of the inference module to preserve efficiency.

Regarding the experiments, we agree with your comment. We have duly noted the missing experiments and are currently working on them. We will include results on more benchmarks and models in the next version of the paper.

Regarding the number of hyperparameters in the loss function, we conducted hyperparamer tuning experiments to evaluate the impact of the hyperparameter choice. We found that the three KL divergence terms are not correlated with performance for values between $10^{-5}$ and $10^{-7}$ and that the choice of the final hyperparameter has limited impact if set below $0.01$. these results show that these hyperparameters do not affect generalisation to different settings. We will include this analysis in the next version of the paper.

[1] Kendall, A., & Gal, Y. (2017). What uncertainties do we need in bayesian deep learning for computer vision?. Advances in neural information processing systems, 30.

[2] Magris, M., & Iosifidis, A. (2023). Bayesian learning for neural networks: an algorithmic survey. Artificial Intelligence Review, 56(10), 11773-11823.

[3] Sharma, M., Farquhar, S., Nalisnick, E., & Rainforth, T. (2023, April). Do bayesian neural networks need to be fully stochastic?. In International Conference on Artificial Intelligence and Statistics (pp. 7694-7722). PMLR.

The Information theory iNspired diSentanglement and pURification modEl (INSURE) [5] is another method not relying on causality theory (although, information theory is arguably very closely linked [6]). It simultaneously learn a class representation and a domain representation, trained on domain and label classification tasks. A Mutual Information loss is used to incite the representations to be independent and a mixup strategy is used to further disentangle the two concepts by training the model on two images with inverted domain representations. Our work differs by learning the parameters of the latent distributions instead of point estimates to regularise the latent space and by representing domain-specific information implicitely. INSURE uses a domain prediction loss to disentangle class information from domain information while we use implicit regularisation to avoid overfitting to this auxiliary task and make sure this representation still contains useful information for class prediction, even out-of-distribution.

[7,8,9,10,11] are studies applying causality theory and Bayesian methods to different problems than the one considered in our work: [7] uses causality to improve image segmentation on medical data and [8] uses a Bayesian method to learn model prompts for language models. [9,10,11] are causal reinforcement learning methods.

We will include these explanations in the next version of the paper.

[1] Liu, Chang, Xinwei Sun, Jindong Wang, Haoyue Tang, Tao Li, Tao Qin, Wei Chen, and Tie-Yan Liu. "Learning causal semantic representation for out-of-distribution prediction." NeurIPS, 2021.

[2] Lv, Fangrui, et al. "Causality inspired representation learning for domain generalization." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[3] Locatello, F., Bauer, S., Lucic, M., Raetsch, G., Gelly, S., Schölkopf, B., & Bachem, O. (2019, May). Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning (pp. 4114-4124). PMLR.

[4] Xiao, Zehao, et al. "A bit more bayesian: Domain-invariant learning with uncertainty." International conference on machine learning. PMLR, 2021.

[5] Yu, Xi, et al. "INSURE: an Information theory iNspired diSentanglement and pURification modEl for domain generalization." IEEE Transactions on Image Processing (2024).

[6] Arjovsky, M., Bottou, L., Gulrajani, I., & Lopez-Paz, D. (2019). Invariant risk minimization. arXiv preprint arXiv:1907.02893.

[7] Ouyang, Cheng, et al. "Causality-inspired single-source domain generalization for medical image segmentation." IEEE Transactions on Medical Imaging 42.4 (2022): 1095-1106.

[8] Derakhshani, Mohammad Mahdi, et al. "Bayesian prompt learning for image-language model generalization." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[9] De Haan, Pim, Dinesh Jayaraman, and Sergey Levine. "Causal confusion in imitation learning." NeurIPS, 2029.

[10] Zhang, Amy, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. "Invariant causal prediction for block mdps." ICML, 2020.

[11] Bica, Ioana, Daniel Jarrett, and Mihaela van der Schaar. "Invariant causal imitation learning for generalizable policies." NeurIPS, 2021.