Title: Should We Learn Most Likely Functions or Parameters?
Venue: NeurIPS 2023
Accepted: yes
Keywords: Function-Space Modeling

Function-Space Regularization

Maximum A Posteriori Estimation

Generalization
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?

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.

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 we use to train our models. Minimizing standard regularized loss functions, including cross-entropy for classification and mean-squared or mean-absolute error for regression, with ℓ 1 or ℓ 2 regularization, exactly corresponds to maximizing a posterior distribution over model parameters [2, 19]. This standard procedure is known in probabilistic modeling as maximum a posteriori (MAP) parameter estimation. However, parameters have no meaning independent of the functional form of the models they parameterize. In particular, our models f θ ( 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 to 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 1a, we visualize the posterior over a mixture coefficient θ R in a Gaussian mixture regression model in both parameter and function space, using the same Gaussian prior for the mixture coefficients (corresponding to ℓ 2 regularization in parameter space). We see that each distribution is maximized by a different parameter θ R , leading to very different learned functions in Figure 1b. Moreover, we can 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.

5

0

R

0.0

0.2

0.4

0.6

p

(

f

|

)

p

(

|

)

(a) Distinct Posterior Densities

1

0

1

x

0.00

0.25

0.50

0.75

1.00

y

Training Data

FS-MAP

PS-MAP

(b) Distinct Learned Functions

On one hand, we might expect the answer to be a clear-cut 'We should learn most likely functions!'. In Section 3, we present a well-defined function-space MAP objective through a generalization of the change of variables formula for probability densities. In addition to 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 flat 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 it difficult to scale to modern neural networks. We also show that function-space MAP will not necessarily be closer than parameter-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 a 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 to improve our understanding of what it means to 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 a complex picture and unearths distinct benefits and drawbacks of learning most likely functions instead of parameters. We conclude with a discussion of practical considerations around function-space MAP estimation, its relationship to flat minima and generalization, and the practical relevance of parameterization invariance.

Our code is available at https://github . com/activatedgeek/function-space-map .

2 Preliminaries

We consider supervised learning problems with a training dataset D = ( x D , y D ) = { x ( i ) D , y ( i ) D } N i =1 of inputs x ∈ X and targets y ∈ Y with input space X ⊆ R D and output space Y ⊆ R K . For example, the inputs x could correspond to times, spatial locations, tabular data, images, and the targets y to regression values, class labels, etc.

To create a model of the data D , one typically starts by specifying a function f θ : X → Y parametrized by a vector θ ∈ R P which maps inputs to outputs. f θ could be a neural network, a polynomial model, a Fourier series, and so on. Learning typically amounts to estimating model parameters θ . To this end, we can relate the function f θ to the targets through an observation model , p ( y | x, f θ ) . For example, in regression, we could assume the outputs y = f θ ( x ) + ϵ , where ϵ ∼ N (0 , σ 2 ) is additive Gaussian noise with variance σ 2 . Equivalently, p ( y | x, θ ) = N ( y | f θ ( x ) , σ 2 ) . Alternatively, in classification, we could specify p ( y | x, θ ) = Categorical(softmax( f θ ( x ))) . We then use this observation model to form a likelihood over the whole dataset p ( y D | x D , θ ) . In each of these example observation models, the likelihood factorizes, as the data points are conditionally independent given model parameters θ , p ( y D | x D , θ ) = ∏ N i =1 p ( y ( i ) D | x ( i ) D , θ ) . We can further express a belief over values of parameters through a prior p ( θ ) , such as p ( θ ) = N ( θ | µ, Σ) . Finally, using Bayes rule, the log of the posterior up to a constant c independent of θ is,

Notably, standard loss functions are negative log posteriors, such that maximizing this parameter posterior, 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 a categorical distribution with a softmax link function, then the log likelihood is negative cross-entropy. If we use a zero-mean Gaussian prior, we recover standard ℓ 2 regularization, also known as weight-decay. If we use a Laplace prior, we recover ℓ 1 regularization, also known as LASSO [26].

Once we maximize the posterior to find

we can condition on these parameters to form our function f ˆ θ PS-MAP to make predictions. However, as we saw in Figure 1, f ˆ θ PS-MAP is not in general the function that maximizes the posterior over functions, p ( f θ | y D , x D ) . In other words, if

then generally f ˆ θ PS-MAP = f ˆ θ FS-MAP . Naively, one can write the log posterior over f θ up to the same constant c as above as

̸

where p ( y D | x D , f θ ) = p ( y D | x D , θ ) , but just written in terms of the function f θ . The prior p ( f θ ) however is a different function from p ( θ ) , because we incur an additional Jacobian factor in this change of variables, making the posteriors also different.

We must take care in interpreting the quantity p ( f θ ) since probability densities in infinite-dimensional 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 a finite number of evaluation points, we provide a more general objective that enables the use of infinitely many evaluation points to construct a more informative prior and makes the relevant design choices of function-space MAP estimation more interpretable.

3 Understanding Function-Space Maximum A Posteriori Estimation

Function-space MAP estimation seeks to answer a fundamentally different question than parameterspace 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 a function-space MAP objective that generalizes the objective considered by prior work and analyze its properties both theoretically and empirically.

3.1 The Finite Evaluation Point Objective

Starting from Equation (4), Wolpert [30] proposed to instead find the MAP estimate for f θ (ˆ x ) , the function evaluated at a finite set of points ˆ x = { x 1 , ..., x M } , where M < ∞ can be chosen to be arbitrarily large so as to capture the behavior of the function to arbitrary resolution. This choice then yields the FS-MAP optimization objective

where p ( f θ (ˆ x )) is a 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 θ , and K the number of function outputs. Assuming the set of evaluation points is sufficiently large so that MK ≥ 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 f (ˆ x ) is given by

where J ( θ ; ˆ x ) is a P -byP matrix defined by and J θ (ˆ x ) ˙ = ∂f θ (ˆ x ) /∂θ is the MK -byP Jacobian of f θ (ˆ x ) , viewed as an MK -dimensional vector, with respect to the parameters θ . Substituting Equation (6) into Equation (5) the function-space MAP objective as a function of the parameters θ can be expressed as

where we are allowed to condition directly on f θ instead of f θ (ˆ x ) because by assumption ˆ x 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 ˆ x, this objective makes it unclear how we should select ˆ x and how that choice can be interpreted or justified in a principled way.

3.2 Deriving a More General Objective and its Interpretation

We now derive a more general class of function-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 ∈ R n to u ∈ R m via an injective map u = φ ( v ) , their probability densities relate as p ( u )d µ ( u ) = p ( v )d ν ( v ) , where µ and ν define respective volume measures. Suppose we let d µ ( u ) = √ det( g ( u ))d m u, the volume induced by an M × M metric tensor g, and d ν ( v ) = d n v , the Lebesgue measure, then we can write d µ ( u ) = √ det(˜ g ( v ))d ν ( v ) , where the N × N metric ˜ g ( v ) = J ( v ) ⊤ g ( u ) J ( v ) is known as the pullback of g via φ and J is the Jacobian of φ [6]. As a result, we have p ( u ) = p ( v )det -1 / 2 ( J ( v ) ⊤ g ( u ) J ( v )) . Applying this argument and identifying u with f θ (ˆ x ) and v with θ, Wolpert [30] thereby establishes p ( f θ (ˆ x )) = p ( θ ) det -1 / 2 ( J θ (ˆ x ) ⊤ J θ (ˆ x )) .

However, an important implicit assumption in the last step is that the metric in function space is Euclidean. That is, g = I and the squared distance between f θ (ˆ x ) and f θ (ˆ x ) + d f θ (ˆ x ) is d s 2 = ∑ M i =1 d f θ ( x i ) 2 , rather than the general case d s 2 = ∑ M i =1 ∑ M j =1 g ij d f θ ( x i )d f θ ( x j ) . To account for a generic metric g, we therefore replace J θ (ˆ x ) ⊤ J θ (ˆ x ) with J θ (ˆ x ) ⊤ gJ θ (ˆ x ) . For simplicity, we assume the function output is univariate ( K = 1 ) and only consider g that is constant i.e., independent of f θ (ˆ x ) and diagonal, with g ii = g ( x i ) for some function g : X → R + . For a discussion of the general case, see Appendix B.2. To better interpret the choice of g, we rewrite

where we suggestively defined the alias ˜ p X ( x ) ˙ = g ( x ) and J θ ( x j ) is the K -byP -dimensional Jacobian evaluated at the point x j . Generalizing from the finite and discrete set ˆ x to the possibly infinite entire domain X ⊆ R D and further dividing by an unimportant normalization constant Z , we obtain

where p X = ˜ p X /Z is a normalized probability density function, with normalization Z -which is independent of θ -only appearing as an additive constant in log det J ( θ ; p X ) . We include a further discussion about this limit in Appendix B.4.

Under this limit, we can identify log p ( θ ) -1 2 log det J ( θ ; p X ) with log p ( f θ ) up to a constant, and thereby express log p ( f θ |D ) , the function-space MAP objective in Equation (4) as

where the choice of p X corresponds to a choice of the metric g in function space. Equation (11) (and approximations thereof) will be the primary object of interest in the remainder of this paper.

Evaluation Distribution as Function-Space Geometry. From this discussion and Equations (10) and (11), it is now clear that the role of the metric tensor g is to impose a distribution p X over the evaluation points. Or conversely, a given distribution over evaluation points implies a metric tensor, and as such, specifies the geometry of the function space. This correspondence is intuitive: points x with higher values of g ( x ) contribute more to 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, f θ is an image classifier for which we only care about its outputs on set of natural images I when comparing it to another image classifier. The metric g ( · ) and therefore the evaluation distribution p X only needs support in I , and the FS-MAP objective is defined only in terms of the Jacobians evaluated at natural images x ∈ I .

Finite Evaluation Point Objective as a Special Case. Consequently, the finite evaluation point objective in Equation (8) can be arrived at by specifying the evaluation distribution p X ( x ) to be ˆ p X ( x ) ˙ = 1 M ∑ x ′ ∈ ˆ x δ ( x -x ′ ) , where δ is the Dirac delta function and ˆ x = { x 1 , ..., x M } with M < ∞ , as before. It is easy to see that J finite ( θ ; ˆ x ) ∝ J ( θ ; ˆ p X ) . Therefore, the objective proposed by Wolpert [30] is a 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, we consider the class of models f θ ( x ) = ∑ P i =1 σ ( θ i ) φ i ( x ) with domain X = [ -1 , 1] , where { φ i } P i =1 is a fixed set of basis functions and σ is 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 J ( θ ; p X ) has a simple closed form, such that where Φ is a constant matrix with elements Φ ij = E p X [ φ i ( X ) φ j ( X )] . Therefore, Φ can be precomputed once and reused throughout training. In this experiment, we use the set of Fourier features { cos( k i · ) , sin( k i · ) } 100 i =1 where k i = iπ and set σ = tanh . We generate training data by sampling x train ∼ Uniform( -1 , 1) , θ i ∼ N (0 , α 2 ) with α = 10 , evaluating f θ ( x train ) , and adding Gaussian noise with standard deviation σ ∗ = 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 ( θ ) and the likelihood to correspond to the data-generating process. For FS-MAP, we specify p X = Uniform( -1 , 1) , the ground-truth distribution of test inputs, which conveniently results in Φ = I/ 2 .

FS-MAP Finds More Probable Functions. Figure 2a shows the improvement in log p ( f θ |D ) when using FS-MAP over PS-MAP. As expected, FS-MAP consistently finds functions with much higher posterior probability p ( f θ |D ) .

10 2

10 3

# Train samples

10 3

6×10 3

(a) Log Posterior Increase

10 2

10 3

# Train samples

10

3

10

4

PS-MAP

FS-MAP

(b) Loss Curvature

10 2

10 3

# Train samples

10

2

10

1

(c) Test RMSE

10 2

10 3

# Train samples

10

1

3×10 2

(d) Train RMSE

10

50

200

0.04

0.06

Test RMSE

FS-MAP

PS-MAP

(a) Number of evaluation points

U(

1, 1)U(

0.1, 0.1)

N/A

pX

0.00

0.02

0.04

0.06

Test RMSE

(b) Evaluation distribution

FS-MAP Prefers Flat Minima. In Figure 2b, we compare the curvature of the minima found by FS-MAP and PS-MAP, as measured by the average eigenvalue of the Hessian of the mean squared error loss, showing that FS-MAP consistently finds flatter minima. As the objective of FS-MAP favors small singular values for the Jacobian J θ , 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 ∇ 2 L ( θ ) . If the model fits the training data well, we can apply the Gauss-Newton approximation: ∇ 2 θ L ( θ ) ≈ 1 |D| J θ ( x D ) ⊤ J θ ( x D ) , which is identical to J ( θ ; p X ) if p X is chosen to be the empirical distribution of the training inputs. More generally, a distribution p X with high density over likely inputs will assign high density to the training inputs, and hence minimizing J ( θ ; p X ) will similarly reduces the magnitude of J θ ( x D ) . Therefore, the FS-MAP objective explicitly encourages finding flatter minima, which have been found to correlate with generalization [10, 17, 18], robustness to data perturbations and noisy activations [3, 12] for neural networks.

FS-MAP can Achieve Better Generalization. Figure 2c shows that FS-MAP achieves lower test RMSE across a wide range of sample sizes. It's worth noting that this synthetic example satisfies two important criteria such that FS-MAP is likely to improve generalization over PS-MAP. First, the condition outlined in Section 4 for a well-behaved FS-MAP objective is met, namely that the set of partial derivatives { ∂ θ i f j θ ( · ) } P i =1 are linearly independent. Specifically, the partial derivatives are given by { sech( θ i ) sin( k i · ) , sech( θ i ) cos( k i · ) } P i =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-generating process. Therefore, FS-MAP seeks the most likely function according to its true probability, without incorrect modeling assumptions.

FS-MAP is Less Prone to Overfitting. As shown in Figure 2d, FS-MAP tends to have a higher train RMSE than PS-MAP as a result of the additional log determinant regularization. While PS-MAP achieves near-zero train RMSE with a small number of samples by overfitting to the noise, FS-MAP's train RMSE is consistently around σ ∗ = 0 . 1 , the true standard deviation of the observation noise.

Performance Improves with Number of Evaluation Points. We compare FS-MAP estimation with p X = Uniform( -1 , 1) and with a finite number of equidistant evaluation points in [ -1 , 1] , where the former corresponds to the latter with infinitely many points. In Figure 3a, we show that the test RMSE of the FS-MAP estimate (evaluated at 400 training samples) decreases monotonically until the number of evaluation points reaches 200. The more evaluation points, the more the finite-point FS-MAP objective approximates its infinite limit and the better it captures the behavior of the function. Indeed, 200 points is the minimum sampling rate required such that the Nyquist frequency reaches the maximum frequency in the Fourier basis, explaining the saturation of the performance.

Choice of Evaluation Distribution is Important. In Figure 3b, we compare test RMSE for 400 training samples for the default choice of p X = Uniform( -1 , 1) , p X = Uniform( -0 . 1 , 0 . 1) -a distribution that does not reflect inputs at test time-and PS-MAP ( p X = N / A ) for reference. The result shows that specifying the evaluation distribution to correctly reflect the distribution of inputs at test time is required for FS-MAP to achieve optimal performance in this case.

4 Limitations and Practical Considerations

We now discuss the limitations of function-space MAP estimation and propose partial remedies, including a scalable approximation that improves its applicability to modern deep learning.

4.1 FS-MAP Does not Necessarily Generalize Better than PS-MAP

There is no guarantee that the most likely function is the one that generalizes the best, especially since our prior can be arbitrary. For example, FS-MAP can under-fit the data if the log determinant introduces excessive regularization, which can happen if our probabilistic model is misspecified by overestimating the observation noise. Returning to the setup in Section 3.3, in Figure 4, we find that FS-MAP (with 400 training samples) is sensitive to the observation noise scale σ in the likelihood. As σ deviates from the true noise scale σ ∗ , the test RMSE of FS-MAP can change dramatically. At σ/σ ∗ = 1 / 3 , the likelihood dominates in both the FS-MAP and PS-MAP objectives, resulting in similar test RMSEs. At σ/σ ∗ = 10 , the likelihood is

1

1/3

10

/

*

0.2

0.4

Test RMSE

PS-MAP

FS-MAP

overpowered by the log determinant regularization in FS-MAP and causes the model to under-fit the data and achieve a high test error. By contrast, the performance of the PS-MAP estimate is relatively stable. Finally, even if our prior exactly describes the data-generating process, the function that best generalizes won't necessarily be the most likely under the posterior.

4.2 Pathological Solutions

In order for function-space MAP estimation to be useful in practice, the prior 'density' p ( f θ ) = p ( θ ) det -1 2 ( J ( θ ; p X )) must not be infinite for any allowed values of θ , since if it were, those values would constitute global optima of the objective function independent of the actual data. To ensure that there are no such pathological solutions, we require the matrix J ( θ ; p X ) to be non-singular for any allowed θ . We present two results that help determine if this requirement is met.

Theorem 1. Assuming p X and J θ = ∂ θ f θ are continuous, J ( θ ; p X ) is non-singular if and only if the partial derivatives { ∂ θ i f θ ( · ) } P i =1 are linearly independent functions over the support of p X .

̸

Proof.

Theorem 2. If there exists a non-trivial symmetry S ∈ R P × P with S = I such that f θ = f Sθ for all θ , then J ( θ ∗ ; p X ) is singular for all fixed points θ ∗ where Sθ ∗ = θ ∗ .

Theorem 1 is analogous to A ⊤ A having the same null space as A for any matrix A . It suggests that to avoid pathological optima the effect of small changes to each parameter must be linearly independent. Theorem 2 builds on this observation to show that if the model exhibits symmetries in its parameterization, FS-MAP will necessarily have pathological optima. Since most neural networks at least possess permutation symmetries of the hidden units, we show that these pathological FS-MAP solutions are almost universally present, generalizing the specific cases observed by Wolpert [30].

A Simple Remedy to Remove Singularities. Instead of finding a point approximation of the function space posterior, we can perform variational inference under a variational family { q ( ·| θ ) } θ , where q ( ·| θ ) is localized around θ with a small and constant function-space entropy h . Ignoring constants and O ( h ) terms, the variational lower bound is given by

Similar to how convolving an image with a localized Gaussian filter removes high-frequency components, the expectation E q ( θ ′ | θ ) removes potential singularities in log det J ( θ ′ ; p X ) . This effect can be approximated by simply adding a small diagonal jitter ϵ to J ( θ ; p X ) :

Alternatively, we know E q ( θ ′ | θ ) [log det J ( θ ′ ; p X )] must be finite because the variational lower bound cannot exceed the log marginal likelihood. ˆ L ( θ ; p X ) can be used as a minimal modification of FS-MAP that eliminates pathological optima. We provide full derivation for these results in Appendix B.7.

4.3 Does the Function-Space MAP Better Approximate the Bayesian Model Average?

The Bayesian model average (BMA), f BMA , of the function f θ is given by the posterior mean, that is:

This expectation is the same when computed in parameter or function space, and has theoretically desirable generalization properties [19, 29]. Both PS-MAP and FS-MAP provide point approximations of f BMA . However, FS-MAP seeks the mode of the distribution p ( f θ |D ) , the mean of which we aim to compute. By contrast PS-MAP finds the mode of a distinct distribution, p ( θ |D ) , which can markedly diverge from p ( f θ |D ) , depending on the parameterization. Consequently, it is reasonable to anticipate that the FS-MAP objective generally encourages finding a superior estimate of the BMA.

̸

Consider the large data regime where the posterior in both parameter and function space follows a Gaussian distribution, in line with the Bernstein-von Mises theorem [27]. In Gaussian distributions, the mode and mean coincide, and therefore f ˆ θ FS-MAP = f BMA . However, even in this setting, where ˆ θ PS-MAP = E p ( θ |D ) [ θ ] , generally f ˆ θ PS-MAP = f BMA , because the expectation does not distribute across a function f that is non-linear in its parameters θ : f E p ( θ |D ) [ θ ] ( · ) = E p ( θ |D ) [ f θ ( · )] .

̸

However, we find that whether FS-MAP better approximates the BMA depends strongly on the problem setting. Returning to the setup in Figure 1, where we have a Gaussian mixture regression model, we compare the BMA function with functions learned by FS-MAP and PS-MAP under the prior p ( θ ) = N (0 , α 2 ) with several settings of α. We observe in Figure 5 that FS-MAP only approximates the BMA function better than PS-MAP at larger α values. To understand this behavior, recall the height h R for the right Gaussian bump is given by exp( θ R ) , which has a lognormal(0 , α 2 ) prior. As we increase α, more prior mass is assigned to h R with near-zero value and therefore to functions with small values within x ∈ [0 , 1] . While the lognormal distribution also has a heavy tail at large h R , the likelihood constrains the posterior p ( h R |D ) to only place high mass for small h R . These two effects combine to make the posterior increasingly concentrated in function space around functions described by h L ≈ 1 and h R ≈ 0 , implying that the mode in function space should better approximate its mean. In contrast, as we decrease α, both the prior and posterior become more concentrated in parameter space since the parameter prior p ( θ ) = N (0 , α 2 ) approaches the delta function at zero, suggesting that PS-MAP should be a good approximation to the BMA function. By varying α , 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 J ( θ ; p X ) = E p X [ J θ ( X ) ⊤ J θ ( X )] is almost never analytically tractable due to the integral over X . We show in Appendix B.9 that a simple Monte Carlo estimate for J ( θ ; p X ) with S samples of X can yield decent accuracy. For large neural networks, this estimator is still prohibitively expensive: each sample will require K backward passes, taking a total of O ( SKP ) time. In addition, computing the resulting SK -by P determinant takes time O ( SKP 2 ) (assuming P ≥ 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 ϵ ≫ λ i for all eigenvalues λ i of J ( θ ; p X )) . To first order in max i λ i /ϵ, we have log det( J ( θ ; p X )) + ϵI ) = 1 2 ∆ ψ d ( θ, θ + ψ )) ∣ ∣ ψ =0 where d ( θ, θ ′ ) ˙ = E p X [ ∥ f θ ( X ) -f θ ′ ( X ) ∥ 2 ] and ∆ = ∑ P i =1 ∂ 2 θ i denotes the Laplacian operator. Exploiting the identity 1 2 ∆ ψ d ( θ, θ + ψ )) ∣ ∣ ψ =0 = 1 β 2 E ψ ∼N (0 ,β 2 I ) [ d ( θ, θ + ψ )] + O ( β 2 ) and choosing β small enough, we obtain an accurate Monte Carlo estimator for the Laplacian

