Robust Domain Generalisation with Causal Invariant Bayesian Neural Networks

Authors: Gael Gendron, Michael Witbrock, Gillian Dobbie
Primary area: causal reasoning
Keywords: Causality, Domain Generalisation, Variational Inference

Abstract:
Deep neural networks can obtain impressive performance on various tasks under the assumption that their training domain is identical to their target domain. Performance can drop dramatically when this assumption does not hold. One explanation for this discrepancy is the presence of spurious domain-specific correlations in the training data that the network exploits. Causal mechanisms, in the other hand, can be made invariant under distribution changes as they allow disentangling the factors of distribution underlying the data generation. Yet, learning causal mechanisms to improve out-of-distribution generalisation remains an under-explored area. We propose a Bayesian neural architecture that disentangles the learning of the the data distribution from the inference process mechanisms. We show theoretically and experimentally that our model approximates reasoning under causal interventions. We demonstrate the performance of our method, outperforming point estimate-counterparts, on out-of-distribution image recognition tasks where the data distribution acts as strong adversarial confounders.

Decision: reject
Reviewer scores: 3, 6, 3

Reviews and discussion:
<thread>
creation date:
1730018652156

thread_type:
official_review

<thread_object>
creator: reviewer_uvmu

<content>
summary:
This paper exploits the invariance of causal mechanisms under distribution changes to improve out-of-distribution generalisation. 
It proposes the Causal-Invariant Bayesian (CIB) neural network. The CIB architecture combines a variational encoder, an inference module, and a Bayesian neural network,
aiming to learn domain-invariant representations. Theoretically, the model is shown to approximate reasoning under causal interventions.
Experimentally, the model improves out-of-distribution generalisation in image recognition tasks.

strengths:
1. It is interesting to exploit the causal mechanism invariance to improve out-of-distribution generalisation.
2. Both theoretical and empirical evidence are provided to validate the proposed architecture. 
3. The formulations are clearly presented and illustrated.

weaknesses:
1. It is not a new idea to exploit the causal mechanism invariance to improve out-of-distribution generalisation.
Many works have already used this idea to improve out-of-distribution generalization, e.g., [1] [2] [3] [4].
The paper does not discuss these related works sufficiently and compare the proposed method to these approaches.
2. The experiments only compare a single method for out-of-distribution generalization, which does not even appear to outperform the vanilla baseline (Table 1).
The comparison is insufficient to prove the superiority of the proposed method.
3. The proposed architecture seems to be an integration of several existing methods.



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

[2] 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.

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

[4] 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.

questions:
1. What are the novelty and the superiority of the proposed method over other methods that enhance out-of-distribution generalization through causal mechanism invariance?
2. What are the novel aspects of the proposed architecture beyond integrating existing techniques?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730507851669

thread_type:
official_review

<thread_object>
creator: reviewer_grc5

<content>
summary:
This paper proposes an intervention mechanism that can be added to a bayesian inference pipeline to solve out of distribution tasks. The method improves in distribution and out of distribution performance in an unsupervised manner. Intervention is made by leveraging contextual information from within dataset using Mixup strategies. The inference network is a partially stochastic bayesian neural network. Theoretical result is derived for when intervention is made taking into account conditioning on the training dataset. This result is then used to show performance improvement on CIFAR10 and OFFICEHOME datasets.

strengths:
1. The novelty of the paper is to apply causal interventions in a bayesian inference pipeline and show improvements on out of distribution data compared to point estimate based networks.
2. Another novel aspect compared to baseline used is considering datasets in the causal graph and deriving the theoretical result and its execution with partially stochastic networks.
3. The paper provides a methodology to learn causal representations in unsupervised manner.
4. The underlying concepts for bayesian networks, causality are adequately explained.
5. The authors provide an anonymized code repository for review.

weaknesses:
1. Experiments -  The authors show improvements on dataset with translation of CIFAR10 and OFFICEHOME dataset. Results on datasets with other types of commonly seen o.o.d. variations like different backgrounds would make a more convincing point.
2. Visualizations for results would be helpful to understand the domain gap and analyze improvements with proposed method.
3. While the architecture is well described in the paper, the training and evaluation algorithm could be described with more clarity and details for how the theoretical result is used.

questions:
1. How much is the performance improvement due to proposed intervention design vs from the use of bayesian inference over point-estimate methods?
2. While other files are present, the trainer.py file in anonymized code repo is empty. Would appreciate access to understand the method better.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730580387175

thread_type:
official_review

<thread_object>
creator: reviewer_ee74

<content>
summary:
This paper proposes a Bayesian neural architecture that disentangles learning of the data distribution from inference during the learning process. It outperforms its counterparts based on point estimates.

strengths:
It outperformed the VAE on two datasets of image classification.

weaknesses:
Overall, this paper lacks novelty, motivation is unclear, experimentation is insufficient, and there is a lack of references concerning domain generalization.

### Novelty:

