ablation_experiments.tex (13345B)
1 \setchapterpreamble[u]{\margintoc} 2 \chapter{Ablation Experiments} 3 \labch{ablation-experiments} 4 5 \section{Hessian-Vector Products}[Hessian-Vector Products] 6 7 In this project, we propose performing the Laplace approximation by only storing the computational graph for the implicit Hessian-vector product, thereby avoiding the explicit computation and storage of the quadratically-scaling Hessian matrix. 8 How significant of an effect does this have, though? 9 To test this, we perform Hessian-vector products \(\bm K \bm v\) and inverse Hessian-vector products \(\bm K^{-1} \bm v\), where \(\bm K \in \reals^{N \times N}\), and compare their performance. 10 Specifically, we apply a manual instantiation and inverse (or solve) of the Hessian matrix with a random vector \(\bm v\) and compare this to the performance of the JAX efficient Hessian-vector products (which we will call HVPs) and GGN-vector products (which we will call GVPs). 11 In reality, JAX computes these products implicitly by performing two Jacobian-vector products (for the HVP) and one Jacobian-vector product (for the GVP). 12 To compare performance with regard to memory usage and runtime, we perform these products for an increasing problem dimensionality \(N \in \{10^1, 10^2, \ldots, 10^9\}\). 13 14 It is difficult to benchmark memory consumption in JAX, since the just-in-time compiler optimises computations using XLA. 15 This then prevents us from stopping the program and measuring memory usage at a single point in time, as many computations will be performed at once. 16 Ideally, it would be possible to measure the \emph{peak} memory usage by JAX during the running of the program. 17 This is not trivial, though, as it would require querying the memory usage from the GPU asynchronously at a very high frequency and ``hope'' the peak usage point was captured. 18 Instead, we simply measure the value of \(N\) at which the program fails from an out-of-memory error instead, as a proxy for memory usage. 19 This experiment should then answer two questions: 20 \begin{enumerate} 21 \item How much larger problems can we solve if we do not need to instantiate the whole Hessian? 22 \item How much faster is it to compute the GVP versus, say, the HVP or the explicit Hessian-vector product? 23 \end{enumerate} 24 25 To perform this experiment, we define the the loss function composition \(\mathcal{L}(\bm x) \equiv (g \circ f) (\bm x) \coloneqq \sum^N_{i=1} (2 x_i + 0.5)^2\) as 26 \(f(\bm x) \coloneqq 2 \bm x + 0.5, g(\bm x) \coloneqq \sum^N_{i=1} x_i^2\). 27 This is equivalent to the sum-of-square-error loss \(g\) for a simple linear model \(f\). 28 Since \(f(x)\) is a linear function, the GGN matrix will exactly equal the actual Hessian.\sidenote{This allows us to compare our implementation of the GGN-vector product to the Hessian-vector product. While this test is not perfect (since it only holds for linear functions), it is a good sanity check. We find that the two are indeed the same in these experiments.} 29 We then initialise \(\bm x\) to a random vector and compute the matrix-vector product \(\bm K \bm v\) where \(\bm v\) is a random vector using three Hessian-vector products: the JAX efficient HVP and GVP functions and a manual Hessian instantiation and multiplication. 30 We then measure the total wall-clock time for computing this Hessian product by \(\bm v\). 31 32 \begin{marginfigure} 33 \centering 34 \includegraphics{hessian_profile_broken.pdf} 35 \caption[Comparison of Hessian-vector product performance.]{Top: comparison of the performance for manual Hessian-vector products, HVPs, and GVPs for increasing numbers of parameters \(D\). Bottom: comparison of the performance of inverse-vector products for the same methods. Implicit inverse is performed via conjugate gradient for HVP and GVP. We can see that the explicit inversion fails for \(D > 10 k\) and that explicit Hessian-vector products fail for \(D > 1 M\) (and, for \(D = 1 M\), the manual product is extremely slow). Implicit products, however, can be computed so long as the output vector fits in memory, while implicit inversion fails for \(D = 1 B\). Notice the broken axis and thus, in particular, the performance difference between the explicit and implicit inversion methods.} 36 \label{fig:hessian-profile} 37 \end{marginfigure} 38 39 Furthermore, since we are looking to compute some variation of the matrix inverse of \(\bm K\), and inverting a high-dimensional matrix can also be very expensive, we also compare the performance of the \emph{inverse} Hessian-vector product \(\bm K^{-1} \bm v\). 40 To compute the inverse for the efficient implicit HVP and GVP functions, we use the conjugate gradient method, which only requires the use of matrix-vector multiplications to compute this inverse product. 41 To compute the explicit inverse, we use the \(\texttt{jax.scipy.linalg}\) package, which performs the inverse using LAPACK (on CPU) or cuSOLVER (on GPU). 42 These experiments were run on a single NVIDIA A100 GPU. 43 44 % RESULTS 45 The results for this experiment have been visualised in \cref{fig:hessian-profile}. 46 Additionally, the raw results for this experiment have been summarised in the appendix, in \cref{tab:hessian-profile}. 47 We can see that the explicit Hessian-vector product fails for \(D > 1 M\) parameters, while the explicit inversion fails for \(D > 10 k\) parameters. 48 This is not surprising, as the explicit Hessian is a \(D \times D\) matrix, and so the memory requirements for this matrix depend on the \(D^2\) parameters. 49 Additionally, while the wall clock time for computing the Hessian-vector product is sublinear with \(D\), explicit Hessian-vector product is over an order of magnitude slower for \(D = 1 \mathrm{M}\) than for \(D = 100 \mathrm{k}\). 50 This may be due to the explicit Hessian, at this size, being too large to fit in the cache of the GPU. 51 52 The largest number of parameters tested is 1\,B. 53 Above 1\,B parameters, even implicit methods fail, although this failure occurs in the instantiation of a single random vector of that size (so it would fail on this hardware for any method, since we can not even store the \(\bm v\) nor the product \(\bm K \bm v\)). 54 Notably, however, these implicit Hessian-vector product methods allowed for computation of matrix products for problems three orders of magnitude larger, and computation of solves for problems four orders of magnitude larger, than the explicit methods. 55 This is a significant improvement in scalability, and is a key reason why we use these methods in our experiments. 56 Additionally the implicit inverse products are significantly faster than explicit inversion, even for small problems. 57 Performance also appears identical regardless of whether the explicit inverse is computed by actual inversion or by solving a system of linear equations. 58 59 \begin{marginfigure} 60 \centering 61 \includegraphics{hessian_profile_ratios.pdf} 62 \caption[Comparison of speedup from HVP to GVP.]{Comparison of the speedup from computation of the Hessian-vector product (HVP) versus the GGN-vector product (GVP) (top) and inverse HVP versus inverse GVP using the conjugate gradient method (bottom). Speedup is calculated as the ratio of the wall-clock time of the HVP to the GVP. We can see that the two methods have the same time complexity (up to a linear factor).} 63 \label{fig:hessian-profile-ratios} 64 \end{marginfigure} 65 66 However, due to the scale of the figure, it is not clear how the performance of the Hessian-vector product compares to the performance of the GGN-vector product. 67 We can then plot the speedup from going from the HVP to the GVP as a function of the number of model parameters. 68 This can be calculated as the ratio of the wall-clock time of the HVP to the wall-clock time of the GVP. 69 Thus, larger values suggest greater performance gains for the GVP, while values closer to one suggest there is no performance gain. 70 These results can be seen in \cref{fig:hessian-profile-ratios}. 71 The number of parameters on the \(x\)-axis are exponentially increasing. 72 We can see that the computation of the GVP is always faster than the computation of the HVP, and that the speedup is approximately linear in the number of parameters. 73 This is in line with what we expected (see \cref{sec:practical-ggn}). 74 We can also see a slight decrease in the speedup as the number of parameters increases for the Hessian-vector products. 75 However, it is not clear why this is the case. 76 77 \section{Sampling Ablation} 78 79 \begin{marginfigure} 80 \centering 81 \includegraphics{ablation/mnist_ablation.pdf} 82 \caption[Ablation over quadrature points for MNIST.]{Ablation over quadrature points for MNIST. We can see that sampling fails unambiguously for \(Q < 5\) and succeeds for \(Q \geq 10\).} 83 \label{fig:mnist-ablation} 84 \end{marginfigure} 85 86 Contour integral quadrature (CIQ) converges to the exact solution for \(\lim_{Q \to \infty} \bm s^{\mathrm{CIQ}}_{Q} = \bm K^{-1/2} \bm v\). 87 However, CIQ is typically close to the exact solution for low values of \(Q\).\sidenote{\textcite{pleiss2020fast} find that \(Q=20\) is approximately sufficient for most problems.} 88 We now attempt to determine the number of quadrature points \(Q\) required to compute an adequate posterior samples. 89 90 We vary \(Q \in \{1, \ldots, 20\}\) to determine the number of quadrature points required to effectively sample from the Laplace approximation. 91 We then compute the Mahalanobis distance between the samples and the true posterior mean, and compare this to the chi-squared distribution with \(D\) degrees of freedom. 92 By comparing the empirical CDF of the Mahalanobis distances to the theoretical \(\chi^2_D\) CDF, we can perform the Kolmogorov-Smirnov test for each value of \(Q\). 93 This p-value is then used to determine whether the null hypothesis that the samples are multivariate normal can be rejected based on some significance threshold. 94 Sample evaluation is explained in more detail in \cref{sec:sampling-evaluation}. 95 The results of this experiment are shown in \cref{tab:mnist-ablation,fig:mnist-ablation}. 96 Quantile--quantile plots for each value of \(Q\) are shown in the appendix, in \cref{fig:ablation-chisq}. 97 These experiments were, again, run on a single NVIDIA A100 GPU. 98 99 \begin{margintable} 100 \centering 101 \caption[Ablation over quadrature points for MNIST.]{Sampling time and p-values for Kolmogorov-Smirnov with varying quadrature points \(Q\) on MNIST with prior precision \(\alpha = 0.1\), for 200 posterior samples. Lower p-value indicates a higher probability of rejecting normality. Bold indicates sampling was successful based on visual inspection of quantile--quantile plots. Since, for values of \(Q \geq 5\), samples appear visually to be multivariate normal, we can conclude that we should choose a significance threshold of 0.01 or lower.} 102 \label{tab:mnist-ablation} 103 \include{tables/mnist_ablation.tex} 104 \end{margintable} 105 106 The p-values in \cref{tab:mnist-ablation} suggest that sampling is successful for \(Q \geq 10\) with a significance threshold of \(0.05\). 107 However, significance testing can be sensitive to the number of samples used, and, for large numbers of samples, these tests can excessively reject the null hypothesis. 108 In our case, the number of samples is generally fixed as a function of the number of parameters in the model and the amount of memory available. 109 As such, it is important to select an appropriate significance threshold for the test based on the number of posterior samples that are required. 110 Inspection of \cref{fig:ablation-chisq} in the appendix suggests that there is not a significant difference between using 5, 10, or 20 quadrature points. 111 Using less than 4 quadrature points, however, results in sampling failure, while using exactly 4 quadrature points results in a borderline success. 112 For 200 posterior samples, a more appropriate significance threshold of \(0.01\) or lower would reject sampling for \(Q \leq 4\), since the cases where sampling fails completely have p-values that are below single floating point precision (i.e., they are effectively zero). 113 114 In taking as much as 20 to 30 minutes to compute 200 posterior samples for an MNIST model with 15\,k parameters, the time required to compute the samples is not negligible. 115 In practice, it is likely that only a small number of samples will be obtained (possibly around 5--15), since for each sample it is necessary to store a parameter vector of size \(D\). 116 However, the sampling procedure still takes over 20 minutes to compute 10 samples. 117 As such, the time required to compute the samples may be a significant factor in the overall time required to compute the Laplace approximation. 118 In particular, for models with a very large number of parameters, the time required to compute the samples may be prohibitive when compared when the near-instant time required to sample from the diagonal Laplace approximation. 119 Even if the sampling time scales sublinearly with the number of parameters, it would still prevent the application in Laplace methods that require the Laplace approximation to be computed at each iteration from being used in practice. 120 However, this sampling time constraint is also significant for MCMC methods, which are the most common method for sampling from the posterior distribution of a BNN. 121 122 \marginnote{Results for the experiments in this chapter run on CPU are shown in the appendix, in \cref{fig:hessian-profile-cpu}.}