1

0

1

x

0.0

0.5

1.0

y

Training Data

FS-MAP

PS-MAP

BMA

(a)

α

= 1

.

2

1

0

1

x

0.0

0.5

1.0

y

Training Data

FS-MAP

PS-MAP

BMA

(b)

α

= 10

1

0

1

x

0.0

0.5

1.0

y

Training Data

FS-MAP

PS-MAP

BMA

(c)

α

= 20

using only forward passes. The resulting objective which we refer to as Laplacian Regularized MAP (L-MAP) is given by

We use a single sample of ψ and S samples of evaluation points for estimating d ( θ, θ + ψ ) at each step, reducing the overhead from O ( SKP 2 ) to only O ( SP ) . A more detailed exposition is available in Appendix B.10.

4.5 Experimental Evaluation of L-MAP

We now evaluate L-MAP applied to neural networks on various commonly used datasets. We provide full experimental details and extended results in Appendix C. Unlike the synthetic task in Section 3.3, in applying neural networks to these more complex datasets we often cannot specify a prior that accurately models the true data-generating process, beyond choosing a plausible architecture whose high-level inductive biases align with the task (e.g. CNNs for images) and a simple prior p ( θ ) favoring smooth functions (e.g. an isotropic Gaussian with small variances). Therefore, we have less reason to expect L-MAP should outperform PS-MAP in these settings.

