FSP-Laplace: Function-Space Priors for the Laplace Approximation in Bayesian Deep Learning

The markdown interpreter does not seem to work correctly in math mode. We apologise for this and can provide a PDF with the proof upon request.

Proposition 2. Let Assumptions A.1 to A.3 hold.

Let $\{\lambda_n\}_{n \in N} \subset R_{> 0}$ with $\lambda_n \to 0$, and $\{ f^\star_n \}_{n \in \N} \subset H_{\Sigma}$ such that $f^\star_n$ is a minimizer of $\R_\text{FSP}^{\lambda_n}$.

Then

$\{f_n^\star\}_{n \in \mathbb{N}}$

has an $H_{\Sigma}$-weakly convergent subsequence with limit

$f^\star \in F$.

Our proof makes use of ideas from (Dashti et al., 2013) and (Lambley, 2023).

Proof. Without loss of generality, we assume $\mathbf{\mu} = \mathbf{0}$. By Assumption A.2, there are constants $K, \alpha > 0$ such that

$R_\text{FSP}(\mathbf{f}) \ge K + \alpha \lVert \mathbf{f} \rVert_{H_{\mathbf{\Sigma}}}^2$

for all $\mathbf{f} \in H_{\mathbf{\Sigma}}$ (Lambley, 2023, Section 4.1). Now fix an arbitrary $\mathbf{f} \in H_{\mathbf{\Sigma}} \cap F$. Then

$R_\text{FSP}(\mathbf{f}) = R_\text{FSP}^{\lambda_n}(\mathbf{f})$

$\ge R_\text{FSP}^{\lambda_n}(\mathbf{f}^\star_n)$

$= R_\text{FSP}(\mathbf{f}^\star_n) + \lambda_n^{-1} d_{B}(\mathbf{f}^\star_n, F)$

$\ge K + \alpha \lVert \mathbf{f}^\star_n \rVert_{H_{\mathbf{\Sigma}}}^2 + \lambda_n^{-1} d_{B}(\mathbf{f}^\star_n, F)$

$\ge K + \alpha \lVert \mathbf{f}^\star_n \rVert_{H_{\mathbf{\Sigma}}}^2,$

and hence

$\lVert \mathbf{f}^\star_n \rVert_{H_{\mathbf{\Sigma}}}^2 \le \frac{1}{\alpha} \left( R_\text{FSP}(\mathbf{f}) - K \right),$

i.e. the sequence $\{f^\star_n\}_{n \in \mathbb{N}} \subset H_{\Sigma}$ is bounded.

By the Banach-Alaoglu theorem (Conway, 1997, Theorems V.3.1 and V.4.2(d)) and the Eberlein-Šmulian theorem (Conway, 1997, Theorem V.13.1), there is a weakly convergent subsequence $\{ \mathbf{f}^\star_{n_k} \}_{k \in \N}$ with limit $\mathbf{f}^\star \in H_{\mathbf{\Sigma}}$.

It remains to show that $\mathbf{f}^\star \in F$. From the inequality above, it follows that

$0 \le d_{B}(\mathbf{f}^\star_{n_k}, F) \le \lambda_{n_k} \left( R_\text{FSP}(\mathbf{f}) - K - \alpha \lVert \mathbf{f}^\star_n \rVert_{H_{\mathbf{\Sigma}}}^2 \right),$

where the right-hand side converges to 0 as $k \to \infty$.

Hence, $\lim_{k \to \infty} d_{B}(\mathbf{f}^\star_{n_k}, F) = 0$.

The embedding $\iota \colon H_{\mathbf{\Sigma}} \to B$ is compact (Bogachev, 1998, Corollary 3.2.4) and hence $\{\iota[\mathbf{f}^\star_{n_k}] \}_{k \in \mathbb{N}} \subset B$ converges ($B$-strongly) to $\iota[f^\star] \in B$ (Conway, 1997, Proposition 3.3(a)).

The continuity of $d_{B}(\,\cdot\,, F)$ implies

$d_{B}(\mathbf{f}^\star, F)= d_{B}(\iota[\mathbf{f}^\star], F)$

$= d_{B} \left( \lim_{k \to \infty} \iota[\mathbf{f}^\star_{n_k}], F \right)$

$= \lim_{k \to \infty} d_{B}(\iota[\mathbf{f}^\star_{n_k}], F)$