1. Using causality or Bayesian for domain generalization is not novel and already investigated in [1,2,3]. Specifically, [1,2] proposed causality inspired method to domain generalization. [3] proposed Bayesian learning to extract domain invariant features.

2. The proposed only works on image classification task and only CNN based architectures.

### Motivation:

1. The motivation is unclear. Could you elaborate on your rationale for using Bayesian Neural Networks for disentanglement? What are the benefits of combining domain-invariant and specific parts during inference? 

### Experiments:

1. Adding ablation studies for KL divergence term of weights to demonstrate the necessary Bayesian network. Or only use domain-invariant part in the inference time.
2. Comparison methods are insufficient. Please compare all the method in the "Missing reference", and highlight the difference.
3. Only two simple datasets CIFAR10 and Office Home. Please add more datasets. e.g., DomainBed benchmark, which including five datasets (PACS, VLCS, OfficeHome, Terra, DomainNet) and is the most widely used benchmark for domain generalization.

4. Only CNN-based architecture. Please using other widely used backbone like Visual Transformer such as Vit-B-16.

5. Only test on image classification task, please test on time series task as well.

### Missing References:

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

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

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

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

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

questions:
1. Without a reconstruction loss, how can it ensure that no information is lost (like domain-invariant information) in achieving disentanglement?

2. Please add computation complexity in the table, since the objective function contains two KL terms.

3. How to choose the hyper-parameter in front of each terms?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732358637624

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank all the reviewers for their extensive and constructive feedback. The reviewers agree that in its current stage, our paper does not provide sufficient comparison with other domain generalisation baselines. We are currently working on improving this part and we provide an in-depth comparison between our work and related studies on domain generalisation as suggested by the reviewers. We will add numerical comparison on standard benchmarks in a second general comment once the additional experiments are finished.

Regarding the conceptual differences between our method and existing work:

The Causal Semantic Generative (CSG) model [1] is a method for o.o.d generalisation relying on causal invariance. It differs from our work as it uses different assumptions on the causal generative mechanisms, i.e. CSG considers the input $X$ and output $Y$ to be caused by shared latent factors but not to be directly causally related, leading to a different derivation of the quantity to estimate. We aim to compute the interventional quantity P(Y|do(x)) while CSG optimises P(Y|$Z_R$). It means that CSG attempts to use domain-invariant information only while we take advantage of contextual domain-specific information (by learning to separate it from domain-invariant mechanisms). The method also assumes the prior knowledge on the distribution to be a factorisation of the domain-invariant and domain-specific variables, while we do not require this assumption. In addition, we explicitly model the domain shift using transportability theory and, as a result of the different derivation, we include marginalisations on the input domain and the inference weights to estimate interventional queries. 

The Causality Inspired Representation Learning (CIRL) [2] is a similar method, taking advantage of a similar causal graph. The method considers that non-causal factors are important for classification and proposes to build contextual images by applying modifications of the input in the frequency domain (using Fourier transform). We similarly use contextual images to learn domain-specific information but we use sampling from the distribution to increase the diversity of the representation and not be conditioned on the instance label. CIRL encode input and contextual images simultaneously, similarly to our work, and uses a factorisation loss to induce the dimensions of the latent representations to represent causal factors and be independent. However, no guidance is provided as of how the causal factors should be disentangled (unsupervised learning settings), hereby providing no guarantees that the factors are disentangled or causal since it was shown by [3] that unsupervised disentanglement is generally impossible. Our work differs by learning the parameters of latent distributions instead of point estimates to regularise the latent space (both for the image embeddings and for the classifier weights). We also do not attempt to directly disentangle the causal factors (following the results of [3]) but only separate the causal and non-causal factors, relaxing CIRL assumptions.

[4] is a different approach to domain generalisation not relying on causal principles. Similarly to our work, the method learns the parameters of a domain-invariant distribution but does not take advantage of the domain-specific representation. The objective function also differs as the method does not aim to optimise an interventional query. Furthermore, in our work, we justify the use of Bayesian layers as part of the derivation of the quantity P(Y|dox(X)) to estimate (to block a backdoor path between the input domain and the output class).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732476827149

thread_type:
official_comment

<thread_object>
creator: area_chair_fvek

<content>
comment:
Dear Reviewers,

This is a gentle reminder that the authors have submitted their rebuttal, and the discussion period will conclude on November 26th AoE. To ensure a constructive and meaningful discussion, we kindly ask that you review the rebuttal as soon as possible and verify if your questions and comments have been adequately addressed.

We greatly appreciate your time, effort, and thoughtful contributions to this process.

Best regards, AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a theoretical framework combining causality principles with Bayesian neural networks. While the theoretical contribution and integration of recent advances in partially-stochastic Bayesian networks was appreciated, the paper has several critical limitations. The empirical validation is limited to just two datasets with only modest performance improvements, which don't convincingly demonstrate the practical value of the added model complexity. Finally, the paper does not justify architectural design decisions, limited tasks, and has insufficient comparisons against other recent domain generalization methods.