UCI Regression. In Table 1, we report normalized test RMSE on UCI datasets [1], using an MLP with 3 hidden layers and 256 units. L-MAP achieves lower error than PS-MAP on 7 out 8 datasets. Since the inputs are normalized and low-dimensional, we use p X = N (0 , I ) for L-MAP.

Image Classification. In Table 2, we compare L-MAP and PS-MAP on image classification using a ResNet-18 [9]. L-MAP achieves comparable or slightly better accuracies and is often better calibrated. We further test the effectiveness of L-MAP with transfer learning with a larger ResNet-50 trained on ImageNet. In Table 3, we show L-MAP also achieves small improvements in accuracy and calibration in transfer learning.

Loss Landscape. In Figure 6 (Left), we show that L-MAP indeed finds flatter minima. Further, we plot the Laplacian estimate in Figure 6 (Right) as the training progresses. We see that the Laplacian is much lower for L-MAP, showing its effectiveness at constraining the eigenvalues of J ( θ ; p X ) .

Distribution of Evaluation Points. In Table 2, we study the impact of the choice of distribution of evaluation points. Alongside our main choice of the evaluation set (KMNIST for FashionMNIST and CIFAR-100 for CIFAR-10), we use two additional distributions - the training set itself and a white noise N (0 , I ) distribution of the same dimensions as the training inputs. For both tasks, we find that using an external evaluation set beyond the empirical training distribution is often beneficial.