$= \lim_{k \to \infty} d_{B}(\mathbf{f}^\star_{n_k}, F)$

$= 0,$

and by Assumption A.3 and Lemma A.2 it follows that $\mathbf{f}^\star \in F$.

We thank the reviewer for their time and feedback. We are happy that they found the paper was "well written" and that our method "performs well".

"Laplace approximations are justified by the Bernstein-von Mises theorem which states that [...] the posterior distribution concentrates around its mode independent of the prior distribution in the large-data limit. [...] However approximating the Hessian, or regularizing it and claiming it's because there's a prior not only does nothing to fix the problem, but perpetuates a line of ill-posed research."

We respectfully disagree. It is true that the Bernstein-von Mises theorem only holds in the large-data limit. However, there is a long line of empirical work suggesting that the linearized Laplace approximation works surprisingly well for neural networks despite the lack of theoretical justification (Papamarkou et al. 2024, Section 3.1 references many such papers). Our paper provides additional empirical evidence supporting this observation. Calling this entire line of research "ill-posed" needs to be measured against the large amount of empirical evidence that clearly demonstrates the value of such research.

More generally, the point of empirical research is precisely to explore regimes in which thoeretical statements are not possible. Rigorous theoretical statements can often only be proven under strong assumptions, but this does not mean that the opposite of the statement is true as soon as assumptions are violated (absence of proof is not proof of absence). Empirical research is indispensable to explore how far away from a perfect limiting case remnants of a theory still apply. For example, the entire field of deep learning relies on optimization methods for which most of the theory applies only to convex functions, and yet methods motivated by this theory are successfully used to train highly non-convex deep neural networks because empirical studys show that they work surprisingly well in certain useful regimes far beyond convexity. This is not an argument against the value of theory. On the contrary: it would have been impossible for us to "guess" the empirically successful approximation scheme that we propose had we not been guided by theory that strictly holds only in a limiting case that is admittedly relatively far away from our applications.

"the work does produce decent empirical results as a quadratic approximation of the posterior inspired by Laplace approximations, however the current reasoning used to get there is ill-considered."

We thank the reviewer for acknowledging our empirical results. We agree that these results were obtained with a method that is inspired by Laplace approximations (see our proofs of Propositions 1 and 2 in the additional comment for clarification on this connection). As with most forms of scalable Bayesian inference, several additional approximations were necessary to make our method usable in practice. We believe that our paper clearly highlights the additional approximations, e.g., by explicitly grouping them into Section 3.2. These additional approximations do not change the fact that the resulting algorithm approximates the posterior with a Gaussian distribution whose mean is the MAP (of the RKHS-regularized neural network) and whose precision matrix is obtained by approximating the curvature at this point. Methods for approximate posterior inference that follow this scheme are generally called Laplace approximations in the literature.

"The dependence of f appearing in Propositions 1 and 2 should be made explicit."

We do not understand your inquiry, would you mind elaborating on what you mean by the "dependence on f appearing in Propositions 1 and 2? From our point of view, there is no dependence on any particular $\mathbf{f}$ in Propositions 1 and 2, since in these propositions $\mathbf{f}$ is an unbound variable denoting an arbitrary element of the sample space $\mathbb{B}$ of the GP prior (or an element of the RKHS $H_{\mathbf{\Sigma}} \subset B$ depending on the context). Moreover, the $\mathbf{f}^\star_n$ in Proposition 2 are explicitly defined as minimizers of a sequence of optimization problems, and $\mathbf{f}^\star$ is also explicitly defined as the weak limit of a certain subsequence of $( \mathbf{f}^\star_n )_{n \in N}$.

"The proofs given in the paper are not sound. In particular, the proof of Proposition 2 [...]. The proof of Proposition 1 is also not clear, and needs to be written out"

We agree that the proof of Propositions 2 is difficult to follow as one has to combine several ideas in the referenced works in a nontrivial manner. Hence, we conducted the proof of Proposition 2 in full detail and attach it as a separate comment below. It will also be included in the appendix in the camera-ready version of the paper.

Under Assumptions A.1 to A.3 detailed in the paper, Proposition 1 is a pretty direct corollary of Theorem 1.1 in (Lambley, 2023). The given proof in the paper verifies the assumptions of Lambley's Theorem 1.1 and is hence complete. However, we will make an effort in the camera-ready version of the paper to add more explanatory comments to make the proof easier to follow.

