SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shikai Qiu*
Tim G. J. Rudner*
Sanyam Kapoor*
Andrew Gordon Wilson
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.
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  $\ell_1$  or  $\ell_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_{\theta}(x)$  are functions given parameters  $\theta$ , 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  $\theta_{R}$  in a Gaussian mixture regression model in both parameter and function space, using the same Gaussian prior for the mixture coefficients (corresponding to  $\ell_2$  regularization in parameter space). We see that each distribution is maximized by a different parameter  $\theta_{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.
(a) Distinct Posterior Densities
(b) Distinct Learned Functions
Figure 1: Illustration of the Difference Between Most Likely Functions and Parameters. The function that is most probable (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 a regression model with two parameters,  $\theta_L$  and  $\theta_R$ , both with a prior  $\mathcal{N}(0,1.2^2)$ . This model is used to learn a mixture of two Gaussians with fixed mean and variance, where the mixture coefficients are given by  $\exp (\theta_L)$  and  $\exp (\theta_R)$ . Both the most probable solution and the posterior density show significant differences when analyzed in function-space versus parameter-space. Since  $\theta_L$  is well-determined, we only plot the posterior as a function of  $\theta_R$ . We normalize the area under  $p(f_\theta|\mathcal{D})$  to 1.
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.
We consider supervised learning problems with a training dataset  $\mathcal{D} = (x_{\mathcal{D}},y_{\mathcal{D}}) = \{x_{\mathcal{D}}^{(i)},y_{\mathcal{D}}^{(i)}\}_{i = 1}^{N}$  of inputs  $x\in \mathcal{X}$  and targets  $y\in \mathcal{V}$  with input space  $\mathcal{X}\subseteq \mathbb{R}^D$  and output space  $\mathcal{V}\subseteq \mathbb{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  $\mathcal{D}$ , one typically starts by specifying a function  $f_{\theta}:\mathcal{X}\to \mathcal{Y}$  parametrized by a vector  $\theta \in \mathbb{R}^P$  which maps inputs to outputs.  $f_{\theta}$  could be a neural network, a polynomial model, a Fourier series, and so on. Learning typically amounts to estimating model parameters  $\theta$ . To this end, we can relate the function  $f_{\theta}$  to the targets through an observation model,  $p(y|x,f_{\theta})$ . For example, in regression, we could assume the outputs  $y = f_{\theta}(x) + \epsilon$ , where  $\epsilon \sim \mathcal{N}(0,\sigma^2)$  is additive Gaussian noise with variance  $\sigma^2$ . Equivalently,  $p(y|x,\theta) = \mathcal{N}(y|f_{\theta}(x),\sigma^2)$ . Alternatively, in classification, we could specify  $p(y|x,\theta) = \mathrm{Categorical}(\mathrm{softmax}(f_{\theta}(x)))$ . We then use this observation model to form a likelihood over the whole dataset  $p(y_{\mathcal{D}}|x_{\mathcal{D}},\theta)$ . In each of these example observation models, the likelihood factorizes, as the data points are conditionally independent given model parameters  $\theta$ ,  $p(y_{\mathcal{D}}|x_{\mathcal{D}},\theta) = \prod_{i = 1}^{N}p(y_{\mathcal{D}}^{(i)}|x_{\mathcal{D}}^{(i)},\theta)$ . We can further express a belief over values of parameters through a prior  $p(\theta)$ , such as  $p(\theta) = \mathcal{N}(\theta|\mu ,\Sigma)$ . Finally, using Bayes rule, the log of the posterior up to a constant  $c$  independent of  $\theta$  is,
$$
\log \overbrace {p (\theta | y _ {\mathcal {D}} , x _ {\mathcal {D}})} ^ {\text {p a r a m e t e r p o s t e r i o r}} = \overbrace {\log p (y _ {\mathcal {D}} | x _ {\mathcal {D}} , \theta)} ^ {\text {l i k e l i h o o d}} + \overbrace {\log p (\theta)} ^ {\text {p r i o r}} + c. \tag {1}
$$
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  $\ell_2$  regularization, also known as weight-decay. If we use a Laplace prior, we recover  $\ell_1$  regularization, also known as LASSO [26].
Once we maximize the posterior to find
$$
\hat {\theta} ^ {\mathrm {P S - M A P}} = \arg \max  _ {\theta} p (\theta | y _ {\mathcal {D}}, x _ {\mathcal {D}}), \tag {2}
$$
we can condition on these parameters to form our function  $f_{\hat{\theta}^{\mathrm{PS - MAP}}}$  to make predictions. However, as we saw in Figure 1,  $f_{\hat{\theta}^{\mathrm{PS - MAP}}}$  is not in general the function that maximizes the posterior over functions,  $p(f_{\theta}|y_{\mathcal{D}},x_{\mathcal{D}})$ . In other words, if
$$
\hat {\theta} ^ {\mathrm {F S - M A P}} = \arg \max  _ {\theta} p \left(f _ {\theta} \mid y _ {\mathcal {D}}, x _ {\mathcal {D}}\right) \tag {3}
$$
then generally  $f_{\hat{\theta}^{\mathrm{PS - MAP}}} \neq f_{\hat{\theta}^{\mathrm{FS - MAP}}}$ . Naively, one can write the log posterior over  $f_{\theta}$  up to the same constant  $c$  as above as
$$
\overbrace {\log p \left(f _ {\theta} \mid y _ {\mathcal {D}} , x _ {\mathcal {D}}\right)} ^ {\text {f u n c t i o n p o s t e r i o r}} = \overbrace {\log p \left(y _ {\mathcal {D}} \mid x _ {\mathcal {D}} , f _ {\theta}\right)} ^ {\text {l i k e l i h o o d}} + \overbrace {\log p \left(f _ {\theta}\right)} ^ {\text {f u n c t i o n p r i o r}} + c, \tag {4}
$$
where  $p(y_{\mathcal{D}}|x_{\mathcal{D}},f_{\theta}) = p(y_{\mathcal{D}}|x_{\mathcal{D}},\theta)$ , but just written in terms of the function  $f_{\theta}$ . The prior  $p(f_{\theta})$  however is a different function from  $p(\theta)$ , 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_{\theta})$  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 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 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_{\theta}(\hat{x})$ , the function evaluated at a finite set of points  $\hat{x} = \{x_1, \dots, x_M\}$ , where  $M < \infty$  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
$$
\mathcal {L} _ {\text {f i n i t e}} (\theta ; \hat {x}) = \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta} (\hat {x})\right) + \log p \left(f _ {\theta} (\hat {x})\right), \tag {5}
$$
where  $p(f_{\theta}(\hat{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  $\theta$ , and  $K$  the number of function outputs. Assuming the set of evaluation points is sufficiently large so that  $MK \geq 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(\hat{x})$  is given by
$$
p \left(f _ {\theta} (\hat {x})\right) = p (\theta) \det ^ {- 1 / 2} \left(\mathcal {J} (\theta ; \hat {x})\right), \tag {6}
$$
where  $\mathcal{J}(\theta ;\hat{x})$  is a  $P$  -by-  $P$  matrix defined by
$$
\mathcal {J} _ {\text {f i n i t e}} (\theta ; \hat {x}) \doteq J _ {\theta} (\hat {x}) ^ {\top} J _ {\theta} (\hat {x}) \tag {7}
$$
and  $J_{\theta}(\hat{x}) \doteq \partial f_{\theta}(\hat{x}) / \partial \theta$  is the  $MK$ -by- $P$  Jacobian of  $f_{\theta}(\hat{x})$ , viewed as an  $MK$ -dimensional vector, with respect to the parameters  $\theta$ . Substituting Equation (6) into Equation (5) the function-space MAP objective as a function of the parameters  $\theta$  can be expressed as
$$
\mathcal {L} _ {\text {f i n i t e}} (\theta ; \hat {x}) = \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta}\right) + \log p (\theta) - \frac {1}{2} \log \det  (\mathcal {J} (\theta ; \hat {x})), \tag {8}
$$
where we are allowed to condition directly on  $f_{\theta}$  instead of  $f_{\theta}(\hat{x})$  because by assumption  $\hat{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  $\hat{x}$ , this objective makes it unclear how we should select  $\hat{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 \in \mathbb{R}^n$  to  $u \in \mathbb{R}^m$  via an injective map  $u = \varphi(v)$ , their probability densities relate as  $p(u)\mathrm{d}\mu(u) = p(v)\mathrm{d}\nu(v)$ , where  $\mu$  and  $\nu$  define respective volume measures. Suppose we let  $\mathrm{d}\mu(u) = \sqrt{\operatorname*{det}(g(u))}\mathrm{d}^m u$ , the volume induced by an  $M \times M$  metric tensor  $g$ , and  $\mathrm{d}\nu(v) = \mathrm{d}^n v$ , the Lebesgue measure, then we can write  $\mathrm{d}\mu(u) = \sqrt{\operatorname*{det}(\tilde{g}(v))}\mathrm{d}\nu(v)$ , where the  $N \times N$  metric  $\tilde{g}(v) = J(v)^{\top}g(u)J(v)$  is known as the pullback of  $g$  via  $\varphi$  and  $J$  is the Jacobian of  $\varphi$  [6]. As a result, we have  $p(u) = p(v)\mathrm{det}^{-1/2}(J(v)^{\top}g(u)J(v))$ . Applying this argument and identifying  $u$  with  $f_{\theta}(\hat{x})$  and  $v$  with  $\theta$ , Wolpert [30] thereby establishes  $p(f_{\theta}(\hat{x})) = p(\theta)\mathrm{det}^{-1/2}(J_{\theta}(\hat{x})^{\top}J_{\theta}(\hat{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_{\theta}(\hat{x})$  and  $f_{\theta}(\hat{x}) + \mathrm{d}f_{\theta}(\hat{x})$  is  $\mathrm{ds}^2 = \sum_{i=1}^{M} \mathrm{d}f_{\theta}(x_i)^2$ , rather than the general case  $\mathrm{ds}^2 = \sum_{i=1}^{M} \sum_{j=1}^{M} g_{ij} \mathrm{d}f_{\theta}(x_i) \mathrm{d}f_{\theta}(x_j)$ . To account for a generic metric  $g$ , we therefore replace  $J_{\theta}(\hat{x})^\top J_{\theta}(\hat{x})$  with  $J_{\theta}(\hat{x})^\top g J_{\theta}(\hat{x})$ . For simplicity, we assume the function output is univariate ( $K = 1$ ) and only consider  $g$  that is constant i.e., independent of  $f_{\theta}(\hat{x})$  and diagonal, with  $g_{ii} = g(x_i)$  for some function  $g: \mathcal{X} \to \mathbb{R}^+$ . For a discussion of the general case, see Appendix B.2. To better interpret the choice of  $g$ , we rewrite
$$
J _ {\theta} (\hat {x}) ^ {\top} g J _ {\theta} (\hat {x}) = \sum_ {j = 1} ^ {M} \tilde {p} _ {X} \left(x _ {j}\right) J _ {\theta} \left(x _ {j}\right) ^ {\top} J _ {\theta} \left(x _ {j}\right), \tag {9}
$$
where we suggestively defined the alias  $\tilde{p}_X(x) \doteq g(x)$  and  $J_{\theta}(x_j)$  is the  $K$ -by- $P$ -dimensional Jacobian evaluated at the point  $x_j$ . Generalizing from the finite and discrete set  $\hat{x}$  to the possibly infinite entire domain  $\mathcal{X} \subseteq \mathbb{R}^D$  and further dividing by an unimportant normalization constant  $Z$ , we obtain
$$
\mathcal {J} (\theta ; p _ {X}) \doteq \frac {1}{Z} \int_ {\mathcal {X}} \tilde {p} _ {X} (x) J _ {\theta} (x) ^ {\top} J _ {\theta} (x) d x = \mathbb {E} _ {p _ {X}} \left[ J _ {\theta} (X) ^ {\top} J _ {\theta} (X) \right], \tag {10}
$$
where  $p_X = \tilde{p}_X / Z$  is a normalized probability density function, with normalization  $Z$  —which is independent of  $\theta$  —only appearing as an additive constant in  $\log \det \mathcal{J}(\theta; p_X)$ . We include a further discussion about this limit in Appendix B.4.
Under this limit, we can identify  $\log p(\theta) - \frac{1}{2}\log \det \mathcal{J}(\theta ;p_X)$  with  $\log p(f_{\theta})$  up to a constant, and thereby express  $\log p(f_{\theta}|\mathcal{D})$ , the function-space MAP objective in Equation (4) as
$$
\mathcal {L} (\theta ; p _ {X}) = \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta}\right) + \log p (\theta) - \frac {1}{2} \log \det  \mathcal {J} (\theta ; p _ {X}), \tag {11}
$$
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_\theta$  is an image classifier for which we only care about its outputs on set of natural images  $\mathcal{I}$  when comparing it to another image classifier. The metric  $g(\cdot)$  and therefore the evaluation distribution  $p_X$  only needs support in  $\mathcal{I}$ , and the FS-MAP objective is defined only in terms of the Jacobians evaluated at natural images  $x \in \mathcal{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  $\hat{p}_X(x) \doteq \frac{1}{M} \sum_{x' \in \hat{x}} \delta(x - x')$ , where  $\delta$  is the Dirac delta function and  $\hat{x} = \{x_1, \dots, x_M\}$  with  $M < \infty$ , as before. It is easy to see that  $\mathcal{J}_{\mathrm{finite}}(\theta; \hat{x}) \propto \mathcal{J}(\theta; \hat{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_{\theta}(x) = \sum_{i=1}^{P} \sigma(\theta_i) \varphi_i(x)$  with domain  $\mathcal{X} = [-1,1]$ , where  $\{\varphi_i\}_{i=1}^P$  is a fixed set of basis functions and  $\sigma$  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  $\mathcal{J}(\theta; p_X)$  has a simple closed form, such that
$$
\mathcal {J} _ {i j} (\theta ; p _ {X}) = \mathbb {E} _ {p _ {X}} \left[ \partial_ {\theta_ {i}} f _ {\theta} (X) \partial_ {\theta_ {j}} f _ {\theta} (X) \right] = \sigma^ {\prime} \left(\theta_ {i}\right) \sigma^ {\prime} \left(\theta_ {j}\right) \Phi_ {i j}, \tag {12}
$$
where  $\Phi$  is a constant matrix with elements  $\Phi_{ij} = \mathbb{E}_{p_X}[\varphi_i(X)\varphi_j(X)]$ . Therefore,  $\Phi$  can be precomputed once and reused throughout training. In this experiment, we use the set of Fourier features  $\{\cos (k_i\cdot),\sin (k_i\cdot)\}_{i = 1}^{100}$  where  $k_{i} = i\pi$  and set  $\sigma = \tanh$ . We generate training data by sampling  $x_{\mathrm{train}}\sim \mathrm{Uniform}(-1,1)$ ,  $\theta_{i}\sim \mathcal{N}(0,\alpha^{2})$  with  $\alpha = 10$ , evaluating  $f_{\theta}(x_{\mathrm{train}})$ , and adding Gaussian noise with standard deviation  $\sigma^{*} = 0.1$  to each observation. We use 1,000 test points sampled from  $\mathrm{Uniform}(-1,1)$ . To train the model, we set the prior  $p(\theta)$  and the likelihood to correspond to the data-generating process. For FS-MAP, we specify  $p_X = \mathrm{Uniform}(-1,1)$ , the ground-truth distribution of test inputs, which conveniently results in  $\Phi = I / 2$ .
FS-MAP Finds More Probable Functions. Figure 2a shows the improvement in  $\log p(f_{\theta}|\mathcal{D})$  when using FS-MAP over PS-MAP. As expected, FS-MAP consistently finds functions with much higher posterior probability  $p(f_{\theta}|\mathcal{D})$ .
(a) Log Posterior Increase
(b) Loss Curvature
Figure 2: FS-MAP exhibits desirable properties. On a 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.
(c) Test RMSE
(d) Train RMSE
(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 probabilistic model.