Number of Evaluation Point Samples. In Figure 7(a), we compare different Monte Carlo sample sizes S for estimating the Laplacian. Overall, L-MAP is not sensitive to this choice in terms of accuracy. However, calibration error [20] sometimes monotonically decreases with S .

2.5

0.0

2.5

0

2

4

Loss

×10 1

Method

L-MAP

PS-MAP

0

10 20 30 40 50

Epoch

0

2

4

×10 4

(a) FashionMNIST

2.5

0.0

2.5

0

2

4

Loss

×10 1

Method

L-MAP

PS-MAP

0

50

100 150 200

Epoch

0.0

2.5

5.0

7.5

×10 4

(b) CIFAR-10

Robustness to Label Noise. In Figure 7(b), we find L-MAP is slightly more robust to label noise than PS-MAP on CIFAR-10.

Main Takeaways. L-MAP shows qualitatively similar properties as FS-MAP such as favoring flat minima and often provides better calibration. However, it achieves comparable or only slightly better accuracy on more complex image classification tasks. In line with our expectations, these results suggest that accounting for the precise difference between FS-MAP and PS-MAP is less useful without a sufficiently well-motivated prior.

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 it is natural to assume that FS-MAP more closely approximates a Bayesian model average, we clearly demonstrate how there can be a significant discrepancy. Moreover, while PS-MAP is not invariant to re-parametrization, which can be seen as a fundamental pathology, we show FS-MAP has its own failure modes such as pathological optima, as well as practical challenges around scalability. In general, our results suggest the benefits of FS-MAP will be greatest when the prior is sufficiently well-motivated.

Our findings engage with and contribute to active discussions across the literature. For example, several works have argued conceptually-and found empirically-that solutions in flatter regions of the loss landscape correspond to better generalization [8, 10, 11, 12, 25]. On the other hand, Dinh et al. [5] argue that the ways we measure flatness, for example, through Hessian eigenvalues, are not parameterization invariant, making it possible to construct striking failure modes. Similarly, PS-MAP estimation is not parameterization invariant. However, our analysis and comparison to FS-MAP estimation raise the question to what extent lack of parametrization invariance is actually a significant practical shortcoming-after all, we are not reparametrizing our models on the fly, and we have evolved our models and training techniques conditioned on standard parameterizations.

