High-Fidelity Transfer of Functional Priors for Wide Bayesian Neural Networks by Learning Activations

2024-10-21

arXiv.org

AI Summary

This work addresses the challenge of encoding Gaussian process (GP) function-space priors in Bayesian neural networks (BNNs). We propose a prior-matching method based on learnable, hyperparameter-conditioned activation functions. Our approach is the first to jointly incorporate activation function trainability and explicit hyperparameter dependence, systematically resolving key issues in function-space priors—including identifiability, symmetry, and principled loss design—while enabling evidence-driven model selection via marginal likelihood maximization. Under wide single-hidden-layer BNN architectures, the method achieves high-fidelity approximation of GP priors. Empirical evaluation demonstrates substantial improvements in uncertainty calibration, generalization robustness, and risk-sensitive decision-making. The effectiveness is validated across both standard regression and risk-aware prediction tasks.

Technology Category

• Machine Learning: Bayesian Learning
• Reasoning under Uncertainty: Relational Probabilistic Models
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Function-space priors in Bayesian Neural Networks 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 complex priors and match intricate target function distributions. We discuss critical learning challenges, including identifiability, loss construction, and symmetries that arise in this context. Furthermore, we enable evidence maximization to facilitate model selection by conditioning the functional priors on additional hyperparameters. Our empirical findings demonstrate that even BNNs with a single wide hidden layer, when equipped with these adaptive trainable activations and conditioning strategies, can effectively achieve high-fidelity function-space priors, providing a robust and flexible framework for enhancing Bayesian neural network performance.

Problem

• Mapping Gaussian Processes to Bayesian Neural Networks
• Importance of trainable activations in BNNs
• Ensuring global stationarity with periodic activations

Innovation

• Trainable activations optimize BNNs
• Closed-form 2-Wasserstein distance enhances efficiency
• Trainable periodic activations ensure global stationarity