"The claim that the term appearing in line 195 is negligible should be confirmed numerically."

We provide evidence that the term $P_0 \Lambda P_0^\top$ is negligible in four different configurations. First, considering the regression setup described in Appendix B.1, we have:

• RBF kernel: $\frac{||P_0 \Lambda P_0^\top||_F}{||\Lambda||_F} = 1.026 \times 10^{-7}$
• Matern-1/2: $\frac{||P_0 \Lambda P_0^\top||_F}{||\Lambda||_F} = 4.305 \times 10^{-6}$

Considering the classification setup described in Appendix B.1, we have:

• RBF kernel: $\frac{||P_0 \Lambda P_0^\top||_F}{||\Lambda||_F} = 3.548 \times 10^{-4}$
• Matern-1/2: $\frac{||P_0 \Lambda P_0^\top||_F}{||\Lambda||_F} = 1.299 \times 10^{-2}$

We thank the reviewer for their time and positive feedback. We are glad that they found our paper was "well-written", "easy to follow" and had a "smooth flow".

"The comparative experiments are not sufficient. Various efficient LA methods should be compared. [...] It would be better to include some functional BNN as well." "Can you please add comparative results with functional BNN?"

We agree with the reviewer that a comparison with a functional BNN would be useful. Therefore, we ran our experiments on FVI (Sun et al., 2019) and present the results in the general rebuttal (see Figures D.1 and D.2 and Tables D.1, D.2 and D.3). We find that our method generally outperforms FVI.

We use the full Laplace posterior covariance for regression, Mauna Loa, and ocean current modeling; due to computational limitations, we use only the KFAC approximation for the classification experiments and for Bayesian optimization. KFAC is standard in Laplace approximations for neural network (Ritter et al. 2018, Immer et al. 2021) and has shown to perform very competitively. The methods proposed by the reviewer offer more scalable alternatives when the full GGN is not available. While these are interesting in themselves, the point of this paper is to discuss the influence of a function-space prior and not of the GGN approximation. Exploring the use of the methods proposed by the reviewer to make our method more scalable sounds like an interesting direction for future work.

"Can you please explain the equation in Line148 in more details?"

The equation follows from inserting the definition of $f^\text{lin}$ (Eq. 2.2) into the left-hand side and completing the square. We note that the $\dagger$ symbol designates the pseudo-inverse. In detail:

<!--$$\frac{1}{2} ||f^{\text{lin}}(\,\cdot\ , \mathbf{w}) - \mathbf{\mu}||_{\mathbb{\mathbf{\Sigma}}}^2 = \frac{1}{2} (\mathbf{w} - \mathbf{w}^*)^\top \mathbb{\mathbf{\Sigma}}_{\mathbf{w}^*}^\dagger(\mathbf{w} - \mathbf{w}^*) + \mathbf{v}\mathbf{v}^\top (\mathbf{w} - \mathbf{w}^*) + \frac{1}{2} ||f(\,\cdot\ , \mathbf{w}^*) - \mathbf{\mu}||_\mathbb{\mathbf{\Sigma}}^2$$ -->

from which we complete the square to obtain the right-hand side of line 148.

"What is the meaning of 'rigorous probabilistic interpretation' in Line 128? Does that mean the new P_B(df) is a properly defined functional posterior?"

We mean that $\frac{1}{\lambda} d_B(f, F)$ can be interpreted as a negative log likelihood encoding that we observe $d_B(f, F) = 0$ with Laplace distributed noise (with parameter $\lambda$). The proposition shows that the optimization problem is the corresponding MAP estimation problem.

"The objective function is similar to the kernel ridge regression. Can you please explain the difference?"

Good observation! The kernel ridge regression (KRR) can be seen as the maximum a-posteriori of a GP. Thus, the first part of our algorithm, which learns an RKHS-regularized neural network, bears some similarity with KRR. However, KRR is a nonparametric method wherease the first part of our algorithm optimizes the parameters of a neural network (this is motivated by the generalization performance of deep neural networks, as mentioned in the abstract). Further, the second part of our algorithm estimates uncertainties of the neural network parameters, which has no analog in KRR.

We thank the reviewer for their time and constructive feedback. We are glad that they found that the "quality of the manuscript is [...] high", that they "like the spirit of the work" and that our method has "good performance".

