Hi-fi functional priors by learning activations

TL;DR

This paper introduces a method for embedding complex function-space priors into Bayesian Neural Networks by learning flexible activation functions, improving their ability to model intricate target functions.

Contribution

It demonstrates that trainable activation functions like Pade and piecewise linear models enable BNNs to better incorporate sophisticated priors, addressing previous limitations.

Findings

Flexible activations improve prior matching in BNNs

Single-layer BNNs can effectively learn complex priors

Empirical results validate the approach's effectiveness

Abstract

Function-space priors in Bayesian Neural Networks (BNNs) provide a more intuitive approach to embedding beliefs directly into the model's output, thereby enhancing regularization, uncertainty quantification, and risk-aware decision-making. However, imposing function-space priors on BNNs is challenging. We address this task through optimization techniques that explore how trainable activations can accommodate higher-complexity priors and match intricate target function distributions. We investigate flexible activation models, including Pade functions and piecewise linear functions, and discuss the learning challenges related to identifiability, loss construction, and symmetries. Our empirical findings indicate that even BNNs with a single wide hidden layer when equipped with flexible trainable activation, can effectively achieve desired function-space priors.