SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[18] Wesley J. Maddox, Gregory W. Benton, and Andrew Gordon Wilson. Rethinking parameter counting in deep models: Effective dimensionality revisited. CoRR, abs/2003.02139, 2020.
[19] Kevin P. Murphy. Machine learning: a probabilistic perspective. MIT Press, Cambridge, Mass. [u.a.], 2013. ISBN 9780262018029 0262018020.
[20] Mahdi Pakdaman Naeini, Gregory Cooper, and Milos Hauskrecht. Obtaining well calibrated probabilities using bayesian binning. In Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015.
[21] Michael Peskin. An introduction to quantum field theory. CRC press, 2018.
[22] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do CIFar-10 classifiers generalize to CIFar-10? 2018.
[23] Tim G. J. Rudner, Zonghao Chen, Yee Whye Teh, and Yarin Gal. Tractable Function-Space Variational Inference in Bayesian Neural Networks. In Advances in Neural Information Processing Systems 35, 2022.
[24] Tim G. J. Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Teh, and Yarin Gal. Continual Learning via Sequential Function-Space Variational Inference. In Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2022.
[25] Tim G. J. Rudner, Sanyam Kapoor, Shikai Qiu, and Andrew Gordon Wilson. Function-Space Regularization in Neural Networks: A Probabilistic Perspective. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2023.
[26] Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the royal statistical society series b-methodological, 58:267-288, 1996.
[27] Aad W Van der Vaart. Asymptotic statistics, volume 3. Cambridge university press, 2000.
[28] Steven Weinberg. Gravitation and cosmology: principles and applications of the general theory of relativity. 1972.
[29] Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020.
[30] David H. Wolpert. Bayesian backpropagation over i-o functions rather than weights. In J. Cowan, G. Tesauro, and J. Alspector, editors, Advances in Neural Information Processing Systems, volume 6. Morgan-Kaufmann, 1993.
[31] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
A PS-MAP is not Invariant to Reparametrization 14
B Mathematical Details & Derivations 15
B.1 Prior Distributions over Functions 15
B.2 The General Case 15
B.3 Why Should a Constant Metric Affect FS-MAP? 16
B.4 Comments on the Infinite Limit 16
B.5 Condition for a Non-Singular  $\mathcal{J}(\theta ;p_X)$  16
B.6 Architectural Symmetries Imply Singular  $\mathcal{J}(\theta ;p_X)$  16
B.7 Addressing the Infinite Prior Density Pathology with a Variational Perspective 17
B.8 Approximating the Expectation 18
B.9 Monte Carlo Estimator for the Log Determinant 18
B.10 Laplacian Regularized MAP Objective 18
C Further Empirical Results and Experimental Details 20
C.1 Training details for Synthetic Experiments 20
C.2 Hyperparameters for UCI Regression 20
C.3 Hyperparameters for Image Classification Experiments 20
C.4 Verifying the L-MAP Approximation 20
C.5 Effective Eigenvalue Regularization using the Laplacian Regularizer 20
C.6 Visualizing the Effect of Laplacian Regualrization 21
C.7 Further Neural Network Experiments 22
Appendix A PS-MAP is not Invariant to Reparametrization
Parameter-space MAP estimation has conceptual and theoretical shortcomings stemming from the lack of reparameterization-invariance of MAP estimation. Suppose we reparameterize the model parameters  $\theta$  with  $\theta' = \mathcal{R}(\theta)$  where  $\mathcal{R}$  is an invertible transformation. The prior density on the reparameterized parameters is obtained from the change-of-variable formula that states  $p'(\theta') = p(\theta)|\mathrm{det}^{-1}(\mathrm{d}\theta'/\mathrm{d}\theta)|$ . In the language of differential geometry, the prior density  $p(\theta)$  is not a scalar but a scalar density, a quantity that is not invariant under coordinate transformations [28]. The probabilistic model is fundamentally unchanged, as we are merely viewing the parameters in a different coordinate system, but the MAP objective is not invariant to reparameterization. Specifically, the reparameterized parameter-space MAP objective becomes
$$
\mathcal {L} ^ {\prime \mathrm {M A P}} \left(\theta^ {\prime}\right) = \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, \theta^ {\prime}\right) + \log p ^ {\prime} \left(\theta^ {\prime}\right) = \mathcal {L} ^ {\mathrm {M A P}} (\theta) - \underbrace {\log | \det  ^ {- 1} \left(\mathrm {d} \theta^ {\prime} / \mathrm {d} \theta\right) |} _ {\text {n e w t e r m}}. \tag {A.1}
$$
Since  $\mathcal{L}^{\mathrm{MAP}}(\theta')$  and the original objective  $\mathcal{L}^{\mathrm{MAP}}(\theta)$  differ by a term which is non-constant in the parameters if  $T$  is non-linear, the maxima  $\theta^{\prime \mathrm{MAP}}\doteq \arg \max_{\theta'}\mathcal{L}^{\mathrm{MAP}}(\theta')$  and  $\theta^{\mathrm{MAP}}\doteq \arg \max_{\theta}\mathcal{L}^{\mathrm{MAP}}(\theta)$  under the two objectives will be different. Importantly, by "different" we don't just mean  $\theta^{\prime \mathrm{MAP}}\neq \theta^{\mathrm{MAP}}$  but  $\theta^{\prime \mathrm{MAP}}\neq \mathcal{R}(\theta^{\mathrm{MAP}})$ . That is, they are not simply each other viewed in a different coordinate system but actually represent different functions. More accurately, therefore, one should say MAP estimation is not equivariant (rather than invariant) under reparameterization. As a result, when using a parameter-space MAP estimate to make predictions at test time, the predictions can change dramatically simply depending on how the model is parameterized. In fact, one can reparameterize the model so that PS-MAP will return an arbitrary solution  $\theta_0$  irrespective of the observed data, by choosing a reparameterization  $\theta^{\prime} = \mathcal{R}(\theta)$  such that  $\log |\operatorname*{det}^{-1}(\mathrm{d}\theta' / \mathrm{d}\theta)|_{\theta = \theta_0} = -\infty$ . One such choice is  $\theta^{\prime} = (\theta -\theta_{0})^{3}$ , where the exponent is taken element-wise.
As a less extreme example, consider the reparameterization  $\theta' = 1/\theta$ . A Gaussian prior on the original parameters  $p(\theta) \propto \exp(-\theta^2/2)$  translates to a prior  $p'(\theta') \propto \exp(-1/2\theta'^2)/\theta'^2 = \theta^2 \exp(-\theta^2/2)$  on the inverted parameters. Note the transformed prior, when mapped onto the original parameters  $\theta$ , assigns a higher weight on non-zero values of  $\theta$  due to the quadratic factor coming from the log determinant. We illustrate this effect with a simple experiment where a linear model with RBF features is trained with a Gaussian likelihood and a Gaussian prior  $\mathcal{N}(0, I)$  using PS-MAP to fit noisy observations from a simple function. Figure 8 shows the predictions and learned weights when PS-MAP is performed both in the original parameterization and in the inverted parameterization. While PS-MAP favors small weights in the original parameterization, it favors non-zero weights in the inverted one, learning a less smooth function composed of a larger number of RBF bases.
(a) Learned functions
(b) Parameter distributions
Figure 8: PS-MAP is not invariant to reparameterization. Simply inverting the parameterization of the weights removes the preference for small coefficients in a linear model, leading to a less smooth learned function, even though the underlying probabilistic model is unchanged.
Appendix B Mathematical Details & Derivations
B.1 Prior Distributions over Functions
Since the function evaluations  $f_{\theta}(\hat{x})$  are related to the parameters  $\theta$  through the map  $T: \theta \mapsto f_{\theta}(\hat{x})$ , the prior density  $p(f_{\theta}(\hat{x}))$  is simply the push forward of the density  $p(\theta)$  through  $T$ , which can be expressed without loss of generality using the Dirac delta function as
$$
p \left(f _ {\theta} (\hat {x})\right) = \int_ {\mathbb {R} ^ {P}} p \left(\theta^ {\prime}\right) \delta \left(T \left(\theta^ {\prime}\right) - f _ {\theta} (\hat {x})\right) \mathrm {d} \theta^ {\prime}. \tag {B.2}
$$
This generic expression as an integral over the parameter space is difficult to evaluate and thus not directly useful. By assuming that  $T$  is injective (one-to-one) modulo some global symmetries in parameter space (e.g., permutation of the neurons), which holds for many neural network architectures [2], Wolpert [30] established the simplification in Equation (6). Note that since the map  $T$  is not guaranteed to be surjective (onto), especially when there are more evaluation points than parameters  $(|\hat{x}| > P)$ ,  $p(f_{\theta}(\hat{x}))$  does not generally have support everywhere and is to be interpreted as a surface density defined over the image of  $T$  [30]. If the model parameterization has continuous symmetries, such as scaling symmetries in MLPs with ReLU activations [5], then the log determinant in Equation (6) should be replaced by the log pseudo-determinant to account for the fact that the dimension of the image of  $T$  is less than  $P$ , the number of parameters.
B.2 The General Case
Here we present a more general result for which the result in Section 3.2 is a special case, without assuming the function output is univariate ( $K = 1$ ) or that the metric is diagonal. To simplify notations, we suppress all  $\theta$ -dependence and write  $f_{\theta}^{k}(x_{i})$ , the  $k$ -th output of the function at  $x_{i}$  as  $f_{i}^{k}$  for some  $\hat{x} = \{x_{i}\}_{i=1}^{M}$ . We use Einstein notation to imply summation over repeated indices and leave everything in un-vectorized forms. In this more general case, the metric is given by  $\mathrm{d}s^2 = g_{ikj\ell} \mathrm{d}f_i^k \mathrm{d}f_j^\ell$ , where  $i, j$  are input indices and  $k, \ell$  are output indices. Denote the un-vectorized  $M$ -by- $K$ -by- $P$  Jacobian  $J_{ika} \doteq \partial_{\theta_a} f_i^k$ . The  $P$ -by- $P$  matrix  $\mathcal{I}$  appearing inside the determinant in Equation (6) is now given by
$$
\mathcal {J} _ {a b} = J _ {i k a} g _ {i k j \ell} J _ {j \ell b} = g _ {i k j \ell} \partial_ {\theta_ {a}} f _ {i} ^ {k} \partial_ {\theta_ {b}} f _ {j} ^ {\ell} \tag {B.3}
$$
In the limit of infinitely many evaluation points uniformly populating the domain  $\mathcal{X}$ , the sum over  $i, j$  becomes a double integral<sup>2</sup>
$$
\mathcal {J} _ {a b} = \int_ {\mathcal {X} \times \mathcal {X}} \partial_ {\theta_ {a}} f _ {\theta} ^ {k} (x) \partial_ {\theta_ {b}} f _ {\theta} ^ {\ell} \left(x ^ {\prime}\right) g _ {k \ell} \left(x, x ^ {\prime}\right) d x d x ^ {\prime}, \tag {B.4}
$$
where we have re-written the metric  $g_{ikj\ell}$  as  $g_{k\ell}(x_i,x_j)$  where  $g:\mathcal{X}\times \mathcal{X}\to \mathbb{R}^{K\times K}$  is the metric represented now as a matrix-valued function. After dividing by a constant  $Z$ , we can similarly write  $\mathcal{J}_{ab}$  as the expectation
$$
\mathcal {J} _ {a b} = \mathbb {E} _ {p _ {X X ^ {\prime} C C ^ {\prime}}} \left[ \partial_ {\theta_ {a}} f _ {\theta} ^ {C} (X) \partial_ {\theta_ {b}} f _ {\theta} ^ {C ^ {\prime}} \left(X ^ {\prime}\right) \operatorname {s g n} \left(g _ {C C ^ {\prime}} \left(X, X ^ {\prime}\right)\right) \right], \tag {B.5}
$$
where  $p_{XX^{\prime}CC^{\prime}}(x,x^{\prime},c,c^{\prime}) = |g_{cc^{\prime}}(x,x^{\prime})| / Z$  for some normalization constant  $Z$  that only shifts the log determinant by a  $\theta$ -independent constant. Note Equation (10) is recovered if the metric is diagonal, that is  $g_{c,c^{\prime}}(x,x^{\prime})\neq 0$  only if  $c = c^{\prime}$  and  $x = x^{\prime}$ . We can further remove the assumption that  $g$  is constant (independent of  $f$ ), but at the cost of letting the normalization constant  $Z$  depend on  $f$  (or equivalently  $\theta$ ), making it invalid to ignore  $Z$  during optimization. While more general than the result presented in Section 3.2, specifying a (possibly  $\theta$ -dependent) non-diagonal metric adds a significant amount of complexity without being particularly well-motivated. Therefore we did not further investigate this more general objective.
2To be more accurate, an integral of the form  $\int \phi (x)\mathrm{d}x$  is the limit of  $\sum_{i = 1}^{M}\phi (x_{i})\Delta x$  with the extra factor  $\Delta x = x_{i + 1} - x_i(\forall i)$ , as  $M\to \infty$ . However, accounting for this factor does not affect the final result in our case because again  $\Delta x$  is independent of  $\theta$ .
B.3 Why Should a Constant Metric Affect FS-MAP?
One may object that a constant metric  $g$  should have no effect in the MAP estimate for  $f_{\theta}$ . However, the subtlety arises from the fact that  $p(f_{\theta})$  does not have support everywhere but is, in fact, a surface density defined on the image of  $T$  (as defined in Appendix B.1), a curved manifold in function space which gives rise to locally preferred directions such that even a global linear transformation will change the density in a non-homogeneous way. If  $T$  were instead surjective, then in Section 3.2  $J_{\theta}(\hat{x})$  would be a square matrix and we could write  $\log \det \big(J_{\theta}(\hat{x})^{\top}gJ_{\theta}(\hat{x})\big) = \log \det \big(J_{\theta}(\hat{x})^{\top}J_{\theta}(\hat{x})\big) + \log \det (g)$  and conclude that a constant metric indeed has no effect on the MAP estimate. Similarly, if the image of  $T$  is not curved, meaning  $J_{\theta}(\hat{x})$  is constant in  $\theta$ , then  $\log \det \big(J_{\theta}(\hat{x})^{\top}gJ_{\theta}(\hat{x})\big)$  would be constant in  $\theta$  regardless of  $g$  and the metric would have no effect.
B.4 Comments on the Infinite Limit
In Section 3.2, we generalized the finite evaluation point objective from Equation (8) by considering the limit as the number of evaluation points  $\hat{x}$  approaches infinity and the points cover the domain  $\mathbb{R}^P$  uniformly and densely. This technique, known as the continuum limit [21], is commonly used in physics to study continuous systems with infinite degrees of freedom, by first discretizing and then taking the limit as the discretization scale approaches zero, thereby recovering the behavior of the original system without directly dealing with intermediate (possibly ill-defined) infinite-dimensional quantities. Examples include lattice field theories where the continuous spacetime is discretized into a lattice, and numerical analysis where differential equations are discretized and solved on a mesh, where the solution often converges to the continuous solution as the mesh size goes to zero.
In a similar vein, we utilize this technique to sidestep direct engagement with the ill-defined quantity  $p(f_{\theta})$ . It may be possible to assign a well-defined value to  $p(f_{\theta})$  through other techniques, though we expect the result obtained would be consistent with ours, given the continuum limit has historically yielded consistent results with potentially more refined approaches in other domains.
B.5 Condition for a Non-Singular  $\mathcal{J}(\theta ;p_X)$
Proof. Recall  $\mathcal{J}(\theta; p_X) = \mathbb{E}_{p_X}\left[J_\theta(X)^\top J_\theta(X)\right]$ . Suppose  $\mathcal{J}(\theta; p_X)$  is singular. Then there is a vector  $v \neq 0$  for which  $0 = v^\top \mathcal{J}(\theta; p_X)v = \mathbb{E}_{p_X}\left[\|J_\theta(X)v\|_2^2\right] = \int_{\mathcal{X}} \|J_\theta(X)v\|_2^2 p_X(x)\mathrm{d}x = 0$ . Since the integrand is non-negative and continuous by assumption, it is zero everywhere. Therefore, we have  $\sum_{i=1}^{P} v_i \partial_{\theta_i} f_\theta(x) = 0$  for all  $x \in \operatorname{supp}(p_X)$  for a non-zero  $v$ , showing  $\{\partial_{\theta_i} f_\theta(\cdot)\}_{i=1}^P$  are linearly dependent functions over the support of  $p_X$ .
Conversely, if  $\{\partial_{\theta_i}f_\theta (\cdot)\}_{i = 1}^P$  are linearly dependent functions over the support of  $p_X$ , then  $\sum_{i = 1}^{P}v_{i}\partial_{\theta_{i}}f_{\theta}(x) = J_{\theta}(x)v = 0$  for all  $x\in \operatorname {supp}(p_X)$  for some  $v\neq 0$ . Hence  $\mathcal{J}(\theta ;p_X)v = \mathbb{E}_{p_X}\big[J_\theta (X)^\top J_\theta (X)v\big] = 0$ , showing  $\mathcal{J}(\theta ;p_X)$  is singular.
B.6 Architectural Symmetries Imply Singular  $\mathcal{J}(\theta ;p_X)$
Proof. Differentiating with respect to  $\theta$ , we have  $J_{\theta} = \partial_{\theta}f_{\theta} = \partial_{\theta}f_{S\theta} = J_{S\theta}S$  for all  $\theta$ . Suppose there exists  $\theta^{*}$  that is invariant under  $S$ ,  $S\theta^{*} = \theta^{*}$ , then  $J_{\theta^{*}} = J_{\theta^{*}}S$  and  $J_{\theta^{*}}(I - S) = 0$ . Since, by assumption  $S \neq I$ , we have shown that  $J_{\theta^{*}}$  has a non-trivial nullspace. Because the nullspace of  $\mathcal{J}(\theta^{*};p_{X})$  contains the nullspace of  $J_{\theta^{*}}$ , we conclude  $\mathcal{J}(\theta;p_{X})$  is also singular.
As an example, consider a single hidden layer MLP  $f_{\theta}(x) = w_{2}^{\top}\sigma (w_{1}x)$ ,  $\theta = \{w_{1}\in \mathbb{R}^{2},w_{2}\in \mathbb{R}^{2}\}$ , where we neglected the bias for simplicity. The permutation  $P_{12}\bigoplus P_{12}$  is a symmetry of  $f_{\theta}$ , where  $\bigoplus$  is the direct sum and  $P_{12}$  is the transposition  $\begin{pmatrix} 0 & 1\\ 1 & 0 \end{pmatrix}$ . The nullspace of  $J_{\theta^{*}}$  contains the image of
$$
\left( \begin{array}{c c} 1 & - 1 \\ - 1 & 1 \end{array} \right) \oplus \left( \begin{array}{c c} 1 & - 1 \\ - 1 & 1 \end{array} \right) \tag {B.6}
$$
with a basis  $\{(1, -1, 0, 0), (0, 0, 1, -1)\}$ , for any  $\theta^*$  of the form  $(a, a, b, b)$ . It's easy to check that perturbing the parameters in directions  $(1, -1, 0, 0)$  and  $(0, 0, 1, -1)$  leaves the function output unchanged, due to symmetry in the parameter values and the network topology.
Continuous symmetries, among other reasons, can also lead to singular  $\mathcal{J}(\theta; p_X)$  and are not covered by the above analysis, which is specific to discrete symmetries. One example is scaling symmetries in MLPs with ReLU activations [5]. However, these scaling symmetries cause  $\mathcal{J}(\theta; p_X)$  to be singular for all  $\theta$  in a manner that does not introduce any pathological optimum. This becomes clear when noting that the space of functions implemented by an MLP with  $P$  parameters using ReLU activations lies in a space with dimension lower than  $P$  due to high redundancies in the parameterization. The resolution is to simply replace all occurrences of log det  $\mathcal{J}(\theta; p_X)$  with  $\log |\mathcal{J}(\theta; p_X)|_+$ , where  $|\cdot|_+$  represents the pseudo-determinant. This step is automatic if we optimize Equation (14) instead of Equation (11), where adding a small jitter will automatically compute the log pseudo-determinant (up to an additive constant) when  $\mathcal{J}(\theta; p_X)$  is singular. By contrast, symmetries such as permutation symmetries do create pathological optima because they only make the Jacobian singular at specific settings of  $\theta$ , thereby assigning to those points infinitely higher prior density compared to others.
B.7 Addressing the Infinite Prior Density Pathology with a Variational Perspective
We now describe the remedy that leads to the objective in Equation (13). As shown in Section 4.2, almost all commonly-used neural networks suffer from the pathology such that a singular Jacobian leads infinite prior density at the singularities. Such singularities arise because FS-MAP optimize for a point approximation for the function space posterior, equivalent to variational inference with the family of delta functions  $\{q(f_{\theta^{\prime}}|\theta) = \delta (f_{\theta^{\prime}} - f_{\theta})\}_{\theta}$ .
We can avoid the singularities if we instead choose a variational family where each member is a distribution over  $f$  localized around  $f_{\theta}$  for some  $\theta$ . Namely, we consider the family  $\mathcal{Q} \doteq \{ q(f_{\theta'} | \theta) \}_{\theta \in \Theta}$  parameterized by a mean-like parameter  $\theta$  and a small but fixed entropy in function space  $H(q(f_{\theta'} | \theta)) = \mathbb{E}_{q(f_{\theta'} | \theta)}[-\log q(f_{\theta'} | \theta)] = h$  for all  $\theta \in \Theta$ . For convenience, we overload the notation and use  $q(\theta' | \theta)$  to denote the pullback of  $q(f_{\theta'} | \theta)$  to parameter space. The resulting variational lower bound is
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {V L B}} (\theta ; p _ {X}) \\ = \mathbb {E} _ {q (f _ {\theta^ {\prime}} | \theta)} \left[ \log \frac {q (f _ {\theta^ {\prime}} | \theta)}{p (f _ {\theta} | \mathcal {D})} \right] \\ = \mathbb {E} _ {q (f _ {\theta^ {\prime}} | \theta)} [ - \log p (f _ {\theta^ {\prime}} | \mathcal {D}) ] - \overbrace {H (q (f _ {\theta^ {\prime}} | \theta))} ^ {\text {c o n s t .}} \\ = \mathbb {E} _ {q \left(f _ {\theta^ {\prime}} | \theta\right)} \left[ - \log p \left(\mathcal {D} \mid f _ {\theta^ {\prime}}\right) \right] + \mathbb {E} _ {q \left(f _ {\theta^ {\prime}} | \theta\right)} \left[ - \log p \left(f _ {\theta^ {\prime}}\right) \right] + \text {c o n s t .} \tag {Bayes' rule} \\ = \mathbb {E} _ {q \left(\theta^ {\prime} \mid \theta\right)} \left[ - \log p \left(\mathcal {D} \mid \theta^ {\prime}\right) \right] + \mathbb {E} _ {q \left(\theta^ {\prime} \mid \theta\right)} \left[ - \log p \left(\theta^ {\prime}\right) \right] + \frac {1}{2} \mathbb {E} _ {q \left(\theta^ {\prime} \mid \theta\right)} \left[ \log \det  \left(\mathcal {J} \left(\theta^ {\prime}; p _ {X}\right)\right) \right] \quad (\text {E q u a t i o n (1 1)}) \\ = - \log p (\mathcal {D} | f _ {\theta}) - \log p (\theta) + \frac {1}{2} \mathbb {E} _ {q (\theta^ {\prime} | \theta)} [ \log \det (\mathcal {J} (\theta^ {\prime}; p _ {X})) ] + \mathcal {O} (h) + \mathrm {c o n s t .}, \\ \end{array}
$$
where in the last line we used the assumption that  $q(\theta'|\theta)$  is localized around  $\theta$  to write the expectation of  $\log p(\mathcal{D}|\theta)$  and  $\log p(\theta)$  as their values at  $\theta' = \theta$  plus  $\mathcal{O}(h)$  corrections (which we will henceforth omit since by assumption  $h$  is tiny), because they vary smoothly as a function of  $\theta$ . More care is required to deal with the remaining expectation, since when  $\mathcal{J}(\theta; p_X)$  is singular at  $\theta$ ,  $\log \det(\mathcal{J}(\theta; p_X))$  is infinite, but its expectation,  $\mathbb{E}_{q(\theta'|\theta)}[\log \det(\mathcal{J}(\theta'; p_X))]$ , must be finite (assuming  $p(\theta)$  and  $p(\mathcal{D}|\theta)$  are finite), given that  $\mathcal{L}_{\mathrm{VLB}}(\theta; p_X)$  lower bounds the log marginal likelihood  $\log p(\mathcal{D}) = \log \int p(\mathcal{D}|\theta)p(\theta)\mathrm{d}\theta < \infty$ . As we will show in Appendix B.8, the effect of taking the expectation of the log determinant is similar to imposing a lower limit  $\epsilon$  on the eigenvalue of  $\mathcal{J}(\theta; p_X)$ , which can be approximated by adding a jitter  $\epsilon$  when computing the log determinant. This effect is similar to applying a high-frequency cutoff to an image by convolving it with a localized Gaussian filter.
With this approximation, we arrive at a simple objective inspired by such a variational perspective, equivalent to adding a jitter inside the log determinant computation in the FS-MAP objective,
$$
\hat {\mathcal {L}} (\theta ; p _ {X}) = \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, \theta\right) + \log p (\theta) - \frac {1}{2} \log \det  (\mathcal {J} (\theta ; p _ {X}) + \epsilon I). \tag {B.7}
$$
Note that while using a jitter is common in practice for numerical stability, it arises for an entirely different reason here.
B.8 Approximating the Expectation