"Context points in section 3.2 are kind of an obscure part of the Algorithm to me. [...] I think a bit of related work and connections in this point could be needed. In general, it reminds me of inducing points in GPs and kind of pseudo-coresets. I also see that some limitations or issues with this sampling and the number of context points could appear, and somehow it is not really considered in the manuscript (in my opinion)."

Methods for regularizing neural networks in function space always rely on context points to evaluate the regularizer (Sun et al. 2019, Tran et al. 2022, Rudner et al. 2023). Popular strategies include uniform sampling from a distribution over a bounded subset of the feature space (Tran et al. 2022, Rudner et al. 2023), from the data (Sun et al., 2019), and from additional (possibly unlabeled) datasets (Rudner et al. 2023).

The goal of the context points is to regularize the neural network on a finite set of input locations which includes the training data and any point where we would like to evaluate the model. During the MAP estimation phase of our method (see Algorithm 2), context points are different from inducing points in GPs: context points are resampled at each update, unlike inducing points of a GP, which are optimized or kept fixed. When computing the posterior covariance, however, context points indeed share some similarity with inducing points in a GP: at this point, they are fixed and define where we regularize the model. Unlike pseudo-coresets, context points do not aim to summarize the training data but rather to cover a finite set of input locations that includes the training data and any point where we wish to evaluate the model.

