SlowGuess/NeurIPS_2023_OCR huggingface dataset:
FS-MAP Prefers Flat Minima. In Figure 2b, we compare the curvature of the minima found by FS-MAP and PS-MAP, as measured by the average eigenvalue of the Hessian of the mean squared error loss, showing that FS-MAP consistently finds flatter minima. As the objective of FS-MAP favors small singular values for the Jacobian  $J_{\theta}$ , when multiple functions can fit the data, the FS-MAP objective will generically prefer functions that are less sensitive to perturbations of the parameters, leading to flatter minima. To make this connection more precise, consider the Hessian  $\nabla^2\mathcal{L}(\theta)$ . If the model fits the training data well, we can apply the Gauss-Newton approximation:  $\nabla_{\theta}^{2}\mathcal{L}(\theta)\approx \frac{1}{|\mathcal{D}|}J_{\theta}(x_{\mathcal{D}})^{\top}J_{\theta}(x_{\mathcal{D}})$ , which is identical to  $\mathcal{J}(\theta ;p_X)$  if  $p_X$  is chosen to be the empirical distribution of the training inputs. More generally, a distribution  $p_X$  with high density over likely inputs will assign high density to the training inputs, and hence minimizing  $\mathcal{J}(\theta ;p_X)$  will similarly reduce the magnitude of  $J_{\theta}(x_{\mathcal{D}})$ . Therefore, the FS-MAP objective explicitly encourages finding flatter minima, which have been found to correlate with generalization [10, 17, 18], robustness to data perturbations and noisy activations [3, 12] for neural networks.
FS-MAP can Achieve Better Generalization. Figure 2c shows that FS-MAP achieves lower test RMSE across a wide range of sample sizes. It's worth noting that this synthetic example satisfies two important criteria such that FS-MAP is likely to improve generalization over PS-MAP. First, the condition outlined in Section 4 for a well-behaved FS-MAP objective is met, namely that the set of partial derivatives  $\{\partial_{\theta_i}f_\theta^j (\cdot)\}_{i = 1}^P$  are linearly independent. Specifically, the partial derivatives are given by  $\{\mathrm{sech}(\theta_i)\sin (k_i\cdot),\mathrm{sech}(\theta_i)\cos (k_i\cdot)\}_{i = 1}^P$ . Since sech is non-zero everywhere, the linear independence follows from the linear independence of the Fourier basis. As a result, the function space prior has no singularities and FS-MAP is thus able to learn from data. The second important criterion is the well-specification of the probabilistic model, which we have defined to precisely match the true data-generating process. Therefore, FS-MAP seeks the most likely function according to its true probability, without incorrect modeling assumptions.
FS-MAP is Less Prone to Overfitting. As shown in Figure 2d, FS-MAP tends to have a higher train RMSE than PS-MAP as a result of the additional log determinant regularization. While PS-MAP achieves near-zero train RMSE with a small number of samples by overfitting to the noise, FS-MAP's train RMSE is consistently around  $\sigma^{*} = 0.1$ , the true standard deviation of the observation noise.
Performance Improves with Number of Evaluation Points. We compare FS-MAP estimation with  $p_{X} = \mathrm{Uniform}(-1,1)$  and with a finite number of equidistant evaluation points in  $[-1,1]$ , where the former corresponds to the latter with infinitely many points. In Figure 3a, we show that the test RMSE of the FS-MAP estimate (evaluated at 400 training samples) decreases monotonically until the number of evaluation points reaches 200. The more evaluation points, the more the finite-point FS-MAP objective approximates its infinite limit and the better it captures the behavior of the function. Indeed, 200 points is the minimum sampling rate required such that the Nyquist frequency reaches the maximum frequency in the Fourier basis, explaining the saturation of the performance.
Choice of Evaluation Distribution is Important. In Figure 3b, we compare test RMSE for 400 training samples for the default choice of  $p_X = \mathrm{Uniform}(-1,1)$ ,  $p_X = \mathrm{Uniform}(-0.1,0.1)$  a distribution that does not reflect inputs at test time—and PS-MAP  $(p_X = \mathrm{N / A})$  for reference. The result shows that specifying the evaluation distribution to correctly reflects the distribution of inputs at test time is required for FS-MAP to achieve optimal performance in this case.
4 Limitations and Practical Considerations
We now discuss the limitations of function-space MAP estimation and propose partial remedies, including a scalable approximation that improves its applicability to modern deep learning.
4.1 FS-MAP Does not Necessarily Generalize Better than PS-MAP
There is no guarantee that the most likely function is the one that generalizes the best, especially since our prior can be arbitrary. For example, FS-MAP can under-fit the data if the log determinant introduces excessive regularization, which can happen if our probabilistic model is misspecified by overestimating the observation noise. Returning to the setup in Section 3.3, in Figure 4, we find that FS-MAP (with 400 training samples) is sensitive to the observation noise scale  $\sigma$  in the likelihood. As  $\sigma$  deviates from the true noise scale  $\sigma^{*}$ , the test RMSE of FS-MAP can change dramatically. At  $\sigma / \sigma^{*} = 1/3$ , the likelihood dominates in both the FS-MAP and PS-MAP objectives, resulting in similar test RMSEs. At  $\sigma / \sigma^{*} = 10$ , the likelihood is
overpowered by the log determinant regularization in FS-MAP and causes the model to under-fit the data and achieve a high test error. By contrast, the performance of the PS-MAP estimate is relatively stable. Finally, even if our prior exactly describes the data-generating process, the function that best generalizes won't necessarily be the most likely under the posterior.
Figure 4: FS-MAP can underperform PS-MAP with a misspecified probabilistic model.
4.2 Pathological Solutions
In order for function-space MAP estimation to be useful in practice, the prior "density"  $p(f_{\theta}) = p(\theta)\operatorname*{det}^{-\frac{1}{2}}(\mathcal{J}(\theta ;p_X))$  must not be infinite for any allowed values of  $\theta$ , since if it were, those values would constitute global optima of the objective function independent of the actual data. To ensure that there are no such pathological solutions, we require the matrix  $\mathcal{J}(\theta ;p_X)$  to be non-singular for any allowed  $\theta$ . We present two results that help determine if this requirement is met.
Theorem 1. Assuming  $p_X$  and  $J_{\theta} = \partial_{\theta}f_{\theta}$  are continuous,  $\mathcal{I}(\theta ;p_X)$  is non-singular if and only if the partial derivatives  $\{\partial_{\theta_i}f_\theta (\cdot)\}_{i = 1}^P$  are linearly independent functions over the support of  $p_X$ .
Proof. See Appendix B.5
Theorem 2. If there exists a non-trivial symmetry  $S \in \mathbb{R}^{P \times P}$  with  $S \neq I$  such that  $f_{\theta} = f_{S\theta}$  for all  $\theta$ , then  $\mathcal{J}(\theta^{*}; p_{X})$  is singular for all fixed points  $\theta^{*}$  where  $S\theta^{*} = \theta^{*}$ .
Proof. See Appendix B.6
Theorem 1 is analogous to  $A^\top A$  having the same null space as  $A$  for any matrix  $A$ . It suggests that to avoid pathological optima the effect of small changes to each parameter must be linearly independent. Theorem 2 builds on this observation to show that if the model exhibits symmetries in its parameterization, FS-MAP will necessarily have pathological optima. Since most neural networks at least possess permutation symmetries of the hidden units, we show that these pathological FS-MAP solutions are almost universally present, generalizing the specific cases observed by Wolpert [30].
A Simple Remedy to Remove Singularities. Instead of finding a point approximation of the function space posterior, we can perform variational inference under a variational family  $\{q(\cdot |\theta)\}_{\theta}$  where  $q(\cdot |\theta)$  is localized around  $\theta$  with a small and constant function-space entropy  $h$ . Ignoring constants and  $\mathcal{O}(h)$  terms, the variational lower bound is given by
$$
\mathcal {L} _ {\mathrm {V L B}} (\theta ; p _ {X}) = \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta}\right) + \log p (\theta) - \frac {1}{2} \mathbb {E} _ {q \left(\theta^ {\prime} \mid \theta\right)} \left[ \log \det  \mathcal {J} \left(\theta^ {\prime}; p _ {X}\right) \right]. \tag {13}
$$
Similar to how convolving an image with a localized Gaussian filter removes high-frequency components, the expectation  $\mathbb{E}_{q(\theta'|\theta)}$  removes potential singularities in log det  $\mathcal{J}(\theta'; p_X)$ . This effect can be approximated by simply adding a small diagonal jitter  $\epsilon$  to  $\mathcal{J}(\theta; p_X)$ :
$$
\hat {\mathcal {L}} (\theta ; p _ {X}) \doteq \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta}\right) + \log p (\theta) - \frac {1}{2} \log \det  \left(\mathcal {J} \left(\theta ; p _ {X}\right) + \epsilon I\right). \tag {14}
$$
Alternatively, we know  $\mathbb{E}_{q(\theta'|\theta)}[\log \det \mathcal{J}(\theta'; p_X)]$  must be finite because the variational lower bound cannot exceed the log marginal likelihood.  $\hat{\mathcal{L}}(\theta; p_X)$  can be used as a minimal modification of FS-MAP that eliminates pathological optima. We provide full derivation for these results in Appendix B.7.
4.3 Does the Function-Space MAP Better Approximate the Bayesian Model Average?
The Bayesian model average (BMA),  $f_{\mathrm{BMA}}$ , of the function  $f_{\theta}$  is given by the posterior mean, that is:
$$
f _ {\mathrm {B M A}} (\cdot) \dot {=} \mathbb {E} _ {p (\theta | \mathcal {D})} [ f _ {\theta} (\cdot) ] = \mathbb {E} _ {p (f _ {\theta} | \mathcal {D})} [ f _ {\theta} (\cdot) ].
$$
This expectation is the same when computed in parameter or function space, and has theoretically desirable generalization properties [19, 29]. Both PS-MAP and FS-MAP provide point approximations of  $f_{\mathrm{BMA}}$ . However, FS-MAP seeks the mode of the distribution  $p(f_{\theta}|\mathcal{D})$ , the mean of which we aim to compute. By contrast PS-MAP finds the mode of a distinct distribution,  $p(\theta |\mathcal{D})$ , which can markedly diverge from  $p(f_{\theta}|\mathcal{D})$ , depending on the parameterization. Consequently, it is reasonable to anticipate that the FS-MAP objective generally encourages finding a superior estimate of the BMA.
Consider the large data regime where the posterior in both parameter and function space follows a Gaussian distribution, in line with the Bernstein-von Mises theorem [27]. In Gaussian distributions, the mode and mean coincide, and therefore  $f_{\hat{\theta}^{\mathrm{FS - MAP}}} = f_{\mathrm{BMA}}$ . However, even in this setting, where  $\hat{\theta}^{\mathrm{PS - MAP}} = \mathbb{E}_{p(\theta | \mathcal{D})}[\theta]$ , generally  $f_{\hat{\theta}^{\mathrm{PS - MAP}}} \neq f_{\mathrm{BMA}}$ , because the expectation does not distribute across a function  $f$  that is non-linear in its parameters  $\theta$ :  $f_{\mathbb{E}_{p(\theta | \mathcal{D})}[\theta]}(\cdot) \neq \mathbb{E}_{p(\theta | \mathcal{D})}[f_{\theta}(\cdot)]$ .
However, we find that whether FS-MAP better approximates the BMA depends strongly on the problem setting. Returning to the setup in Figure 1, where we have a Gaussian mixture regression model, we compare the BMA function with functions learned by FS-MAP and PS-MAP under the prior  $p(\theta) = \mathcal{N}(0,\alpha^2)$  with several settings of  $\alpha$ . We observe in Figure 5 that FS-MAP only approximates the BMA function better than PS-MAP at larger  $\alpha$  values. To understand this behavior, recall the height  $h_R$  for the right Gaussian bump is given by  $\exp (\theta_R)$ , which has a lognormal  $(0,\alpha^2)$  prior. As we increase  $\alpha$ , more prior mass is assigned to  $h_R$  with near-zero value and therefore to functions with small values within  $x\in [0,1]$ . While the lognormal distribution also has a heavy tail at large  $h_R$ , the likelihood constrains the posterior  $p(h_{R}|\mathcal{D})$  to only place high mass for small  $h_R$ . These two effects combine to make the posterior increasingly concentrated in function space around functions described by  $h_L\approx 1$  and  $h_R\approx 0$ , implying that the mode in function space should better approximate its mean. In contrast, as we decrease  $\alpha$ , both the prior and posterior become more concentrated in parameter space since the parameter prior  $p(\theta) = \mathcal{N}(0,\alpha^2)$  approaches the delta function at zero, suggesting that PS-MAP should be a good approximation to the BMA function. By varying  $\alpha$ , we can interpolate between how well PS-MAP and FS-MAP approximate the BMA.
4.4 Scalable Approximation for Large Neural Networks
In practice, the expectation  $\mathcal{J}(\theta; p_X) = \mathbb{E}_{p_X}[J_{\theta}(X)^{\top} J_{\theta}(X)]$  is almost never analytically tractable due to the integral over  $X$ . We show in Appendix B.9 that a simple Monte Carlo estimate for  $\mathcal{J}(\theta; p_X)$  with  $S$  samples of  $X$  can yield decent accuracy. For large neural networks, this estimator is still prohibitively expensive: each sample will require  $K$  backward passes, taking a total of  $\mathcal{O}(SKP)$  time. In addition, computing the resulting  $SK$ -by  $P$  determinant takes time  $\mathcal{O}(SKP^2)$  (assuming  $P \geq SK$ ). However, we can remedy the challenge of scalability by further leaning into the variational objective described in Equation (13) and consider the regime where  $\epsilon \gg \lambda_i$  for all eigenvalues  $\lambda_i$  of  $\mathcal{J}(\theta; p_X)$ . To first order in  $\max_i \lambda_i / \epsilon$ , we have  $\log \det(\mathcal{J}(\theta; p_X)) + \epsilon I) = \left.\frac{1}{2} \Delta_{\psi} d(\theta, \theta + \psi)\right|_{\psi=0}$  where  $d(\theta, \theta') \doteq \mathbb{E}_{p_X}\left[\|f_{\theta}(X) - f_{\theta'}(X)\|^2\right]$  and  $\Delta = \sum_{i=1}^{P} \partial_{\theta_i}^2$  denotes the Laplacian operator. Exploiting the identity  $\left.\frac{1}{2} \Delta_{\psi} d(\theta, \theta + \psi)\right|_{\psi=0} = \frac{1}{\beta^2} \mathbb{E}_{\psi \sim \mathcal{N}(0, \beta^2 I)}[d(\theta, \theta + \psi)] + \mathcal{O}(\beta^2)$  and choosing  $\beta$  small enough, we obtain an accurate Monte Carlo estimator for the Laplacian
(a)  $\alpha = 1.2$
(b)  $\alpha = 10$
(c)  $\alpha = 20$
Figure 5: FS-MAP does not necessarily approximate the BMA better than PS-MAP. A Gaussian mixture regression model, where the right Gaussian has weight  $\exp (\theta_R)$  with prior  $p(\theta_R) = \mathcal{N}(0,\alpha^2)$ . As we increase  $\alpha$ , we interpolate between PS-MAP and FS-MAP approximating the BMA.
using only forward passes. The resulting objective which we refer to as Laplacian Regularized MAP (L-MAP) is given by
$$
\mathcal {L} _ {\mathrm {L - M A P}} (\theta ; p _ {X}) \doteq \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta}\right) + \log p (\theta) - \frac {1}{2 \epsilon \beta^ {2}} \mathbb {E} _ {\psi \sim \mathcal {N} \left(0, \beta^ {2} I\right)} [ d (\theta , \theta + \psi) ]. \tag {15}
$$
We use a single sample of  $\psi$  and  $S$  samples of evaluation points for estimating  $d(\theta, \theta + \psi)$  at each step, reducing the overhead from  $\mathcal{O}\bigl( SKP^2\bigr)$  to only  $\mathcal{O}(SP)$ . A more detailed exposition is available in Appendix B.10.
4.5 Experimental Evaluation of L-MAP
We now evaluate L-MAP applied to neural networks on various commonly used datasets. We provide full experimental details and extended results in Appendix C. Unlike the synthetic task in Section 3.3, in applying neural networks to these more complex datasets we often cannot specify a prior that accurately models the true data-generating process, beyond choosing a plausible architecture whose high-level inductive biases align with the task (e.g. CNNs for images) and a simple prior  $p(\theta)$  favoring smooth functions (e.g. an isotropic Gaussian with small variances). Therefore, we have less reason to expect L-MAP should outperform PS-MAP in these settings.
UCI Regression. In Table 1, we report normalized test RMSE on UCI datasets [1], using an MLP with 3 hidden layers and 256 units. L-MAP achieves lower error than PS-MAP on 7 out 8 datasets. Since the inputs are normalized and low-dimensional, we use  $p_{X} = \mathcal{N}(0,I)$  for L-MAP.
Table 1: Normalized test RMSE (↓) on UCI datasets. We report mean and standard errors over six trials.
METHOD	BOSTON	CONCRETE	ENERGY	NAVAL	POWER	PROTEIN	WINERED	WINEWHITE
PS-MAP	.329±.033	.272±.016	.042±.003	.032±.005	.219±.006	.584±.005	.851±.029	.758±.013
L-MAP	.352±.040	.261±.013	.041±.002	.018±.002	.218±.005	.580±.005	.792±.031	.714±.017
Image Classification. In Table 2, we compare L-MAP and PS-MAP on image classification using a ResNet-18 [9]. L-MAP achieves comparable or slightly better accuracies and is often better calibrated. We further test the effectiveness of L-MAP with transfer learning with a larger ResNet-50 trained on ImageNet. In Table 3, we show L-MAP also achieves small improvements in accuracy and calibration in transfer learning.
Loss Landscape. In Figure 6 (Left), we show that L-MAP indeed finds flatter minima. Further, we plot the Laplacian estimate in Figure 6 (Right) as the training progresses. We see that the Laplacian is much lower for L-MAP, showing its effectiveness at constraining the eigenvalues of  $\mathcal{J}(\theta; p_X)$ .
Distribution of Evaluation Points. In Table 2, we study the impact of the choice of distribution of evaluation points. Alongside our main choice of the evaluation set (KMNIST for FashionMNIST and CIFAR-100 for CIFAR-10), we use two additional distributions - the training set itself and a white noise  $\mathcal{N}(0, I)$  distribution of the same dimensions as the training inputs. For both tasks, we find that using an external evaluation set beyond the empirical training distribution is often beneficial.
Number of Evaluation Point Samples. In Figure 7(a), we compare different Monte Carlo sample sizes  $S$  for estimating the Laplacian. Overall, L-MAP is not sensitive to this choice in terms of accuracy. However, calibration error [20] sometimes monotonically decreases with  $S$ .
Table 2: We report the accuracy (ACC.), negative log-likelihood (NLL), expected calibration error [20] (ECE), and area under selective prediction accuracy curve [7] (SEL. PRED.) for FashionMNIST [31] ( $\mathcal{D}' =$  KMNIST [4]), and CIFAR-10 [15] ( $\mathcal{D}' =$  CIFAR-100). A ResNet-18 [9] is used. Std. errors are reported over five trials.
METHOD	FASHIONMNIST	CIFAR-10
ACC.↑	SEL. PRED.↑	NLL↓	ECE↓	ACC.↑	SEL. PRED.↑	NLL↓	ECE↓
PS-MAP	93.9%±.1	98.6%±.1	.26±.01	4.0%±.1	95.4%±.1	99.4%±.0	.18±.00	2.5%±.1
L-MAP pX=N(0,I)	94.0%±.0	99.2%±.0	.25±.01	4.0%±.2	95.3%±.1	99.4%±.0	.20±.00	3.0%±.1
pX=TRAIN	93.8%±.1	99.2%±.1	.27±.01	4.3%±.2	95.6%±.1	99.5%±.0	.18±.01	2.6%±.0
pX=D'	94.1%±.1	99.2%±.1	.26±.01	4.1%±.1	95.5%±.1	99.5%±.0	.16±.01	1.4%±.1
Table 3: Transfer learning from ImageNet with ResNet-50 on CIFAR-10 can lead to slightly better accuracy and improved calibration.
METHOD	ACC. ↑	SEL. PRED. ↑	NLL ↓	ECE ↓
PS-MAP	96.3%±0.1	99.5%±0.1	0.18±0.01	2.6%±0.2
L-MAP	96.4%±0.1	99.5%±0.1	0.15±0.01	2.1%±0.1
Robustness to Label Noise. In Figure 7(b), we find L-MAP is slightly more robust to label noise than PS-MAP on CIFAR-10.
Main Takeaways. L-MAP shows qualitatively similar properties as FS-MAP such as favoring flat minima and often provides better calibration. However, it achieves comparable or only slightly better accuracy on more complex image classification tasks. In line with our expectations, these results suggest that accounting for the precise difference between FS-MAP and PS-MAP is less useful without a sufficiently well-motivated prior.
Figure 6: Empirical evidence that L-MAP finds flatter minima. (Left) For various step sizes in 20 randomly sampled directions starting at the minima, we compute the training loss averaged over all directions to summarize the local loss landscape. We use filter normalization for landscape visualization [16]. L-MAP visibly finds flatter minima. (Right) We plot the Laplacian estimate throughout training, showing L-MAP is indeed effective at constraining the eigenvalues of  $\mathcal{J}(\theta; p_X)$ .
Figure 7: (Left) Calibration error can be reduced by increasing the number of Monte Carlo samples for the Laplacian estimator in Equation (15). (Right) L-MAP is slightly more robust to training with label noise. For each level of noise, we replace a fraction of labels with uniformly random labels.
While we typically train our models through PS-MAP, we show that FS-MAP has many appealing properties in addition to re-parametrization invariance. We empirically verify these properties, including the potential for better robustness to noise and improved calibration. But we also reveal a more nuanced and unexpected set of pros and cons for each approach. For example, while it is natural to assume that FS-MAP more closely approximates a Bayesian model average, we clearly demonstrate how there can be a significant discrepancy. Moreover, while PS-MAP is not invariant to re-parametrization, which can be seen as a fundamental pathology, we show FS-MAP has its own failure modes such as pathological optima, as well as practical challenges around scalability. In general, our results suggest the benefits of FS-MAP will be greatest when the prior is sufficiently well-motivated.
Our findings engage with and contribute to active discussions across the literature. For example, several works have argued conceptually—and found empirically—that solutions in flatter regions of the loss landscape correspond to better generalization [8, 10, 11, 12, 25]. On the other hand, Dinh et al. [5] argue that the ways we measure flatness, for example, through Hessian eigenvalues, are not parameterization invariant, making it possible to construct striking failure modes. Similarly, PS-MAP estimation is not parameterization invariant. However, our analysis and comparison to FS-MAP estimation raise the question to what extent lack of parametrization invariance is actually a significant practical shortcoming—after all, we are not reparametrizing our models on the fly, and we have evolved our models and training techniques conditioned on standard parameterizations.
Acknowledgements. We thank Marc Finzi, Pavel Izmailov, and Micah Goldblum for helpful discussions. This work is supported by NSF CAREER IIS-2145492, NSF I-DISRE 193471, NSF IIS-1910266, BigHat Biosciences, Capital One, and an Amazon Research Award.