masters-thesis

introduction.tex (10439B)

---

 \setchapterpreamble[u]{\margintoc}
 \chapter{Introduction}
 \labch{intro}
 
 \section{Bayesian Deep Learning}
 
 Bayesian methods allow for estimates of uncertainty which enable more efficient usage of data (e.g., via active learning) and avoid overfitting.
 They accomplish this by using Bayes' theorem to model a distribution on the model parameters, which is then used to make predictions.
 From here we can sample from this distribution to make predictions, thereby yielding uncertainty estimates, thus enabling the assessment of model predictive confidence.
 In doing so, they allow for uncertainty estimates, which is important for many applications in machine learning, such as those involving safety-critical decisions which rely on understanding risk, where traditional deterministic neural networks are overconfident or fail to provide uncertainty estimates at all.
 The estimation of model confidence is particularly important in domains where silent failures can have catastrophic consequences~\cite{jospin2022hands}, such as autonomous driving~\cite{rao2018deep}, medical diagnosis~\cite{ker2017deep}, and financial markets~\cite{cavalcante2016computational}.
 These methods have been successfully applied to a wide range of classical problems in statistics, though attempts to apply them to deep learning have had limited success.
 For deep learning, the Bayesian approach is often limited by the computational cost of Bayesian inference.
 This is because Bayesian inference requires computing the posterior distribution of the model parameters given the data, which, due to the nonlinear complexity of the neural network architecture, is often intractable.
 
 Typically, the Bayesian approach to deep learning is to use a prior distribution over the model parameters and then compute the posterior distribution of the model parameters given the data, where the posterior is maximised to obtain the optimal model parameters for the available data.
 This posterior distribution can then be sampled from to make predictions.
 The optimisation can, depending on the choices of likelihood and prior, often be done analytically or numerically.
 However, computing and sampling from the posterior is often difficult, and so approximate Bayesian methods are used instead.
 These approaches depend greatly on the choice of prior, which is non-trivial.
 Often, a prior is chosen which is simple and has a convenient analytical form, such as a Gaussian prior.
 However, even in this case, the choice of prior precision is still difficult, often being estimated by cross-validation.
 Instead, maximising the marginal likelihood allows for optimisation of hyperparameters in the likelihood and prior, which can be done using gradient-based optimisation algorithms.
 However, this is computationally infeasible, and so approximate Bayesian methods are used instead.
 
 \section{Current Methods}
 
 Currently, approximate Bayesian methods are either expensive to compute (Markov chain Monte Carlo) or are limited in their Bayesian interpretation (e.g., Monte Carlo dropout~\cite{gal2016dropout}, deep ensembles~\cite{lakshminarayanan2017simple}, SWAG~\cite{maddox2019simple}).
 Of these, MC dropout and Markov chain Monte Carlo are the most widely used, though both have significant drawbacks.
 As such, there is demand for a method which exhibits the same computational cost as the optimization algorithms for deterministic neural networks while providing accurate posterior approximations and working out-of-the-box for any given architecture.
 
 The Laplace approximation~\cite{daxberger2021laplace,laplace1774memoire} is a simple yet theoretically well-supported posterior approximation suitable for Bayesian modeling.
 The Laplace approximation has been applied in, for example, the prediction of earthquake hypocenters~\cite{izzatullah2022laplace}.
 The primary advantages of the Laplace approximation are that it is simple to implement, effective at out-of-distribution detection~\cite{miani2022laplacian} and can be performed post-hoc on a pre-trained model.
 This last point is important, since it also means that the Laplace approximation can be as cheap to train as the optimization algorithms for deterministic neural networks, while only requiring a small amount of additional computation to compute the posterior samples for inference.
 
 The Laplace approximation corresponds to performing a second-order Taylor expansion of the log-posterior around the \emph{maximum a posteriori} (MAP) estimate of the model parameters as
 \begin{align*}
     \log p(\bm \theta\,|\,\bm y)
     \approx{} & \log p(\bm \theta_{\textsc{map}}\,|\,\bm y) + \frac{1}{2} (\bm \theta - \bm \theta_{\textsc{map}})\T \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y) (\bm \theta - \bm \theta_{\textsc{map}}).
 \end{align*}
 This is equivalent to approximating the posterior as a Gaussian distribution with mean and precision given by the MAP estimate and the Hessian matrix of the log-posterior respectively, i.e.,
 \begin{align*}
     p(\bm \theta\,|\,\bm y) \approx{} & \normal\left(\bm \theta\,\Big|\,\bm \theta_{\textsc{map}}, \left(- \nabla^2_{\bm \theta} \log p(\bm \theta\,|\,\bm y)\right)^{-1}\right).
 \end{align*}
 To compute the posterior distribution of the model parameters, we thus need to compute the Hessian of the negative log-posterior loss with respect to the model parameters.
 Typically, this is done by storing the posterior precision matrix explicitly and computing the inverse of this matrix.
 Since the full Hessian matrix is a \(D \times D\) square matrix with \(D^2\) elements, where \(D\) is the number of parameters, this is intractable for large models, which can contain millions or billions of parameters, thus limiting the applicability of full Laplace outside of toy problems of a few thousand parameters.
 To overcome this, we can use a crude approximation, such as only storing the posterior precision matrix's diagonal~\cite{lecun1989optimal,denker1990transforming}, its Kronecker factorisation~\cite{heskes2000natural,martens2015optimizing,botev2017practical}, and other low-rank approximations.
 Many of these approximations to the posterior precision have successfully been applied to the Laplace approximation~\cite{ritter2018scalable,ritter2018online}, though are limited by the quality of the approximation.
 To overcome this, we propose a method which performs the full Laplace approximation without requiring the precision to be stored explicitly.
 
 \section{Large-Scale Laplace}
 
 We propose an implementation of the Laplace approximation which does not require instantiation of the full posterior precision matrix.
 In general, the Laplace approximation has two steps.
 First, we would like to learn a posterior distribution over the model parameters by maximising either the posterior or the marginal likelihood.
 Maximising the posterior during training is trivial, since it has a cheap closed-form gradient.
 However, maximising the marginal likelihood is more difficult.
 
 \begin{marginfigure}
     \centering
     \input{figures/overview.tex}
     \caption[Overview of the Laplace approximation.]{Overview of the Laplace approximation (LA) for Bayesian neural networks. We highlight in green the computations which we approximate in this project using only Jacobian-vector products, namely the log-determinant (see \cref{ch:training}) and inverse square root (see \cref{ch:sampling}) of the posterior precision. Our contributions allow us to forgo the typical crude approximations of the precision by performing the LA using the full posterior precision without instantiating it explicitly.}
     \label{fig:overview}
 \end{marginfigure}
 
 As we will prove later, the Laplace log-marginal likelihood is given by, ignoring constant terms,
 \begin{align*}
     \log p(\bm y) \stackrel{\textsc{la}}{\approx}{} & - \frac{1}{2} \rho \sum_{i=1}^N \norm{\bm y_i - f_{\bm \theta}(\bm x_i)}^2 - \frac{1}{2} \alpha \norm{\bm \theta}^2 - \frac{1}{2} \log \det \bm \Lambda,
 \end{align*}
 where \(\bm \Lambda\) is the posterior precision.
 The first two terms are easy to compute (and differentiate), but computing the log-determinant of very large matrices is generally infeasible, since it typically has \(O(D^3)\) complexity or requires storing the matrix in memory.
 To overcome this, we thus need to estimate the log-determinant of the posterior precision matrix without instantiating it.
 
 The second step of the Laplace approximation is to perform inference and make predictions from the posterior predictive distribution.
 We estimate the posterior predictive via Monte Carlo by sampling from the posterior distribution and averaging the model predictions over the posterior samples.
 Since the posterior distribution is approximated by a Gaussian, we sample from the posterior as
 \begin{align*}
     \bm \varepsilon = \bm \theta_{\textsc{map}} + \bm \Lambda^{-1/2} \bm \varepsilon_0,
 \end{align*}
 which effectively only requires computing the inverse square root of the posterior precision matrix.
 This computation is also intractable for large models, since it usually involves computing the Cholesky decomposition of the precision, which suffers from the same performance issues as the computation of the log-determinant.
 
 As stated before, we will present a method which performs the full Laplace approximation without explicitly storing the posterior precision matrix.
 Our method approximates the posterior precision's inverse square root and log-determinant using only implicit Jacobian-vector products (JVPs).
 These JVPs compute products of the posterior precision matrix with a vector without storing the matrix itself in memory.
 Since this product outputs a vector, the memory requirements are linear in the number of parameters (rather than quadratic), allowing us to perform the full Laplace approximation with significantly larger models.
 We use JAX~\cite{jax2018github}, a Python library for GPU-accelerated automatic differentiation and computational graph compilation, to efficiently compute Jacobian-vector products while only storing the implicit computational graph.
 JAX uses XLA~\cite{sabne2020xla}, a compiler for linear algebra, to compile the JVP computation into efficient machine code.
 The computation of the log-determinant and inverse square root of the precision will thus enable us to perform training and inference using the Laplace approximation.
 An overview of the Laplace approximation and our contributions is shown in \cref{fig:overview}.