Relying on a set of context points is fine for datasets with low-dimensional features, which are precisely the cases for which we have prior beliefs in the form of a GP prior (GP priors are much harder to specify with high-dimensional data). Our FSP-Laplace method was precisely designed with these scenarios in mind and the scalable matrix-free linear algebra methods can perfectly cope with a very large number of context points (> 20'000). For high-dimensional data, using context points is manageable if we have prior knowledge about where the training data lies (for example, a manifold in the ambient feature space). We can then use a set of context points that has approximately the same support, as shown in our image classification examples where we use a set of context point from the KMNIST dataset. Nevertheless, Table D.2 in the PDF linked in the general rebuttal also shows that our method remains competitive on MNIST image classification when using context points drawn from a uniform distribution during training (RND). Due to limited time during the rebuttal period, results for Fashion MNIST are not yet available (N.As in Table D.2) and we will update the reviewers during the discussion phase with the results. Another example where prior beliefs can often be formulated in the form of a GP is scientific data (see, e.g., our ocean current experiment). For many scientific data, context points could also be generated by a simulation that does not need to be very accurate (since it is only meant to generate points in the vicinity of the true data distribution).

A sufficient number of context points is necessary to capture the beliefs specified by the GP prior (see Figure D.3 and D.4 in the PDF link to the general rebuttal). If the number of context points is too small, the beliefs are not accurately captured, and the function draws will not look like the prior but the uncertainty should not collapse (see Figures D.3 and D.4 in PDF in the general rebuttal).

In the case of FSP-Laplace, it is not clear to me if the work is also fully applicable to both classification and regression problems in an easy way, to which likelihood models, and if doing that makes things more difficult on the FSP-Laplace evaluation part or the Laplace approximation side.

FSP-Laplace works in exactly the same settings as the standard Laplace approximation, e.g., for classification and regression settings using Gaussian and Categorical likelihoods, respectively (see regression and classification examples in Section 4.1). Just like Laplace, we require that the likelihood function is twice differentiable (we need to compute its Hessian). Note that for the classification setting with $c$ classes, we place a GP prior on every logit such that we have $c$ priors. Therefore, $M$ is the $p \times rc$ matrix containing the concatenation of the $c$ $(J_i^T L_i)_{i=1}^c$ matrices in line 4 of Algorithm 1. We will make this clearer in the camera-ready version.

"In L201, it is mentioned that there is an observation on the fact that the trace of the posterior is 'always' upper-bounded by the trace of the prior covariance. Due to this, a little condition is imposed as a remedy. Is this an empirical observation? it is always true? or is there a proof for this?"

This follows from the expression of the covariance of the posterior in a Gaussian process, which is the prior covariance minus a term corresponding to the conditioning on the training data. The latter is always non-negative.

"I don't see exactly how feasible could be to compute the potential \phi and Eq. (3.1)"

Intuitively, the potential can be understood as the negative log-likelihood.

We thank the reviewer for their useful feedback and time. We are glad that they found our "method is well-motivated", our "techniques [...] are solid" and that the paper was "very well-written" and "easy-to-follow".

"My main concern is the relation between this paper and Sun et al. (2019) [8]. Sun et al. also consider Bayesian inference of function posteriors with a functional prior directly defined in the function spaces (e.g., Gaussian processes). The main difference seems to be that Sun et al. perform fully Bayesian inference, while this work aims to obtain a MAP estimate followed by a Laplace approximation to obtain the posterior. Sun et al.'s work should be included as a baseline, and more discussion on the relation between the two works is needed."

While Sun et al. (2019) and our method both specify GP priors on the function represented by the neural network, Sun et al. use variational inference (VI), whereas our method uses the Laplace approximation. VI approximates the posterior by a parametric distribution that maximizes the ELBO (Eq. 3 in Sun et al.), while the Laplace approximation approximates the posterior by a Gaussian centered at the MAP estimate of the parameters. Note that the Laplace approximation and VI are both parametric approximate posterior inference methods that are commonly used in the machine learning and statistics literature. Neither one of them can be considered more "fully Bayesian" than the other. We agree that a comparison with Sun et al. is interesting, and we therefore include experiments using this method (labeled "FVI") in our overall rebuttal above. Our method FSP-Laplace generally outperforms FVI. This can possibly be explained by a complication in the FVI method: evaluating the ELBO requires access to the density of the variational distribution, which is not available in function space. FVI addresses this issue with implicit score function estimators, which introduce another approximation and makes FVI difficult to use in practice (Ma and Hernández-Lobato, 2021). On a more theoretical level, note that the KL divergence between measures (Eq. 4 in Sun et al.) is actually infinite when using GP priors (Burt et al., 2020) and the function-space ELBO (Eq. 3 in Sun et al.) is undefined. Our Laplace approximation does not suffer from either of these issues and approximates a well-defined objective.

"Similar to the difficulty in evaluating the functional KL divergence in Sun et al., when evaluating the 'prior probability' (i.e., RKHS norm) of $f$, a set of context points must be specified to approximate it. My second concern is that this sampling strategy could be inefficient and impractical for high-dimensional datasets. For 1D or 2D datasets, the context points can be sampled uniformly or fixed to grid points. However, for higher-dimensional cases, this approach becomes less effective. This is a critical issue for the proposed method, as in high-dimensional spaces, there is little chance that the context points will adequately cover the support of the data distributions. For example, using KMNIST as the context points for MNIST may not always be feasible in general cases."

Methods that regularize a neural network in function space (both for approximate Bayesian inference or regularized empirical risk minimization) always use a set of context points to evaluate the regularizer (Sun et al. 2019, Tran et al. 2022, Rudner et al. 2023). As you say, relying on a set of context points is fine for datasets with low-dimensional features, which are precisely the cases for which we have prior beliefs in the form of a GP prior (GP priors are much harder to specify with high-dimensional data). Our FSP-Laplace method was precisely designed with these scenarios in mind and the scalable matrix-free linear algebra methods can perfectly cope with a very large number of context points (> 20'000). For high-dimensional data, using context points is manageable if we have prior knowledge about where the training data lies (for example, a manifold in the ambient feature space). We can then use a set of context points that has approximately the same support, as shown in our image classification examples where we use a set of context points from the KMNIST dataset. Nevertheless, Table D.2 in the PDF linked in the general rebuttal also shows that our method remains competitive on MNIST image classification when using context points drawn from a uniform distribution during training. Due to limited time during the rebuttal period, results for Fashion MNIST are not yet available (N.As in Table D.2) and we will update the reviewers during the discussion phase with the results. Another example where prior beliefs can often be formulated in the form of a GP is scientific data (see, e.g., our ocean current experiment). For many scientific data, context points could also be generated by a simulation that does not need to be very accurate (since it is only meant to generate points in the vicinity of the true data distribution). We will add more details in the camera-ready paper.

"What will happen if the same context points are used for both training and approximating the Hessian?"

This is fine. It is only important that the set of context points covers a finite set of input locations containing all the training data as well as any points where we would like to evaluate the model. During training, we can amortize the coverage by resampling the context points at each update step (step 5 in Algorithm 2). When computing the covariance, we can no longer amortize and must use a set of points that covers the space well, such as a low discrepancy sequence (Latin hypercube sampling, for example).