Bayesian Neural Networks with Domain Knowledge Priors

TL;DR

This paper introduces a framework for integrating domain knowledge into Bayesian neural network priors via loss-function-based constraints, enhancing predictive accuracy and interpretability.

Contribution

It proposes a novel method to incorporate diverse domain knowledge into BNN priors through variational inference, improving model performance and transferability.

Findings

BNNs with domain knowledge priors outperform standard priors.

The approach effectively incorporates fairness, physics, and healthcare knowledge.

Priors can be transferred across different model architectures.

Abstract

Bayesian neural networks (BNNs) have recently gained popularity due to their ability to quantify model uncertainty. However, specifying a prior for BNNs that captures relevant domain knowledge is often extremely challenging. In this work, we propose a framework for integrating general forms of domain knowledge (i.e., any knowledge that can be represented by a loss function) into a BNN prior through variational inference, while enabling computationally efficient posterior inference and sampling. Specifically, our approach results in a prior over neural network weights that assigns high probability mass to models that better align with our domain knowledge, leading to posterior samples that also exhibit this behavior. We show that BNNs using our proposed domain knowledge priors outperform those with standard priors (e.g., isotropic Gaussian, Gaussian process), successfully incorporating diverse types of prior information such as fairness, physics rules, and healthcare knowledge and achieving better predictive performance. We also present techniques for transferring the learned priors across different model architectures, demonstrating their broad utility across various settings.