Learning Active Subspaces for Effective and Scalable Uncertainty
  Quantification in Deep Neural Networks
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.
Sanket Jantre, Nathan M. Urban, Xiaoning Qian, Byung-Jun Yoon
Learning Active Subspaces for Effective and Scalable Uncertainty Quantification in Deep Neural Networks

###### 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.

Sanket Jantre\({}^{\star}\) Nathan M. Urban\({}^{\star}\) Xiaoning Qian\({}^{\star\dagger}\) Byung-Jun Yoon\({}^{\star\dagger}\)\({}^{\star}\) Computational Science Initiative, Brookhaven National Laboratory, Upton, NY

\({}^{\dagger}\) Department of Electrical & Computer Engineering, Texas A&M University, College Station, TX Active subspace, Bayesian deep learning, subspace inference, uncertainty quantification (UQ)

## 1 Introduction

Neural networks (NN) are highly flexible and good approximators of complex functions at the expense of overparameterization, where the number of learnable parameters far exceeds the available training data. This increases computational demands and elevates the risk of overfitting [1], where the model starts to memorize noise rather than capture meaningful patterns. To avoid overconfidence and miscalibration, uncertainty quantification (UQ) in neural network predictions is crucial, especially in fields involving critical decision-making, such as clinical diagnostics or autonomous driving [2].

Addressing these challenges, Bayesian modeling presents a principled way to quantify the model prediction uncertainty [3]. Moreover, the Bayesian framework demonstrates heightened resilience against noise and adversarial perturbations due to their inherent probabilistic predictive capabilities [4]. To this effect, the confluence of deep neural networks and Bayesian inference in the form of Bayesian neural networks (BNN) has significantly advanced probabilistic machine learning. However, exact posterior inference is intractable in neural networks due to the extremely high dimensionality. Instead, approximation methods, such as mean-field variational inference [5], provide a computationally feasible way to perform posterior inference in BNNs. Nonetheless, this approximation severely limits the expressiveness of the inferred posterior, ultimately degrading the quality of the uncertainty estimates [6].

To this end, Bayesian inference over a low-dimensional subspace of neural network weights offers an elegant solution for improving accuracy, robustness, and uncertainty quantification [7]. Alternatively, active subspace methods [8] developed for inverse problems involving computer models perform dimension reduction by constructing a linear subspace of inputs with directions accounting for most variation in the function's output. In this paper, we demonstrate the utility of active subspace inference for Bayesian deep learning. Accordingly, we identify a low-dimensional subspace embedded in a high-dimensional parameter space to capture most of the variability in the neural network output. Our main contributions:

* We propose two active subspace methods: output-informed (**AS**) and likelihood-informed (**LIS**) for scalable Bayesian inference in deep learning. The compact subspace facilitates the use of otherwise intractable posterior approximation methods.
* We empirically demonstrate the superior uncertainty quantification offered by our approach compared to Bayesian inference over the full network and an existing subspace inference method.

**Related work.** The framework for interpretable inference through _effective dimensionality_ of the parameter space was provided by [9] in the context of BNNs. [10] demonstrated that near local optima, numerous directions in parameter space have low impact on neural network predictions. [11] performed inference over low-dimensional subspace spanned by the SGD iterates of NN weights. [7] further substantiated that posterior inference with a compact subspace of the entire parameter space is as effective as with the full network.

Alternately, active subspace finds its roots in the computer model UQ literature [12], where dimension reduction is performed on the high-dimensional input space using model outputs. The identified active subspace is then exploited for cheap approximate modeling of the computationally expensive simulators [13, 14]. The combination of active subspace meth