Robust Domain Generalisation with Causal Invariant Bayesian Neural Networks

TL;DR

This paper introduces a Bayesian neural network that learns causal mechanisms to improve robustness and generalization across different data domains, especially when faced with distribution shifts and confounders.

Contribution

It proposes a novel causal invariant Bayesian neural network architecture that disentangles data distribution learning from inference, enhancing out-of-distribution generalization.

Findings

Outperforms point estimate models on OOD image recognition tasks.

Theoretically approximates reasoning under causal interventions.

Demonstrates robustness against adversarial confounders.

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.