10 1

10 2

S

)

1

2

3

4

Calibration Error

FashionMNIST

CIFAR-10

0.0

0.2

0.4

0.6

0.8

0

20

40

60

80

100

Method

L-MAP

PS-MAP

Acknowledgements. We thank Marc Finzi, Pavel Izmailov, and Micah Goldblum for helpful discussions. This work is supported by NSF CAREER IIS-2145492, NSF I-DISRE 193471, NSF IIS-1910266, BigHat Biosciences, Capital One, and an Amazon Research Award.

References

Arthur Asuncion and David Newman. Uci machine learning repository, 2007.

Christopher M. Bishop. Pattern recognition and machine learning (information science and statistics). 2006.

Pratik Chaudhari, Anna Choroma´ nska, 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.

Tarin Clanuwat, Mikel Bober-Irizar, Asanobu Kitamoto, Alex Lamb, Kazuaki Yamamoto, and David Ha. Deep learning for classical japanese literature, 2018.

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.

Christopher Terence John Dodson and Timothy Poston. Tensor geometry: the geometric viewpoint and its uses , volume 130. Springer Science & Business Media, 2013.

Ran El-Yaniv and Yair Wiener. On the foundations of noise-free selective classification. Journal of Machine Learning Research , 11(53):1605-1641, 2010.

Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur. Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations , 2020.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016 , pages 770-778. IEEE Computer Society, 2016. doi: 10 . 1109/CVPR . 2016 . 90.

Sepp Hochreiter and Jürgen Schmidhuber. Flat minima. Neural Computation , 9:1-42, 1997.

P Izmailov, AG Wilson, D Podoprikhin, D 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.

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.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Leo Klarner, Tim G. J. Rudner, Michael Reutlinger, Torsten Schindler, Garrett M. Morris, Charlotte 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.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Hao Li, Zheng Xu, Gavin Taylor, and Tom Goldstein. Visualizing the loss landscape of neural nets. In Neural Information Processing Systems , 2017.

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 in Neural Information Processing Systems , volume 31. Curran Associates, Inc., 2018.

Wesley J. Maddox, Gregory W. Benton, and Andrew Gordon Wilson. Rethinking parameter counting in deep models: Effective dimensionality revisited. CoRR , abs/2003.02139, 2020.

Kevin P. Murphy. Machine learning : a probabilistic perspective . MIT Press, Cambridge, Mass. [u.a.], 2013. ISBN 9780262018029 0262018020.

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.

Michael Peskin. An introduction to quantum field theory . CRC press, 2018.

Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do cifar-10 classifiers generalize to cifar-10? 2018.

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.

Tim G. J. Rudner, Freddie Bickford Smith, Qixuan Feng, Yee Whye Teh, and Yarin Gal. Continual Learning via Sequential Function-Space Variational Inference. In Proceedings of the 39th International Conference on Machine Learning , Proceedings of Machine Learning Research. PMLR, 2022.

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.

Robert Tibshirani. Regression shrinkage and selection via the lasso. Journal of the royal statistical society series b-methodological , 58:267-288, 1996.

Aad W Van der Vaart. Asymptotic statistics , volume 3. Cambridge university press, 2000.

Steven Weinberg. Gravitation and cosmology: principles and applications of the general theory of relativity. 1972.

Andrew Gordon Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, MariaFlorina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

David H. Wolpert. Bayesian backpropagation over i-o functions rather than weights. In J. Cowan, G. Tesauro, and J. Alspector, editors, Advances in Neural Information Processing Systems , volume 6. Morgan-Kaufmann, 1993.

Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.

Table of Contents

Appendix

Appendix A PS-MAP is not Invariant to Reparametrization

Parameter-space MAP estimation has conceptual and theoretical shortcomings stemming from the lack of reparameterization-invariance of MAP estimation. Suppose we reparameterize the model parameters θ with θ ′ = R ( θ ) where R is an invertible transformation. The prior density on the reparameterized parameters is obtained from the change-of-variable formula that states p ′ ( θ ′ ) = p ( θ ) | det -1 (d θ ′ / d θ ) | . In the language of differential geometry, the prior density p ( θ ) is not a scalar but a scalar density, a quantity that is not invariant under coordinate transformations [28]. The probabilistic model is fundamentally unchanged, as we are merely viewing the parameters in a different coordinate system, but the MAP objective is not invariant to reparameterization. Specifically, the reparameterized parameter-space MAP objective becomes

̸

̸

As a less extreme example, consider the reparameterization θ ′ = 1 /θ . AGaussian prior on the original parameters p ( θ ) ∝ exp ( -θ 2 / 2 ) translates to a prior p ′ ( θ ′ ) ∝ exp ( -1 / 2 θ ′ 2 ) /θ ′ 2 = θ 2 exp ( -θ 2 / 2 ) on the inverted parameters. Note the transformed prior, when mapped onto the original parameters θ , assigns a higher weight on non-zero values of θ due to the quadratic factor coming from the log determinant. We illustrate this effect with a simple experiment where a linear model with RBF features is trained with a Gaussian likelihood and a Gaussian prior N (0 , I ) 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, it favors non-zero weights in the inverted one, learning a less smooth function composed of a larger number of RBF bases.

Since L ′ MAP ( θ ′ ) and the original objective L MAP ( θ ) differ by a term which is non-constant in the parameters if T is non-linear, the maxima θ ′ MAP ˙ = arg max θ ′ L ′ MAP ( θ ′ ) and θ MAP ˙ = arg max θ L MAP ( θ ) under the two objectives will be different. Importantly, by "different" we don't just mean θ ′ MAP = θ MAP but θ ′ MAP = R ( θ MAP ) . That is, they are not simply each other viewed in a different coordinate system but actually represent different functions. More accurately, therefore, one should say MAP estimation is not equivariant (rather than invariant) under reparameterization. As a result, when using a parameter-space MAP estimate to make predictions at test time, the predictions can change dramatically simply depending on how the model is parameterized. In fact, one can reparameterize the model so that PS-MAP will return an arbitrary solution θ 0 irrespective of the observed data, by choosing a reparameterization θ ′ = R ( θ ) such that log | det -1 (d θ ′ / d θ ) | θ = θ 0 = -∞ . One such choice is θ ′ = ( θ -θ 0 ) 3 , where the exponent is taken element-wise.

0.0

0.2

0.4

0.6

0.8

1.0

x

5

0

5

10

y

PS-MAP

PS-MAP Reparam.

True function

Observations

(a) Learned functions

2.0

1.5

1.0

0.5

0.0

0.5

1.0

1.5

2.0

10 0

10 1

10 2

Count

PS-MAP

PS-MAP Reparam.

(b) Parameter distributions

Appendix B Mathematical Details & Derivations

B.1 Prior Distributions over Functions

