Should We Learn Most Likely Functions or Parameters?
Authors: Shikai Qiu
Tim G. J. Rudner
Sanyam Kapoor
Andrew Gordon Wilson
Venue: NeurIPS 2023 poster (2023)
Keywords: Function-Space Modeling
Function-Space Regularization
Maximum A Posteriori Estimation
Generalization
TL;DR: This paper investigates the implications of directly learning the most likely function, instead of the most likely parameters, implied by a model and the data.

Abstract: Standard regularized training procedures correspond to maximizing a posterior distribution over parameters, known as maximum a posteriori (MAP) estimation. However, model parameters are of interest only insomuch as they combine with the functional form of a model to provide a function that can make good predictions. Moreover, the most likely parameters under the parameter posterior do not generally correspond to the most likely function induced by the parameter posterior. In fact, we can re-parametrize a model such that any setting of parameters can maximize the parameter posterior. As an alternative, we investigate the benefits and drawbacks of directly estimating the most likely function implied by the model and the data. We show that this procedure leads to pathological solutions when using neural networks and prove conditions under which the procedure is well-behaved, as well as a scalable approximation. Under these conditions, we find that function-space MAP estimation can lead to flatter minima, better generalization, and improved robustness to overfitting.

Should We Learn Most Likely Functions or Parameters?
Shikai Qiu"
Tim G. Rudner"
Sanyam Kapoor*
Andrew Gordon Wilson
New York University
Abstract
Standard regularized training procedures correspond to maximizing posterior distribution over parameters, known as maximum a posteriori (MAP) estimation. However, model parameters are of interest only insomuch as they combine with the functional form of model to provide function that can make good predictions. Moreover, the most likely parameters under the parameter posterior do not generally correspond to the most likely function induced by the parameter posterior. In fact, we can p--rameetrrzee a model such that any setting of parameters can maximize the parameter posterior. As an alternative, directly estimating the most likely function implied by the model and the data. We show that this procedure leads to pathological solutions when using neural networks and prove conditions under which the procedure bell-beeaved, as well scalable approximation. Under these conditions, find that function-space MAP estimation can lead to flatter minima, better generalization, and improved robustness overfitting.
investigate the benefits and drawbacks
1 Introduction
Machine learning has matured to the point where we often take key design decisions for granted. One of the most fundamental such decisions is the loss function use to train models. Minimizing standard regularized loss functions, including cross-entropy for classification and mean-squared or mean-absolute error for regression, with & or & regularization, exactly corresponds to maximizing posterior distribution over model parameters [2, 19). This standard procedure is known in probabilistic modeling as maximum posteriori (MAP) parameter estimation. However, parameters have no meaning independent of the functional form of the models they parameterize. In particular, our models fz(x) are functions given parameters which map inputs x (e.g., images, spatial locations, etc.) to targets (e.g., softmax probabilities, regression outputs, etc.). We are typically only directly interested in the function and its properties, such as smoothness, which we use make predictions.
Alarmingly, the function corresponding to the most likely parameters under the parameter posterior does not generally correspond to the most likely function under the function posterior. For example, in Figure la, we visualize the posterior over a mixture coefficient On a Gaussian mixture regression model in both parameter and function space, using the same Gaussian prior for the mixture coefficients (corresponding to E regularization in parameter space). We see that each distribution is maximized by different parameter OR. leading to very different learned functions in Figure 1b. Moreover, we can re- re-parametrize the functional form of any model such that any arbitrary setting of parameters maximizes the parameter posterior (we provide further discussion and an example in Appendix A).
Should we then be learning the most likely functions or parameters? As we will see, the nuanced pros and cons of each approach are fascinating and often unexpected.
Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

