Learning Active Subspaces for Effective and Scalable Uncertainty Quantification in Deep Neural Networks

TL;DR

This paper introduces a method to identify low-dimensional active subspaces in neural network parameters, enabling scalable Bayesian inference and improved uncertainty quantification in deep learning models.

Contribution

It proposes a novel active subspace construction technique that reduces parameter space dimensionality, making Bayesian inference computationally feasible for deep neural networks.

Findings

Effective uncertainty quantification in regression tasks

Scalable Bayesian inference via reduced subspaces

Reliable predictions with robust uncertainty estimates

Abstract

Bayesian inference for neural networks, or Bayesian deep learning, has the potential to provide well-calibrated predictions with quantified uncertainty and robustness. However, the main hurdle for Bayesian deep learning is its computational complexity due to the high dimensionality of the parameter space. In this work, we propose a novel scheme that addresses this limitation by constructing a low-dimensional subspace of the neural network parameters-referred to as an active subspace-by identifying the parameter directions that have the most significant influence on the output of the neural network. We demonstrate that the significantly reduced active subspace enables effective and scalable Bayesian inference via either Monte Carlo (MC) sampling methods, otherwise computationally intractable, or variational inference. Empirically, our approach provides reliable predictions with robust uncertainty estimates for various regression tasks.