Since the function evaluations f θ (ˆ x ) are related to the parameters θ through the map T : θ ↦→ f θ (ˆ x ) , the prior density p ( f θ (ˆ x )) is simply the push forward of the density p ( θ ) through T, which can be expressed without loss of generality using the Dirac delta function as

This generic expression as an integral over the parameter space is difficult to evaluate and thus not directly useful. By assuming that T is injective (one-to-one) modulo some global symmetries in parameter space (e.g., permutation of the neurons), which holds for many neural network architectures [2], Wolpert [30] established the simplification in Equation (6). Note that since the map T is not guaranteed to be surjective (onto), especially when there are more evaluation points than parameters ( | ˆ x | > P ) , p ( f θ (ˆ x )) does not generally have support everywhere and is to be interpreted as a surface density defined over the image of T [30]. If the model parameterization has continuous symmetries, such as scaling symmetries in MLPs with ReLU activations [5], then the log determinant in Equation (6) should be replaced by the log pseudo-determinant to account for the fact that the dimension of the image of T is less than P, the number of parameters.

B.2 The General Case

Here we present a more general result for which the result in Section 3.2 is a special case, without assuming the function output is univariate ( K = 1 ) or that the metric is diagonal. To simplify notations, we suppress all θ -dependence and write f k θ ( x i ) , the k -th output of the function at x i , as f k i for some ˆ x = { x i } i M =1 . We use Einstein notation to imply summation over repeated indices and leave everything in un-vectorized forms. In this more general case, the metric is given by d s 2 = g ikjℓ d f k i d f ℓ j , where i, j are input indices and k, ℓ are output indices. Denote the un-vectorized M -byK -byP Jacobian J ika ˙ = ∂ θ a f k i . The P -byP matrix J appearing inside the determinant in Equation (6) is now given by

In the limit of infinitely many evaluation points uniformly populating the domain X , the sum over i, j becomes a double integral 2

where we have re-written the metric g ikjℓ as g kℓ ( x i , x j ) where g : X × X → R K × K is the metric represented now as a matrix-valued function. After dividing by a constant Z, we can similarly write J ab as the expectation

̸

where p XX ′ CC ′ ( x, x ′ , c, c ′ ) = | g cc ′ ( x, x ′ ) | /Z for some normalization constant Z that only shifts the log determinant by a θ -independent constant. Note Equation (10) is recovered if the metric is diagonal, that is g c,c ′ ( x, x ′ ) = 0 only if c = c ′ and x = x ′ . We can further remove the assumption that g is constant (independent of f ), but at the cost of letting the normalization constant Z depend on f (or equivalently θ ), making it invalid to ignore Z during optimization. While more general than the result presented in Section 3.2, specifying a (possibly θ -dependent) non-diagonal metric adds a significant amount of complexity without being particularly well-motivated. Therefore we did not further investigate this more general objective.

B.3 Why Should a Constant Metric Affect FS-MAP?

One may object that a constant metric g should have no effect in the MAP estimate for f θ . However, the subtlety arises from the fact that p ( f θ ) does not have support everywhere but is, in fact, a surface density defined on the image of T (as defined in Appendix B.1), a curved manifold in function space which gives rise to locally preferred directions such that even a global linear transformation will change the density in a non-homogeneous way. If T were instead surjective, then in Section 3.2 J θ (ˆ x ) would be a square matrix and we could write log det ( J θ (ˆ x ) ⊤ gJ θ (ˆ x ) ) = log det ( J θ (ˆ x ) ⊤ J θ (ˆ x ) ) + log det( g ) and conclude that a constant metric indeed has no effect on the MAP estimate. Similarly, if the image of T is not curved, meaning J θ (ˆ x ) is constant in θ, then log det ( J θ (ˆ x ) ⊤ gJ θ (ˆ x ) ) would be constant in θ 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 ˆ x approaches infinity and the points cover the domain R P uniformly and densely. This technique, known as the continuum limit [21], is commonly used in physics to study continuous systems with infinite degrees of freedom, by first discretizing and then taking the limit as the discretization scale approaches zero, thereby recovering the behavior of the original system without directly dealing with intermediate (possibly ill-defined) infinite-dimensional quantities. Examples include lattice field theories where the continuous spacetime is discretized into a lattice, and numerical analysis where differential equations are discretized and solved on a mesh, where the solution often converges to the continuous solution as the mesh size goes to zero.

In a similar vein, we utilize this technique to sidestep direct engagement with the ill-defined quantity p ( f θ ) . It may be possible to assign a well-defined value to p ( f θ ) 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 J ( θ ; p X )

̸

Proof. Recall J ( θ ; p X ) = E p X [ J θ ( X ) ⊤ J θ ( X ) ] . Suppose J ( θ ; p X ) is singular. Then there is a vector v = 0 for which 0 = v ⊤ J ( θ ; p X ) v = E p X [ ∥ J θ ( X ) v ∥ 2 2 ] = ∫ X ∥ J θ ( X ) v ∥ 2 2 p X ( x )d x = 0 . Since the integrand is non-negative and continuous by assumption, it is zero everywhere. Therefore, we have ∑ P i =1 v i ∂ θ i f θ ( x ) = 0 for all x ∈ supp( p X ) for a non-zero v, showing { ∂ θ i f θ ( · ) } P i =1 are linearly dependent functions over the support of p X .

̸

Conversely, if { ∂ θ i f θ ( · ) } P i =1 are linearly dependent functions over the support of p X , then ∑ P i =1 v i ∂ θ i f θ ( x ) = J θ ( x ) v = 0 for all x ∈ supp( p X ) for some v = 0 . Hence J ( θ ; p X ) v = E p X [ J θ ( X ) ⊤ J θ ( X ) v ] = 0 , showing J ( θ ; p X ) is singular. ■

B.6 Architectural Symmetries Imply Singular J ( θ ; p X )

̸

Proof. Differentiating with respect to θ, we have J θ = ∂ θ f θ = ∂ θ f Sθ = J Sθ S for all θ . Suppose there exists θ ∗ that is invariant under S, Sθ ∗ = θ ∗ , then J θ ∗ = J θ ∗ S and J θ ∗ ( I -S ) = 0 . Since, by assumption S = I , we have shown that J θ ∗ has a non-trivial nullspace. Because the nullspace of J ( θ ∗ ; p X ) contains the nullspace of J θ ∗ , we conclude J ( θ ; p X ) is also singular. ■

As an example, consider a single hidden layer MLP f θ ( x ) = w ⊤ 2 σ ( w 1 x ) , θ = { w 1 ∈ R 2 , w 2 ∈ R 2 } , where we neglected the bias for simplicity. The permutation P 12 ⊕ P 12 is a symmetry of f θ , where

⊕ is the direct sum and P 12 is the transposition ( 0 1 1 0 ) . The nullspace of J θ ∗ contains the image of

with a basis { (1 , -1 , 0 , 0) , (0 , 0 , 1 , -1) } , for any θ ∗ of the form ( a, a, b, b ) . It's easy to check that perturbing the parameters in directions (1 , -1 , 0 , 0) and (0 , 0 , 1 , -1) leaves the function output unchanged, due to symmetry in the parameter values and the network topology.