0.6
1.00
p(f-[D)
Trining Data
p(()2))
0.75
FS-MAP
0.4
0.50
0.2
0.25
0.0
0.00 -1
-5
0
0R
X
(a) Distinct Posterior Densities
(b) Distinct Learned Functions
Figure I: Illustration of the Difference Between Most Likely Functions and Parameters. The function that is most prohable (denoted as FS-MAP) under the posterior distribution over functions can substantially diverge from the function represented by the most probable parameters (denoted as PS-MAP) under the posterior distribution over parameters. We illustrate this fact with regression model with two parameters, 0,, and OR. both with a prior (0, 1.2*). This model is used to learn mixture of Gaussians with fixed mean and variance, where the mixture coefficients are given by exp(01. and exp(OR). Both the most probable solution and posterior density show significant differences when analyzed in function-space sepciion-ppcc versus parameter-span Since Oz is
we only plot the posterior as function of 0g. We normalize the area under p(fa (D) to 1.
On one hand, might expect the answer to clear-cut "We should learn most likely functions!". In Section 3, we present a well-defined function-space MAP objective through generalization of the change of variables formula for probability densities. addition functions being the direct quantity of interest and the function-space MAP objective being invariant to the parametrization of our model, we show that optimization of the function-space MAP objective indeed results in more probable functions than standard parameter-space MAP, and these functions often correspond to flar minima [10, 12] that are more robust to overfitting.
On the other hand, function-space MAP is not without its own pathologies and practical limitations. We show in Section 4 that the Jacobian term in the function-space MAP objective admits trivial solutions with infinite posterior density and can require orders of magnitude more computation and memory than parameter-space MAP, making difficult to scale to modern neural networks. We also show that function-space MAP will not necessarily be closer than parameter-spp space MAP to the posterior expectation over functions, which in Bayesian inference forms the mean of the posterior predictive distribution and has desirable generalization properties [2. 19, 29]. To help address the computational limitations, we provide in Section 4 scalable approximation to the function-space MAP objective applicable to large neural networks using Laplacian regularization, which we refer to as L-MAP. We show empirically that parameter-space MAP is often able to perform on par with L-MAP in accuracy, although L-MAP tends to improve calibration.
The aim of this paper is improve our understanding of what it means learn most likely functions instead of parameters. Our analysis, which includes theoretical insights as well as experiments with both carefully designed tractable models and neural networks, paints complex picture and unearths distinct benefits and drawbacks of learning most likely functions instead of parameters. We conclude with discussion of practical considerations around canetion-space MAP estimation, its relationship to flat minima and generalization, and the practical relevance of parameterization invariance.
Our code is available http::///thhha..cc//////tttteeee
Preliminaries
We consider supervised learning problems with training dataset D = (ID, (10 Up I.v of inputs x E X' and targets y E Y with input space X C RR and output space y C RR For example, the inputs x could correspond to times, spatial locations, tabular data, images, and the targets y to regression values, class labels, etc.
(1)

---

To create a model of the data D, one typically starts by specifying a function fo : X - ) parametrized by a vector 0 E Iº which maps inputs outputs. fo could be a neural network, polynomial model, a Fourier series, and so on. Learning typically amounts to estimating model parameters 0. To this end, we can relate the function fo to the targets through an observation model, ply fo). For example, in regression, we could assume the outputs y = (x) + where ~ ((0,, o²) is additive Gaussian noise with variance o². Equivalently, p(y) æ 0) = NNy (x)), o²). Alternatively, in classification, we could specify p(y (x, 0) Categorical certoriica(sootmaaxx foftmax(ff(e (x))) We then use this observation model to form a likelihood over the whole dataset p(y a 0). In each of these example observation models, the likelihood factorizes, as the data points are conditionally independent given model parameters
0, p(yD IDD o) = II." p(yD ( 0). We can further express a belief over values of parameters through prior p(), such p(0) (O0u,, 2.) Finally, using Bayes rule, the log of the posterior up to a constant e independent of is,
(6)
parameter posterior
likelihood
prior
log p(0 TD) =
0) + logg
log (()) +
Notably, standard loss functions negative log posteriors, such that maximizing this parameter pos- terior, which we refer to as parameter-space maximum a posteriori (PS-MAP) estimation, corresponds to minimizing standard loss functions [19]. For example, if the observation model is regression with Gaussian noise or Laplace noise, the negative log likelihood is proportional to mean-square or mean-absolute error functions, respectively. If the observation model is categorical distribution with softmax link function, then the log likelihood is negative cross-entropy. If we use zero-mean Gaussian prior, we recover standard f2 regularization, also known as weight-decay. If use Laplace prior, we recover l regularization, also known as LASSO [26].
Once we maximize the posterior to find
MSS-MPP
arg max, p (Olyy,>D);
we can condition on these parameters to form our function fyrsmar to make predictions. However, as we saw in Figure 1, Sar not general the function that maximizes the posterior over functions, p( fo (D) In other words, if
G
arg maxg ((JolyD,
then generally fersuur # feesnar Naively, one can write the log posterior over fo up to the same constant e above as
function posterior
likelibood
function prior
log Dogr(folDD.zg =
log pogg((DDDDD fo) + log p(fo) +
(4)
where p(yp|æp pp, fo) = p(yplupp IDD,0, 0), but just written in terms of the function fa. The prior p(fo) however is different function from p(0). because ineur additional Jacobian factor in this change of variables, making the posteriors also different.
We must take care in interpreting the quantity p(fo) since probability densities in dimensioad vector spaces are generally ill-defined. While prior work [14, 23, 24, 25, 30] avoids this problem by considering a prior only over functions evaluated at finite number of evaluation points, we provide a more general objective that enables the use of infinitely many evaluation points to construct more informative prior and makes the relevant design choices of function-space MAP estimation more interpretable.
Understanding Function-Space Maximum A Posteriori Estimation
Function-space MAP estimation seeks to answer a fundamentally different question than parameter- space MAP estimation, namely, what is the most likely function under the posterior distribution over functions implied by the posterior distribution over parameters, rather than the most likely parameters under the parameter posterior.
To better understand the benefits and shortfalls of function-space MAP estimation, we derive function-space MAP objective that generalizes the objective considered by prior work and analyze its properties both thcoretically and empirically.
3

---

3.1 The Finite Evaluation Point Objective
Starting from Equation (4), Wolpert [30] proposed to instead find the MAP estimate for fo()), the function evaluated at a finite set of points à {ai,..., TM where M < oo can be chosen to be arbitrarily large as capture the behavior of the function to arbitrary resolution. This choice then yields the FS-MAP optimization objective
Lanite (0; i:) = Σ
(8) 1=1 log logp(u?'as 1- fo(()) + log log
(5)
where (i)) is well-defined but not in general analytically tractable probability density function. (See Appendix B.1 for further discussion.) Let P be the number of parameters A and K the number of function outputs. Assuming the set of evaluation points is sufficiently large so that MK 2 P, using a generalization of the change of variable formula that only assumes injectivity rather than bijectivity, Wolpert [30] showed that the prior density over (() is given by
p(fo())) = p(0) det dct- (T(;;)),
(6)
where J(0; B) is P-by-P P matrix defined by
Tanite (0: +) Jo( Jo(
(7)
is the MK-by-P Jacobian of fo()), viewed as an M -idiimensoonal vector, with respect to the parameters O. Substituting Equation (6) into Equation (5) the function-space MAP objective as a function of the parameters / can be expressed as
and O(a())
.V Σ logp(1 (1) [ fo) + p() log (0;;
Lanite (0;
8
where we are allowed condition directly on fo instead of (Ø) because by assumption è is large enough to uniquely determine the function. In addition to replacing the function, our true quantity of interest, with its evaluations at a finite set of points é this objective makes unclear how we should select & and how that choice can be interpreted or justified in principled way.
3.2 Deriving More General Objective and its Interpretation
We now derive a more general class of function-spact space MAP objectives that allow us to use effectively infinitely many evaluation points and meaningfully interpret the choice of those points.
In general, when performing a change of variables from v E R to u E IR" via an injective map p())du(v); where μ and define p(t), their probability densities relate p(u)dı((u)
the volume induced by an M x M metric tensor g, and dv(c)) = d' the Lebesgue measure, then we can write dı(u)) where the N x N metric ā(v) = J(v) (()) is known the pullback of via oo and J is the Jacobian of y [6]. As a result, we have p(u) ()? (v)). Applying this argument and identifying u with fo( and with 0, Wolpert [30] thereby establishes p(ff p(0) det (/J Jz)))
respective volume measures. Suppose we let du(u) =v (del(g
/
However, an important implicit assumption in the last step is that the metric function space is Euclidean. That is, I and the squared distance between fo(t) and So(ă) + dfa(E) is d EMI account for a generic metric y, we therefore replace Ja())) Jo( with we assume the function output is univariate (K 1) and only consider g that constant i.c.. independent of fa( and diagonal, with Iü = g(x;)) for some function g X + Rt. For a discussion of the general case, see Appendix B.2. To better interpret the choice of g, we rewrite
To M dff(z.)², rather than the general case ds² EM E (f ( For simplicity.
Jo(3) gJ(() = Σ (x(x;) Jo(zz),
where we suggestively defined the alias px - g(x) and Jo(x;) is the T-by--P-dimensional Jaco- bian evaluated at the point x.. Generalizing from the finite and discrete set å to the possibly infinite entire domain X' C R and further dividing by an unimportant normalization constant Z, we obtain
1 [(0:px) = Z / (x(x))o(1) Jo(x) dx Epp [o(())To(X))
(10)

---

where px Px is a normalized probability density function, with normalization Z 2-which is independent of 0 onolyy appearing as an additive constant in log det [[:: px). We include a further discussion about this limit in Appendix B.4.
Under this limit, can identify log p(0) 2 log det J(0: with log p(fo) up a constant, and thereby express log p [D) the function-space MAP objective in Equation (4) as
px): Σ i=1 logp(yD (1) mp fo) + (()
log det [(:ppx):
(11)
where the choice of px corresponds to a choice of the metric y in function space. Equation (11) (and approximations thereof) will be the primary object of interest in the remainder of this paper.
Evaluation Distribution as Funetion-Space Geometry. From this discussion and Equations (10) and (II), it is now clear that the role of the metric tensor is to impose distribution px over the evaluation points. Or conversely, given distribution over evaluation points implies metric tensor, and as such, specifies the geometry of the function space. This correspondence is intuitive: points with higher values of g(x) contribute more defining the geometry in function space and therefore should be assigned higher weights under the evaluation distribution when maximizing function space posterior density. Suppose, for example, So is an image classifier for which we only care about its outputs on set of natural images Z when comparing it to another image classifier. The metric g(.) and therefore the evaluation distribution px only needs support in I, and the FS-MAP objective is defined only in terms of the Jacobians evaluated at natural images E I.
Finite Evaluation Point Objective as Special Case. Consequently, the finite evaluation point objective in Equation (8) can be arrived at by specifying the evaluation distribution Px (x) to be px M Ex'es 6(x x'), where is the Dirac delta function and fæ1,..,, Tas) with M < oo, before. It is easy see that Tanite (0; 2) J p0:px. Therefore, the objective proposed by Wolpert [30] is special case of the more general class of objectives.
3.3 Investigating the Properties of Function-Space MAP Estimation
To illustrate the properties of FS-MAP, consider the class of models fa(a) = = (a))
with domain X' [-1, where { is fixed set of basis functions and a a non-linear function to introduce a difference between function-space and parameter-space MAP. The advantage of working with this class of models is that I (0;px) has simple closed form, such that
[1x(0:px) = Epx So So (x, [o()] = = (0.))'(00)
(12)
where & is a constant matrix with elements P.j = Epx puted once and reused throughout training. In this experiment, (kos(()))), (;;)) ] 100 where k; = in and set a = tanh. We generate training data by sampling 10, evaluating fo (.train)) and adding Gaussian noise with standard deviation o" = 0.1 to each observation. We use 1,000 test points sampled from Uniform(-1,, 1). To train the model, we set the prior p(0) and the likelihood correspond to the data-generating process. For FS-MAP, we specify Px Uniform (11, 1), the ground-truth distribution of test inputs, which conveniently results in T I/2.
Therefore, + can be precom- the set of Fourier features
Twain ~ Uniform( 1). 02 ~O, a²) with
FS-MAP Finds More Probable Functions. Figure 2a shows the improvement in p(fo (D) when using FS-MAP over PS-MAP. As expected, FS-MAP consistently finds functions with much higher posterior probability p(fol (D).
6x 10
PS: MAP FS-MAP
10
10-2
10a
10-2 10
10a
10a
10
10a # Train samples
10a
10a
10a
10 # Train samples
# Train samples
# Train samples
(a) Log Posterior Increase
(b) Loss Curvature
(c) Test RMSE
(d) Train RMSE
Figure 2: FS-MAP exhibits desirable properties. On non-linear regression problem, FS-MAP empirically (a)
learns more probable functions, (b) finds flatter minima, (c) improves generalization, and (d) is less prone to overfitting. The plot shows means and standard deviations computed from 3 random seeds.
5

---

FS
Test RMSE
PS-MAP
Test RMSE
0.06
0.06
0.04
0.04
0.02
10
50
200
0.00
Vl-1, 1.00-0
NAA
Evall Points
Px
(a) Number of evaluation points
(b) Evaluation distribution
Figure 3: Effect of important hyperparameters in FS-MAP. FS-MAP improves with the number of evaluation points and requires carefully specifying the evaluation distribution and the prohabilistic model.
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 Jg. 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 D²) If the model fits the training data well, we can apply the Gauss-Newton approximation: V2.(0)= x [D] Jø (=D)) Jg (#D).
which is identical to I (0; px) if px is chosen to be the empirical distribution of the training inputs. More generally, a distribution px with high density over likely inputs will assign high density to the training imputs, and hence minimizing J (0: px) ) will similarly reduces the magnitude of Ja (TD). 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 well-behaved FS-MAP objective is met, namely that the set of partial derivatives (o are linearly independent. Specifically, the partial derivatives are given by {sech(O:) sin(k:-), sech(O:) (ks(k:;))]==1 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-gent genteaannn process. Therefore, FS-MAP seeks the most likely function according to its true probability, without incorrect modeling assumptions.
MAP is Less Prone Overfitting. As shown in Figure 2d, FS-MAP tends have higher train RMSE than PS-MAN 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 o* 0.1, the true standard deviation of the observation noise.
Performance Improves with Number of Evaluation Points. We compare FS-MAP estimation with Px = Uniform( 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 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 Px = Uniform( (-1, 1), px = Uniform( 100.1 0.1) - distribution that does not refleet inputs at test time and PS-MAP (px N/A) for reference. The result shows that specifying the evaluation distribution to correctly reflect the distribution inputs at test time is required for FS-MA FS-MAP to achieve optimal performance in this case.

---

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 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 intro- duces excessive regularization, which can happen our probabilistic model is misspecified by overestimating the observation noise. Re- turning to the setup in Section 3.3, in Figure we find that FS-MAP (with 400 training samples) is sensitive to the observation noise scale a in the likelihood. As a deviates from the true noise scale o², the test RMSE of FS-MAP can change dramatically. At 0//o" 1/3, the likelihood dominates in both the FS-MAP and PS- PS-MAP objectives, resulting in similar test RMSEs. At a 10, the likelihood is overpowered by the log determinant regularization in FS-MAP and causes the model to under-fit the data and achieve high test error. By contrast, the performance of the PS-MAP estimate is relatively stable. Finally, even if our prior exactly describes the mata-generating process, the function that best generalizes won't necessarily be the most likely under the posterior.
Test RMSE
0.4
PS-MAP
FS-MAP
0.2
1/3
ala
Figure 4: FS-MAP can underper- form 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" =(Ja)=p(t det (1 p0;))) must not be infinite for any allowed values of 0, since if it were, those values would constitute global optima of the objective function independent of the actual data. To ensure that there are such pathological solutions, we require the matrix J(0: px) to be non-singular for any allowed 0. We present two results that help determine if this requirement is met.
Theorem 1. Assuming px and Jo fo are continuous, J(0: Px) is non oonguir if and only if the partial derivatives (o [-) are linearly independent functions over the support of px-
Proof. See Appendix B.5
Theorem 2. If there exists a non-trivial symmetry S E PxP with S # such that fo = fse for all 0, then I (0" :p)) is singular for all fixed points Or where S0 0*
Proof. See Appendix B.6
Theorem 1 is analogous to A A having the same null space as A for any matrix A. It suggests that avoid pathological optima the effect of small changes each parameter must be linearly independent. Theorem 2 builds on this observation to show that if the model exhibits symmetries in its parameterization, PA 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].
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 {g(-10)},, where g(-10) is localized around 0 with a small and constant function-space entropy h. Ignoring constants and Ohh) terms, the variational lower bound given by
(0;p)(;px) Σ
log log((v (9) [ fo) - log p(0) Bycoroe) [log det J :x):
(13)
Similar to how convolving an image with localized Gaussian filter removes high-frequency compo- nents, the expectation Ey(0 removes potential singularities in log det J (Px) This effect can be approximated by simply adding a small diagonal jitter to [(::px):
L(B: [(0:px)= ΣxTh log p(v fo) log p(0) log det (J(0: + cT).
(1) (6)
1
(14) Alternativel,, we know Egyorrn (0) [log det J(B': px)] must be finite because the variational lower bound cannot exceed the log marginal likelihood. L(:; px) 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.
7

---

4.3 Does the Function-Space MAP Better Approximate the Bayesian Model Average?
The Bayesian model average (BMA), SIMA, of the function fo is given by the posterior mean, that is: fBMA (-) = ()p(rr))ff(()) = (E((..
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 SBMA. However, FS-MAP seeks the mode of the distribution p p(f))D), the mean of which we aim to compute. By contrast PS-MAP finds the mode of a distinct distribution, p(0 which can markedly diverge from p(fol DD), 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 fgrswar SBMA- However, even in this setting, where @'S-MAP ( p) [0]. generally fgrsww # fISMA, because the expectation does not distribute across a function f that is non-linear in its parameters A: |0 (-) o Egp(o D) [fo()-))-
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 Gaussian mixture regression model, we compare the BMA function with functions learned by FS-MAP and PS-MAP under the prior p(0) A/0 a²) with several settings of a. We observe in Figure 5 that FS-MAP only approximates the BMA function better than PS-MAP at larger a values. To understand this behavior, recall the height hR for the right Gaussian bump is given by exp(OR), which has a lognormal(O, a²) prior. As we increase at more prior mass is assigned to hR with near-zero value and therefore to functions with small values within x E [o, 1j. While the lognormal distribution also has a heavy tail at large hR, the likelihood constrains the posterior p(h (D) to only place high mass for small hg. These two effects combine to make the posterior increasingly concentrated in function space around functions described by hg R 1 and R A 0, implying that the mode in function space should better approximate its mean. In contrast, as we decrease a, both the prior and posterior become more concentrated in parameter space since the parameter prior p(0)= = (0, a²) approaches the delta function at zero, suggesting that PS-MAP should be good approximation to the BMA function. By varying a, 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 I (0; = Epx (X) due to the integral over X. We show in Appendix B.9 that a simple Monte Carlo estimate for J(0: px) with S samples of X can yield decent accuracy. For large neural networks, this estimator is still prohibitively expensive: each sample will require backward passes, taking a total of O(SKP) time. In addition, computing the resulting S SK-byP determinant takes time O(SK (assuming P& Z 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 >> X for all eigenvalues X; of Apd(0 A + ))) - where d(0, 8') = Epp (x) fo- ( and
(oo)TTT(X))))))) Jo(X)) is almost never analytically tractable
px)) + el) = To first order in max, da/e, we have log det [[:; (0:
Σ 0g denotes the Laplacian operator. Exploiting the identity 1 d 0 +b)] = ( ga 110.0110 |d(0, 0 }) +
(B2) and choosing B small enough, we obtain an accurate Monte Carlo estimator for the Laplacian
1.0
1.0
1.0
Training Data
Faining Deta
Tairing Data
HAP
PS-MP
HS-MAP
PS MAP
PS HAP BMA
PS NAP
0.5
BMM
0.5
> 0.5
RMA
0.0
0.0
0.0
0
X
(a) a = 1.2
(b) a = 10
(C) a = 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(OR) with prior p(On) = N1(0, a As we increase a, we interpolate between PS-MAP and FS-MAP approximating the BMA.
8

---

using only forward passes. The resulting objective which we refer to Laplacian Regularized MAP (L-MAP) is given by
G.:ma)(Opx)) Σ. log plyp (8) I .o) + +ogg p(0) 20B2 N00 +)]- (15)
We use single sample of d and S samples of evaluation points for estimating d(0, 0 v) each step, reducing the overhead from (S P2) only OSSP). more detailed exposition is available in Appendix B.10.
4.5 Experimental Evaluation 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 prior that accurately models the true mata--ennraaing process, beyond choosing plausible architecture whose high-level inductive biases align with the task (c.g. CNNs for images) and a simple prior p(0) favoring smooth functions (c.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, 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 out 8 datasets. Since the inputs are normalized and low-dimensional, we use px (o0, for L-MAP.
Table 1: Normalized test RMSE (1) on UCI datasets. We report mean and standard errors over six trials.
METHOD
BOSTON CONCRETE ENERGY
NAVAL
POWER
PROTEIN
WINERED WINEWHITE
PS-MAP
.3294.08
2722.016
.0422.003
.0324 005
.219 006
.584
.851
.758
L-MAP
.352=.04
.6112.018 .041 .02
.018 002
2188 005
.5802.005 .7922.081
.714=.017
Image Classification. In Table 2, we compare L-MAP and PS-MAP on image classification using ResNet-18 [9]. L-MAP achieves comparable or slightly better accuracies and 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.-AP also achieves small improvements accuracy and calibration in transfer learning.
Figure (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 much lower for L-MAP, showing its effectiveness constraining the eigenvalues of J (0: px)).
Loss Landscape.
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 (O, 7) distribution of the same dimensions the training inputs. For both tasks, find that using an external evaluation set beyond the empirical training distribution is often beneficial.
Distribution of Evaluation Points.
Number of Evaluation Point Samples. In Figure 7(a), compare different Monte Carlo sample sizes S for estimating the Laplacian. Overall, L-MAP not sensitive to this choice in terms of accuracy. However, calibration error [20] sometimes monotonically decreases with S.
Table 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] (D' KMNIST [4]), and CIFAR- [15] (D' CIFAR-100). A ResNet- 18 [9] is used. Std. errors reported over five trials.
FASHIONMNIST
CIFAR-10
METHOD
Acc.t
SEL. PRED. NLLL ECE,
Acc.t
SEL. PRED. NLL, ECE,
PS-MAP
93.9%+..
98.6%
262. 4.0%
95.4%...
99.4%...
.18 2.5%
L-MAP px =V(00 /)
94.0%++o
99.2%+..
25+.01 4.0% +
+..3%++.
99.4%+.o
.20+.00 +.0%+.1
Px TRAIN
13.8%++1
99.2%++.1
27
4.3% +
95.6%+.1
99.5%...
.18
2.6%+.o
Px =D=
94.1%
99.2%++.1
.26
4.1% +
95.5%+.1
99.5%+..
.16+.01 1.4%+.1

---

Table Transfer learning from ImageNet with ResNet-50 on CIFAR-10 can lead slightly better accuracy and improved calibration.
METHOD
Acc. f
SEL, PRED. 1
NLL
ECE
PS MAP
96.3%
99.5%=o.1 99.5%=z.1
+.18 0.1540.01
2.6%
1.-APP
96.1%
2.1 1
x10
x10
x101
x10*
Method
Laplacian Est.
Method
L-MAP
L-MAP PS-MAP
Loss
PS-MAP
5.0
2.5
0.0
-2.5
0.0
2.5
-2.5
2.5
100 150 200 Epoch
Step Size
Epoch
Step Size
(a) FashionMNIST
(b) CIFAR-10
Figure 6: Empirical evidence that L-MAP finds flatter minima. (Left) For various step sizes in 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-MAA is indeed effective at constraining the eigenvalues of px:
Robustness to Label Noise. In Figure 7(b), we find L-MAP is slightly more robust to label noise than PS-MAP CIFAR-IO.
100
Method
Calloration ratton Errr
L-HAP
Accuracy
PSMAP
Fashamm9as1 1
CIRAR-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 only slightly better accuracy on more complex image classifi cation tasks. In line with our expectations, these results suggest that accounting for the precise dif- ference between FS-MAP and PS-MAP is less use- ful without a sufficiently well-motivated prior.
10a
0.0
0.4
0.8
Samples (5)
Noise
Figure 7: (Left) Calibration error can be reduced by increasing the number of Monte Carlo samples for the Laplacian estimator Equation (15). (Right) L-MAP is slightly more robust training with label noise, For each level of noise, we replace a fraction of labels with uniformly random lahels.
5
Discussion
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 is natural to assume that FS- MAP more closely approximates a Bayesian model average, we clearly demonstrate how there can be significant discrepancy. Moreover, while PS-MAP is not invariant to re- parametriizatinn which can be seen as fundamental pathology, we show FS-MAI 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-MAU estimation 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 -fter all, we are not reparametrizing our models on the fly, and we have evolved our models and training techniques conditioned on standard parameterizations.
10

---

Acknowledgements. We thank Mare Finzi, Pavel Izmailov, and Micah Goldblum for helpful discussions. This work supported by NSF CAREER IIS-2145492, NSF I-DISRE 193471, NSF I18-1910266, BigHat Biosciences, Capital One, and an Amazon Research Award.
References
[1] Arthur Asuncion and David Newman. Uci machine learning repository, 2007.
[2] Christopher M. Bishop. Pattern recognition and machine learning (information science and statistics). 2006.
[3] Pratik Chaudhari, Anna Choromańska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer T. Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-sgd: biasing gradient descent into wide valleys. Journal of Statistical Mechanics: Theory and Experiment, 2019, 2016.
[4] Tarin Clamuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature, 2018.
[5] Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pages 1019-1028, 2017.
[6] Christopher Terence John Dodson and Timothy Poston. Tensor geometry: the geometric viewpoint and its uses, volume 130. Springer Science & Business Media, 2013.
[7] Ran Yaniv and Yair Wiener. On the foundations of noise-free selective classification. Journal of Machine Learning Research, 1153):1605-1641, 2010.
[8] Pierre Foret, Ariel Kleiner, Hossein Mobahi. and Behnam Neyshabur. Sharpness-aware min- imization for efficiently improving generalization. In International Conference on Learning Representations, 2020.
[9] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Puttern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 770-778. IEEE Computer Society, 2016. doi: 10. PR910CVR..2206.990
[10] Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation, 9:1-42, 1997.
[11] Izmailov, AG Wilson, D Podoprikhin, Vetrov, and T Garipov. Averaging weights leads to wider optima and better generalization. In 34th Conference on Uncertainty in Artificial Intelligence 2018 UAI 2018, pages 876-885, 2018
[12] Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Peter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2016.
[13] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
[14] Leo Klarner, Tim G. J. Rudner, Michael Reutlinger, Torsten Schindler, Garrett M. Morris, Char- lotte Deane, and Yee Whye Teh. Drug Discovery under Covariate Shift with Domain-Informed Prior Distributions over Functions. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2023.
[15] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[16] Hao Li, Zheng Xu, Gavin Taylor, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Neural Information Processing Systems, 2017.
[17] Hao Li, Zheng Xu, Gavin Taylor, Christoph Studer, and Tom Goldstein. Visualizing the loss landscape of neural nets. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018

---

[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 026201 8020.
[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 quantum field theory. CRC press, 2018.
[22] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do cifar- 10 classifiers generalize to cifar-107 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. 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, Mare' Aurelio Ranzato, Raia Hadsell, Maria- Florina Balcan, and Hswan- Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12 12, 2020, virtual, 2020.
[30] David H. Wolpert. Bayesian backpropagation over i-o functions rather than weights. In Cowan, G. Tesauro, and Alspector, editors, Advances in Neural Information Processing Systems, volume 6. Morgan-Kauffann 1993.
[31] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnin novel image dataset for benchmarking machine learning algorithms, 2017.
12

---

Appendix
Table of Contents
A
PS-MAP is not Invariant to Reparametrization
14
Mathematical Details & Derivations
15
B.I 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 Non-Singular [(0:px)
16
B.6 Architectural Symmetries Imply Singular J (0:
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
13

---

Appendix A PS-MAP is not Invariant to Reparametrization
Parameter-space MAP estimation has conceptual and theoretical shortcomings stemming from the lack of reparameterization-invee
of MAP estimation. Suppose we reparameterize the model parameters 0 with 6' = (0) where R is an invertible transformation. The prior density on the reparameterized parameters is obtained from the change-of-variable formula that states y/ 3/(0) = p(0) Jdet (d0' /do)l. In the language of differential geometry, the prior density p(0) is not scalar but 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 different coordinate system, but the MAP objective is not invariant to reparameterization. Specifically, the reparameterized parameter-space MAP objective becomes
C (0') Σ
(8) log pluD T 0') + logz'(0') = L (0) log |det (do' /d0)
(A.1)
new term
Since ["MAR(6) (0") and the original objective (0) differ by term which is non-constant in the pa- rameters if I' is non-linear, the maxima @ = arg maxg. (CMAF(0') and @MAP - arg maxo CMAP(O) under the two objectives will different. Importantly, by "different" we don't just mean @TMAP # gMAP but ( * R R(MM))). That is, they are not simply each other viewed in different coordinate sys- tem but actually represent different functions. More accurately, therefore, one should say MAP estimation is not equivariant (rather than invariant) under reparameterization. As 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 model so that PS-MAP will return an arbitrary solution On irrespective of the observed data, by choosing a reparameterization 8' R(() such that log |det (det- (df'/d0) /d0)lo=0o = -Cc. One such choice is 6' = (# 0u)², where the exponent taken element-wise.
As less extreme example, consider the reparameterizzatio 0' 1/00. A Gaussian prior on the original
exp(-0² (0pp-) /2) parameters p(0) x exp(-0²/2) translates to a prior p' (00) exp(-1/20')) /0²2 = 6
on the inverted parameters. Note the transformed prior, when mapped onto the original parameters 0, assigns a higher weight on non-zero values of 0 due to the quadratic factor coming from the log determinant. We illustrate this effect with simple experiment where linear model with RBF features is trained with a Gaussian likelihood and Gaussian prior N(O,, 7) using PS-MAP to fit noisy observations from a simple function. Figure 8 show 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, favors non-zero weights the inverted one, learning a less smooth function composed of a larger number of RBF bases.
PS-NAP Reparam.
Hepar
Thue function Deservaliins
Count
0.0
0.2
0.4
0.8
1.0
2.0
1.5
1.0
0.5
0.0
0.5
1.0
1.5
2.0
(a) Learned functions
(b) Parameter distributions
Figure 8: PS-MAP is not invariant to reparameterizatiin. 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.
14

---

Appendix B Mathematical Details & Derivations
B.1
Prior Distributions over Functions
Since the function evaluations fo( are related to the parameters through the map T: 0 + (a), the prior density p(f((ā)) is simply the push forward of the density p(0) through T, which can be expressed without loss of generality using the Dirac delta function as
p(f(()) - =/e-
-
(B.2)
This generic expression 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 7 is not guaranteed be surjective (onto), especially when there
evaluation points than parameters (l+) > P). p(f) 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 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 present more general result for which the result in Section 3.2 is special case, without assuming the function output is univariate (K = 1) or that the metric is diagonal. To simplify notations, suppress all 6-dependence and write f (arz), the k-th output of the function wi, repeated indices
as f for some å {r;; ) /= We use Einstein notation imply summation and leave everything in un-vectorized forms. In this more general case, the metric is given by ds² df dff, where i.j input indices and k, are output indices. Denote the un-vectorized M-by-K --y--- Jacobian Jika = f.t The P-by-P matrix J appearing inside the determinant in Equation (6) is now given by
Job JikagikseJ = f
(B.3)
In the limit of infinitely many evaluation points uniformly populating the domain X, the sum i becomes a double integral?
{ (z) (nel(, x') dr
Jab
(B.4)
where we have re-ttitten the metric Jskje as gkz( 2.)) where g X x X + R is the metric represented now as matrix- -aluux function. After dividing by a constant Z, we can similarly write Tab the expectation
Jab = E Exxxico
(X)0s, (X))sg((cco
(B.5)
where CCC (I,, x', , (,,,,,) = (x, x' 1/2 for some normalization constant that only shifts the log determinant by f-independent constant. Note Equation (10) is recovered if the metric is and z = x'. We can further remove the assumption
diagonal, that is 9e.c' (z, x²) # only if that g is constant (independent of f), but at the cost of letting the normalization constant Z depend on f (or equivalently 0), making invalid to ignore Z during optimization. While more general than the result presented in Section 3.2, specifying a (possibly 0-dependent) non-diagonal metric adds significant amount of complexity without being particularly well-motivated. Therefore we did not further investigate this more general objective.
"To be more accurate, an integral of the form f dxx) da is the limit of M o(x,) (.,))An Ar with the extra factor Ar æ+41 2, (Vi), as M- + ec. However, accounting for this factor does not affect the final result in our case because again Ar is independent of 0.
15

---

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 fo. However, the subtlety arises from the fact that p(fe) does not have support everywhere but is, in fact, surface density defined on the image of I (as defined in Appendix B.1), curved manifold in function space which gives rise to locally preferred directions such that even a global linear transformation will change the density in non-homogeneous way. I were instead surjective, then in Section 3.2 Jo (a) would be a square matrix and we could write log det (J)( (J) = log det (ā) Jo - log (g) and conclude that constant metric indeed has effect the MAP estimate. Similarly, if the image of T is not curved, meaning Je (2) is constant in 0, then log det det (J)(B)) gJo( would be constant in 9 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 è approaches infinity and the points cover the domain R" 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) 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 mesh, where the solution often converges to the continuous solution the mesh size goes to zero.
In similar vein, we utilize this technique to sidestep direct engagement with the ill-defined quantity p(fe). It may be possible to assign well- value to p(fe) 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 p(:
Proof. Recall px) Epx [Js(X) (00X))Jo(XX)) Jx( Suppose [(::Px) is singular. Then there is a vector v * 0 for which 0 = v J(0: = Epx [X) 2 = fx (xx = 0. Since the integrand is nen-negative and continuous by assumption, it is zero everywhere. Therefore, we have Σ v;dp, fe(x)) for all x E supp(px) for a non-zero v, showing fo are linearly dependent functions over the support of px-
Conversely, if Ei So(x) = Jo(x)u 0 for x E supp(px) for some e + 0. Hence J(0: px)e Epx [Jo(X) (o(X)))((Xou) = 0, showing J (0:px)) is singular.
are linearly dependent functions over the support of px, then
B.6 Architectural Symmetries Imply Singular [(0:px)
Proof. Differentiating with respect to 0, we have Jø -aff = Oefse S for all 0. Suppose
there exists 8 that is invariant under S, SS* = 0", then Jø- = Jo- S and Jn. (T -7--5 S) = 0. Since, by assumption S - T, we have shown that Jo- has non-trivial nullspace. Because the nullspace of J(0": px) contains the nullspace of Jo-, we conclude J(e: px) is also singular.
As an example, consider a single hidden layer MLP fo(x) = wi a (w, x)) 0 = (w) E R ERR uy E R2), where we neglected the bias for simplicity. The permutation P1a + P12 is symmetry of fe, where
(0
1 The nullspace of Jg- contains the image of
is the direct sum and Piz is the transposition
(1
( )
-1 1 +
(B.6)
with a basis ((1, -1,-110,,0 10,0, for any 0" of the form (a, a, b, b). It's easy to check that
perturbing the parameters in directions (1, -,, 0) and (0,0, 1, -1) leaves the function output unchanged, due to symmetry in the parameter values and the network topology.
16

---

Continuous symmetries, among other reasons, can also lead to singular J [(:; px) and are not covered by the above analysis, which is specific to discrete symmetries. One example is scaling symmetries in MLPs with ReL.U activations [5]. However, these scaling symmetries cause J (0; px) to be singular for all 0 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 ReI.U 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 J (0; px) with log (0; where |- I+ represents the pseudo-determinant This step is automatic if we optimize Equation (14) instead of Equation (11), where adding small jitter will automatically compute the log detuddddttemmnnnt (up to an additive constant) when J [(: (8: is singular. By contrast, symmetries such as permutation symmetries do create pathological optima because they only make the Jacobian singular at specific settings of 0, 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 decribe 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( (0) ö(fe - fo
We can avoid the singularities if we instead choose a variational family where each member is a dis- tribution over localized around fo for some O. Namely, we consider the family Q = {q( (o) beee parameterized by a mean-like parameter and a small but fixed entropy in function space H (g(For (0)) = Euifg [0) log g(for [0)] h for all 0 E O. For convenience, we overload the notation and use ((01)) to denote the pullback of q(fer (0) to parameter space. The resulting varia-
tional lower bound
(v)
log ((frr 10)
= -Ba(srr (8)
p(folD)
Eu(sorlo) - log log [D)] (()
log Ea(s (oloo/)))-oop((D))) (D))f()
+ Eg 1)| log (ferr }1 const.
(Bayes' rule)
Repeoo (olgg((((0)) + | log log ((0)) -
2 T (og pxx)) (Equation (11))
= =ologp(D(.a) - (0) + E
+ (hh) + const.,
where in the last line we used the assumption that (0) is localized around 0 write the ex- pectation of log p(D) (DD|0 ()) and l(g)(0) as their values at 0' = 0 plus O(h) corrections (which we will henceforth omit since by assumption h is tiny), because they vary smoothly as a function of 0. More care is required to deal with the remaining expectation, since when J [0:px is singular at 0, log de (O; px) is infinite, but its expectation, (ggoor0) 00) [log detdde(([[] (0'; :x ))}, must be finite (assuming p(0) and p (D10) are finite). given that Lvi.n lower bounds the log marginal likeli- hood log P(D) = log ʃ p(D)0))p(OdO < co. 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 on the eigenvalue of I [(:; p.).. which can be approximated by adding a jitter e when computing the log determinant. This effect is similar to applying a high-frequency cutoff an image by convolving it with a localized Gaussian filter.
With this approximation, arrive at simple objective inspired by such a variational perspective, equivalent to adding a jitter inside the log determinant computation in the FS-MAP objective,
[(0:xx) = Σ, log 1
(8) (0) I 0) + p(0) log (T + +e))
(B.7)
Note that while using a jitter is common in practice for numerical stability, it arises for an entirely different reason here.
17

---

B.8
Approximating the Expectation
Rewriting log (T (xx;))) in terms the eigenvalues Ao(0"), we have
Becoro0 [ :xx]]] - P
(B.8)
Consider each term (0)1)) y) [log A (6') in isolation. Since the distribution ((11)) is highly localized by assumption, there are only two regimes. If (0) * O, then the variation of the log eigenvalue is small and the expectation is well approximated by log (0). Otherwise, if (, = O, then log (o) -00 is no longer valid
log (0)) has singularity 0' 0 and (Eg(0110) /) [log A, (0)) ] approximation since we' ve shown the expectation is finite. Therefore, in this case, the expectation must evaluate to large but finite quantity, say log(c) for some small c. where the exact value of с depends on the details of ((('0)) and (f"). Consequently, taking the expectation can be approximated by adding jitter cl to [60'pp when computing the log determinant to clip its eigenvalues from below at some threshold > 0.
B.9 Monte Carlo Estimator for the Log Determinant
Monte Carlo Estimator. Exactly computing px) Epx [Jo(X)T [Jx())TT{(X)) and its log deter- minant is often intractable. Instead, we can use a simple Monte Carlo estimator to approximate the expectation inside the log determinant:
log det (Epx [Jo(X)To(XX))) = log det
(((60))To(((0))) + +el). E
S
(B.9)
3
where we add a small amount of jitter prevent the matrix inside the determinant from be- coming singular when SK < Since 1/2 5 Jo(tli)) Ja (z((x) ( ( where is {2(1){ and JØ(és) is an SK-by-P matrix,
compute the determinant of s Jo(()s)_ + el through the product of squared singular values of Jo( in Omin(SK P2) time (with zero singular values replaced with /t), much faster than ((p)) for computing the original P-by-P determinant if S is small enough, without ever storing the matrix J(e:
Accuracy of the Estimator. Due to the nonlinearity in taking the log determinant, the estimator is not unbiased. To test its accuracy, we evaluate the estimator with S (800, 400, 200} for neural network with two inputs, two outputs, four hidden layers, and 898 parameters, using tanh activations. The network is chosen small enough so that computing the exact log determinant is feasible. Figure we compare exact log determinant for px = Nे Σ. 55; with M = 1, 600, and its estimate with S Monte Carlo samples. Here (x;) are linearly spaced in the region [-5, 5)2. We observe that using S 800 can almost perfectly approximate the exact log determinant. With = 400 and S = 200, the approximation underestimates the exact value because the number of zero singular values increases for Jo T Jo (Øs), but it still maintains a strong correlation and monotonic relation with the exact value. We see that it is possible to approximate the log dog-determinant with sufficient accuracy with simple Monte Carlo using fraction of the compute and memory for exact evaluation, though there can be a larger discrepancy as we continue to decrease S.
B.10 Laplacian Regularized MAP Objective
For that large enough compared to the eigenvalues of ((::)))): log det(T((; + el) can be sufficiently accurate. Expanding to first order in p= max, X//e, have
)): a first order approximation to
= log (et(de)) + = N x+
1 = - + O(p2), +c(p)) (B.10)
where c. = Plog(c) is independent of 0. Defining d(0,6') = (1ff((X - (x)(1) we have
d(e, 0 + x+) Epx [|| (X) 0 = T 1 O
(B.11)
18

---

05
X12
4124
Estimated lop
-40
Exact det
Exact log det
Exact
(a) 800
(h) 400
(c) 200
Figure 9: Approximating the log determinant via an S-sample simple Monte Carlo estimator and
jitter. Each dot represents the exact and estimated log determinant evaluated random parameters sampled by randomly initializing the network followed by scaling by factor ~ loguniform 10) include parameters of different magnitudes. The dashed line shows x = y.
which shows [(0:px)) =
30 1:-00 Therefore,
Tr(J(0;px)) d( A0,0000 +b)) m=)0
(B.12)
where A is the Laplacian operator.
To estimate the Laplacian, consider the following expectation:
Eu [d(0, +b)]
(B.13)
Ep [v J [(; +
(B.14)
=E 36 [(::px)) 0(*))
(B.15)
= T((:x))) + O(8))
(B.16)
where we used
E J(0:px))
E [1,10 8²
(B.17)
Therefore, we have
Tr([(O:px)) 1 A d(0,0 +1)1 za 0o.. +)] 0
(B.18)
Combining Equation (0. and Equation (B.18), have shown Equation (14) reduces to the more efficiently computable L-MAP objective for large enough e and small enough 3
LL-MAP CL:MAr(0:p2)) (0:px) Σ. log = fo) + log 2 Es |d(,, v)]. (B.19)
The entire objective is negated and divided by N to yield the loss function
LE-Mar (e:MA)(::px)=
1 N Σ log poyp 2 fo) logp(!) +A ( E.t Ex.- +)
(i)
(1)
Standard regularized loss
Laplacian regularization R(0:9)
(B.20)
2xN Therefore, using L-MAP amounts to simply adding regularization XR(0; (32) to standard regularized training (PS-MAP), is a tolerance parameter for approximating the Laplacian and can simply be fixed to small value such as 10-3
where we have absorbed the 1 1/N factor into the hyperparameter = 1
19

---

Appendix C Further Empirical Results and Experimental Details
C.1
Training details for Synthetic Experiments
For both PS-MAP and FS-MAP, we train with the Adam [13] optimizer with a learning rate 0.1 for 2,500 steps to maximize the respective log posteriors. For FS-MAP, precompute the t matrix and reuse it throughout training. When px is not Uniform (),,1), Φ does not have a simple form and we use 10,000 Monte Carlo samples evaluate it.
C.2 Hyperparameters for UCT Regression
We use an MLP with hidden layers, 256 units, and ReLU activations. We train it with the Adam optimizer for 10,000 steps with learning rate of 10-3. For each dataset, we tune hyperparameters based on validation RMSE, where the validation set is constructed by holding out 10% of training data. We tune both the weight decay (corresponding to prior precision) and L-MAP'S À over the choices {10 5 1004, 4, 10 - 10 10-1). We then report the mean and standard error of test RMSE across 6 runs using the best hyperparameters selected this way. In each dataset, the inputs and outputs are standardized based on training statistics.
C.3 Hyperparameters for Image Classification Experiments
Both the weight decay scale, corresponding to the variance of Gaussian prior, and the L-MAP hyperparameter À in Equation (B.20) is tuned over the range [10 10010-11 using randomized grid search. For PS-MAP, À is set to O. The parameter variance in the Laplacian estimator is fixed to 33 = 10-6, In addition, we clip the gradients to unit norm. We use learning rate of 0.1 with SGD and cosine decay schedule over 50 epochs for MNSH and 200 epochs for CIFAR-10. The mini-batch size is fixed 128.
C.4 Verifying the L.-MAP Approximation
To verify that the I.-MAP objective is good approximation to Equation B.7 when is large enough compared to the eigenvalues of [0:;pu px), we compare log det( px) + cT) Plog(() with the Laplacian estimate Ze37 Eo [d( + v)]] used by L-MAP in Appendix B. as its first order
approximation in max; > (0)/ee. We reuse the same network architecture, evaluation distribution Px. and sampling procedure for the parameters from Appendix B.9 to perform the comparison in Figure 10 and color each point by X(0)/e, the average eigenvalue of J [(; px) divided by The smaller this value, the better the approximation should be. We observe that indeed as e increases and X(0)/e approaches O, the Laplacian estimate becomes more accurate approximation of the log determinant. Interestingly, even when c < A()) and the approximation is not accurate, there still appears to be monotonic relation between the estimate and the log determinant, suggesting that the Laplacian estimate will continue to produce a qualitatively similar regularization effect as the log determinant in that regime.
X10
Japlacan estimate
sasconnn cstimane 59033
12
l0g det
det litter
lgg del with jiter
1og det with Jitee
(a) 6 10-3
cb)==1
(c)==
(d) 6 10
Figure 10: Log determinant v.s. the Laplacian estimate different values c. The r-axis shows log det(( (T(: px) + +c) - Plog(()) while the y-axis shows its first order approximation Zesa BM (0,6 v)], estimated with 10 samples of vb
C.5 Effective Figenvalue Regularization using the Laplacian Regularizer
20

---

The Laplacian regularizer in Equation (B.19) only uses sample- based estimator that is unbiased only in the limit that + 0. keep the computational overhead at minimam so that M. can scale to large neural networks, we only take sample of per gradient step. To test the effectiveness of I.-AP in regularizing the eigenvalues of J (0; px) under this practical setting, we train an MLP with hidden layers, 16 units, and tanh activations on the Two Moons dataset (gen- erated with sklearn.. sklearn.datasets wwwemmaooeaaaeeeaaamaamma mmons (n_ sasssamplese es=200, shuffle-Trun
PS.MAP
PS MAP
Figure 11: L-MAP is effective at minimizing the cigenvalues.
noise=0.2 2, sannommstate--eee
for 10* steps
with the Adam optimizer and a learning rate of 0.001. Here we choose px M Ein dz, with M = 1, 600 and { linearly spaced in the region [-5, 5] In Fig- 11, we compare the of eigenvalues of J px; for L-MAP and FS-MAP with different levels of jitter e. Both FS-MAP and L-MAP significantly reduce the eigenvalues compared to PS-MAP with small values of c, corresponding to stronger regularization.
C.6 Visualizing the Effect of Laplacian Regualrization
The hyperparameter is inversely related to the strength of the Laplacian regularization, We visualize the effect of in Figure 12, showing the L-MAP solution varies smoothly with evolving from near-zero function that severely underfits the data to one that fits the data perfectly.
0.5032
0.92
0.5021
0.80
0..011
0.68
0.5008
0,5000
00.36
-0.9992
0.44
-2
0.4984
0.32
0.4976
0.4968
0.20
0.4960
0.08
(a) c 10
(b) c = 10
4
00900
4
0.90
0.75
0.75
0.60 bd
0.60
0.45
0.45
-2
0..00
2
0.30
0.15
-..15
0.00
0.00
-4
-2
4
(c) =
(d) 10
Figure 12: Visualization of L-MAP solustions at various e.
21

---

C.7 Further Neural Network Experiments
In Table 4, we report the detailed results for varying the number of Monte Carlo samples as reported by Section 4.5 and Figure 7. As noted previously, the performance is fairly robust to this hyperparameter S in terms of accuracy, selective accuracy, and negative log- likkeiihoodd. However, we find that increasing S can sometimes lead to significant improvement in calibration, as seen for MNSST [31].
Table 4: This table provides detailed quantitative performance for the results plotted Figure 7, ablating the choice of the number of samples used for evaluation of the Laplacian estimator S. The standard deviations are reported over five trials.
FASHIONMNIST
CIFAR-
# SAMPLES (S)
Acc.
SEL. PRED. 1 NLL + ECE +
Acc. 1
SEL. PRED. t NLL +
ECE +
94.0%
99.2%1..0
0.28 4.3% 10.6
99.5%..0.
0.17:0.01 20.
94.0%
99.2%1100
0.27=0.01 4.2% 10.6%
99.5%110.
0.17:2.000 10.
16
14.1%
99.2%.10.
0.27 4.1%%to 95.6%.0.2
99.5%.10.0
0.17:1.000 00 2.1
32
93.8%+o.1
99.2%+0.o
0.28=0.01 2. 95.5%+o.1
99.5%+...
0.16#..00 1.8%+0.2
91.0%%+..1
99.2%+4.o
0.27 1..% 95.7%+b.2
99.5%+..
0.164a.00
1.7%+0.2
128
91.1%+o.1
99.2%++.1
0.26=#0.1 4.1 95.5%+o.1
99.5%+0.0
0.16# 1.4
256
94.0%
99.2%
0.27 4.. 95.5%
99.5%610.
0.16110.0
1.2%
512
93.9%
99.1%11..
0.26 4.2 95.5% 99.5%
0.16 1..1
Distribution Shift. We additionally evaluate our L-MAP trained models assess their performance under covariate shift. Specifically, for models trained on CIFAR-10 [15], we use the additional test set from CIFAR-10.1 [22] which contains additional images collected after the original data, mimicing a covariate shift. As reported Table L-MAP tends to improve calibration, while retaining performance in terms of accuracy.
Table 5: We evaluate the performance of CIFAR-10. 10.1 [22] using the models trained on CIFAR-I 10 [15]. L-MAP tends to improve the data fit in terms of the log likelihood and is better calibrated while retaining the same performance as PS- MAP. The standard deviations are reported over five trials.
METHOD
Acc. f
SEL. PRED.
NLL +
ECE 4
PS-MAP
89. 20.44%
97.8%++o.2
0.43=2.03
6.2%+4..
L-MAP
89.2%=u.7
97.9%+o.z
0.40=a.02
4.1%
Transfer Learning. Using ResNet-50 trained on ImageNet, we the effectiveness of L-MAP with transfer learning. We choose hyperparameters as in Appendix C.3, except lower learning rate of 10 3 and a smaller batch size of 64 due computational constraints. Table reports all the results.
22


Reviews:
[{"review_id": "XI9VhP7txe", "reviewer": "Reviewer_NXLr", "summary": "This paper explores the question of learning most likely functions vs most likely parameters in the context of MAP estimation, a common setting including minimizing log likelihood with L1 or L2 regularization, as these regularizations correspond to imposing a prior. Differences in parameterization affect the MAP through change of measure for the prior, and potentially in very unhelpful ways. Because of these potential problems with parameterizations, authors explore doing MAP not on parameters as we usually do (referred to as PS-MAP), but on function outputs for a set of inputs, an idea from Wolpert (1993). Authors refer to this as function-space MAP (FS-MAP). \n\nThe contributions are as follows:\n- While Wolpert considered function outputs on a finite set of inputs, authors generalize to function outputs on a measure on input space. Authors emphasize that this measure of input space can be specified according to the task at hand (e.g. restricting to natural images, rather than all images). \n- Authors empirically compare FS-MAP and PS-MAP on a simple example and empirically find that FS-MAP finds flatter minima, has better generalization in some settings, and produces less overfitting. Authors also discuss how FS-MAP and PS-MAP compare with the posterior mean in variations on this simple setting. Also, FS-MAP’s test performance depends on the measure of inputs that are considered important. \n- Authors discuss problems with FS-MAP: FS-MAP can be computationally very expensive, and can behave in pathological ways under certain conditions. Authors also provide conditions under which it is well-behaved, as well as a scalable approximation to FS-MAP, called L-MAP, that also behaves ok even when FS-MAP would behave pathologically. Authors empirically compare FS-MAP and PS-MAP on neural networks for image classification. L-MAP performs similarly to PS-MAP, but is better calibrated in one experimental setting, and also finds flatter minima. ", "questions": "1. On line 135-136, “the distance between two function evaluations is measured locally by [...]” : what two function evaluations are being referred to?\n2. Why is it important to have infinitely many evaluation points? I assume that’s not the appeal but that’s how it was advertised. (Lines 127-128)\n3. What is p_X=N/A in Figure 3b?\n4. Can you provide intuition for why FS-MAP archives better generalization, in the settings where it does achieve better generalization? Can you also provide intuition for why it is less prone to overfitting? Can you discuss when this wouldn’t be the case? (Section 3.3)\n5. The numbers in Table 1 are all very similar. What do you think is the reason for this? Should we be surprised by the similarities between the different choices of p_X, given how much p_X mattered in Section 3.3?\n6. “For example, while it is 331 natural to assume that FS-MAP more closely approximates a Bayesian model average”: Why is it natural to make this assumption? (lines 330-331)\n7. Could you briefly discuss why the posterior mean is particularly important (line 254)? The citation that follows is a textbook. \n8. Why do we want to constrain eigenvalues of the Jacobian? I understand that you can run into problems if they are quite large. Is that the only reason? Is there a risk of over-constraining? \n9. Why is calibration better with FS-MAP for CIFAR-10 and p_X=CIFAR-100 compared to PS-MAP, but not other settings? \n10. Can you explain further why L-MAP is more robust to training with label noise? The explanation in line 325 is Section 3.3, which appears to be about input space, not about label noise\n11. Should I expect these comparisons between FS-MAP and PS-MAP in other settings? In what settings should these results hold vs not (beyond satisfying the assumptions)? \n12. While Wolpert’s evaluation on finite sets of points is a special case of the authors’ formulation, are they not effectively the same in practical settings? \n13. Authors discuss a measure on which to consider function outputs. I am curious about the connections to distribution shift. ", "limitations": "The experiments and discussion around experiments could be more thorough. See questions. It's not clear to me which experimental results should generalize, and under what conditions. ", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Authors raise and discuss an interesting question that I had not thought about. \n2. There is a substantive discussion of PS-MAP and FS-MAP, including weaknesses of both approaches. \n3. Authors provide a computationally feasible, well-behaved approximation to a computationally too-expensive and poorly-behaved objective (L-MAP, vs FS-MAP). ", "weaknesses": "1. There is a tension between statements like “parameters have no meaning independent of the functional form of the models they parameterize” (lines 22-23) and the fact that the difference between best parameters vs best functions is a matter of parameterization of the prior, and the prior is a prior on the parameters (Equation 1). This tension is central to the question of parameterization for MAP and is not addressed. \n2. This tension continues into the experiments, many of which emphasize flat minima. Authors cite Dinh’s “Sharp minima can generalize for deep nets”, which points out that the flatness / sharpness of minima is a function of parameterization. \n3. Function-space MAP and priors on functions are defined much later in the paper than ideal. For example, in Figure 1, I do not know how FS-MAP is calculated. p_X (important to FS-MAP) is also not specified. \n4. The terminology of “most likely functions” and “more probable functions” obscures the dependence of these terms on the choice of p_X (e.g. line 37, 184). \n5. The connection to variational inference does not seem particularly strong. MAP as VI would involve choosing the best q, while the construction in Section B.9 does not do that. It’s also unclear to me why we care about these constructions; e.g. why do we care about q, besides that it ultimately leads to (13) and (B.13)?\n6. Also, (B.8) looks like a KL divergence, not a variational lower bound / ELBO. Maybe I’m missing something here. \n7. The experiments and discussion around experiments could be more thorough. See questions. \n8. If someone told me they were looking at most likely functions, rather than most likely parameters for MAP, I would be curious about priors on functions in place of priors on parameters. (This isn't a weakness / question / limitation; it's more of a comment.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We're glad you appreciated the paper and found it well-written. We address your comments below.\n\nPlease let us know if you have any remaining questions.\n\n1. **When and why FS-MAP can outperform PS-MAP and reduce over-fitting**\nThank you for asking these important questions at the heart of this paper. While in general there is little reason to expect either FS-MAP or PS-MAP will perform better, we provide a thorough discussion on important special cases in the general response that explains our empirical findings. In terms of reducing overfitting, a simple reason why FS-MAP can reduce overfitting is the introduction of the additional regularization term which can prevent the model from fitting the data. In addition, fitting noise often requires precise settings of certain parameters to account for unnatural noisy patterns, which would incur a large Jacobian determinant penalty.\n\n2. **Infinitely many evaluation points**\nWhen using only a finite set of evaluation points $\\hat{x}$ as in Section 3.1, one can only find the MAP estimate for the function evaluated at a finite set of points $f\\_\\theta(\\hat{x})$, which is different from the MAP estimate for the function $f\\_\\theta(\\cdot),$ equivalent to an infinite dimensional vector. As our goal is to find the MAP estimate for the whole function $f\\_\\theta(\\cdot),$ we seek to use an infinite set of evaluation points. In Figure 3a, we also empirically show that FS-MAP performs monotonically better as we use more evaluation points and achieves optimal performance with an infinite set.\n\n3. **Distinction from Wolpert's formulation**\nOur formulation is different from Wolpert's formulation even when using finite samples, since a random finite set of points is sampled at every gradient step and the underlying $p\\_X$ can have support on infinitely many points. \n\n4. **Benchmark Performance**: \nOur prior in FashionMNIST and CIFAR-10 experiments was non-informative, as a result we do not expect FS-MAP / L-MAP to outperform PS-MAP though it did improve calibration in many cases. In contrast, the improvements in Section 3.3 stemmed from using a prior matching the data-generating process. We additionally evaluate L-MAP on UCI dataset and transfer learning on CIFAR-10 and present the results in the the general response.\n\n5. **Dependence of mostly likely functions on $p_X$**\nWe agree that this is an important subtlety. Indeed, revealing this subtlety and establishing the connection to choosing a metric in function space is a key contribution of this work and we will emphasize it more. We additionally note that such dependence on an underlying metric applies to any MAP estimation problems in continuous spaces, including PS-MAP, as the probability density always depends on a choice of base measure.\n\n5. **FS-MAP approximating BMA** \nIt is natural to assume FS-MAP approximates BMA better due to arguments we presented in line 259-263. Namely, with sufficient amount of data, both $p(\\theta | \\mathcal{D})$ and $p(f_\\theta | \\mathcal{D})$ are approximately Gaussian. Given Gaussian posteriors, the BMA function coincides with its posterior mode $f_{\\theta}$, i.e., FS-MAP. But since PS-MAP seeks the posterior mode of $\\theta$ rather than $f_\\theta,$ it learns a different function compared to FS-MAP and thus BMA.\n\n6. **Posterior mean importance**: \nThe posterior predictive mean function combines all possible model hypotheses weighted by their posterior likelihood, rather than a single point estimate. It often achieves better generalization and uncertainty estimates as a result [1].\n\n7. **On model parameterization**\nWe emphasize that given any prior $p(\\theta)$ and parameterization $f_\\theta$, PS-MAP will learn different functions upon reparameterization as we show in Appendix A, whereas FS-MAP will always learn the same function. So while the precise difference between PS-MAP and FS-MAP depends on PS-MAP's parameterization, the existence of this difference is universal. Similarly, while flatness metrics such as Hessian eigenvalues are not parameterization-invariant, FS-MAP's preference for flat minima is universal when compared with PS-MAP as evident from Equation 11. As such, we believe there is no tension between these observations but would appreciate further clarifications.\n\n8. **On the L-MAP derivation**\nThank you for catching our typo. Eqn. (B.8) is indeed a **negative** evidence lower bound. Regarding motivation for considering variational inference (VI), unlike in conventional VI, the goal of the formulation described in Appendix B.9 is not primarily to learn an accurate approximation to the posterior but to obtain a non-pathological objective in a principled way. In addition, we do choose the best $q$ within the variational family defined in Appendix B.9 through learning the optimal mean parameter $\\theta$.\n\n9. **Numerical Complexity**: \nL-MAP requires an additional forward pass on $ S $ evaluation points. It's faster than FS-MAP and only slightly slower than PS-MAP. For a Two Moons dataset experiment with an MLP, L-MAP was 33 times faster than FS-MAP and only 1.4 times slower than PS-MAP. In scalable experiments, L-MAP takes roughly twice the time of PS-MAP.\n\n| Dataset | Method | Gradient Step (ms) | Epoch (s) |\n| ------- | -------- | -------- | -------- |\n| FashionMNIST | PS-MAP     |  54     |  24     |\n|  | L-MAP     |  109     |  45    |\n| CIFAR-10 | PS-MAP     |  64     |  27     |\n|  | L-MAP     |  126     |  52    |\n\n11. **Calibration**: \nWe believe L-MAP's calibration improvement on CIFAR-10 with $ p_X = $ CIFAR-100 stem from the evaluation distribution covering relevant regions of the input space without duplicating the training set.\n\n12. **Notation**: \nWe'll clarify unclear notations. Specifically, $ p_X = N/A $ corresponds to PS-MAP in Figure 3b and in line 135-136 the the functions are $f$ and $f+df$ where $df$ is an infinitesimal separation.\n\n\n\n[1] Bayesian deep learning and a probabilistic perspective of generalization, Wilson et al., NeurIPS 2020\n"}, {"review_id": "oXhRB3szUk", "reviewer": "Reviewer_VTcx", "summary": "The paper compares the widely-used parameter estimation of machine learning against function estimation for the maximum a posterior (MAP) estimation. The authors provide detailed analysis and mathematical variation why there are significant difference in results for function vs. paramter MAP, and introduce conditions where the function MAP can avoid its failure modes. \n\nTo be honest, I am not an expert in the area, and cannot make an assessment how novel the proposed method is against relevant literature. I would rely on other reviewers on the aspect of the evaluation.", "questions": "- This might be simply my ignorance. But what is the difference between N samples ($x_\\mathcal{D}^{(i)}$) vs. M samples ($\\hat{x}$) used in Equation 5? Is it necessary to derive different samples from the space, or can we use the same samples? Roughly how big number would you choose for N and M, respectively? The paper writes condition based on M ($MK \\geq P$). But is there any condition on N?", "limitations": "The authors clearly stated limitations in extensive aspects. However, it would be informative to have a detailed numbers on the increased numerical complexity or time/ memory requirement for using FS-MAP (or LS-MAP) over PS-MAP for some of famous architecture.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is fairly well-written in most of cases, and grounds strong motivation and intuition for the proposed idea. Even though I am not expert in the area, I enjoyed reading through the paper.\n\n", "weaknesses": "- The paper is heavy on math. Even though they authors did a decent job to make the equations approachable, some of important definitions are missing, and many of details are deferred to the supplementary material. It is suggested to double-check the completeness of the formulation in the main paper. If it cannot be self-contained, I am afraid the publication should head to another venue with a longer length limit (maybe a journal). For example, P was never introduced, but could be only inferred as the parameter dimension. Also, $\\theta_R$ and $\\theta_L$ was not clearly defined in Figure 1, which is not desirable considering that it should serve as a strong motivating example. \n\n- The results are not very competitive. Especially for Table 1 for MNIST and CIFAR-10. Is there any intuitions why it is the case? On the other hand, graphs in Figures appear quite promissing. I would appreciate more discussion in the paragraph on benchmark performance (line 297-298).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your thoughtful and constructive questions and suggestions!\n\nWe were pleased that you **\"enjoyed reading the paper\"** and that you found it to be **\"well-written\"**.\n\nWe address your questions and comments below. Please let us know if you have any remaining questions.\n\n---\n\n> Even though they authors did a decent job to make the equations approachable, some of important definitions are missing. [...] For example, P was never introduced, but could be only inferred as the parameter dimension. Also, $\\theta_R$ and $\\theta_L$ was not clearly defined in Figure 1, which is not desirable considering that it should serve as a strong motivating example.\n\nThank you for suggesting clearer definitions of notations. We will clarify that $P$ is the number of parameters and $\\theta \\in \\mathbb{R}^P$ in the preliminaries. We will also clarify that $\\theta\\_L$ and $\\theta\\_R$ refers to the parameters controlling the height of the left and right Gaussian, respectively. If you think it would be helpful, we will include a list of definitions in the appendix.\n\n\n---\n\n\n> The results are not very competitive. Especially for Table 1 for MNIST and CIFAR-10. Is there any intuitions why it is the case? On the other hand, graphs in Figures appear quite promising. I would appreciate more discussion in the paragraph on benchmark performance (line 297-298).\n\nThank you for pointing out that further discussion of the results would be useful. We included such a discussion in our general response and will add it to the manuscript. As discussed in our general response, we were not expecting FS-MAP or L-MAP to lead to significant improvements in these settings where we do not have a highly informative prior. In contrast, we use a prior that exactly corresponds to the data-generating process in Section 3.3, leading to significant gains by performing FS-MAP.  Here, since our Gaussian prior is rather non-informative about the data-generating process in the FashionMNIST and CIFAR-10 experiments, we were pleasantly surprised to find that L-MAP achieves significant reduction in ECE on CIFAR-10 with $p\\_X=$CIFAR-100 and noticeable improvement in selective prediction accuracy on FashionMNIST with all three choices for $p\\_X$, though the accuracy is comparable with PS-MAP.\n\n\n\n---\n\n\n> This might be simply my ignorance. But what is the difference between $\\mathrm{N}$ samples $(x_{\\mathcal{D}}^{(\\imath)})$ vs. $M$ samples $(\\hat{x})$ used in Equation 5? Is it necessary to derive different samples from the space, or can we use the same samples? Roughly how big number would you choose for $N$ and $M$, respectively? The paper writes conditions based on $M$ ($MK > P$). But is there any condition on N?\n\n$N$ is the number of points in the training set $x_{\\mathcal{D}}$ and there is no condition based on $N$. $M$ is the number of points in the set of evaluation points $\\hat{x}$ for computing the prior $p(f\\_{\\theta}(\\hat{x}))$. $\\hat{x}$ and $x_{\\mathcal{D}}$ ar generally not the same since $\\hat{x}$, being a special case of specifying an evaluation distribution $p_X,$ corresponds to a choice of metric in function space as we show in Section 3.2. Technically, using the same samples is possible, and we investigate setting $\\hat{x} = x\\_{D}$ in our experiments (see $p\\_X = Train$ in Table 1). More generally, however, we wish to set $M$ to be as large as possible to appropriately account for the behavior of the function throughout the entire input space, and Figure 3a shows that FS-MAP indeed can benefit from using larger $M.$ While it is generally intractable to use exceedingly large $M$ as discussed in Section 4.3, there we also present an efficient approximation via L-MAP that doesn't suffer from this intractability by allowing unbiased Monte-Carlo estimates for objectives defined with arbitrarily large $M$.\n\n\n---\n\n> The authors clearly stated limitations in extensive aspects. However, it would be informative to have a detailed numbers on the increased numerical complexity or time/ memory requirement for using FS-MAP (or LS-MAP) over PS-MAP for some of famous architecture.\n\nCompared to PS-MAP, L-MAP only requires an additional forward pass on $S$ evaluation points to evaluate the Laplacian regularization term, where $S$ is the number of Monte Carlo samples. This is much faster compared to the exact FS-MAP, as we discuss in Section 4.3 line 291-293.\n\nTo showcase concrete run times, we train an MLP for 10,000 steps on the Two Moons datsets with a batch size of $200$ and $1600$ evaluation points. The MLP has only 2 hidden layers and 16 units each so that FS-MAP is tractable. L-MAP, finishing in 31 seconds, is 33 times faster than FS-MAP, which took 16 minutes, and only 1.4 times slower than PS-MAP, which took 22 seconds. We also show the time per gradient step and peak memory consumption for ResNet-18:\n\n\nFor the scalable experiments with neural networks in Section 4.3, we report the approximate wall clock times between standard PS-MAP and our proposed approximation L-MAP in the table below. All experiments were run with batch size $128$ and $S=128$ Monte Carlo samples for the evaluation points. As expected, L-MAP takes approximately twice the amount of time due to the additional forward pass. Note it is not feasible to run FS-MAP in this case as it requires too much memory.\n\n\n| Dataset | Method | Gradient Step (ms) | Epoch (s) |\n| ------- | -------- | -------- | -------- |\n| FashionMNIST | PS-MAP     |  54     |  24     |\n|  | L-MAP     |  109     |  45    |\n| CIFAR-10 | PS-MAP     |  64     |  27     |\n|  | L-MAP     |  126     |  52    |\n\nAll times are reported based on an NVIDIA TITAN RTX 24 GB GPU, where each run takes approximately 7GB of GPU memory in JAX.\n\n\n---\n\nPlease let us know if you have any further questions!\n"}, {"review_id": "B5TxYviqgP", "reviewer": "Reviewer_PZwV", "summary": "This paper analyzes the question: should we find the parameter that maximizes $p(\\theta|D)$ or $p(f_\\theta|D)$, given a data distribution $D$.\nThe former is the classic MAP, or PS-MAP (\"parameter space\") in this work, and the latter is called FS-MAP (\"function space\").\n\nThe paper shows that neither is universally more desirable than the other, and describes pros and cons of both methods:\n- FS-MAP has the advantage of directly optimizing the quantity we care about (i.e. function value), is invariant to reparameterization, and can be more robust to label noise.\n- On the other hand, PS-MAP tends to be more stable, and while it is not invariant to reparameterization, it's unclear how much this would be a disadvantage in practice. Moreover, PS-MAP can be closer to the Bayesian model average than FS-MAP.\n\nThe findings are supported by both theory and experiments.", "questions": "- Sec 3.3: the choice of prior is important; is there guidance on how to choose a prior in practice?\n- Is there guidance on how to choose between FS-MAP and PS-MAP given a practical problem? How much would the choice matter (the experiments seem to show that the choice doesn't matter much)?\n- Fig 2b: why is there an increase in curvature a the number of samples increases? \n    - also, is there a reason for measuring flatness as the average eigenvalue (rather than e.g. max)?\n", "limitations": "The paper clearly discusses the limitations.\n\nThere is no direct societal impact.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper has extensive comparisons on the limitations and applicability of both methods.\n- The paper addresses practical considerations:\n  - The problem of having a (near) singular Jacobian is addressed by adding perturbations.\n  - For scalability, the problem of requiring two forward steps is addressed by zeroth-order approximation, which requires 1 forward pass only.\n", "weaknesses": "While I appreciate the discussions, I'm not sure the evidences are clear enough.\n- Fig 2 shows that there's no difference between FS-MAP and PS-MAP when the number of samples is around 10k, a sample size that most practical settings can afford.\n- In Table 1, the results from both methods are very close, often within the margin of standard deviation.\n\nOverall, it seems that both methods compare similarly, and it's unclear how much we would lose by simply always choosing PS-MAP.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful and constructive questions and suggestions!\n\nWe were pleased that you highlighted the **\"extensive comparison\"** of limitations and applicability of PS-MAP and FS-MAP and that the paper addresses **\"practical considerations\"**.\n\nWe address your questions and comments below. Please let us know if you have any remaining questions.\n\n---\n\n> In Table 1, the results from both methods are very close, often within the margin of standard deviation.\n\nThank you for pointing out that further discussion of the results would be useful. Indeed, as we discussed in our general response, we were not expecting FS-MAP or L-MAP to lead to significant improvements in these settings where we do not have a highly informative prior. In contrast, we use a prior that exactly corresponds to the data-generating process in Section 3.3, leading to significant gains by performing FS-MAP.  Here, since our Gaussian prior is rather non-informative about the data-generating process in the FashionMNIST and CIFAR-10 experiments, we were pleasantly surprised to find that L-MAP achieves significant reduction in ECE on CIFAR-10 with $p\\_X=$CIFAR-100 and noticeable improvement in selective prediction accuracy on FashionMNIST with all three choices for $p\\_X$, though the accuracy is comparable with PS-MAP.\n\n---\n\n> Fig 2 shows that there's no difference between FS-MAP and PS-MAP when the number of samples is around 10k, a sample size that most practical settings can afford.\n\nThank you for the observation. This early convergence to low test error for both methods is because the data is generated exactly from our prior, greatly simplifying the task and reducing the sample complexity to a much lower level compared to many real-world applications. Relative to the complexity of this task, FS-MAP still achieves a noticeable improvement in efficiency compared to PS-MAP.\n\n---\n\n> Sec 3.3: the choice of prior is important; is there guidance on how to choose a prior in practice?\n\nIndeed, Sec 3.3 shows that well-specification of the prior as well as $p_X$ was important to achieving the best performance with FS-MAP. Choosing a good prior often involves exploiting problem-specific inductive biases (e.g., convolutional networks induce a prior that favors translation equivariant functions, suitable for problems in computer vision) and is a fundamental question in machine learning and outside the scope of this work. On the other hand, our work does provide guidance on choosing $p_X$ by revealing its connection to the metric in function space (Section 3.2) and shows choosing $p_X$ to closely approximate the input distribution at test time often works better (Figure 3).\n\n---\n\n> Is there guidance on how to choose between FS-MAP and PS-MAP given a practical problem? How much would the choice matter (the experiments seem to show that the choice doesn't matter much)?\n\nIn general, despite the flatness-favoring properties of FS-MAP, there is no strong reason to expect that either FS-MAP or PS-MAP will generalize better in the absence of a good prior that is highly informative of the data-generating process. When such a prior is available, we expect that the choice between FS-MAP and PS-MAP to be more important and FS-MAP to more likely outperform PS-MAP. We provide a comprehensive discussion on potentially useful criteria we can rely on to choose between FS-MAP or PS-MAP in the general response.\n\n---\n\n\n> Fig 2b: why is there an increase in curvature a the number of samples increases?\n\nThanks for the intriguing observation! There is indeed a slight increase in curvature for FS-MAP as we increase training samples. We hypothesize this is because as we add more training data, the volume of solutions that can fit the data decreases since more constraints need to be satisfied. As a result, it becomes more difficult to not increase the loss as we move away from the minimum.\n\n---\n\n> is there a reason for measuring flatness as the average eigenvalue (rather than e.g. max)?\n\nWe didn't have any strong reason against using max, though taking the mean captures more information about the local geometry as it is aggregated over all dimensions of the high dimensional parameter space, whereas the max eigenvalue only measures the loss curvature in a single direction.\n\n\n---\n\nPlease let us know if you have any further questions!"}, {"review_id": "LEdHQNxTRY", "reviewer": "Reviewer_Zoj6", "summary": "This work investigates the task of performing inference over model functions rather than model parameters. In particular, the authors propose an objective function *function-space MAP (FS-MAP)*, which intuitively is the usual MAP objective where the prior term is taken over functions rather than parameters. Defining this function-space prior is fundamentally challenging, and to do so, the authors generalize the techniques of Wolpert (1993) to consider functions observed on (possibly) infinite sets.\n\nThe authors perform an in-depth empirical evaluation of this proposed objective on several synthetic and real-world benchmarks. In addition, a scalable approximation to the objective is proposed in Section 4.3. ", "questions": "- Under what conditions is Equation (10) well-defined? It was not clear that the integral in this expression is finite without e.g. further assumptions on the model class. (It is straightforward to give examples where the integral is infinite)\n- Can the infinite-observation objective (Equation (11)) be justified as an MAP estimate? Equation (8) follows from the change-of-variables formula, but it was not clear to me if (or how) Equation 11 was being justified in the same manner, or merely as a heuristic.", "limitations": "The authors appropriately discuss the limitations of their work within the submission.\n", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is overall very well-written and clear in its aims and conclusions. \n- The key ideas of the paper are likely of interest to the community and the findings have the potential to be impactful for future work on function-space inference.\n- I found the discussion relating the function-space geometry to the sampling distribution in Section 3.2 fascinating, and this was a novel connection that I had not considered before. \n- The experimental evaluation is generally sound and the conclusions drawn are supported by the evidence provided throughout the paper.\n- The proposed objective is a non-trivial extension of previous work in this area, and the scalable approximation in Section 4.3 is similarly non-trivial. ", "weaknesses": " - There are several places in the paper lacking rigor and where I believe there to be undressed latent fundamental issues (see the questions section). As noted on Line 96, a fundamental difficulty in function-space inference is defining an appropriate notion of a prior density over infinite dimensional spaces. The methods of this work are obtained via discretization of the functions followed by a heuristic passage to the limit, but it is unclear how this relates to the underlying function-space distributions, which are the true objects of interest. \n- The empirical results on real-world data (Table 1) do not show benefits of using the FS-MAP objective over the standard PS-MAP objective. \n\t- Note that I do not view this necessarily as a weakness of the work itself but rather a potential limitation of the methodology itself, and an honest evaluation of the proposed methodology is valuable.\n- The paper needs to be better contextualized within the literature, e.g. there is no related work section. See e.g. [1-4] for some potentially relevant works on function-space inference which are potentially of interest.\n\n\n[1] Function-Space Regularization in Neural Networks: A Probabilistic Perspective, Rudner et al., ICML 2023\n\n[2] Understanding Variational Inference in Function-Space, Burt et al., AABI 2020\n\n[3] Generalized Variational Inference in Function Spaces: Gaussian Measures meet Bayesian Deep Learning, Wild et al., NeurIPS 2022\n\n[4] Tractable Function-Space Variational Inference in Bayesian Neural Networks, Rudner et al., NeurIPS 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful and constructive questions and suggestions!\n\nWe were pleased that find that our submission is **\"of interest to the community\"** and that our findings **\"have the potential to be impactful for future work on function-space inference\"**. We were also happy to see that you found the manuscript to be **\"very well-written\"** and **\"clear in its aims and conclusions\"**.\n\n\nWe address your questions and comments below. Please let us know if you have any remaining questions.\n\n---\n> [...] The methods of this work are obtained via discretization of the functions followed by a heuristic passage to the limit, but it is unclear how this relates to the underlying function-space distributions\n\nYou correctly noted that the main difficulty in establishing any connection between our result and the so-called true object of interest involving the density evaluation $p(f_\\theta)$ is due to technical challenges in defining the infinite-dimensional probability density function $p(f_\\theta)$ in the first place.\n\nSection 3.2 establishes the limiting behavior of the finite-point objective derived by Wolpert (which exactly corresponds to the value of the prior probability density function of a distribution over finite function evaluations computed at a given parameterized function $f_\\theta$) and thereby avoids explicitly defining $p(f_\\theta)$. We expect the evaluation of $p(f_\\theta)$ obtained using the infinite-limit expression to agree with an evaluation that could be obtained using a suitably-defined base measure over the function space. A justification for this is provided in Appendix B.2. The technique used in Section 3.2 is referred to as taking the continuum limit and has previously been applied to analogous problems in physics and applied mathematics.\n\nFor example, in numerical analysis, the solutions of discretized differential equations often converge to the solutions of the continuous equations as the mesh size goes to zero. Similarly, lattice field theories in physics, where the continuous spacetime is discretized into a lattice, has provided consistent results with the continuum approach in many cases.\n\n---\n> The empirical results on real-world data (Table 1) do not show benefits of using the FS-MAP objective over the standard PS-MAP objective.\n\nThank you for pointing out that further discussion of the results would be useful. Indeed, as we discussed in our general response, we were not expecting FS-MAP or L-MAP to lead to significant improvements in these settings where we do not have a highly informative prior. In contrast, we use a prior that exactly corresponds to the data-generating process in Section 3.3, leading to significant gains by performing FS-MAP. Here, since our Gaussian prior is rather non-informative about the data-generating process in the FashionMNIST and CIFAR-10 experiments, we were pleasantly surprised to find that L-MAP achieves a significant reduction in ECE on CIFAR-10 with $p\\_X=$CIFAR-100 and noticeable improvement in selective prediction accuracy on FashionMNIST with all three choices for $p\\_X$, though the accuracy is comparable with PS-MAP.\n\n---\n> Under what conditions is Equation (10) well-defined? It was not clear that the integral in this expression is finite without e.g. further assumptions on the model class.\n\nIt is true that Equation (10) is not finite for any model. However, for most models, including commonly used neural networks, we can show it is indeed finite. To show the expectation in Equation (10) is finite, it is sufficient to show that the integrand $J_\\theta(x)^\\top J_\\theta(x)$ is bounded for any $x$ in the support of $p_X,$ a sufficient condition for which is that the function $f_\\theta$ has bounded Lipschitz constants over the support of $p_X$. Consequently, Equation (10) is finite for any $\\theta$ for a wide range of models, e.g., any MLP with standard activation functions (ReLU, tanh). Note this is not a necessary condition, and Equation (10) can still be finite depending on the specifics of $p_X$ and $J_\\theta(x).$ \n\n---\n> Can the infinite-observation objective (Equation (11)) be justified as an MAP estimate? [...] it was not clear to me if (or how) Equation 11 was being justified in the same manner, or merely as a heuristic.\n\nAs mentioned above, the main difficulty in establishing any connection between our result and the true mode of $p(f_\\theta | \\mathcal{D})$ lies in technical challenges in the prior density $p(f_\\theta)$ in the first place. On the other hand, our result by taking the continuum limit is not merely a heuristic but follows an established approach with a history of successful applications in other areas of mathematics and physics that enable tractable computational methods when the continuum approach presents challenges.\n\n---\n> The paper needs to be better contextualized within the literature, e.g. there is no related work section.\n\nFunction-space regularization in NNs is a burgeoning field, and we appreciate the opportunity to further contextualize our work:\n\n[1] discuss function-space VI in NNs. They explain in which cases the function-space variational objective proposed in [2] is not well-defined. This pathology is conceptually similar to the pathology identified in our work.\n\n[3] propose approximations to make the variational objective proposed by [2] tractable and identify a prior over functions for which the objective is well-defined. Unlike our approach, the sampling distribution in [3] is not part of the probabilistic model but of the approximation.\n\n[4] also perform variational inference but they minimize a Wasserstein instead of a KL divergence, and they consider Gaussian process models instead of stochastic NNs.\n\n[5] specify a prior over parameters that induce a desirable prior over functions but perform inference in parameter space instead of deriving a function-space objective to avoid the challenges described in our submission.\n\n---\nPlease let us know if you have any further questions."}, {"review_id": "7yG1eScvnA", "reviewer": "Reviewer_wSt3", "summary": "This paper investigates an important yet neglected question in machine learning: should we train the model in the parameter space or function space? The authors show the benefits and shortage of function space MAP estimation which could provide a guide for practitioners.  Last but not least, the authors propose a scalable approximation for function space MAP estimation for deep learning.", "questions": "See weakness above", "limitations": "See weakness above", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- Thorough analysis of the pros and cons of parameter space MAP estimation and function space MAP estimation\n- Propose a scalable approximation for large neural networks which extends the application area of function space MAP estimation\n- Well written and easy to follow", "weaknesses": "- I wonder how universal are the pros and cons demonstrated in the paper, considering they're being shown in a relatively simple toy data set\n- I understand the author is not claiming L-MAP works well in all cases, but I am still interested to see the performance of L-MAP on tasks other than classification", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your thoughtful and constructive questions and suggestions!\n\nWe were pleased that you found our manuscript to be **\"well-written and easy to follow\"** and that you highlighted our **\"thorough analysis\"** of the pros and cons of PS-MAP and FS-MAP.\n\nWe address your questions and comments below. Please let us know if you have any remaining questions.\n\n---\n\n> I wonder how universal are the pros and cons demonstrated in the paper, considering they're being shown in a relatively simple toy data set\n\nThank you for this important question. The goal of our experiments on synthetic datasets was precisely to probe the pros and cons of FS-MAP in a well-controlled setting and the results are highly in line with the expectation from our theoretical analysis. We also conducted extensive evaluations with neural networks on datasets commonly used in deep learning. Moreover, we have run new experiments on UCI regression and transfer learning to broaden our evaluations, which can be found in our general response. Finally, we added a discussion on a comprehensive set of criteria for when FS-MAP and L-MAP should or shouldn't be expected to outperform PS-MAP in more general settings.\n\n---\n\n\n> I understand the author is not claiming L-MAP works well in all cases, but I am still interested to see the performance of L-MAP on tasks other than classification\n\nWe appreciate the suggestions for additional evaluations and ran additional experiments to compare L-MAP and PS-MAP on UCI regression datasets and transfer learning on CIFAR-10. \n\n### UCI regression\nWe found that L-MAP outperformed PS-MAP on 7 out of 8 datasets according to normalized test RMSE, showing L-MAP can also benefit generalization on many regression tasks.\n\n| Dataset   | L-MAP             | PS-MAP            |\n|:----------|:------------------|:------------------|\n| Boston    | 0.352 ± 0.040     | **0.329 ± 0.033** |\n| Concrete  | **0.261 ± 0.013** | 0.272 ± 0.016     |\n| Energy    | **0.041 ± 0.002** | 0.042 ± 0.003     |\n| Naval     | **0.018 ± 0.002** | 0.032 ± 0.005     |\n| Power     | **0.218 ± 0.005** | 0.219 ± 0.006     |\n| Protein   | **0.580 ± 0.005** | 0.584 ± 0.005     |\n| Winered   | **0.792 ± 0.031** | 0.851 ± 0.029     |\n| Winewhite | **0.714 ± 0.017** | 0.758 ± 0.013     |\n\nWe adopt the following setup: we use a 3 hidden layer MLP with 256 units and tune the weight decay (prior variance) and Laplacian regularization strengths on a validation set. We standardize both the inputs and targets and report the mean and standard error of test RMSE across six different runs. For L-MAP we set $p_X = \\mathcal{N}(0, I)$ since the inputs were standardized and fairly low-dimensional.\n\n### Transfer Learning from ImageNet to CIFAR-10\n\nUsing a ResNet18 pre-trained with ImageNet, we tested the performance on further fine-tuning on CIFAR-10 with samples from CIFAR-100 as the evaluation set for 5 epochs. As with training from scratch, we still use a batch size of 128 and an evaluation points size of 128.\n\n| Method | Acc. | Sel. Acc. | Avg. NLL | ECE |\n| ----- | ----- | ----- | ----- | ----- |\n| PS-MAP | 95.3% | 99.5% | 0.14 | 1.2% |\n| L-MAP | 95.4% | 99.5% | 0.14 | 1.0% |\n\nWe find L-MAP is able to perform marginally better while improving the calibration of the classifier slightly.\n\n\n---\n\nPlease let us know if you have any further questions."}]

Author rebuttal:
## General Response to All Reviewers

We thank all reviewers for their feedback, support, and unanimous recognition that our paper is interesting, well-written, and thorough in presenting the pros and cons of the considered approaches. We start by providing a general response aimed at addressing the most common questions from the reviewers. 

We first present additional experiments and find that L-MAP outperforms PS-MAP on UCI Regression and transfer learning from ImageNet to CIFAR-10 and then present a comprehensive set of criteria on when FS-MAP and L-MAP should or shouldn't be expected to outperform PS-MAP in general settings, explaining why FS-MAP significantly outperforms PS-MAP in some of our experiments but no others. We hope our response addresses the key points and inquiries raised by the reviewers and can be taken into account in the final assessment.

## Additional Experiments on UCI Regression and Transfer Learning

We provide additional regression and transfer learning experiments with L-MAP, as suggested by some reviewers. The results are consistent with the results in the manuscript and show that L-MAP tends to perform as well or better than PS-MAP.

### UCI

We found that L-MAP outperformed PS-MAP on 7 out of 8 datasets according to normalized test RMSE, showing L-MAP can also benefit generalization on many regression tasks.

| Dataset   | L-MAP             | PS-MAP            |
|:----------|:------------------|:------------------|
| Boston    | 0.352 ± 0.040     | **0.329 ± 0.033** |
| Concrete  | **0.261 ± 0.013** | 0.272 ± 0.016     |
| Energy    | **0.041 ± 0.002** | 0.042 ± 0.003     |
| Naval     | **0.018 ± 0.002** | 0.032 ± 0.005     |
| Power     | **0.218 ± 0.005** | 0.219 ± 0.006     |
| Protein   | **0.580 ± 0.005** | 0.584 ± 0.005     |
| Winered   | **0.792 ± 0.031** | 0.851 ± 0.029     |
| Winewhite | **0.714 ± 0.017** | 0.758 ± 0.013     |

We adopt the following setup: we use a 3 hidden layer MLP with 256 units and tune the weight decay (prior variance) and Laplacian regularization strengths on a validation set. We standardize both the inputs and targets and report the mean and standard error of test RMSE across six different runs. For L-MAP we set $p_X = \mathcal{N}(0, I)$ since the inputs were standardized and fairly low-dimensional.

### Transfer Learning from ImageNet to CIFAR-10

Using a ResNet18 pre-trained with ImageNet, we tested the performance on further fine-tuning on CIFAR-10 with samples from CIFAR-100 as the evaluation set for five epochs. As with training from scratch, we still use a batch size of 128 and an evaluation points size of 128.

| Method | Acc. | Sel. Acc. | Avg. NLL | ECE |
| ----- | ----- | ----- | ----- | ----- |
| PS-MAP | 95.3% | 99.5% | 0.14 | 1.2% |
| L-MAP | 95.4% | 99.5% | 0.14 | 1.0% |

We find L-MAP is able to perform marginally better while improving the calibration of the classifier slightly.

## When FS-MAP should or shouldn't outperform PS-MAP

Many reviewers noted that while FS-MAP leads to noticeably better performance in the synthetic example in Section 3.3, less improvement is observed in neural network experiments on image classification in Section 4.4. We note that these results are exactly in line with our theoretical observation and empirical findings in Section 3.3 that FS-MAP's superior performance depends on a well-specified probabilistic model (prior and likelihood). As one of our key points in the paper, in general, there is no strong reason to expect either FS-MAP or PS-MAP will generalize better since our prior can be arbitrarily different from the true data-generating process. Nevertheless, we now summarize our current best understanding of when FS-MAP should or shouldn't outperform PS-MAP.

The example in Section 3.3 exemplifies four criteria that generally favor FS-MAP:

1) The Jacobian is non-singular everywhere and therefore, FS-MAP does not lead to pathological solutions as described in Section 4.1. 
2) The likelihood and the prior exactly correspond to the data-generating process. Consequently, decision theory states that Bayesian Model Average (BMA) is the optimal predictor in terms of expected test RMSE.
3) The prior $p(\theta) = \mathcal{N}(\theta | 0, 10^2)$ is diffuse in parameter space but fairly concentrated in function space because each coefficient $\tanh(\theta_i)$ follows a highly bimodal distribution at $-1$ and $1$ as a result of the prior standard deviation being much larger compared to the scale at which $\tanh$ saturates (which is approximately 2).
4) Therefore, FS-MAP will better approximate the Bayesian model average than PS-MAP, following a similar argument in Section 4.2 line 271-274, and therefore achieve better generalization.

Correspondingly, there are situations where FS-MAP will likely underperform PS-MAP:

1) When the prior is poorly specified. For example, in a regression setting, if the observation noise is extremely over-estimated, FS-MAP may severely underfit the data because the log determinant regularization can overpower the likelihood. We demonstrate this effect in Appendix B.7, Figure 8.
2) When the log determinant contains singularities. As discussed in Section 4.1, in this case, FS-MAP leads to pathological solutions that ignore the data, though L-MAP mitigates this issue.

We hope this response addresses the potential concern that FS-MAP doesn't always outperform PS-MAP, and our thorough and transparent demonstration of the pros and cons of both approaches is viewed as a strength.

---

**Thank you for reviewing our work!**



Meta-review: This paper studies learning most likely functions vs. parameters in the context of MAP estimation. Differences in parameterization affect the parameter space, and the authors provide an intellectual discussion on contrasting this approach to function space MAP. Reviewers have pointed out several dimensions of improvement to maximize clarity and I advise the authors to carefully address them in the camera-ready version. In particular, I find the limitation of the work to be inadequately discussed. For example, the algorithmic takeaway from the discussion in the paper is not clear as multiple reviewers have noted. This is fine within the scope of a single paper but is worth highlighting for future research.
Decision: Accept (poster)