SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rewriting  $\log \det (\mathcal{J}(\theta^{\prime};p_X))$  in terms of the eigenvalues  $\lambda_{i}(\theta^{\prime})$ , we have
$$
\mathbb {E} _ {q \left(\theta^ {\prime} \mid \theta\right)} \left[ \log \det  \left(\mathcal {J} \left(\theta^ {\prime}; p _ {X}\right)\right) \right] = \sum_ {i = 1} ^ {P} \mathbb {E} _ {q \left(\theta^ {\prime} \mid \theta\right)} \left[ \log \lambda_ {i} \left(\theta^ {\prime}\right) \right]. \tag {B.8}
$$
Consider each term  $\mathbb{E}_{q(\theta'|\theta)}[\log \lambda_i(\theta')]$  in isolation. Since the distribution  $q(\theta'|\theta)$  is highly localized by assumption, there are only two regimes. If  $\lambda_i(\theta) \neq 0$ , then the variation of the log eigenvalue is small and the expectation is well approximated by  $\log \lambda_i(\theta)$ . Otherwise, if  $\lambda_i(\theta) = 0$ , then  $\log \lambda_i(\theta')$  has a singularity at  $\theta' = \theta$  and  $\mathbb{E}_{q(\theta'|\theta)}[\log \lambda_i(\theta')] \approx \log \lambda_i(\theta) = -\infty$  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(\epsilon)$  for some small  $\epsilon$ , where the exact value of  $\epsilon$  depends on the details of  $q(\theta'|\theta)$  and  $\lambda_i(\theta')$ . Consequently, taking the expectation can be approximated by adding a jitter  $\epsilon I$  to  $\mathcal{J}(\theta'; p_X)$  when computing the log determinant to clip its eigenvalues from below at some threshold  $\epsilon > 0$ .
B.9 Monte Carlo Estimator for the Log Determinant
Monte Carlo Estimator. Exactly computing  $\mathcal{J}(\theta; p_X) = \mathbb{E}_{p_X} [J_\theta(X)^\top J_\theta(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:
$$
\log \det  \left(\mathbb {E} _ {p _ {X}} \left[ J _ {\theta} (X) ^ {\top} J _ {\theta} (X) \right]\right) \approx \log \det  \left(\frac {1}{S} \sum_ {j = 1} ^ {S} J _ {\theta} \left(x ^ {(j)}\right) ^ {\top} J _ {\theta} \left(x ^ {(j)}\right) + \epsilon I\right), \tag {B.9}
$$
where we add a small amount of jitter to prevent the matrix inside the determinant from becoming singular when  $SK < P$ . Since  $\frac{1}{S} \sum_{j=1}^{M} J_{\theta}(x^{(j)})^{\top} J_{\theta}(x^{(j)}) = \frac{1}{S} J_{\theta}(\hat{x}_S)^{\top} J_{\theta}(\hat{x}_S)$  where  $\hat{x}_S = \{x^{(j)}\}_{j=1}^S$  and  $J_{\theta}(\hat{x}_S)$  is an  $SK$ -by-  $P$  matrix, one can compute the determinant of  $\frac{1}{S} J_{\theta}(\hat{x}_S)^{\top} J_{\theta}(\hat{x}_S) + \epsilon I$  through the product of squared singular values of  $J_{\theta}(\hat{x}_S) / \sqrt{S}$  in  $\mathcal{O}(\min(SKP^2, S^2K^2P))$  time (with zero singular values replaced with  $\sqrt{\epsilon}$ ), much faster than  $\mathcal{O}(P^3)$  for computing the original  $P$ -by-  $P$  determinant if  $S$  is small enough, without ever storing the  $P \times P$  matrix  $\mathcal{J}(\theta; 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} = \frac{1}{M}\sum_{i=1}^{M}\delta_{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  $\frac{1}{S}J_{\theta}(\hat{x}_S)^{\top}J_{\theta}(\hat{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  $\epsilon$  that is large enough compared to the eigenvalues of  $\mathcal{J}(\theta; p_X)$ , a first order approximation to  $\log \det(\mathcal{J}(\theta; p_X) + \epsilon I)$  can be sufficiently accurate. Expanding to first order in  $\rho \doteq \max_i \lambda_i / \epsilon$ , we have
$$
\log \det  (\mathcal {J} (\theta ; p _ {X}) + \epsilon I) = \frac {1}{\epsilon} \sum_ {i = 1} ^ {P} \lambda_ {i} + c _ {\epsilon} + \mathcal {O} \left(\rho^ {2}\right) = \frac {1}{\epsilon} \operatorname {T r} (\mathcal {J} (\theta ; p _ {X})) + c _ {\epsilon} + \mathcal {O} \left(\rho^ {2}\right), \tag {B.10}
$$
where  $c_{\epsilon} = P\log (\epsilon)$  is independent of  $\theta$ . Defining  $d(\theta ,\theta^{\prime})\doteq \mathbb{E}_{p_X}\left[\| f_\theta (X) - f_{\theta^{\prime}}(X)\| ^2\right]$ , we have
$$
d (\theta , \theta + \psi) = \mathbb {E} _ {p _ {X}} \left[ \| J _ {\theta} (X) \psi \| ^ {2} \right] + \mathcal {O} \left(\psi^ {4}\right) = \psi^ {\top} \mathcal {J} (\theta ; p _ {X})) \psi + \mathcal {O} \left(\psi^ {4}\right), \tag {B.11}
$$
(a)  $S = 800$
(b)  $S = 400$
Figure 9: Approximating the log determinant via an  $S$ -sample simple Monte Carlo estimator and jitter. Each dot represents the exact and estimated log determinant evaluated at random parameters sampled by randomly initializing the network followed by scaling by a factor  $s \sim \mathrm{loguniform}(0.1, 10)$  to include parameters of different magnitudes. The dashed line shows  $x = y$ .
(c)  $S = 200$
which shows  $\mathcal{J}(\theta ;p_X)) = \frac{1}{2}\nabla_{\psi}^2 d(\theta ,\theta +\psi)\big|_{\psi = 0}$ . Therefore,
$$
\left. \operatorname {T r} (\mathcal {J} (\theta ; p _ {X})) = \frac {1}{2} \operatorname {T r} \left(\nabla_ {\psi} ^ {2} d (\theta , \theta + \psi)\right) \mid_ {\psi = 0} = \frac {1}{2} \Delta_ {\psi} d (\theta , \theta + \psi)) \right| _ {\psi = 0}, \tag {B.12}
$$
where  $\Delta$  is the Laplacian operator.
To estimate the Laplacian, consider the following expectation:
$$
\mathbb {E} _ {\psi \sim \mathcal {N} \left(0, \beta^ {2} I\right)} [ d (\theta , \theta + \psi) ] \tag {B.13}
$$
$$
= \mathbb {E} _ {\psi \sim \mathcal {N} (0, \beta^ {2} I)} \left[ \psi^ {\top} \mathcal {J} \left(\theta ; p _ {X}\right)) \psi + \mathcal {O} \left(\psi^ {4}\right) \right] \tag {B.14}
$$
$$
= \mathbb {E} _ {\psi \sim \mathcal {N} (0, \beta^ {2} I)} [ \psi^ {\top} \mathcal {J} (\theta ; p _ {X})) \psi ] + \mathcal {O} (\beta^ {4}) \tag {B.15}
$$
$$
= \beta^ {2} \operatorname {T r} \left(\mathcal {J} \left(\theta ; p _ {X}\right)\right) + \mathcal {O} \left(\beta^ {4}\right), \tag {B.16}
$$
$$
\mathbb {E} _ {\psi} \left[ \psi^ {\top} \mathcal {J} \left(\theta ; p _ {X}\right)\right) \psi ] = \sum_ {i j} \mathbb {E} _ {\psi} \left[ \psi_ {i} \psi_ {j} \right] \mathcal {J} _ {i j} (\theta ; p _ {X}) = \beta^ {2} \sum_ {i j} \delta_ {i j} \mathcal {J} _ {i j} (\theta ; p _ {X}) = \beta^ {2} \operatorname {T r} \left(\mathcal {J} \left(\theta ; p _ {X}\right)\right). \tag {B.17}
$$
$$
\left. \operatorname {T r} (\mathcal {J} (\theta ; p _ {X})) = \frac {1}{2} \Delta_ {\psi} d (\theta , \theta + \psi)) \right| _ {\psi = 0} = \frac {1}{\beta^ {2}} \mathbb {E} _ {\psi \sim \mathcal {N} (0, \beta^ {2} I)} [ d (\theta , \theta + \psi) ] + \mathcal {O} \left(\beta^ {2}\right). \tag {B.18}
$$
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  $\epsilon$  and small enough  $\beta$ :
$$
\mathcal {L} _ {\mathrm {L - M A P}} (\theta ; p _ {X}) \doteq \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta}\right) + \log p (\theta) - \frac {1}{2 \epsilon \beta^ {2}} \mathbb {E} _ {\psi \sim \mathcal {N} (0, \beta^ {2})} [ d (\theta , \theta + \psi) ]. \tag {B.19}
$$
The entire objective is negated and divided by  $N$  to yield the loss function
$$
L _ {\mathrm {L - M A P}} (\theta ; p _ {X}) \doteq \underbrace {- \frac {1}{N} \sum_ {i = 1} ^ {N} \log p \left(y _ {\mathcal {D}} ^ {(i)} \mid x _ {\mathcal {D}} ^ {(i)}, f _ {\theta}\right) - \frac {1}{N} \log p (\theta)} _ {\text {S t a n d a r d r e g u l a r i z e d l o s s}} + \lambda \underbrace {\left(\frac {1}{\beta^ {2}} \mathbb {E} _ {\psi \sim \mathcal {N} (0 , \beta^ {2})} [ d (\theta , \theta + \psi) ]\right)} _ {\text {L a p l a c i a n r e g u l a r i z a t i o n R (\theta ; \beta)}}, \tag {B.20}
$$
where we have absorbed the  $1 / N$  factor into the hyperparameter  $\lambda = \frac{1}{2\epsilon N}$ . Therefore, using L-MAP amounts to simply adding a regularization  $\lambda R(\theta; \beta^2)$  to standard regularized training (PS-MAP).  $\beta$  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  $\Phi$  matrix and reuse it throughout training. When  $p_{X}$  is not  $\mathrm{Uniform}(-1,1)$ ,  $\Phi$  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  $\lambda$  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  $\lambda$  in Equation (B.20) is tuned over the range  $[10^{-1}, 10^{-10}]$  using randomized grid search. For PS-MAP,  $\lambda$  is set to 0. The parameter variance in the Laplacian estimator is fixed to  $\beta^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  $\epsilon$  is large enough compared to the eigenvalues of  $\mathcal{J}(\theta; p_X)$ , we compare  $\log \det(\mathcal{J}(\theta; p_X) + \epsilon I) - P\log(\epsilon)$  with the Laplacian estimate  $\frac{1}{2\epsilon\beta^2}\mathbb{E}_{\psi \sim \mathcal{N}(0,\beta^2)}[d(\theta, \theta + \psi)]$  used by L-MAP in Appendix B.10 as its first order approximation in  $\max_i \lambda_i(\theta)/\epsilon$ . 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  $\bar{\lambda}(\theta)/\epsilon$ , the average eigenvalue of  $\mathcal{J}(\theta; p_X)$  divided by  $\epsilon$ . The smaller this value, the better the approximation should be. We observe that indeed as  $\epsilon$  increases and  $\bar{\lambda}(\theta)/\epsilon$  approaches 0, the Laplacian estimate becomes a more accurate approximation of the log determinant. Interestingly, even when  $\epsilon \ll \bar{\lambda}(\theta)$  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.
(a)  $\epsilon = 10^{-3}$
(b)  $\epsilon = 1$
Figure 10: Log determinant v.s. the Laplacian estimate at different values of  $\epsilon$ . The  $x$ -axis shows  $\log \det(\mathcal{J}(\theta; p_X) + \epsilon I) - P\log(\epsilon)$  while the  $y$ -axis shows its first order approximation  $\frac{1}{2\epsilon\beta^2}\mathbb{E}_{\psi \sim \mathcal{N}(0,\beta^2)}[d(\theta, \theta + \psi)]$ , estimated with 10 samples of  $\psi$ .
(c)  $\epsilon = 10$
(d)  $\epsilon = 10^{3}$
C.5 Effective Eigenvalue Regularization using the Laplacian Regularizer
The Laplacian regularizer in Equation (B.19) only uses a sample-based estimator that is unbiased only in the limit that  $\sigma \rightarrow 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  $\psi$  per gradient step. To test the effectiveness of L-MAP in regularizing the eigenvalues of  $\mathcal{J}(\theta; 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
choose  $p_{X} = \frac{1}{M}\sum_{i = 1}^{M}\delta_{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  $\mathcal{J}(\theta ;p_X)$  for L-MAP and FS-MAP with different levels of jitter  $\epsilon$ . Both FS-MAP and L-MAP significantly reduce the eigenvalues compared to PS-MAP with small values of  $\epsilon$ , corresponding to stronger regularization.
Figure 11: L-MAP is effective at minimizing the eigenvalues.
C.6 Visualizing the Effect of Laplacian Regualrization
The hyperparameter  $\epsilon$  is inversely related to the strength of the Laplacian regularization. We visualize the effect of  $\epsilon$  in Figure 12, showing the L-MAP solution varies smoothly with  $\epsilon$ , evolving from a near-zero function that severely underfits the data to one that fits the data perfectly.
Figure 12: Visualization of L-MAP solutions at various  $\epsilon$ .
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].
Table 4: This table provides detailed quantitative performance for the results plotted in Figure 7, abating the choice of the number of samples used for evaluation of the Laplacian estimator  $S$ . The standard deviations are reported over five trials.
# SAMPLES (S)	FASHIONMNIST	CIFAR-10
ACC.↑	SEL. Pred.↑	NLL↓	ECE↓	ACC.↑	SEL. Pred.↑	NLL↓	ECE↓
4	94.0%±0.1	99.2%±0.0	0.28±0.01	4.3%±0.2	95.6%±0.1	99.5%±0.0	0.17±0.01	2.4%±0.1
8	94.0%±0.2	99.2%±0.0	0.27±0.01	4.2%±0.3	95.6%±0.1	99.5%±0.0	0.17±0.00	2.2%±0.1
16	94.1%±0.1	99.2%±0.0	0.27±0.01	4.1%±0.1	95.6%±0.2	99.5%±0.0	0.17±0.00	2.1%±0.1
32	93.8%±0.1	99.2%±0.0	0.28±0.01	4.4%±0.2	95.5%±0.1	99.5%±0.0	0.16±0.00	1.8%±0.2
64	94.0%±0.1	99.2%±0.0	0.27±0.0	4.2%±0.1	95.7%±0.2	99.5%±0.0	0.16±0.00	1.7%±0.2
128	94.1%±0.1	99.2%±0.1	0.26±0.01	4.1%±0.1	95.5%±0.1	99.5%±0.0	0.16±0.00	1.4%±0.1
256	94.0%±0.2	99.2%±0.0	0.27±0.0	4.3%±0.1	95.5%±0.1	99.5%±0.0	0.16±0.00	1.2%±0.2
512	93.9%±0.1	99.1%±0.0	0.26±0.01	4.2%±0.1	95.5%±0.1	99.5%±0.0	0.16±0.00	1.2%±0.1
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, mimicking a covariate shift. As reported in Table 5, L-MAP tends to improve calibration, while retaining performance in terms of accuracy.
Table 5: We evaluate the performance of CIFAR-10.1 [22] using the models trained on CIFAR-10 [15]. L-MAP tends to improve the data fit in terms of the log likelihood and is better calibrated while retaining the same performance as PS-MAP. The standard deviations are reported over five trials.
METHOD	ACC. ↑	SEL. PRED. ↑	NLL ↓	ECE ↓
PS-MAP	89.2%±0.4	97.8%±0.2	0.43±0.03	6.2%±0.3
L-MAP	89.2%±0.7	97.9%±0.2	0.40±0.02	4.1%±0.5
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.