Continuous symmetries, among other reasons, can also lead to singular J ( θ ; p X ) and are not covered by the above analysis, which is specific to discrete symmetries. One example is scaling symmetries in MLPs with ReLU activations [5]. However, these scaling symmetries cause J ( θ ; p X ) to be singular for all θ in a manner that does not introduce any pathological optimum. This becomes clear when noting that the space of functions implemented by an MLP with P parameters using ReLU activations lies in a space with dimension lower than P due to high redundancies in the parameterization. The resolution is to simply replace all occurrences of log det J ( θ ; p X ) with log |J ( θ ; p X ) | + , where | · | + represents the pseudo-determinant. This step is automatic if we optimize Equation (14) instead of Equation (11), where adding a small jitter will automatically compute the log pseudo-determinant (up to an additive constant) when J ( θ ; p X ) is singular. By contrast, symmetries such as permutation symmetries do create pathological optima because they only make the Jacobian singular at specific settings of θ, 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 ( f θ ′ | θ ) = δ ( f θ ′ -f θ ) } θ .

We can avoid the singularities if we instead choose a variational family where each member is a distribution over f localized around f θ for some θ . Namely, we consider the family Q ˙ = { q ( f θ ′ | θ ) } θ ∈ Θ parameterized by a mean-like parameter θ and a small but fixed entropy in function space H ( q ( f θ ′ | θ )) = E q ( f θ ′ | θ ) [ -log q ( f θ ′ | θ )] = h for all θ ∈ Θ . For convenience, we overload the notation and use q ( θ ′ | θ ) to denote the pullback of q ( f θ ′ | θ ) to parameter space. The resulting variational lower bound is

where in the last line we used the assumption that q ( θ ′ | θ ) is localized around θ to write the expectation of log p ( D| θ ) and log p ( θ ) as their values at θ ′ = θ plus O ( h ) corrections (which we will henceforth omit since by assumption h is tiny), because they vary smoothly as a function of θ . More care is required to deal with the remaining expectation, since when J ( θ ; p X ) is singular at θ , log det( J ( θ ; p X )) is infinite, but its expectation, E q ( θ ′ | θ ) [log det( J ( θ ′ ; p X ))] , must be finite (assuming p ( θ ) and p ( D| θ ) are finite), given that L VLB ( θ ; p X ) lower bounds the log marginal likelihood log p ( D ) = log ∫ p ( D| θ ) p ( θ )d θ < ∞ . 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 J ( θ ; p X ) , which can be approximated by adding a jitter ϵ when computing the log determinant. This effect is similar to applying a high-frequency cutoff to an image by convolving it with a localized Gaussian filter.

With this approximation, we arrive at a simple objective inspired by such a variational perspective, equivalent to adding a jitter inside the log determinant computation in the FS-MAP objective,

Note that while using a jitter is common in practice for numerical stability, it arises for an entirely different reason here.

B.8 Approximating the Expectation

Rewriting log det( J ( θ ′ ; p X )) in terms of the eigenvalues λ i ( θ ′ ) , we have

̸

Consider each term E q ( θ ′ | θ ) [log λ i ( θ ′ )] in isolation. Since the distribution q ( θ ′ | θ ) is highly localized by assumption, there are only two regimes. If λ i ( θ ) = 0 , then the variation of the log eigenvalue is small and the expectation is well approximated by log λ i ( θ ) . Otherwise, if λ i ( θ ) = 0 , then log λ i ( θ ′ ) has a singularity at θ ′ = θ and E q ( θ ′ | θ ) [log λ i ( θ ′ )] ≈ log λ i ( θ ) = -∞ is no longer a valid approximation since we've shown the expectation is finite. Therefore, in this case, the expectation must evaluate to a large but finite quantity, say log( ϵ ) for some small ϵ, where the exact value of ϵ depends on the details of q ( θ ′ | θ ) and λ i ( θ ′ ) . Consequently, taking the expectation can be approximated by adding a jitter ϵI to J ( θ ′ ; p X ) 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 J ( θ ; p X ) = E p X [ J θ ( X ) ⊤ J θ ( X ) ] and its log determinant is often intractable. Instead, we can use a simple Monte Carlo estimator to approximate the expectation inside the log determinant:

where we add a small amount of jitter to prevent the matrix inside the determinant from becoming singular when SK < P . Since 1 S ∑ M j =1 J θ ( x ( j ) ) ⊤ J θ ( x ( j ) ) = 1 S J θ (ˆ x S ) ⊤ J θ (ˆ x S ) where ˆ x S = { x ( j ) } S j =1 and J θ (ˆ x S ) is an SK -byP matrix, one can compute the determinant of 1 S J θ (ˆ x S ) ⊤ J θ (ˆ x S ) + ϵI through the product of squared singular values of J θ (ˆ x S ) / √ S in O (min( SKP 2 , S 2 K 2 P )) time (with zero singular values replaced with √ ϵ ), much faster than O ( P 3 ) for computing the original P -byP determinant if S is small enough, without ever storing the P × P matrix J ( θ ; p X ) .

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 a 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. In Figure 9, we compare the exact log determinant for p X = 1 M ∑ M i =1 δ x i with M = 1 , 600 , and its estimate with S Monte Carlo samples. Here { x i } are linearly spaced in the region [ -5 , 5] 2 . We observe that using S = 800 can almost perfectly approximate the exact log determinant. With S = 400 and S = 200 , the approximation underestimates the exact value because the number of zero singular values increases for 1 S J θ (ˆ x S ) ⊤ J θ (ˆ x 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-determinant with sufficient accuracy with simple Monte Carlo using a 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 is large enough compared to the eigenvalues of J ( θ ; p X )) , a first order approximation to log det( J ( θ ; p X ) + ϵI ) can be sufficiently accurate. Expanding to first order in ρ ˙ = max i λ i /ϵ, we have

where c ϵ = P log( ϵ ) is independent of θ . Defining d ( θ, θ ′ ) ˙ = E p X [ ∥ f θ ( X ) -f θ ′ ( X ) ∥ 2 ] , we have

4

3

2

1

Exact log det

×10 4

4.5

4.0

3.5

3.0

2.5

2.0

1.5

1.0

0.5

×10 4

(a)

S

= 800

4

3

2

1

Exact log det

×10 4

4.5

4.0

3.5

3.0

2.5

2.0

1.5

1.0

0.5

×10 4

(b)

S

= 400

5

4

3

2

1

Exact log det

×10 4

5

4

3

2

1

×10 4

(c)

S

= 200

which shows J ( θ ; p X )) = 1 2 ∇ 2 ψ d ( θ, θ + ψ ) ∣ ∣ ψ =0 . Therefore,

where ∆ is the Laplacian operator.

To estimate the Laplacian, consider the following expectation:

where we used

Therefore, we have

Combining Equation (B.10) and Equation (B.18), we have shown Equation (14) reduces to the more efficiently computable L-MAP objective for large enough ϵ and small enough β :

The entire objective is negated and divided by N to yield the loss function

where we have absorbed the 1 /N factor into the hyperparameter λ = 1 2 ϵN . Therefore, using L-MAP amounts to simply adding a regularization λR ( θ ; β 2 ) to standard regularized training (PS-MAP). β is a tolerance parameter for approximating the Laplacian and can simply be fixed to a small value such as 10 -3 .

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, we precompute the Φ matrix and reuse it throughout training. When p X is not Uniform( -1 , 1) , Φ does not have a simple form and we use 10,000 Monte Carlo samples to evaluate it.

C.2 Hyperparameters for UCI Regression

We use an MLP with 3 hidden layers, 256 units, and ReLU activations. We train it with the Adam optimizer for 10,000 steps with a 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 , 10 -4 , 10 -3 , 10 -2 , 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 a Gaussian prior, and the L-MAP hyperparameter λ in Equation (B.20) is tuned over the range [10 -1 , 10 -10 ] using randomized grid search. For PS-MAP, λ is set to 0. The parameter variance in the Laplacian estimator is fixed to β 2 = 10 -6 . In addition, we clip the gradients to unit norm. We use a learning rate of 0 . 1 with SGD and a cosine decay schedule over 50 epochs for FashionMNIST and 200 epochs for CIFAR-10. The mini-batch size is fixed to 128.

C.4 Verifying the L-MAP Approximation

To verify that the L-MAP objective is a good approximation to Equation B.7 when ϵ is large enough compared to the eigenvalues of J ( θ ; p X ) , we compare log det( J ( θ ; p X ) + ϵI ) -P log( ϵ ) with the Laplacian estimate 1 2 ϵβ 2 E ψ ∼N (0 ,β 2 ) [ d ( θ, θ + ψ )] used by L-MAP in Appendix B.10 as its first order approximation in max i λ i ( θ ) /ϵ . We reuse the same network architecture, evaluation distribution p X , and sampling procedure for the parameters from Appendix B.9 to perform the comparison in Figure 10 and color each point by ¯ λ ( θ ) /ϵ , the average eigenvalue of J ( θ ; p X ) divided by ϵ. The smaller this value, the better the approximation should be. We observe that indeed as ϵ increases and ¯ λ ( θ ) /ϵ approaches 0, the Laplacian estimate becomes a more accurate approximation of the log determinant. Interestingly, even when ϵ ≪ ¯ λ ( θ ) and the approximation is not accurate, there still appears to be a 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.

10 2

10 4

10 6

10 2

10 4

10 6

Laplacian estimate

1

2

3

4

/

×10 4

(a)

ϵ

= 10

-

3

10 1

10 3

10 0

10 1

10 2

10 3

10 4

Laplacian estimate

1

2

3

4

/

×10 1

(b)

ϵ

= 1

10 0

10 1

10 2

10 3

10 0

10 1

10 2

10 3

Laplacian estimate

1

2

3

4

/

(c)

ϵ

= 10

10

2

10 0

10

3

10

2

10

1

10 0

10 1

Laplacian estimate

1

2

3

4

/

×10

2

(d)

ϵ

= 10

3

C.5 Effective Eigenvalue Regularization using the Laplacian Regularizer

The Laplacian regularizer in Equation (B.19) only uses a samplebased estimator that is unbiased only in the limit that σ → 0 . To keep the computational overhead at a minimum so that L-MAP can scale to large neural networks, we only take 1 sample of ψ per gradient step. To test the effectiveness of L-MAP in regularizing the eigenvalues of J ( θ ; p X ) under this practical setting, we train an MLP with 2 hidden layers, 16 units, and tanh activations on the Two Moons dataset (generated with sklearn.datasets.make_moons(n_samples=200, shuffle=True, noise=0.2, random_state=0) ) for 10 4 steps with the Adam optimizer and a learning rate of 0 . 001 . Here we

10

-

27

10

-

20

10

-

13

10

-

6

10

1

0

25

50

75

100

125

Sum Eigenvalue

L-MAP

FS-MAP

PS-MAP

/epsilon1

choose p X = 1 M ∑ M i =1 δ x i with M = 1 , 600 and { x i } linearly spaced in the region [ -5 , 5] 2 . In Figure 11, we compare the sum of eigenvalues of J ( θ ; p X ) for L-MAP and FS-MAP with different levels of jitter ϵ. Both FS-MAP and L-MAP significantly reduce the eigenvalues compared to PS-MAP with small values of ϵ, 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 a near-zero function that severely underfits the data to one that fits the data perfectly.

-

4

-

2

0

2

4

-

4

-

2

0

2

4

0

.

4960

0

.

4968

0

.

4976

0

.

4984

0

.

4992

0

.

5000

0

.

5008

0

.

5016

0

.

5024

0

.

5032

E

[

y

|D

;

x

]

(a)

ϵ

= 10

-

2

-

4

-

2

0

2

4

-

4

-

2

0

2

4

0

.

08

0

.

20

0

.

32

0

.

44

0

.

56

0

.

68

0

.

80

0

.

92

E

[

y

|D

;

x

]

(b)

ϵ

= 10

-

1

-

4

-

2

0

2

4

-

4

-

2

0

2

4

0

.

00

0

.

15

0

.

30

0

.

45

0

.

60

0

.

75

0

.

90

E

[

y

|D

;

x

]

(c)

ϵ

= 1

-

4

-

2

0

2

4

-

4

-

2

0

2

4

0

.

00

0

.

15

0

.

30

0

.

45

0

.

60

0

.

75

0

.

90

E

[

y

|D

;

x

]

(d)

ϵ

= 10

C.7 Further Neural Network Experiments

In Table 4, we report the detailed results for varying the number of Monte Carlo samples S 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 -likelihood. However, we find that increasing S can sometimes lead to significant improvement in calibration, as seen for FashionMNIST [31].

Distribution Shift. We additionally evaluate our L-MAP trained models to 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 in Table 5, L-MAP tends to improve calibration, while retaining performance in terms of accuracy.

Transfer Learning. Using ResNet-50 trained on ImageNet, we test the effectiveness of L-MAP with transfer learning. We choose hyperparameters as in Appendix C.3, except a lower learning rate of 10 -3 and a smaller batch size of 64 due to computational constraints. Table 3 reports all the results.

Peer reviews:
comments: 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.

STRENGTHS:
- Thorough analysis of the pros and cons of parameter space MAP estimation and function space MAP estimation
- Propose a scalable approximation for large neural networks which extends the application area of function space MAP estimation
- 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
- 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

QUESTIONS:
See weakness above

LIMITATIONS:
See weakness above
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-22

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\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. (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.\n\nLIMITATIONS:\nThe authors appropriately discuss the limitations of their work within the submission.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe paper clearly discusses the limitations.\n\nThere is no direct societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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).\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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).\n\nWEAKNESSES:\n1. 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.)\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}