Debiasing Mini-Batch Quadratics for Applications in Deep Learning
Authors: Lukas Tatzel
Bálint Mucsányi
Osane Hackel
Philipp Hennig
Venue: ICLR 2025 Poster (2025)
Keywords: quadratic Taylor approximation
mini-batching
second-order optimizers
conjugate gradients
uncertainty quantification
Laplace approximation
stochastic curvature
GGN
KFAC
Primary area: optimization
TL;DR: This paper shows that mini-batching introduces biases in quadratic approximations to deep learning loss functions, discusses their impact on second-order optimization and uncertainty quantification, and proposes debiasing strategies.

Abstract: Quadratic approximations form a fundamental building block of machine learning methods. E.g., second-order optimizers try to find the Newton step into the minimum of a local quadratic proxy to the objective function; and the second-order approximation of a network's loss function can be used to quantify the uncertainty of its outputs via the Laplace approximation. When computations on the entire training set are intractable - typical for deep learning - the relevant quantities are computed on mini-batches. This, however, distorts and biases the shape of the associated *stochastic* quadratic approximations in an intricate way with detrimental effects on applications. In this paper, we (i) show that this bias introduces a systematic error, (ii) provide a theoretical explanation for it, (iii) explain its relevance for second-order optimization and uncertainty quantification via the Laplace approximation in deep learning, and (iv) develop and evaluate debiasing strategies.

Published as a conference paper at ICLR 2025
DEBIASING MINI-BATCH QUADRATICS FOR APPLICATIONS IN DEEP LEARNING
Lukas Tatzel, Bálint Muesányi, Osane Hackel & Philipp Hennig Tübingen AI Center
University of Tübingen Tübingen, Germany
(lukas.tatzel, bal intinnt.an mucsanyi -encesatannenntttuubbbeeeeee
@uni-tuebi www.ueiiinnnnn.ccm de
ABSTRACT
Quadratic approximations form a fundamental building block of machine learning methods. E.g., second-order optimizers try find the Newton step into the min- imum of a local quadratic proxy to the objective function; and the second-order approximation of network's s loss function can be used quantify the uncertainty of its outputs via the Laplace approximation. When computations on the entire training set are intractable intractable-typical for deep learning the relevant quantities are computed on mini-batches. This. however, distorts and biases the shape of the associated stochastic quadratic approximations in an intricate way with detrimen- tal effects applications. In this paper, we (i) show that this bias introduces a systematic error, (ii) provide a theoretical explanation for it, (iii) explain its rele- vance for second-order optimization and uncertainty quantification via the Laplace approximation deep learning, and (iv) develop and evaluate debiasing strategies.
INTRODUCTION
Quadratic approximations of the loss landscape increasingly used by algorithms in deep learning. from pruning methods (Dong et al., 2017; Zeng & Urtasun, 2019) and influence functions (Koh & Liang, 2017) to second-order optimizers (Amari, 1998; Martens, 2010; Martens & Grosse, 2015; Grosse & Martens, 2016; Botev et al., 2017: Zhang et al., 2017; George et al., 2018; Martens et al., 2018; Osawa et al., 2019) and uncertainty quantification via the Laplace approximation (Ritter et al., 2018b;a; Kristiadi et al., 2020; Daxberger et al., 2021; Immer et al., 2021). When the computations are intractable on the entire training set-typical for deep learning the quantities of interest are computed on mini-batches subsampled from the training data. The goal of this work is to highlight that mini-batching systematically biases the shape of a quadratic approximation.
A systematic bias? Figure illustrates the phenomenon. It shows five mini-batch quadratics in their top-curvature 2D subspace for the fully trained ALL-CNN-C model on CIFAR- 100 data, For comparison, the full-batch quadratic, where all quantities are evaluated on the entire training set, is projected into the same 2D subspace. Within that subspace, the two quadratics are quite different: The mini-batch quadratic is much "narrower" (exhibits larger curvature) than the full-batch version. Given that the full-batch quadratic the "right" object to serve as the basis for, e.g., a Newton step or Laplace approximation, the mini-batch version is not meaningful surrogate: Its Newton step is overly small and a Laplace approximation yields an overconfident uncertainty estimate.
Contributions. To enable stable and efficient stochastic second-order optimizers, well reliable techniques for uncertainty quantification, we analyze this phenomenon and develop strategies to mitigate it. More specifically, our contributions are as follows: (i) We study mini-batch quadratics empirically and show that their geometry is systematically biased; (ii) we provide an explanation for this phenomenon, explaining the bias as an instance of the classic regression to the mean (directions of extreme steepness/curvatur for one particular mini-bre bann-b are less extreme for other mini-br mani-batches).
We are not concerned with the approximation error arising from the quadratic approximation of the non- quadratic function Cre( D) q( D) (see Equation (2)), but only with the consequences of replacing the full-batch quantities by their mini-batch counterparts q( :DD) +(::3)).

---

Published as a conference paper at ICLR 2025
Top two eigenvectors computed on mini-batch
Bo
83
B2
B3 B4
T1
T1
T1
T1
Figure 1: A systematic bias? We compute five mini-batch quadratics a( :B.)) with batch size
|Bm.| = 512 for the loss landscape of the fully trained ALL--CN--C model on CIFAR-18 100 data around O,- - 0. (shown ). Each mini-batch quadratic defines 2D subspace spanned by the top two cigenvectors u1, U2 of HB... in which we evaluate (i) the quadratic q( 0, + U1 + T2222; Bm) itself (shown in that subspace, the mini-batch quadratic is much "narrower" than the full-batch version which leads to overly small Newton steps and overconfident uncertainty estimates via the Laplace approximation.
and (ii) the full-batch quadratic q(.. + TI u + T2 D) (shown in
In
(iii) we explain the relevance of this bias for second-order optimization and uncertainty quantification via the Laplace approximation, and (iv) develop and evaluate debiasing strategies.
2 NOTATION & BACKGROUND
The regularized loss. We consider general supervised learning problem, where we try to find the optimal parameters 0, arg mingez? Lee (,, D) for the parameterized function So R RCC by
minimizing the regularized loss Lieg N training examples D := {( ynr) TRD R nED:
Logg( () L(0:D) r0) with L(,:D)
1
((fo(an), vn), D (1, (1....., N). mED
Σ
D
(1)
The empirical risk measures the dissimilarity between the model's predictions fo (xxx) and the true outputs Y, via a loss function L RC RC - R. The regularizer r: RRP R. r(0) := 0/2 0 with 3 12 penalizes the "complexity" of the model.
2.1 LOCAL FULL-BATCH MINI-BATCH QUADRATIC APPROXIMATIONS
Full-batch quadratic. A local quadratic approximation of the regularized loss around Og E RP is given by the second-order Taylor expansion
(o 9( } 0 T H 0,) + (0-00)
00) 9D + CD:
where ep == Cigg Co: D), = V V (6: D) and Hp == V Licg D) V2C(eu: D) + 31 is (some approximation of) the Hessian at Og (all derivatives are with respect the parameters 0 unless stated otherwise). As all quantities are evaluated on the entire training set D, we refer to this as the full-batch quadratic. holds that (e H 0)) gD and V (0: D) HD.
Mini-batch quadratic. When the computations are intractable on the entire training set, the quantities Equation (2) are typically computed on mini-batch-a small randomly drawn subset-of the training data B = D, [B] < N, resulting in a stochustic quadratic approximation q( :B) & g( D). As CB. 9B and Hg are unbiased estimates of cp, gD and Hp. this substitution may seem innocent, but, will see in Section 3, it affects the geometry of the quadratic approximation substantially.
Directional slope and curvature. Consider a cut r through the quadratic :::; ;B) from 0, E g² along the normalized direction d, d|| = 1. holds that (derivation in Appendix r()) q(O rd: B) = 2/27dd" V2g(e.;
BB) d (((6::)). rd Vg(O 0.: B) + const. So, as function of r, - is ID d V ( 0.; B) d d (q(O.; B) and ("(r) d V= d.
parabola with derivatives r'( We denote the directional slope and curvature of the quadratic q( BB) at 0. in direction by
da (0: r'( d (q(0.; and (2. r"(0) d (2:
2

---

Published as a conference paper at ICLR 2025
The directional slope and curvature at 0. are thus simply the projections of the quadratic's gradient and Hessian at that location onto the direction.
Eigenvalues directional curvatures. The directional curvature of the quadratic q( B) along one of H3's ss normalized cigenvectors u coincides with the corresponding eigenvalue since 2 ()) u V (0.; B) u Hg Al|za| = X. Thus, in the context of a quadratic, eigenvalue the Hessian H8 has a geometric interpretation as the directional curvature along the respective eigenvector.
GGN & FIM. Next, we discuss second-order optimization methods and the Laplace approximation (LA) for neural networks. Both techniques rely on local quadratic approximations of the regularized loss function and require positive definite curvature matrix Hg. The empirical risk's 's Hessian V²C(00; B) can indefinite and therefore typically replaced by (an approximation of) the positive Ceuss-NNeton matrix (GGN) G or Fisher information matrix (FIM) Fg- In fact, GGN and FIM are often identical (Martens, 2020, Sec. 9.2). The resulting Hg is positive definite if 3 > 0, when damping term 81, ō E R> is added (e.g., in trust-region methods).
semi-definite Generalized
2.2 SECOND-ORDER METHOSS & CONJUGATE GRADIENTS
The Newton step. Due to the simple polynomial form of the quadratic a( (3), its minimum can be derived in closed form and is given by the Newton step arg min, g(; 13) = 0o Hg gs. This serves as the basis for recond-order optimizers like L-BFGS (Liu & Nocedal, 1989; Nocedal. 1980), the Hessian-free approach (Martens, 2010) or K-FAC (Kroneeker-factorred approximate curvature) (Martens & Grosse, 2015; Grosse & Martens, 2016; Martens al., 2018).
Conjugate gradients. We focus on the method of conjugate gradients (CG) (Hestenes & Stiefel, 1952), as a powerful method specifically designed to minimize quadratics with positive definite Hessians effectively (details in Appendix A.2). It is particularly useful in the context of large-scale optimization because only requires access to matrix-vector products with the curvature matrix - Hg v that can be computed in a matrix-fri free manner (Pearlmutter, 1994; Schraudolph, 2002), i.e. without ever materializing the full matrix in memory-ane it has been used successfully for training neural networks (Martens, 2010). Starting Og. CG creates sequence of iterates (fn...., 0r). In each iteration p. two main steps are performed: (i) Given the current position O, and a normalized search direction dy, the algorithm finds the minimum of the quadratic along O, + rdp. i.e.
dd, 1 (0p: 07 (0,;;))
0p+1 = Op + Txdp with
Tp = arg min ((, + Tdp; B)
(3)
reB
In the second step (ii), CG constructs the next search direction dp+1 such that is conjugate to all previous search directions, i.e. d3J+1 Hg d, for E (1.....p).
2.3 LAPLACE APPROXIMATION FOR NEURAL NETWORKS
Laplace approximation. The Laplace approximation (I.A) turns a trained standard neural network into a Bayesian neural network in post-hoc manner (MacKay, 1991; Ritter et al., 2018b; Kristiadi et al., 2020; Daxberger et al., 2021) (details in Appendix A.3). The idea is to reinterpret the regularized loss Lreg the negative unnormalized log posterior log p(O 0 D) of specific Bayesian model. This interpretation identifies the optimal parameters 0, arg ming L (0; D) = arg max p(@ ID) as the mode of the posterior, i.e. the maximum posteriori (MAP) estimate. A second-order approximation of the regularized loss around O,- - 0, then translates into a Gaussian approximation of the posterior- the so-called Laplace approximation (MacKay, 1992), i.e.
Coeg(:D) = ((: = (0 (0)))) ( 0.) const.
p(O ID) R N((;0, Ep), (4)
with Ep :- N-I (or an approximation thereof). We obtain the predictive uncertainty p(yo 2o, D) for some test input x. R by propagating the parameter uncertainty A/(e., Ep) to the model's outputs via the linearized network (Immer et al., 2021; Roy et al., 2024).
Full-batch vs. mini-batch LA. is common practice to compute the LA on the entire training set. Depending the curvature approximation, this comes considerable computational costs. For
3

---

Published as a conference paper at ICLR 2025
example, full-batch LA is prohibitive even for moderately sized models/datasets when low-rank approximation of the Hessian is computed via repeated Hessian- vector products (each of which requires a full pass over the training data); and it is essentially infeasible for model classes like LLMs, which are trained on massive datasets. However, what adds most to the costs is that is standard to tune the prior precision (Daxherger al., 2021). This adds another outer loop that requires the same procedure to be performed multiple times. Being able to emulate the behavior of the full-batch quadratic on mini-batch would thus be useful. We therefore study the mini-batch setting, i.e. we replace q(: : D) by g( : B) in Equation (4). We will see in Section 6.2 that, when mitigating the associated biases, mini-batch LA can be a good proxy for the full-batch LA.
THE SHAPE OF A MINI-BATCH QUADRATIC
This section studies the "shape" of mini-batch quadratic q( -:3) and how it differs from q( :D).
3.1 EMPIRICAL STUDY OF THE DIRECTIONAL SLOPES AND CURVATURES
The high-curvature subspace is relevant for all common cases. In our empirical study, focus on the top-curvature subspace. This subspace is particularly relevant for several reasons:
1. By the Bekart- Young-Mirsky Theorem, a truncated SVD Frobenius norm-optimal, i.e. a low-rank approximation using the top eigenvectors is ideal from theoretical perspective. As the spectrum of the Hessian typically decays quickly (Ghorbani et al., 2019), the bulk of the curvature information is contained in the top-curvature subspace.
2. In the context of optimization, it has been observed that the gradient mainly lives in the high-curvature space (Gur-Ari et al., 2018; Dangel al., 2022). Thus, it makes sense for a second-order method to operate mainly in that space, as steps outside of it can not be expected to reduce the objective function significantly.
3. In the context of the LA, directions of large curvature correspond directions in the weight space with low variance, i.e. these directions capture what we know about the model's parameters. Consequently, Daxberger al. (2021, p. 19) describe low-rank approximation of the Hessian based on its top eigenvectors.
Experimental procedure. We use the same setting as in Figure I: The fully trained ALL-CNN-C model on the CIFAR-100 dataset. To isolate the effect of data subsampling, we eliminate all other sources of noise. Thus, we remove the dropout layers from the model. We use the cross-entropy loss, an reggullarre and train the model with SGD for 350 epochs, see Appendix B. for details.
We then pick mini-batch Bn of size 512 and compute the 100 eigenvectors uI..... U100 to the 100 largest eigenvalues of Hg.,. < GB.. + .I. That the Hessian of the E -ergullrrized where we replaced 2 C(: Bm) with the GGN approximation. Next, we compute the directional slopes and curvatures for all mini-batch quadratics q(. :;; B,,),, mn' {1...., M) along those 100 eigenvectors. For fixed eigenvector Up. the average of those directional slopes/curvatures over all quadratic. i.c.
loss,
mini-batches coincides with the directional slope/curvature of the
M Σ Cu, g(x:: Brr) = iu,9(0.:7 and
1
Σ 201
M
E2 a (.)): E2 ((.;
(5)
M
M
one and many
one
and many
derivation in Appendix A.1. The colored markers in Equation (5) refer to Figure 2. We repeat this procedure for three mini-batches B.., m E {0, 1,
Directional slopes and curvatures are biased. Figure reveals a systematic bias in the directional slopes and, more pronounced, in the directional curvatures: When eigenvectors and directional derivatives are computed on the same mini-batch, curvature is overestimated by roughly one order of magnitude compared to the curvature of the full-batch quadratic g( D). Within the space that actually carries curvature information the top-eigenspace of the quadratic's 's Hessian the curvature magnitude is thus not at all representative of the true underlying loss landscape. For other mini- batches, the directional slopes and curvatures are similar to the full-batch quadratic. This is because the averages in Equation (5) are dominated by these unbiased samples.
4

---

Published as a conference paper at ICLR 2025
Eigenvectors computed mini-bated
B
B1
B3
Directional slope
1.0
Permuted cigenvector index
Permuted cigenvector index
Permuted cigenvector index
Directional curvature
10
10%
25
75
100
25
75
100
25
75
100
Eigenvector index
Eigenvector index
Eigenvector index
Figure 2: Directional slopes and curvatures are biased. We use the CIFAR-100 dataset with the fully trained ALL-CNN-C model and draw three mini-batches Bm, of size [Bm]] = 512 to compute the top 100 cigenvectors u1...., U100- For each mini-batch/column, we show the directional slopes (Top) and curvatures (Bottom) evaluated on g B,.)) (i.e. on the same mini-batch of data) as (ii) @(O;; Bm))) for m' * m (i.e. for all other mini-batches) as and (iii) the full-batch quadratic (the average of the orange and all blue dots, see Equation (5)) as +. For the top panel, we switch the order and sign of the eigenvectors such that the orange dots are all above zero and in descending order. There is a strong, systematic bias, particularly the curvature: Computing the eigenvectors and directional curvatures the same data results in over-estimation by roughly one order of magnitude.
We present additional results for batch size 2048 (instead of 512), the Hessian V CZ. Bmx) (instead of its GGN approximation) and CG search directions (instead of eigenvectors) in Appendix B.4. The bias is present in all settings, but less pronounced for larger batch sizes. For the CG search directions, the bias in directional slope is much more distinct than for the eigenvectors: The biased slope is always negative, while the full-batch quadratic's slope is in fact positive for most CG directions!
3.2 WHERE THE BIASES COME FROM?
In the following, we provide an explanation for the biases in the directional slopes and curvatures (i.c. the gap between the orange and blue dot in Figure 2) from theoretical perspective.
3.2.1 BIAS IN THE DIRECTIONAL SLOPE
Here, we focus on the CG search directions since the bias in the directional slope even more pronounced for those directions than for the eigenvectors. This is not accidental! The CG directions are constructed from gradients-aae particular mini-batch quadratic. For other mini-bated quadratics, the steepness along those directions (i.e. the directional slope is therefore less extreme. We formalize this intuition in Appendix A.4.
gradients are directions that maximize the steepness for one
3.2.2 BIAS IN THE DIRECTIONAL CURVATURE
Next, we consider the bias in the directional curvature along the cigenvectors of the curvature matrix.
Directional curvature along Hg's eigenvectors. Let u1,..., ., up and uI...., üp denote the eigenvectors of two mini-batch Hessians Hg and Hg. respectively. Assume that the corresponding eigenvalues are in descending order, i.e. Xι ... z Ap and λ₁ > Åp. We write the

---

Published as a conference paper at ICLR 2025
directional curvatures on B and B along eigenvector u, (details in Appendix A.5)
P Xo
Cy (0.: u Hgu;
Ship and C ((.:: u Hgu; x
(6)
The weights {Ri.p == (u. ) ) are non-negative and sum to one, i.e. Σp- Ri.p = 1. The bias in the directional curvature thus originates from misalignment of the eigenspaces of the two curvature matrices captured by the weights i and/or systematic difference in their spectra.
Curvature bias not due to differing spectra... Assume that the eigenspaces are perfectly aligned, Le. $2,, = 1 and St = 0 Vp * i. In this case, we obtain C5, ((,; B) = X from Equation (6), so the bias originates exclusively from the differences in the spectra. Figure 2 shows the eigenvalues (as directional curvatures) for three different CIFAR- 00 mini-batches. As they are very similar, this cannot serve as the main explanation for the curvature bias.
but due to misaligned eigenspaces. Now, assume that the spectra of H and Hg are identical, i.e. A, Xp Vpp {I,..., P) which simplifies the unbiased estimate in Equation (6) to é ((,;; ()) Ep- Ap S. Consider the curvature along U1 as an example. If µn = ú there is only one non- zero weight 12 = and the two directional curvatures are identical. However, there is significant overlap with any other eigenvectors (i.e. some weight is distributed also the lower-curvature directions), the resulting curvature ((.:(0) is smaller than 2h ((.; BB). Analogously, for up, the directional curvature on B is larger than on there is significant overlap between up and some of Hg g higher-curvature eigenvectors. This can be formalized by the following inequalities (derivation in Appendix A.5):
2i g B Ce, ((::: and Cup (0.; (..:
In general, if the eigenspaces are not perfectly aligned such that weight is distributed on several eigenvectors, this leads to less extreme directional curvatures on mini-batch B on both ends of the spectrum. Figure 3 shows the weights Stip as pixels. The overlap between the eigenspaces is far from perfect, i.e. eigenvector from B overlaps with several eigenvectors from B to some extent. This explains the curvature overestimation in the top curvature subspace we observe in Figure 2 and identifies the misalignment of the eigenspaces as the dominating factor for the eurvature bias.
3.2.3 SUMMARY OF THE THEORETICAL FINDINGS
The CG search directions and the eigenvectors of the curvature matrix are both designed to be "extreme" in some sense: The CG directions are based on gradients that maximize steepness of the quadratic q( - B), while the top/bottom cigenvectors subsume directions of largest/smallest curvature. However, these directions extreme only for particular mini-batch B. Another mini-batch B has its own extreme directions that typically differ from those of B. Thus, the extreme directions for are less extreme for Projecting both quadratics onto B's directions consequently leads to extreme steepness and curvatures for B, but less extreme values for B. This result an instance of the classic regression to the mean (Galton, 1886): Using algorithm to find the directions of most extreme steepness/curvature in one particular mini-batch, we must expect the steepness/curvature to be less extreme most other batches (and thus also the entire dataset). We can expect this phenomenon to occur for other datasets, models, and curvature proxies as well, since the underlying mechanism is the stochasticity of the geometric information.
DEBIASING MINI-BATCH QUADRATICS FOR APPLICATIONS
We here argue that the biases affect second-order applications, and develop debiasing strategies.
4.1 IMPLICATIONS FOR SECOND-ORDER OPTIMIZATION AND THE LAPLACE APPROXIMATION
Detrimental updates in second-order optimizers. In the context of second-order optimization, both the biases in the directional slopes and curvatures relevant. From Equation (3), the CG

---

Published as a conference paper at ICLR 2025
Bp
B2
Directional magnitude
Bo
B1
B2
Bo
-1-
10
CG search direction index
B1
Figure 4: CG update magnitudes are biased. Same setting Figure 9 (Bottom). We run CG on {Bm me and show the directional up- date magnitudes 71, 71,-..s T10 for the first 10 CG steps using (i) the same mini-batch Bm (as )), (ii) all other mini-batches (as )) and (iii) the entire training set (as +). The magnitudes are given by the negative ratio of the directional slope and curvature (see Equation (3)) and thus inherit the attached biases. Note that most of the update magnitudes that are based on a single mini-batch of data (+) have the wrong sign resulting in detri- mental updates in the wrong direction.
Figure 3: In practice, eigenspaces are mis- aligned. We reuse the setting of Figure 2 and compute the top 100 eigenvectors Um IR Px 100 for (B, me (0. 1.2). The weights D are shown as 100x x 100 greyscale image (color ranges from black for D2 & 10-8 white for Si 1) for m E {0. m' E (0, 2,)) Clearly, the cigenspaces for different mini-batches are not perfectly aligned eigenvectors from Bm over- lap with several eigenvectors from Bm.
update magnitude Ty to reach the minimum of q(O, Tdp; Bm) given by the negative ratio of the directional slope and curvature at the current iterate O, in the direction dh-thht is the ID Newton step along that cut.² As both these quantities are biased, so is Tp as shown in Figure 4: While the update magnitudes for B,, are always positive (a property of CG), minimizing the other-equally valid-quadratics would require negative update magnitudes for most of the CG directions. In this sense, naive CG on single mini-batch of data makes updates in the wrong direction. This can be attributed to the bias in the directional slope since the slope determines the sign of TP. Overestimation of the directional curvature is "accidentally beneficial" in this case as it leads smaller steps.
Unreliable uncertainty quantification with the Laplace approximation. For the LA, only the bias in the directional curvature is relevant. After Section 2.3, the approximate posterior covariance over the network's parameters is given by Es N- H Via the eigendecomposition H
Ep u. we obtain Σg N- Σp u, Due to the biases we describe in Section the directional curvatures X are not representative the true underlying curvature, resulting in deformed posterior covariance Σg. Specifically, the overestimation of the curvature in the top- curvature subspace translates into an underestimation of the uncertainty in the posterior covariance (due to the inversion) with potentially severe consequences in safety-critical applications.
4.2 DEBIASING STRATEGIES
We turn to strategies that mitigate the biases in second-order optimization and the LA. An empirical evaluation of these methods follows in Sections 6.1 and 6.2.
Decoupling directions from magnitudes with two-batch strategy. The approaches propose are simple yet effective. The idea is to commit to the imperfect directions from one mini-batch (as before) but an additional mini-batch estimate the directional derivatives. With this, decouple the mechanism that determines the parameter subspace in which the method operates from the slope and curvature measurements within the space. Effectively, we use the blue dots from Figures and 4 instead of the orange dots and thus obtain much more realistic estimates of the actual loss function's geometry (within the subspace defined by the first mini-batch). Next, we describe in more detail how this strategy can be applied to debias CG and the LA
"The exact same argument can also be made for the eigenvector directions the Newton step be decomposed into ID Newton steps along the eigenvectors.
7

---

Published as a conference paper at ICLR 2025
4.2.1 DEBIASED CONJUGATE GRADIENTS
Debiased approach. For the debiased CG method, we need to implement two processes: (i) The first process applies K & P CG iterations to the mini-batch quadratic g;;; and collects the search directions (d1..... . dx). As before, this defines the subspace in which CG operates. (ii) The second process recomputes the trajectory (0,,... 0k) within that subspace using debiased update
magnitudes that are computed on a different mini-batche B, i.e. we use 0g == 0g and
Ed Ed (d,p(0p: ag 2(pp0
Ed (Opp(0p; Ca, (0p;
Op 0p+1 == O2 ipd with
instead of Tp
(8)
For B- B, (O..... 0x) congruent with the original trajectory ( 0x) from the single- batch CG approach. If the two mani-batches are different, the debiased trajectory will use updates into the directional minima of q( ; B) instead of ;; :)) Processes (i) and (ii) can either be run side by side in an alternating fashion (offering more flexibility regarding e.g. the termination criterion at the cost of small memory overhead, details in Appendix A.2) or one after the other.
4.2.2 DEBIASED LAPLACE APPROXIMATION
Laplace approximation with K-FAC. We briefly explore another popular curvature proxy: Fronecker-Factored Approximate Curvature (K-FAC), which commonly used both in opti- mization (Martens & Grosse, 2015; Martens et al., 2018; Eschenhagen et al., 2023) and un- certainty quantification (Ritter et al., 2018b). It is a block-diagonal approximation to the FIM where each block KS A x (!) is approximated by the Kronecker product of two smaller matrices A and (1), Sampling from the respective LA with covariance Σg N "KK + (3)) can be performed efficiently due to the specific structure of the K-FAC approximation. For details, see Appendix A.3.2.
FB x K := blockdiigg........K
Debiased approach. For the debiased LA, we use two K-FAC approximations K3 and Kg com- puted on different mini-batcher Via the eigendecomposition Kg = UAUT with U (uI...., up) and A = diag( ÀI................. Ap). we can re-write the covariance matrix as Eg = (K' + 31)-1 U + BU)-LU. For the debiased approach, we keep the eigenspace defined by KB but recompute the directional curvatures based on Kg. i.e. we use
Σg N
de) BI) UUT with X u, K Kg instead of Xp u. Kgup-
4.2.3 COMPUTATIONAL COST OF DEBIASING
Both debiasing techniques roughly double runtime compared to their single-batch counterparts: Debiased CG performs one extra matrix- product with H per CG iteration to compute the debiased update magnitude (details in Appendix A.2). Debiased LA requires an additional K-FAC approximation Kg: g: all subsequent operations to compute the debiased K-FAC can be carried out efficiently at Kronecker factor level (details in Appendix A.3.3). In Section 6, we thus use the debiased versions at half the batch size, for fair comparison. We will see that, at the resulting similar computational cost, the debiased versions clearly outperform the single-batch alternatives.
RELATED WORK
We here briefly list other, related forms of bias correction that have been suggested elsewhere.
A different two-batch approach. Other works have proposed to use different (not necessarily disjoint) mini-batches for the gradient and the Hessian (Martens, 2010; Byrd al., 2011) Benzing (2022, Sec. I.5) mentions the idea of using independent mini-batches for the gradient and the Hessian to obtain an, in some sense, unbiased estimate H 9B of the exact Newton step. This does, however, not resolve the biases described in this paper. Via the eigendecomposition of the Hessian
Hg Σx_ up we obtain - H "8 = Dn- 0n, (0.; () q(..; (O) up- While the numerator yields an unbiased estimate of the directional slope (similar to blue dot in the upper
8

---

Published as a conference paper at ICLR 2025
panel of Figure 2), the bias in the denominator remains since eigenvectors and directional curvatures are based on the same mini-bated (as for the orange dots in the bottom panel of Figure 2).
Another debiasing approach. EKFAC (Eigenvalue-correct) ct al., 2018) corrects the cigenvalues of the K-FAC approximation (similar to Equation (9)) which, provably, yields more accurate approximation of the FIM than K-FAC (in Frobenius norm). This correction, however, designed to resolve a different kind of bias that is specific to the K-FAC approximation and does not address the biases described in this work.
Kronecker Factorization) (George
Running averages. Other popular deep learning optimizers aggregate curvature information over multiple steps via exponential moving averages (see e.g. (Martens & Grosse, 2015)). This emulates larger mini-batch sizes and thus reduces the curvature biases. However, when curvature evolves rapidly, obsolete curvature estimates from past steps might slow down training. Aggregating more robust debiased curvature estimates instead might allow for shorter moving average windows and accelerate the optimization progress. We leave it to future work to explore these interactions.
EXPERIMENTS
In this section, we evaluate the effectiveness of the debiasing strategies from Section 4.2. In Appendix B, we provide the experimental details as well additional empirical analyses. For instance, we show how the curvature biases with K-FAC depend on the mini-ban batch size (see Appendix B.8), how the biases evolve over the course training (see Appendix B,9), and that the curvature biases become more pronounced for deeper/wider models (see Appendix B.10).
6.1 DEBIASED CONJUGATE GRADIENTS
We compare the standard single-batch CG method to the debiased version (see Section 4.2. on the fully trained ALL-CNN-C model without curvature damping. For fair comparison, the single-batch approach uses mini-bate batch of size 1024 while the debiased approach uses bini-bbtchhes of size 512, such that a similar amount of data and runtime budget is used. Details in Appendix B.5.
Reg. training loss 4
Reg. test loss J
Training accuracy
Test accuracy
3.4
1.00
60-
Loss
Accuracy
911
3.2
0.95
59
92-
0.904
3.0-
58
30
30
10
CG iteration
CG iteration
CG iteration
CG iteration
[8]
1024
[8] 512 (debissed)
Figure 5: Debiased CG much more stable than the single-batch approach. We compare CG runs without curvature damping (8 0) with K = 30 iterations for the fully trained ALL-CNN-C model on the CIFAR- 100 dataset in terms of training/test loss/accuracy at similar computational cost: The single-batch approach (shown as -) uses one mini-batch of size 1024 while the debiased approach (shown as -) uses two mini-batches of size 512 each. Both approaches use the GGN curvature proxy and are times on different mini-batches. The markers and are placed at peak performance. While the single-batch runs diverge quickly, the debiased CG runs are stable.
Results & discussion. All CG runs in Figure achieve significant improvement in all four performance metrics. The most striking difference between the two approaches is their stability: The single-batch runs quickly reach peak performance and then diverge. In contrast, although the steps tend to be larger (see Figure 4), the debiased CG runs are much more stable (without using any curvature damping). This suggests that the reason for the divergence the single-batch approach is not the missing damping but the misinformed update magnitudes. The peak performance is slightly better for the single-batch runs which can be attributed to its more informative search directions (they computed using double the data). The peak performance of the debiased runs could likely be improved (at the same computational cost) by using a larger batch size for the directions and a smaller one for the update magnitudes (as the former seem harder to estimate, see Section 3.2).

---

Published as a conference paper at ICLR 2025
6.2 DEBIASED LAPLACE APPROXIMATION
Here, we fully trained ALL- model CIFAR- 10 data and compare (i) the vanilla model without LA, (ii) the single-batch K- K-FAC LA approach, (iii) the debiased version (see Section 4.2.2), and (iv) the full-batch approach (where we compute K-FAC on the entire training set) in terms of accuracy, NLL and ECE. Again, we apply the debiased approach at half the batch size of the single- batch approach for a fair comparison. We use prior precisions between 10 and 10. Appendix 6.6 contains the experimental details and additional results on the training and OOD data.
Accuracy in %
NI.L loss 4
ECH 4
60-1
Vanilla model So.
[B]
90
5 10-
[8]] 256
10"
[8]] 1024
10
(dehiased)
128 (dehiased)
3 x 10-1
10
[8] 512 (debiased)
10-2 10 10*
10
10
10ah 10
10
10
10 10
Full-batch
Prior precision 3
Prior precision /
Prior precision
Figure 6: Debiased LA mimics the full-batch LA very well. We compare LAs for the fully trained ALL--CN-CC model on CIFAR-10 10 in terms of accuracy, negative log likelihood loss (NLL.) and expected calibration error (ECE) on the CIFAR-10 test set. For each mini-bated size (lighter color indicating smaller batch size), we draw mini-batches and report the mean performance as dot and min/max as vertical line. We also report the performance of the vanilla model without LA (shown as -) and the full-batch approach based on Kp (shown as +). In contrast to the single-bated approach, the debiased version mimics the behavior of the full-batch approach very well over the entire range of prior precisions.
Results & discussion. Figure 6 shows the results. If the prior precision low (i.e. the LA relies mainly the curvature information without the regularizing diagonal term), the single-batch version acts erratically due to the deformed curvature model-its performance drops drastically. The single- batch approach is also more sensitive to the choice of prior precision (in fact, it suggests much larger prior precision than the full-batch approach). In contrast, the debiased approach achieves good performance over wider range of prior precisions and mimics the full-batch approach (N = =00,000 very well despite using only a tiny fraction of the data.
In order showcase our debiasing strategy's scalability, we provide additional results RESNET-50 and VIT LITTLE on the IMAGENET dataset in Appendix B.6.
Summary of the experimental results. The use of mini-bated quadratics is simple way to keep the costs of second-order optimization and uncertainty quantification manageable. The resulting biases, however, severely degrade their value, requiring large mini-batches or algorithmic add-ons. Our experiments suggest that even simple debiasing strategies can largely mitigate this issue.
7 CONCLUSION
Our main takeaway a general principle: Quadratic approximations to the training loss computed on mini-batches of the training data provide a distorted representation of the true underlying loss landscape. In particular, along the directions of large curvature, the mini-batch quadratics tend to strongly overestimate the curvature of the true loss. Our theoretical analysis shows that these biases can be traced back to the misalignment of the curvature matrices' eigenspaces. These insights are highly relevant for applications: As we demonstrated empirically, the biases in the directional slope and curvature lead to severely misinformed updates in stochastic second-order optimizers, and cause unreliable uncertainty estimates with Laplace approximations. We also proposed simple two-batch strategies to mitigate these biases. Our experiments demonstrate their superiority over the single-batcl approaches in terms of stability and quality at similar computational costs. Our findings reveal a design prerequisite for building better stochastic revvauur--bsee methods, which should be generally considered, and further developed, for all methods using such curvature evaluations.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The authors gratefully acknowledge co-funding by the European Union (ERC, ANUBIS, 101 123955). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council. Neither the European Union nor the granting authority can be held responsible for them. Philipp Hennig is member of the Machine Learning Cluster of Excellence, funded by the Deutsche Forrchungsgemminsschiie
(DFG, German Rescarch Foundation) under Germany's Excellence Strategy EXC number 2064/1 Project number 390727645; The authors further gratefully acknowledge financial support by the DFG through Project HE 71 71145-1 in SPP2298/1; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ:01IS) 8039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württembers Lukas Tatzel and Bálint Muesányi are grateful the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. Further, the authors thank Felix Dangel and Runa Eschenhagen for providing feedback to the manuscript.
REFERENCES
Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 1998.
Frederik Benzing. Gradient descent neurons and its link approximate second-order optimization. In International Conference on Machine Learning (ICML), 2022.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss-Newton optimisation for deep learning. In Proceedings of Machine Learning Research (PMLR), 2017.
Richard H. Byrd, Gillian M. Chin, Will Neveitt, and Jorge Nocedal. On the use of stochastic Hessian information in optimization methods for machine learning. SIAM Journal on Optimization, 2011.
Felix Dangel, Frederik Kunstner, and Philipp Hennig. BackPACK: Packing more into backprop. In International Conference on Learning Representations (ICLR), 2020.
Felix Dangel, Lukas Tatzel, and Philipp Hennig. ViVIT: Curvature access through the generalized gauss 's low-rank structure. Transactions on Machine Learning Research (TMLR). 2022.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and redux-effortee effortless Bayesian deep learning. In Advances in Neural Informa- tion Processing Systems (NeurIPS), 2021.
Philipp Hennig. Laplace redux
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image datahase. In Conference on Computer Vision and Pattern Recognition (CVPR), 2009.
Nicki Detlefsen, Jiff Borovec, Justus Schoek, Ananya Jha, Teddy Koker, Luca Di Liello, Daniel Štanel, Quan Changsheng, Maxim Grechkin, and William Falcon. TorchMetrics measuring reproducibility in PyTorch. The Journal of Open Source Software, 2022.
Xin Dong, Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer wise optimal brain surgeon. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL https: //ppne/iiee......e.. net / forum? i d= iiniiiiddiiiiiiai PicTTd
Runa Eschenhagen, Alexander Immer, Richard E. Turner, Frank Schneider, and Philipp Hennig. Sannecker-Factoree Approximate Curvature for modern neural network architectures. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Francis Galton. Regression towards mediocrity in hereditary stature. The Journal of the Anthropolog-
ical Institute of Great Britain and Ireland, 1886.
11

---

Published as a conference paper at ICLR 2025
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems (NeurlPS), 2018.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural optimization via hessian eigenvalue density, In Proceedings of Machine Learning Research (PMLR), 2019.
Roger Grosse and James Martens. A Kronecker-factoree factored approximate Fisher matrix for convolution layers. International Conference on Machine Learning (ICML), 2016.
Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace. arXiv, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR), 2016.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference Learning Representations (ICLR), 2019.
Magnus Rudolph Hestenes and Eduard Stiefel. Methods of conjugate gradients for solving linear systems, Journal of Research of the National Bureau of Standards, 1952.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions Bayesian neural nets via local linearization. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2021.
Pang Wei Koh and Perey Liang. Understanding black-box predictions via influence functions. In Proceedings of Machine Learning Research (PMLR), 2017.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being Bayesian, even just a bit, fixes overconfidence in ReLU networks. In International Conference on Machine Learning (ICML), 2020.
Alex Krizhevsky. Learning multiple layers of features from tiny images, 2009.
Dong C. Liu and Jorge Nocedal. On the limited memory bígs method for large scale optimization. Mathematical Programming, 1989.
David J. C. MacKay. Bayesian model comparison and backprop nets. In Advances in Neural Information Processing Systems (NeurIPS), 1991.
J. C. MacKay, The evidence framework applied to classification networks. Neural Computation, 1992.
James Martens. Deep learning via Hessian-free optimization. In International Conference on Machine Learning (ICML), 2010.
James Martens. New insights and perspectives on the natural gradient method. arXiv, 2020.
James Martens and Roger Grosse. Optimizing neural networks with Kromecker-factored approximate curvature. In International Conference on Machine Learning (ICML), 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations (ICLR), 2018.
Jorge Nocedal. Updating quasi-newton matrices with limited storage. Mathematics of Computation, 1980.
Jorge Nocedal and Stephen Wright. Numerical optimization. Springer, ed. edition, 2006.
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large- scale distributed second-order optimization using Krotecker-factored approximate curvature for deep convolutional neural networks. In Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
12

---

Published as a conference paper at ICLR 2025
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTTrch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Barak A. Pearlmutter. Fast exact multiplication by the Hessian. Neural Computation, 1994.
Hippolyt Ritter, Aleksandar Botev, and David Barber. Online structured laplace approximations for overcoming catastrophic forgetting. In Advances in Neural Information Processing Systems (NeurIPS), 2018a.
Hippolyt Ritter, Aleksandar Botev, and David Barber. scalable laplace approximation for neural networks. In International Conference Learning Representations (ICLR), 2018b.
Hrittik Roy, Marco Miani, Carl Henrik Ek, Philipp Hennig. Marvin förtner, Lukas Tatzel, and Søren Hauberg, Reparameterization invariance in approximate Bayesian inference. arXiv, 2024.
Robin M Schmidt, Frank Schneider, and Philipp Hennig. Descending through crowded valley benchmarking deep learning optimizers. International Conference on Muchine Learning (ICML), 2021.
Frank Schneider, Lukas Balles, and Philipp Hennig, DeepOBS: A deep learning optimizer benchmark suite. In International Conference on Learning Representations (ICLR), 2019.
Nicol N. Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural Computation, 2002.
Jost T. Springenberg, Alexey Dosovitskiy, Thomas Brox, and Martin Riedmiller. Striving for simplicity: The all convolutional net. arXiv, 2015.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv, 2017.
Wenyuan Zeng and Raquel Urtasun. MLPrune: Multi-layer pruning for automated neural network compression, 2019.
Huishuai Zhang, Caiming Xiong, James Bradbury, and Richard Socher. Block-diagonal Hessian-free optimization for training neural networks. arXiv, 2017.
13

---

Published as a conference paper at ICLR 2025
SUPPLEMENTARY MATERIAL
Below, we provide additional details on the mathematical derivations, describe the experimental setup and present additional results.
A Mathematical details
15
A.I Directional derivatives of quadratic
15
A.2
The method of conjugate gradients (CG)
16
A.3
Laplace approximation for neural networks
17
A.3.1
Derivation of the Laplace approximation for neural networks
17
A.3.2
Sampling from the K-FAC Laplace approximation
19
A.3.3
Debiased K-FAC Laplace approximation
20
A.4 Bias in the directional slope
A.5 Bias in the directional curvature
22
Experimental details
23
B.1
Test problems and training procedures
24
B.2
Matrix-vector products with the curvature matrix
25
B.3 Section I (Introduction): Figure
25
B.4 Section 3 (The shape of mini-batch quadratic): Figures 4
25
B.5 Section 6.1 (Debiased conjugate gradients): Figure 5
27
B.6 Section 6.2 (Debiased Laplace approximation): Figure
32
B.7 Additional experiment: Biases for a WIDERESNET 40-4 model architecture
38
B.8 Additional experiment: K-FAC and the dependence of the biases mini-batch size
38
B.9 Additional experiment: Development of the biases over the course of training
38
B.10 Additional experiment: The curvature bias increases with P
40
14

---

Published as a conference paper at ICLR 2025
A MATHEMATICAL DETAILS
A.1 DIRECTIONAL DERIVATIVES OF QUADRATIC
We claim in Section 2.1 that through the quadratic (::;B) from 0. RRP along the normalized direction d can be written as
r(r) = q(.. + B) = 2 r v ((): d+ rd Vo(0.; B) + const.
(10)
Proof for Equation (10). Here, we provide the derivation for Equation (10). Let 0. TRP be point in parameter space and d E I²' normalized direction, i.c. ||]]] 1. We consider the quadratic approximation q0; B) around 0, and evaluate it along the cut 0. + rd for E IR. We assume that the Hessian (or its approximation) Hg is symmetric and obtain
r(r) q0. 1)
- (0. rd 1 00) HB + rd -dddd-6 (d)-)))
1
00) + (0. rd 00)² gb + cB
2
2 dd Hgd rd Hs 0o) + 20.00 Oo) Hs On)
+ +dd gB + (0. 0o) gs + cB
1 dd Hgd rd) (Hs(0. On) + gs) + const.
2
Since Vg(0.:B) H8 Ou) ge and (e: H8. arrive Equation (10). D
Proof for Equation (5). Next, we show that the average directional slope/curvature over all mini- batches in the training set coincides with the directional slopc/curvature of the full-batch quadratic. Let's assume that all M mini-batches (Bme) M -1 are disjoint, have the same size |B1| and that their union is the training set, i.e.
M |B1| Vm' e (1....,M. M) and UBm
Bm D.
(11)
This implies M MBB1| = |D]. holds
Σ
g (0,g(0. Bin()
Σ m²=1
Up (m((.:Bm)))
with Vg(O.; -s... since O, - O,
M
M
Tr and many
M
V Lreg! (0x; Bm)))
p
M
mn
M Σ
Vr( +
v LO.; Ban)
4
M
M
m =1
M
vr(0.) + M m/=1
Σ
Ve(fo,( ym))
nEB
M
(1) u
Vr(O.) +
Σ
Σ
(x)))
M M
nEB
up I Lreg (O.;
En, ().:
As the mini- batch curvature V ((. Bm)) is also an average over the samples in the mini-batch (this applies to both the actual Hessian and its GGN approximation), the same argument holds for the associated directional curvatures. For K-FAC, the derivation above does not hold as (1/2)(Kg + K B) of K BBB-
15

---

Published as a conference paper at ICLR 2025
A.2 THE METHOD OF CONJUGATE GRADIENTS (CG)
Using CG for minimizing a quadratic. The argmin of a quadratic q( :B) is given by 0p H gB (assuming that H is symmetric and positive definite), see Section 2.2. We can the method of conjugate gradients (see Algorithm 1) for computing the Newton step (or approximation thereof) by setting A Hg g and b gB.
Algorithm 1: Method of conjugate gradients (CG), based on (Nocedal & Wright, 2006, Alg. 5.2)
Input. Access to matrix-vectie vector products - Av, where A E Px symmetric and positive definite, right-hand side b e RR,, convergence tolerance E Rşo, maximum number of iterations Pmax N, Pmax & P
Output. Approximate solution Tp to the linear system Ar
Initialize 3y 0, ro + b and so To for p o,l,..., P. do
if p = Pmax or |rpp12 & e then return (approximate) solution Pp
Termination criteria
4
tp - Asp
Compute the matrix-vector product only once
r rp
Tpp Tp Asp
Op" 8,
Update magnitude ap
1p+1 - Up + 8p8p
Update along search direction 8p Residual ry Axx b computed via recursion
Tp+1 Tp + aptp
r rpt1 pp+1
rp rp 8pt1
10
8pt1 - -1pt1 + Apt18p
Construction of new (conjugate) search direction
Properties of CG & geometric interpretations. In the following, provide some important properties and a geometric interpretation of the quantities involved in the CG method.
Residual Ty- Let x, == 0, 0o. The residual тр can be written as тр = Ax b = Hs(0, 0,) + 9s Vg B)) i.c. it coincides with the gradient of the quadratic 0,- Update magnitude Otp- Consider cut through the quadratic from Op into the direction 8p. i.e. r(T) (0,, + T&P B). Minimizing this 1D quadratic requires its first derivative to vanish. It holds (see Equation (10))
r'(7)
T8 2 9( (0.: B3) 8p + 8 (q(pp:
[g 13)
8 S,
r r.r
P
T
22 @ B3) d,
8 S
A.A
Op:
p
where the equality -8, rp rp rp is due to Equation (5.14a) and Theorem 5.2 in (Nocedal & Wright, 2006). So, the update magnitude op chosen such that minimizes the quadratic along the direction 8p. Note that the update magnitude can also be written ratio of the negative directional slope and directional curvature at 0p. i.c. ID directional Newton step. Let dp == sp/l| spl denote the normalized search direction. It holds
8 ST
d d. = I o (0,; dp
(d,9(0,; B)
2p+1 xx + appp = I, ST. AS
8p = Ip
d Ad
Tp
i.. Tp = apllspl.. Via the shift 0, 0u xp, we arrive Equation (3).
T Search direction sp. CG constructs the search directions to conjugate, i.c. 8p+ 1 A s; for all & p. Note that this property also applies the normalized search directions.
16

---

Published as a conference paper at ICLR 2025
Efficient implementation of the debiased CG approach. For the debiased CG version, need to re-evaluate the update magnitudes Tp for a given set of search directions (di.... dp} on a second mini-batch B (see Equation (8)), i.e.
θ2 (0, B) 22 (0,;:
d ( B)
d. ( 06 60) 9%
V d,
d H dp
Implemented naively, this requires two matrix- vector products-one for the numerator and one for the denominator. However, we can use a recursive formula for the numerator: vg(0p;:)
Hg(O, 0o) + 9
Hg(0p- Tp- 1p-1dp- 1d 00) + +g
Hg(0p-1 0n) 98 8p- H g d
V2(0p-1: B)
(12)
So, if we store the gradient from the previous iteration (g(0p-1; B) and the Hessian vector product with the previous direction H dp-1, the current gradient can be computed without an additional matrix-vector product. At iteration p. we thus (i) compute the gradient Vq(Op: p B) recursively from the cached vectors Vg(0p-1i B) and H g dp-1 via Equation (12) (for p O. we have (g((u; B) gg). (ii) compute the Hessian-vects vector product with the current direction Hg dp. (ii)) store both these vectors and (iv) compute the update magnitude i (both the numerator and the denominator only require a simple dot product of two pre-computed vectors).
A.3 LAPLACE APPROXIMATION FOR NEURAL NETWORKS
A.3.1 DERIVATION THE LAPLACE APPROXIMATION FOR NEURAL NETWORKS
Preliminaries. The softmax function softmax : RC RCC is defined as
exp(z1)
exp(zc)
softmax(2)
exp(ze)
DS_ exp(2ee (ze)
maps an arbitrary vector z E C to a vector whose entries are non-negative and sum up to one. So, the output of the softmax function can be interpreted as probability distribution over C' classes. With this, the cross-entropy loss for a single datum (æ:, y) is given by
{(f(x)) y) = Σ ye (onssofmaxx(fo(z)))))
(13)
Here, we assume that E (0, 1] is a one-hot encoded vector representing the true class label. So, c. e (I,..., C') is the correct class (i.e. yo. 1). the cross-entropy loss is given by the negative log- arithm of the probability the network assigns to this class (.f(()), y) - logg softmax (eetmax(ff@x)))..... Finally, let
Cat(y: p) = Il d
(14)
denote the probability mass function of the categorical distribution, where E TRC is a vector of probabilities (i.e. Pe and Σ. Pc. 1), and y {0, 1] is one-hot encoded vector representing a class label.
Probabilistic interpretation of the regularized loss. Recall from Equation (1) that the regularized loss function is given by
Lag ( D) L(0:D) r(0) with L(,,D)
1 N
Σ ((fo(xn).))) RED
We assume classification problem that uses the cross-entropy loss (see Equation (13)) and an @ regularizer r(0)= 9/2 00 with parameter 3 E R2s. We use one-hot encoded labels yn E {0, 13 (with yne. denote the e-th entry of yx) and obtain
N
N Creg (0;
Σ ((o(xn) yn) D
1011
2
17

---

Published as a conference paper at ICLR 2025
(13)
Σ
Σ
Yne log(softmax(fo((e)))))
No 2
nED
(14)
Σ
log (Cat(gn; (ottma/))))t
1 0 (NB 110)
nED
log
π
Cat (ottmax:(
log
NN(0:0,
N
0:0,
1 N NB
Z,
likelihood p(D
prior p(0)
where P/2 absorbs the normalization constant of the Gaussian prior. The cross entropy loss L for the training set is thus connected to the negative log categorical likelihood and the regularizer can be seen as a negative log Gaussian prior over the parameters. Note that a similar derivation is possible for other loss functions as well, e.g. the MSE loss (which is equivalent to negative log Gaussian likelihood).
The derivation above shows that the (rescaled) regularized loss function can be interpreted as the negative unnormalized log posterior of a Bayesian model
N Leg 0: D) - P(D 0) log ( - - log p(0
(15)
with Gaussian prior and categorical likelihood. With _ we denote equality up to an additive constant.
Training as MAP estimation. Equation (15) allows re-interpreting the training procedure of a neural network as maximum posteriori (MAP) estimation problem since minimizing the regularized loss
arg min Leg(( (0, OERP
arg min (0. DD OERP
(15)
arg max log [
arg max p(0 ID) dek!
OERP
0.
0.
equivalent to maximizing the posterior p(0) D).
Laplace approximation. The idea of the Laplace approximation is to approximate the posterior distribution p(0 D) NN(;; 0., Ex) with a Gaussian at the mode 9, of the posterior, i.c. after training the network. For this, we approximate the log posterior by second-order Taylor expansion around the mode 0g O., i.e.
-N HD) 0.)
log ((0 D) _ - Creg ( D) - ()0 -0-
0.) (N
0.), (16)
where we assumed that 0, is (local) minimum of the regularized loss function (i.e. JD Veee(O,; D) 0) such that the linear term (0 0,) gD in the quadratic approximation van- ishes. The additive constants did not state explicitly in Equation (16) turn into some multiplicative factor (denoted by Z -11) when taking the exponential
- 0.) HD) (-)(0- 0.))
1 (-
p 0
D)
exp
This immediately identifies as the normalization constant of the Gaussian and we obtain the Laplace approximation
1 0 ID) N0 0., Ex) with Ep (N HD) N H
(17)
Mini-batch version of the Laplace approximation. In order to reduce the computational cost of the Laplace approximation, we can replace the full-batch quadratic g0; D) Equation (16) by mini-batch quadratic g0; B), i.e. we obtain Σσ N H
Predictive uncertainty via the Laplace approximation. Ultimately, we want to use the Laplace approximation equip the prediction for unknown test input æ. with uncertainty. Ideally, would compute the expectation of the model likelihood under the approximate posterior, i.e.
p Io, D)
ply) xo. p(O IB) d0
18

---

Published as a conference paper at ICLR 2025
Cat softma((fo(x=))) (0;00., Ex) d0.
One way to approximate this integral is via Monte Carlo sampling, i.c.
1
p Co. DD
Cat softmaxx( (xo)))
where ( ~ NOO., Eg).
(18)
S
As suggested by Immer et al. (2021); Roy et al. (2024), the linearized network f0)(x) (x) + (x) 0.), where V fe, (æ) gCxP is the network's Jacobian at 0..
A.3.2 SAMPLING FROM THE K-FAC LAPLACE APPROXIMATION
The Monte Carlo (MC) approach from Equation (18) requires samples from the Gaussian Noo., Σs) (the following derivations work exactly the same when the full-batch LA based the entire training set is used). However, ΣB E PRP is often too large to be built explicitly in memory and requires inverting the P x P Hessian Hg of the regularized loss function. One approach for that issue is to use the K-FAC curvature approximation Hg <² B) + 31 ~ Kg + 31, i.c. Eg v N -N'(KKK + BI) As we will see in the following, the K-FAC approximation enables to sample efficiently from the corresponding LA.
Leveraging K-FAC's block-diagonal structure. K-FAC block block-diagonal curvature approxi- mation, i.c. KB = blockdiag:- I..... ( and Kg + 3I inherits this structure. The inverse of block-diagonal matrix is also block-diagona, with the block inverses on the diagonal, i.c.
Eg
KB + 37)
N
1 blookdiag
(K + +BI))
N
blockdiag
[ + 3I) -1)
(1)
bloc
To draw sample IR from WoO. Es), we can thus simply sample from the block covariances ()()) (0, stack these samples and add the mean, i.e. v 0, + (v(1), v(())
Leveraging the blocks' Kronecker structure. To sample from .' Σ ()), can exploit the Kronecker structure of the blocks K! A @ B (we omit the layer index I for A and B for brevity). First, we compute the eigendecomposiiioos of the Kronecker factors, i.e. A = UASAU and B = SBS It follows
K0 B + BI = A + BI
U & UBSBU BI
-UA UB) (SA 8 (UAAUU)))) + 3I
=:U
=:
USUT +,3I
U(S + T
U UA & Up forms an orthogonal eigenbasis of the block and the diagonal matrix $ + 3I
SA ② SB 3I contains the corresponding eigenvalues (see e.g. (George al., 2018)). follows
0 1
[ + -11
N
1
= U(S 3I) 1-1--I1-'' -UUT
N
U(S -1)
(S 31) 1/U
19

---

Published as a conference paper at ICLR 2025
VV
So. in order to draw sample v()) ~ (0, E(0), we first draw from standard Gaussian N(O, I) and then transform the sample via vi)) = Vw. The resulting vector w!!) has mean 0 and covariance VIVT y As Gaussians are closed under affine linear transformations, w(() is indeed distributed according to (0, )
The vatrix-vector product w - Vw can be computed efficiently without actually forming in 1/2,
memory. The first step to multiply with (S /I)-1/2 Since this is a diagonal matrix, we can simply multiply the vector element-wise with the inverse square root of the diagonal entries of S + 3I. The second step is to multiply with = UA @ UB. which can be implemented efficiently by using the property (U & BB) vec( W) vec WU)). Finally, we scale by N-1/2,
Summary. Drawing sample from the K-FAC LA NO,, 31)= can be done without ever forming the blocks of the covariance matrix explicitly. Using properties of the Kronecker product, we can efficiently transform a sample from the standard Gaussian to a sample v(!) ~ No E Stacking the samples from all blocks and adding the mean 0, yields sample from the LA due the block-diagonal structure of the covariance matrix.
A.3.3 DEBIASED K-FAC LAPLACE APPROXIMATION
The idea of the debiased K-FAC Laplace approximation is to construct one Kronecker-factored curvature matrix from two mini-batches B and B.
Eigenbasis of a K-FAC block. First, recall that the eigendecomposition of block K from K-FAC matrix can constructed from the eigendecompositions of the Kronecker factors (see e.g.
(!]
(George al., 2018)). Consider the I-th block from both K-FAC approximations K! = A & B and K (1) C (we omit the layer index for A, B, C, and for brevity). Again, let A A=UASAA and B BUBSSU denote the eigendecompossiiinn of the Kronecker factors. It holds
K B
A& = U & UUDSRUD
(UA @ Un) (SA @
=:U
=:
USU
U= & Up forms an orthogonal cigenbasis of the block and the diagonal matrix S = S & SB contains the eigenvalues.
Re-evaluation of the directional curvatures. For the debiased approach, we keep the block's eigenbasis U, but instead using the directional curvatures S, we re-evaluate these measurements on the second mini-batch B. First, consider the projection of the block (1) onto the eigenvectors U, i.c.
U
DDUAAUU) Un) ((
UJCUAOUUUU
The debiased directional curvatures the diagonal of U k U. For a square matrix X Rexn 1et Diag(X( denote the operator that maps the matrix onto its diagonal, i.c. Diag: : gaxn R with Diag( == X for i (l...., n). It holds:
B
Diag(U
k U) Diag (ang(U/CO UUCUAAU B
Diag Diag(U)CCUA) & Diag U
=:
BB
ŠA 8 šB
=: 8.
20

---

Published as a conference paper at ICLR 2025
Construction of debiased block. Now that we have the eigenbasis U and the debiased directional
curvatures 8, construct the debiased block. Let S = diag($ ()). SB diag($8) (&B) and
$ diag($) 8) diag( 5 8B) diag 8 diag 8 S X Sn.
The debiased block is given by USUT. It can be written as the Kronecker product of two matrices A and B since
USU (UA@UB)(SA (B)(UA
=: A =:
This is important as allows for efficient sampling as described in Appendix A.3.2.
Computational cost. Since we need the eigendecomposstions of the Kronecker factors A UASAUA and B = UBSB UB for sampling in any case (see Appendix A.3.2), the com- putational overhead for the debiased K-FAC approximation À & B consists of computing K-FAC approximation C' & D on another mini-batch, re-evaluating the directional curvatures A D and в B UB) and finally computing Á U diag($A) U and = UB diag(āB) -
From block-level to full matrix. So far, have only considered the debiasing of single block. However, correcting the blocks's ' eigenvalues sufficient because they coincide with the eigenvalues of the full matrix (due to the block-diagona structure). Let (!)) denote an eigenvector of the 1-th block ( of some block block-diagonal matrix X blockdiag.........XX (X()) corresponding to the
eigenvalue A. Then, u = (OT, u ,...,0' OT) is an eigenvector of X corresponding to the same eigenvalue X, because
(
0
X u = blookdiag blo. ((x")
X (1)
\u = Au.
==
=
( 0
The eigenvalues of X thus coincide with the eigenvalues of its blocks; and X's eigenvectors can be constructed from the the eigenvectors of the blocks by filling them up with zeros.
Connection Equation (9): The equivalence between Equation (9) and the approach we describe above may not be obvious. We thus show here that the directional curvature of the debiased matrix K along an eigenvector u of K3 indeed coincides with the directional curvature u K g on B,
More concretely, let K blockdiag:.I..... (K (0) denote the debiased K-FAC approximation u (T) OT) denote eigenvector of KB, where (!)) is the i-th eigenvector of KD (i.e. the i-th column of U). It holds
constructed from Kg and K& as described above. Also, let u = (0
u Ku (1) K O (0) UŠUT ( e_ee, 5, u ( (1) ( K u
where e; denotes the i-th eigenvector.
A.4 BIAS IN THE DIRECTIONAL SLOPE
Biased directional slopes along negative gradient directions. Assume that given quadratic q( 13) around the current parameters 0g. We consider the negative normalized gradient direction
Vg(O.; B) (q B)| at some location 0, evaluated mini-batch B. We have
da q(: B da ( B) (V.; [1] cos(a)) > da ((.::
(19)
where a == L(Vg e, B), Vg 0,: B3) and we assumed Vq(0,; B3) = / 0,: BB| which is true least in expectation. Projecting gradient onto its negative direction the right-hand side of the
21

---

Published as a conference paper at ICLR 2025
inequality wwill always result in a directional slope that is & since
V) [v((.: Vg(0.: = - VV. :3)l =
La ( B)) = d' Vg((.,
and = 0 only if (q(00
: B) = 0. Projecting a different gradient (of equal length) onto that direction
that is the left-hand side of the inequality- result in larger (possibly even positive) directional
slope.
Proof of Equation (19). Equation (19) quantifies the bias in the directional slope along d function of the alignment of the two mini-batch gradients. It holds
dd (0,;; B) Ea (0.; B) = d B)) d Vg(0.; B)
d Vi 0.: [3 CoS((T) d [q(0.: cos(?)
= [v. B)][( --s((1)))))))
where = <(d, (Q. B)) and we assumed that VV((..; B)]l (Q) in the last step (which is true, at least, expectation). Next, we re-write ?
7 (d, Vq 0.; B)) VV(O0); Vg(O.; B)) (VV Vo
-:0
follows- -1 cos(1) (7)
cos(T ा (x) cos((x) Substituting this into the expression for
the bias, arrive at
a (00 B) da (0.: B) = VQ(0.::)](
cos(r)) = (-1-cc)
[)V(( 1),
from which Equation (19) follows. D
The first CG search direction. Assume that the current parameters are Ou and we apply CG the local quadratic approximation o( B). The very first search direction is the quadratic's 's normalized negative gradient do above, where 0. is set to O,. Equation (19) thus explains the bias for the first CG direction.
Vy Oo: B) VQ((0; B- at 0g. This is exactly the situation we describe
The subsequent CG search directions. For the subsequent CG search directions (d, with p > 1), the situation is more complex: Each direction d, is linear combination of the current normalized negative residual -r, (note that Tp coincides with the gradient Vg(O,, p B), see Appendix A.2) and the previous search direction dz- 1 (see Appendix A.2) and the previous search direction dz-1 (see Algorithm 1), i.e.
Sp 1p
Pp |spl|
||rp1 |
8pl[sp-11 111
-Tp rpl
dp
rp)
+
8p-1
dy 1
1/1
+ 7/2 d
|8p
|rrp]
||pp||
30
=1pp33
The additional correction term ensures conjugacy. The directional slope along dp thus also splits into two corresponding terms: The slope along the negative gradient direction and the slope along the previous search direction. It holds
ea ()) d, ())
Tp |r,
vg(Op: + 1/2 d Vg(0p:B
= a( (0,; B) + ( B).
The bias in the first term-at explained by Equation (19)~also introduces bias along the search direction dp-
A.5 BIAS IN THE DIRECTIONAL CURVATURE
Derivation for Equation (6). Let Hg = UAUT and H ÜÄÜT denote the eigendecomposi tions of the mini-bate batch Hessians, where U (ui,..., up) U (ū,,..., ur) R. PxP contain
22

---

Published as a conference paper at ICLR 2025
the orthonormal eigenvectors and A diag(A1..... Xp), A diag( (tag() the respective eigenvalues in descending order, i.e. X1 ... Xr and dι along one of H8" ss eigenvectors u, on mini-batch B is given by the corresponding eigenvalue since
Xr. The directional curvature
2 ((.; B) u, Hgu; = u, a = >r|lu; As.
For B, we obtain
P u,
P
P
2 nn
(0, B) = u Hgu; = u
Σ
Up
Σ
u.)²
Σ
Xp
S2ip
p-
Weights sum up to one. The weights Di are non-negative and sum up to one, i.e. > This is because
Slip 1.
P Σ
P
Σ
up u:)²
I0 u. u, ÜÜ u,
(20)
p=1
where we used that Ü is an orthogonal matrix and that u, normalized.
Proof of Equation (7). Here, we assume the simplified case where the spectra of H8 and Hg are identical, i.. Ap Ap Vp P).
First, consider us and assume that and in are not perfectly aligned, i.e. S2 (u, iz)² < 1. It holds
P
a
g(0g
B
6
Σ
-
1 +
Σ A p-2
The sum can be bounded from above by putting all remaining weight 1 52 (see Equation (20)) on the second-largest eigenvalue X2, i.e.
Cn a(u.. B
21 + A2(1 22,1) & 11 A1 11,1) = A1 2 (0. B).
The inequality (*) turns into < if the top two cigenvalues are separated, i.e. A1 d2. The proof for up is similar. We assume that up and up are not perfectly aligned, i.e. S2p, = (ufūp)² < 1. follows
P- Σ Spp Apppp.p
c (0. B)
16 Σ
Spp
This time, we bound the sum from below by putting all remaining weight -52pp (see Equation (20) on the recond-smalless eigenvalue Ap-1, i.c
2n Co. 13) 3 A1-1(1 - 2p.r) Arprp Ar - 2prr) + ApΩpp = Ap = Ga, (6, B).
Again, the inequality (*) turns into > if the bottom two eigenvalues are separated, i.e. Ap-1 > Ap. This concludes the proof of Equation (7).
B EXPERIMENTAL DETAILS
In the following, we provide information on the experimental setup and give detailed instructions on how to replicate all empirical results.
23

---

Published as a conference paper at ICLR 2025
B.1 TEST PROBLEMS AND TRAINING PROCEDURES
Throughout the paper, we use a series of test problems with different models, datasets, and training procedures which we describe in more detail in the following. We use DEEPOBS (Schneider et al., 2019) on top of PYTORCH (Paszke et al., 2019) as our general benchmarking framework as it provides easy access to a variety of datasets and model architectures.
Data. We use the datasets CIFAR-10 (with C = 10 classes) and CIFAR- 100 (with C' = 100 classes) (Krizhevsky, 2009). Each dataset contains 60,000 data points that are split into 40,000 training samples. 10,000 validation samples and 10,000 test samples. For the experiments on out-of- distribution (OOD) data, we create the datasets CIFAR- 10-C and CIFAR- 100-C, each containing 10,000 images, described in (Hendrycks & Dietterich, 2019). In these datasets, each image is corrupted using one out of 15 different corruptions (chosen uniformly at random) at a specific severity level (a number between and 5). We also use the IMAGENET dataset (Deng et al., 2009) which contains images from C' 1000 different classes.
Test problems. All of the following test problems use the cross-entropy loss function.
(A) ALL-CNN-C on CIFAR-100. This test problem uses the ALL-CNN-C model architecture (Springenberg et al., 2015) (where we removed the dropout layers as explained in Section 3. and CIFAR- 100 data. The training hyperparameters are taken from an existing benchmark (Schmidt et al., 2021): We train the model with SGD (learning rate 0. 0.171234) with batch size 256 for 350 epochs. Weight decay 3 = 0.0005 is used on the weights but not the biases of the model.
(B) ALL-CNN-C on CIFAR-10. This test problem is similar to (A) but uses the CIFAR- 10 dataset. The model is trained with SGD (learnnng rate 0.025, momentum 0.9) with batch size 256 for 350 epochs. The learning rate is reduced by a factor of 10 at epochs 200, 250 and 300, as suggested in (Springenberg et al., 2015). Weight decay 8 0.001 is used on the weights but not the biases of the model.
(C) WIDERESNET 40-4 on CIFAR-100. This is a test problem from DEEPOBS. For details on the architecture, see (Zagoruyko & Komodakis, 2017). We train this model with SGD (learning rate 1, momentum 0.9) with batch size 128 for 160 epochs. The learning rate reduced by factor of 5 after 60 and 120 epochs. Weight decay B 0.0005 is used on the non-bias weights of the model.
(D) CONvNET CIFAR-10. This test problem uses simple convolutional neural network with variable depth and width w (denoted by "model d-w" in Figures 7 and 22). The first block of the model consists of convolutional layer (kernel size 5, padding 2) with 3 input and w output channels, a ReLU activation function, and a max-pooling layer (kernel size 2). The last block consists of max-pooling layer (kernel size 2). flatten layer, and dense linear layer. In between those blocks, there are d hidden blocks each consisting of a convolutional layer (kernel size 5, padding 2) with w input and w output channels followed by a ReLU activation function. So, the depth d determines the number of hidden blocks, whereas w controls the number of parameters in the layers and is thus referred to as the width. We use d E {1, 4, 7) and w E {32, 64, 128} (resulting in 9 models) and train each model for 100 epochs on the CIFAR- 10 dataset using ADAM with standard hyperparameters (learning rate 0.001, 31 = 0.9, 82 0.999, e 10 8) with batch size 256. No weight decay is used for this test problem.
(E) RESNET-50 IMAGENET. This test problem uses the RESNET-50 model architecture (He et al., 2016) and the IMAGENET dataset. We use the IMAGENETI K_ V1 weights pre- trained on IMAGENET- 1K using SGD (initial learning rate 0.1, momentum 0.9) with batch size 256 for 90 epochs. For the pre-training, the learning rate was set by multistep scheduler which multiplied the learning rate by factor of 0.1 after every 30 epochs. Additionally, a weight decay of B = 0.0001 was used on all weights apart from biases and the learned batch normalization weights.
(F) VIT LITTLE on IMAGENET. This test problem uses the VIT LITTLE architecture of the VISION TRANSFORMER model family (Dosovitskiy et al., 2021) on the IMAGENET dataset. We use the VIT LITTLE P REG4 GAP 256.sBB IN к weights pre-trained on IMAGENET-LE 1K using NADAMW. For the exact hyperparameters used for the weights, refer to the HuggingFace Model Card.
24

---

Published as a conference paper at ICLR 2025
During training, store the model's parameters 10 checkpoints spaced rigeeeuiiittntly between the first and last epoch. The training metrics for all test problems are shown in Figure 7.
B.2 MATRIX-VECTOR PRODUCTS WITH THE CURVATURE MATRIX
Access to curvature matrices via BACKPACK. In order to compute eigendecomposition of the curvature matrix or apply the CG method to quadratic, we need access to matrix-vector products with the curvature matrix v HB- v. When products with the full-batch curvature matrix are required, we accumulate the mini-batch quantities over manageable chunks of the training set. Our implementation uses BACKPACK (Dangel al., 2020) that provides access to products with the Hessian - (2; B) and GGN e Gg of the empirical risk L as well as the K-FAC curvature approximation Kg.
Eigenvectors of the curvature matrix. Given access to matrix- products with the curvature matrix, we can construct an instance of a scipy- sparse. linalg. LinearOperatee that can be used in scipy. sparse. linall....... e eigh.
B.3 SECTION 1 (INTRODUCTION), FIGURE 1
For the visual abstract in Figure I, we use the fully trained ALL-CNN-C model from test problem (A) (see Appendix B.1). The experimental procedure below is repeated for 5 different mini-batches Bm. m E {0, 1,2, 3, 4] of size [Bm = 512 (we omit the index m in the following).
Experimental procedure. We use the GGN curvature proxy Hg G5 3II and compute its top two eigenvectors uy and U2 (see Appendix B.2). In order to evaluate the quadratic q(f.; B) in the 2D) space spanned by those eigenvectors efficiently, use the following equation:
q(O. Tiul + T222; B)
2 1ilu HB un + T HB u2] + - [u HB HBug]
+ n [u] gs] r2[u]] gs] [cs].
We can compute all the terms on the right-hand side in brackets once, store them, and then evaluate the quadratic for arbitrary values of and 72 basically cost. The derivation for the full-batch quadratic q(,, + Ti u1 + T2u2; D) is analogous.
B.4 SECTION 3 (THE SHAPE OF A MINI- BATCH QUADRATIC): FIGURES TO
For the evaluation of the directional derivatives, we use the fully trained ALL-CNN-C model from test problem (A) (see Appendix B.J).
Three settings. Throughout this section, consider three different settings: (i) Quadraties that use the Hessian of the empirical risk, i.e. Hg V + 3I at batch size 512. (ii) quadratics that use the Hessian's 's GGN approximation, i.c. Hg Gg + 3I batch size 512 and (iii) at batch size 2048. For each setting, the experimental procedure below is repeated for 3 mini- batchss Bms m E {0, 1, 2). The prior precision 3 is set to 0.0005 (the same B was used for training, see Appendix B. and only acts on the weights of the model but not its biases.
Experimental procedure (biases). The experimental procedure consists of two steps: The compu- tations of the directions on mini-batch B,, and the evaluation of the directional derivatives all mini-batches (of the same size) in the training dataset.
1. Directions based on Bm. First, we use the quadratic q(0.; B..) with the given curvature proxy and batch size to compute a set of 100 directions. This is either the top 100 eigen- vectors of Hg computed via scipy sparse dinalgg eigsh (see Appendix B.2) or the first 100 CG search directions (see Appendix B.5). In the case CG, we also store the trajectory of the iterates 0, On, fi...., 0100.
2. Directional derivatives on all curvature for all directions from step 1 on all mini-batches Bm,', m/ (O...., M - 1} in the training dataset (where M denotes the number of mini mini-batches in the training data).
Finally, we evaluate the directional slope and
Directional derivatives along eigenvectors. For an eigenvector u the directional slope and curvature are given by du q (,.; Bmn) and 22 q(.:; : B,n), respectively.
25

---

Published as a conference paper at ICLR 2025
Loss
Accuracy in %
Training loss
1-
Test loss
75-
50-
(A)
25-
Training accuracy
Test accuracy
0-
100
200
300
100
200
300
Epoch
Epoch
Loss
Accuracy in %
Training loss
Test loss
80
2
60
(B)
Training accuracy
20
Test accuracy
100
200
300
100
200
300
Epoch
Epoch
Loss
Accuracy in %
Training loss
Test loss
75-
4-
(C)
50
25-
Training accuracy
Test accuracy
50
100
150
50
100
150
Epoch
Epoch
Training loss
Test loss
Training accuracy in
Test accuracy in %
2.0
80
60
1.5-
60-
(D)
1.0-
1.0
0.5
20-
20
50
100
50
100
0
50
100
0
50
100
Epoch
Epoch
Epoch
Epoch
model 1.32 (P 48,554)
model 7:32 (P 202,346)
model 7 7.64 = 763.082)
model 4-32 (P 125,450)
model 4-64 (P 155,690)
model 4-128 (P 1.730,570)
model (P 148,298)
model 1-128 501,386)
model 7-128 2.939,754)
Figure 7: Training metrics for test problems. The panels show the training/test loss/accuracy during training for all test problems (A) (ALL-CNN-C on CIFAR- 100), (B) (ALL-CNN-CC on CIFAR- 10), (C) (WIDERESNET 40-4 on CIFAR- 100) and (D) (CONvNET on CIFAR- 10). Test problem (E) (RESNET-50 50 IMAGENET) is not shown since pre-trained model is used here.
26

---

Published as a conference paper at ICLR 2025
Directional derivatives along CG search directions. For CG, we evaluate the direc- tional slope and curvature along each search direction dp at the respective iterate 8, of the trajectory, i.e. Cd ((0 Bur) ) for the slope and C2, 7 p0pp Bm)) for the curvature. As the curvature (0q B,n)) H is independent of 0, this small distinction is irrelevant for the directional curvature. For the directional slope, however, does make difference.
Results (biases). The results for setting (i) are shown in Figure 8, for setting (ii) in Figures 2 and 9 and for (iii) in Figure 10. The biases in the directional slopes and curvatures can be observed across all scenarios. The biases in the slopes are even more pronounced along the CG search directions than along the cigenvectors of the curvature matrix. The opposite holds for the curvature biases that tend to be larger along the eigenvectors. This is consistent with our explanations from Section 3.2. Both the biases the slope and curvature decrease with increasing mini-batch size. For setting (i), CG encounters search direction with negative curvature (indicating that, expected, the curvature matrix Hg is indefinite) and thus terminates already after 4 iterations.
Experimental procedure (overlaps). Next, we compute the overlaps between the cigenspaces of the curvature matrices different bani-bbachhes B and B. More specifically, we compute S (u, úp)² where i, p. E {L...., 100). u, is an eigenvector of Hg and Up is an eigenvector of H The values Rh.p are bounded between 0 (u, and Up are orthogonal) and 1 (u, and ü, are identical, up to their sign). We apply log-transform St.p- In Figure and Appendix B.4, values below -8 (i.e. D & 10- are shown in black, values equal to 0 (i.e. R. 10" 1) are shown in white.
Results (overlaps). The results are shown in Appendix B.4. They show that the eigenspaces of the curvature matrices are not perfectly aligned: Eigenvectors from B,, typically overlap with several eigenvectors from B. The eigenspaces are more aligned at the larger batch size 2048 than at batch size 512. This seems reasonable since random vectors in high-dimensional spaces tend to be orthogonal to each other- leess stochasticity (due to a larger batch size) thus leads to better alignment.
B.5 SECTION 6.1 (DEBIASED CONJUGATE GRADIENTS): FIGURE
Here, we describe the experimental details for Figure 5 from Section 6.1. For the derivation and the mathematical details of the standard singlebbach CG method and the debiased approach, see Appendix A.2. The experiment uses the fully trained ALL-CNN-C model from test problem (A) (see Appendix B. and the GGN curvature proxy Hg - GB + 3I with B 0.0005.
Experimental procedure. The experimental procedure consists of two steps: The computation of the CG trajectories and the evaluation of the four performance metrics.
Computation of the CG trajectories. For the single-batch CG approach, we use one mini-batch of size 1024, apply K = 30 CG iterations and store the trajectory 0.
0p,01, .I..... O20- For the debiased approach (details in Appendix A.2), we use two mini- batches of size 512 to construct the trajectory. As the debiased approach uses total of 60 vectvr products (30 for the search directions and 30 for the update magnitudes) at half the computational cost (since the cost for the GGN- vector product scales linearly with the mini- bai-b size), the total cost of both approaches are comparable.
For both approaches, the above procedure is repeated for 5 different mini- menn-baacces/mmii--eee -atch pairs, such that 10 trajectories are computed in total.
We compute two additional trajectories: The standard CG trajectory for the full-batch quadratic q(: D) and the trajectory for debiased full-batch approach, where the training data for the directions and the other half for the update magnitudes.
half
2. Evaluation of performance metrics. For each of the 12 trajectories, we evaluate four performance metries: (i) the regularized loss Leeg (see Equation (1)) on the training set D and (ii) on the test set Prest as well as the accuracy (i.e. the relative number of correctly classified samples) on the (iii) training and (iv) test dataset.
Results. The mini-batch results are shown in Figure 5 and discussed in Section 6.1. The results for the two full-batch CG variants are shown in Figure 12. Surprisingly, the full-batch CG approach diverges after roughly 10 iterations. If we consider the training data as a (large) sample from the data
27

---

Published as a conference paper at ICLR 2025
Biases: Additional results for H8 V2L(0,, + 3I (batch size 512)
Eigenvectors computed on mini-batch
Bg
B2
Directional slope
1.0-
0.5
0.0
Permuted eigenvector index
Permuted eigenvector index
Permuted eigenvector index
Directional eurvature
10
10"
50
75
100
25
75
100
50
100
Eigenvector index
Eigenvector index
Eigenvector index
CG trajectory and search directions computed on mini-batch
Bu
B₁
B₂
Directional slope
2.5
0.0
2.5
CG search direction index
CG search direction index
CG search direction index
Directional curvature
101
10"
CG search direction index
CG search direction index
CG search direction index
Figure 8: Directional slopes and curvatures are biased. The experimental setting is similar to Figure but uses the Hessian of the empirical risk, i.e. Hg 2 L(O.; B)) + ,3I at batch size 512. The upper plot (Top) shows the directional derivatives along the top 100 cigenvectors of Hg. the lower plot (Bottom) shows the directional derivatives along the first 100 CG search directions. For the top panel of the upper plot, we switch the order and sign of the eigenvectors such that the orange dots all above zero and in descending order. There strong systematic biases in the directional slopes and curvatures. The curvature biases are more pronounced along the eigenvectors, whereas the biases in the slope are larger along the CG directions.
28

---

Published as a conference paper at ICLR 2025
Biases: Additional results for Hş G + BI (batch size 512)
Eigenvectors computed on mini-batch
Bo
B
Ba
Directional slope
1.0-
0.5+
0.0-
Permuted eigenvector index
Permuted eigenvector index
Permuted eigenvector index
Directional curvature
10
10
76
100
Figenvector index
Rigenvector index
Eigenvector index
CG trajectory and search directions computed mini-batch
Bp
B:
B2
Directional slope
-2-
CG search direction index
CG search direction index
CG search direction index
Directional curvature
10
10-1_
25 50
100
25
50
100
25
100
CG search direction index
CG search direction index
CG search direction index
Figure 9: Directional slopes and curvatures are biased. The experimental setting the same as in Figure 2: We use the GGN curvature proxy Hg GB 3I at batch size 512. The upper plot (Top) shows the directional derivatives along the top 100 eigenvectors of HB. the lower plot (Bottom) shows the directional derivatives along the first 100 CG search directions. For the top panel of the upper plot, we switch the order and sign of the eigenvectors such that the orange dots are all above zero and in descending order. There are strong systematic biases in the directional slopes and curvatures. The curvature biases are more pronounced along the eigenvectors, whereas the biases in the slope are larger along the CG directions.
29

---

Published as a conference paper at ICLR 2025
Biases: Additional results for H5 GB 31 (batch size 2048)
Eigenvectors computed on mini-batch
Bo
B
B₂
1.0
Directional slope
0.5
0.0-
Permuted eigenvector index
Permuted eigenvector index
Permuted eigenvector index
Directional curvature
10
10a
76
100
Figenvector index
Rigenvector index
Eigenvector index
CG trajectory and search directions computed mini-batch
Bp
B:
B2
Directional slope
2-
04
CG search direction index
CG search direction index
CG search direction index
Directional curvature
10
10%
10-1_
25
50
100
25
50
100
0
25
50
100
CG search direction index
CG search direction index
CG search direction index
Figure 10: Directional slopes and curvatures biased. The experimental setting is similar to Figure 2: We use the GGN curvature proxy Hg GB + 3I but batch size 2048. The upper plot (Top) shows the directional derivatives along the top 100 eigenvectors of HB. the lower plot (Bottom) shows the directional derivatives along the first 100 CG search directions. For the top panel of the upper plot, we switch the order and sign of the eigenvectors such that the orange dots are all above zero and in descending order. There are strong systematic biases in the directional slopes and curvatures. The curvature biases are more pronounced along the eigenvectors, whereas the biases in the slope are larger along the CG directions.
30

---

Published as a conference paper at ICLR 2025
Eigenspace overlaps: Additional result for HB V2L(0,,; BI
Batch size 512
B
Bo
Bi
Ba
Eigenspace overlaps: Additional results for H8 - Gg + 31
Batch size 512
Batch size 248
Bi
B
B2
Bo
B
B2
Bo
Bo
Bi
B
Ba
Ba
Figure 11: In practice, eigenspaces are misaligned. The experimental setting is similar Figure The top plot shows overlaps between three curvature matrices that use the Hessian of the empirical risk, i.e. Hg C²C(0.: 3I batch size 512. The bottom plots use the GGN curvature proxy Hg Gg + 3I at batch size 512 (Left) and 2048 (Right). The weights Ri are shown as a 100 100 greyscale image (color ranges from black for D. & 10-8 to white for 2, = 1) for m, mmm (0, 1, 2). Clearly, the eigenspaces for different mini-batches
perfectly aligned as cigenvectors from B,, overlap with several eigenvectors from B, At the larger batch size, the eigenspaces are more aligned.
31

---

Published as a conference paper at ICLR 2025
distribution, the full-batch CG approach can be seen instance of the mini- mini-batch approach with very large mini-batch size. Thus, might also exhibit the associated biases we describe in Section 3.2. Another explanation would be the approximation error in Equation (2): (- R g(: :D). If the trajectory moves far away from 0.. it leaves the region where the approximation is valid (whereas the debiased full-batch version might take smaller steps and thus remains stable). We leave it to future work to investigate this further.
Additional results. We also provide the results for two other mini-batch sizes [B] (512, 2048) (the debiased approach uses two mini-batches of size 256 and 1024, respectively). The results are shown in Figure 13.
Runtime and memory consumption. Our current implementation uses the naive approach that requires two matrix-vector products with the GGN for computing one debiased update magnitude (see Appendix A.2 for details). As this would skew the runtime comparison, we refrain from providing concrete numbers here. However, as we describe in detail in Appendix A.2, is possible to implement the debiased approach in more efficient way that only requires one additional matrix-vector product for each CG direction to compute the debiased update magnitude. Using the debiased approach at half the mini-batch size, this brings the computational costs down to roughly the same level as the standard approach (assuming that the runtime of a GGN-vector product scales linearly with the mini-bated size).
As explained in detail in Appendix A.2, the memory overhead for the efficiently implemented debiased CG approach is two additional vectors of size P. For the experimental setting used for Figure 5, the number of parameters is P = 1,387, 108, which corresponds to 5.55 MB in single precision. The computational overhead incurred by the debiased approach is thus about 11.1 MB.
B.6 SECTION 6.2 (DEBIASED LAPLACE APPROXIMATION) FIGURE 6
Figure shows a comparison of the vanilla model, the mini- K-FAC LA, the debiased version, and the full-batch LA. Here, we describe the experiment in more detail and present additional results, in particular on OOD data. For the derivation and the mathematical details of these approaches, see Appendix A.3. The experiment uses the fully trained ALL-CNN-C model on CIFAR-10 data from test problem (B) (see Appendix B.
The experimental procedure consists of three steps: The computation of the K-FAC curvature approximations, the evaluation of the corresponding predictive class probabili- ties, and the computation of the performance metrics.
Experimental procedure.
1. K-FAC curvature approximations. We consider a log-equidistant grid of 13 prior preci- sions 3 between 10 + and 10" and add B 10.
For each of those 14 values, we compute the K-FAC curvature approximation via BACK- PACK (Dangel et al., 2020) (see Appendix B.2) using (i) a single batch of size 64, 256, and 1024, (ii) the two-batch debiased LA version (details in Appendix A.3.3) at batch sizes 32, 128, and 512, and (iii) the full-batch LA. For the latter, we accumulate mini-batch K-FAC approximations over the entire training set. As additional baseline, we consider (iv) the vanilla model (without LA), which is independent of the prior precision B. For approaches (i) and (ii), we repeat the experiment for different mini- bww.bbaa/hhhhbbbbiiiiiieeee batch pairs.
2. Evaluation of predictive class probabilities. The first step in the evaluation of the perfor- mance metries is the computation of the predictive uncertainty. For a given test dataset of size No, these can be represented as a matrix P E T N.C, where Pa.c. is the probability that the n-th sample belongs to class c, i.c. rows of P sum to 1. In all experiments, we draw S 40 MC samples {0(*)} | from the weight posterior following the procedure in Appendix A.3.2. The predictive class probabilities are then obtained via Equation (18). In the case of the vanilla model without LA, the sum Equation (18) collapses to a single term corresponding to the MAP model.
The evaluation procedure above can be applied to arbitrary test datasets. We consider the training data, test data and CIFAR- 10-C D-C (i.e. OOD datasets at severity 1 to 5).
3. Performance metrics. We consider the following performance metrics:
Accuracy. The predictive classes are obtained from P by extracting the class with the highest probability. The accuracy is the relative number of correctly classified samples.
32

---

Published as a conference paper at ICLR 2025
Reg. training loss 4
Reg. test loss
Training accuracy
Test accuracy
1
100
Accuracy in %
1.2
3.25
3.00
55-
0.8
80
50-
1
10
30
10
30
10
30
10
30
CG iteration
CG iteration
CG iteration
CG iteration
[8]
[P]
|D] /2 (debinsed)
Figure 12: Comparison of full-batch CG approaches. We use the setting from Figure The full-batch approach (shown -) applies standard CG to q D) while the debiased approach (shown as -) uses one half of the training data for the directions and the other half for the update magnitudes. The markers and are placed at peak performance. Surprisingly, the full-batch run diverges. The debiased CG run stable.
Reg. training loss +
Reg. test loss +
Training accuracy f
Test accuracy t
3.4
1.00
Accuracy in %
60-
Loss
3.2-
0.95
59-
92-
3.0-
58
10
10
30
1
10
30
CG iteration
CG iteration
CG iteration
CG iteration
[8|
88]
(debinsed)
Reg. training loss
+
Reg. test loss 4
Training accuracy 1
Test accuracy 1
3.4
1.00
Accuracy in %
Loss
3.2-
0.95
59-
92-
0.90
3.0-
58
10
10
30
30
CG iteration
CG iteration
CG iteration
CG iteration
[B] = 1024
[8]
512 (debiased)
Reg. training loss
Reg. test loss L
Training accuracy 1
Test accuracy 1
1.00
60
Loss
3.2-
Accuracy
9-1
0.95
59-
0.90
3.0
10
10
30
CG iteration
CG iteration
CG iteration
CG iteration
B 2048
18] 1024 (debiased)
Figure 13: Comparison of CG approaches at different batch sizes. We use the same experimental setting as Figure but consider different mini-batch sizes [B] E {512, 1024, 2048) (from (Top) to (Bottom)). The debiased approach uses two mini-batches of size 256, 512, and 1024, respectively. Across all mini-batch sizes, the single-batch CG runs diverge quickly while our debiased approach maintains stability.
33

---

Published as a conference paper at ICLR 2025
Negative log-likelihood (NLL). For each datum in the test set, compute the negative log-pre probability of the true class. We then average over all samples in the test dataset. This coincides with the empirical risk from Equation (I), evaluated on the test dataset.
Expected calibration error (ECE). The ECE is a measure of the calibration of the model's s predictive probabilities. It groups the classification confidences (i.e. the maxi- mum entry in each of P's rows) into bins, and within these bins, compares the average confidence with the actual accuracy. We use MalticlassCalibratee from torchmet x zics (Detlefsen et al., 2022).
AUROC. For the OOD datasets, we provide the Area Under the Receiver Operating Characteristic curve (AUROC). Our goal is to distinguish in-distribution (ID) and cutoo--distribution (OOD) samples using the model's uncertainty. We use the entropy of the predictive distribution p(ye xx. D) as our uncertainty estimate u(x.) E R for input Eg. The ground-truth binary labels are the ID/OOD indicators of the test inputs. Several thresholds & could be established to turn the scalar uncertainty estimates into binary predictions by using 1 (u(z))=E)- Instead of choosing particular threshold and evaluating its accuracy, the AUROC metric directly evaluates u( by measuring the area under the plot of the true positive rate (TPR) against the false positive rate (FPR) for all possible thresholds.
Additional experimental details.
Uncertainty over weights but not biases. As the prior acts only on the weights of the network but not its biases (see Appendix B.1), we only consider the uncertainty over the weights in the LA. This slightly reduces the size of the covariance matrix as it excludes the bias parameters.
Single vs. double precision. Although the K-FAC factors are positive semi-definite by construction, their numerical eigenvalues can be negative due to numerical inaccuracies. To balance precision and computational cost, we use double precision for all computations until drawing the MC samples and single precision afterwards.
Results. The results presented in Figures 6, 14 and 15. For the single-batch and debiased approach pairs each), we report the average performance as dot and
(where we use 5 minni-btchess/mini-baae the min/max vertical line. The performance of the vanilla model is shown as a horizontal line as its performance does not depend on 3, as explained above.
Across all datasets and performance metrics, the debiased LA mimics the behavior of the full-batch approach much better than the single-batch LA (although both approaches use a comparable amount of computational resources). In particular, for small prior precisions, where the covariance matrix relies almost exclusively on the K-FAC curvature information. the debiased LA maintains good performance in contrast to the single-batol approach.
Additional results.
RESNET-50 on IMAGENET. To showcase the scalability of our debiased LA approach, we repeat the experiment for a RESNET-50 model on the IMAGENET dataset (test problem (E) in Appendix B. The results are shown in Figure 16. Again, the debiased approach behaves similarly to the full-batch LA and maintains stability over the entire spectrum of prior precisions.
VIT LITTLE on IMAGENET. Similar to the previous experiment, we also apply our LA debiasing approach to a VIT LITTLE model on the IMAGENET dataset (test problem (F) in Appendix B. to ensure that it is beneficial for other architectures well. The results shown in Figure 17. The debiased approach mimics the calibration behavior of the full-batch LA remarkably well and remains stable for lower prior precision values as well where the single-batch approach behaves erratically.
Runtime analysis (computational overhead of debiasing). We claim in Section 6.2 that, by using the debiased approach at half the mini-bated size that is used by the single-batch approach, we obtain a fair comparison (more details in Section 4.2.3). Here, we substantiate this claim by providing a detailed runtime comparison of the biased, debiased, and full-batch LA approaches on CIF.R-10.
34

---

Published as a conference paper at ICLR 2025
Performance on the training set
Accuracy f
NLL loss
ECE
Vanilla model fa.
10
[S]]
[ 256
99
10
[3] 1024
10
(debiased)
[ 128 (debiased)
[8] 512 (debiased)
10
10
10'
10
10
10a 10'
10a
10
10 10'
Full-batch
Prior precision 3
Prior precision 8
Prior precision |
Performance on CIFAR-10-C (severity level 1)
Accuracy in %
NL. loss 4
ECE
AUROC 1
85-
10
0.54
84
10
4
10
10"
10
10
10a 10
10
10
10"
10
10
Prior precision B
Prior precision 3
Prior precision 3
Prior precision
Vanilla model fo.
B 61
B
256
B
-0024
(debiased)
8 128 (dehiased)
(debiased)
Performance on CIFAR-10-C (severity level 2)
Accuracy in %
NLL loss
ECF 4
AUROC 1
10
10
81
0.55
10
10
80
10
10a
10
10
10
10 10
10
10
10a 10
10
10
10
10
Prior precision
B
Prior precision 3
Prior precision 3
Prior precision 8
Vanilla model fo.
B
6
256
B 1022
Full butch
B 32 (debiased)
B 128 (dehiased)
15 512 (debiased)
Figure 14: Debiased LA mimics the full-batch LA very well. The experimental setting is the same as in Figure 6, but we report results on additional datasets (training set and OOD datasets at severity level 1 and 2). For the OOD datasets, we report the AUROC metric in addition to the accuracy, NLL and ECE. In contrast to the single-batch approach, the debiased version mimics the behavior of the full-batch approach very well over the entire range of prior precisions and across datasets.
35

---

Published as a conference paper at ICLR 2025
Performance on CIFAR-10-C (severity level 3)
Accuracy %
NLL loss
ECE +
AUROC 1
0.65
78-
10
0.60
77
10
0.55
10 2
10
10
10 10*
10
10
10 10
10
10a
10*
10
10
10a 10
Prior precision |
Prior precision
B
Prior precision 3
Prior precision |
Vanilla model fo.
256
1024
Full batch
B 32 (debiased)
B 128 (dehiased)
18
512 (debiased)
Performance on CIFAR-10-C (severity level
Accuracy in %
NLL loss
ECE
AUROC
1.8 10
0.70
1.6 10
72-
1.1 10
1011.
0.65
10
10
0.55-
4
10
10
10
10-2
10
10"
Prior precision B
Prior precision
Prior precision
Prior precision
Vanilla model fo.
B 61
B
256
B
1024
Full-batch
(debiased)
(dehiased)
(debiased)
Performance on CIFAR-10- (severity level 5)
Accuracy in %
NLL loss |
ECE
AUROC
X 10
10
64
X 10
0.74
1.7 X 10%
1.6 10"
63-
1.4 100
0.6
1.3 x 109
1.2 10
0.5
10
10
10 10.
10 10-2
10 10
10
10"
10"h 10*
10
10
10a 10*
Prior precision /
Prior precision 8
Prior precision 8
Prior precision |
Vanilla model fo.
[8]
256
1024
+
Full-batch
B 32 (debiased)
[8] 128 (dehiased)
18 512 (debiased)
Figure 15: Debiased LA mimics the full-batch LA very well. The experimental setting is the same as Figure 6, but we report results on additional datasets (OOD datasets at severity level 3, 4, and 5). We report the AUROC metric in addition to the accuracy, NLL., and ECE. In contrast the single-batch approach, the debiased version mimics the behavior of the full-batch approach very well over the entire range of prior precisions and across datasets.
36

---

Published as a conference paper at ICLR 2025
Performance on the test set
Accuracy in % f
NLL loss
BCE
Vanilla model fa.
[28]
76.1
0.06
0.97
[28]
256
[S]
1024
76.0
0.04
(debiased)
75.9
0.96
[S] 128 (debiased)
0.03
[S] 512 (debiased)
75.8
0.02-
bull-batch
0.95
0.01
10
10 10
4
2
10
100
Prior precision 3
Prior precision 8
Prior precision 3
The experimental setting is the same in Figure 6, but we use test problem (E) (see Appendix B.1) (RESNET-50 on IMAGENET). The results are consistent with the findings on CIFAR- 100: The debiased approach behaves similarly to the full-batch LA and maintains stability even for small prior precisions.
Figure 16:
Our approach scales to RESNET-50 on IMAGENET.
Performance on the test set
Accuracy f
NLL loss
ECE
Vanilla model fo.
0.25
[28]
79.7
1.03
0.20-
[S]
256
[S]
1024
79.6
1.00
0.15
[S]
(dehiased)
0.95
B 128 (dobiased)
79.5
0.10-
[ 512 (debiased)
0.90
Full-batci
79.4
0.85
10
10
10
10'
10
4
10
10"
10" -4
10
10
10
Prior precision 8
Prior precision B
Prior precision /
Figure 17: Curvature biases are also present for the VIT LITTLE architecture on IMAGENET. The experimental setting the same as in Figure and Figure 16, but we use test problem (F) (see Appendix B.1) (VIT LITTLE on IMAGENET). Again, the debiased approach behaves similarly to the full-batch LA and maintains stability even for small prior precisions (where the covariance is almost exclusively based on the K-FAC curvature information).
Results (runtimes). We report the runtimes of the biased (O), debiased (O), and full-batch (O) K-FAC Laplace approximation schemes in Figure 18. We find that (i) both biased and debiased approaches are orders of magnitude cheaper than the full-batch LA, and (ii) the debiased approach introduces only negligible overhead and can, sometimes, even improve upon the runtime of the biased approach.
Memory consumption. The ALL-CNN-C -C model we use for the LA experiment in Section 6.2 has P = 1,368,488 480 trainable parameters (this amounts to 10.95 MB of memory in double precision). The K-FAC approximation requires storing 11. 11,483,965 (91.87 MB) numbers represent all Kronecker factors. The memory of a K-FAC approximation is thus roughly equivalent to that of 8.4 models. As the debiased approach based on two such approximations in the case of a naive implementation. the overhead of debiasing is another 91.87 MB of memory. However, this can be significantly improved by building up the K-FAC approximations block by block: Having two corresponding blocks available (based on two different mini-batches) already suffices to compute the respective debiased block (see Appendix A.3.3). This way, we only have two blocks (represented by their Kronecker factors) in memory the same time (instead of two entire K-FAC approximations) which greatly reduces the memory overhead of debiasing, down to the largest block. For the ALL-CNN-C model, this is two Kronecker factors of sizes 192 192 and 1728 x 1728 amounting 3,022,848 numbers total, i.e. the equivalent of 2.20 models or 24.18 MB in double precision. The statement that the debiased approach roughly doubles the memory consumption is thus worst-case scenario.
37

---

Published as a conference paper at ICLR 2025
304
800
K-FAC sctup (first)
25
K-FAC setup (second)
Runtime (seconds)
600
A
Dehiased computation
20
MC sampling
Blased
400
Debinsed
Full-hatch
10
200
18]
64
[8] 256
[88
1024
[2]
Batch size |B]
Figure 18: The computational overhead of debiasing is negligible for Laplace approximations. We follow the experimental setup of Section 6.2 and provide detailed runtime comparison of the biased (=), debiased (=), and full-batch (0) K-FAC Laplace approximation schemes on CIFAR- 10. The total runtime includes the construction of K K-FAC's Kronecker factors (on one or two mini-batches for the biased and debiased approach, respectively), the debiasing computations (see Appendix A.3.3), and the Monte Carlo sampling of averages over 5 runs. The error bars for the total runtime cover one standard deviation. The results show that the debiased approach introduces a negligible overhead. Compared to the full-batch variant, both mini-batch approaches orders of magnitude faster. Note that the evaluation (i.e., the computation of the predictive uncertainty) is not included in the comparison, as the time requirements the same for all approaches; however, this step requires the most computational resources.
40 weights. All reported runtimes are
B.7 ADDITIONAL EXPERIMENT: BIASES FOR WIDERESNET 40-4 MODEL ARCHITECTURE
Experimental setting and results. Here, extend our analysis from Section 3.1 (details in Appendix B.4) to test problem (C) (see Appendix B.1): A WIDE RESNET model on CIFAR- 100. We use the GGN curvature proxy Hg GB + 3I at batch size 256 and compute the directional slopes and curvatures along the top 100 eigenvectors of Hg. The results are shown in Figure 19. Along the first few top-curvature directions, the curvature biases are even larger for this test problem than for the ALL-CNN-C model CIFAR-100, but then they decay more quickly.
B.8 ADDITIONAL EXPERIMENT: K- K-FAC AND THE DEPENDENCE OF THE BIASES ON MINI-BATCH SIZE
Experimental setting. Here, extend analysis of test problem (A) (see Appendix B.1) from Section 3.1 to K-FAC and investigate the impact of the mini-batch size on the curvature biases.
Results & discussion. Figure 20 shows the directional curvatures of the K-FAC approximation for four different mini-batch sizes. When we compute the eigenvectors and directional curvatures on the same batni-batch, observe systematic eurvature bias that decreases with increasing mini-batch size. There are two phenomena at play: With increasing mini-batch size, (i) the green crosses move upwards and (ii) the orange dots move downwards. Our intuition for this is as follows: With increasing mini-batch size, the eigenvectors become more meaningful such that, on other data, they also exhibit large curvature- ttis explains (i). Similarly, with increasing mini-batch size, it gets harder to find directions of extreme curvature these directions have to exhibit large curvature on all data points within the mini-batch- -hhs explains (ii).
B.9 ADDITIONAL EXPERIMENT: DEVELOPMENT OF THE BIASES OVER THE COURSE OF TRAINING
Experimental setting and results. Here, we investigate how the biases in the slope and curvature develop the course of training for test problem (A) (see Appendix B.I). We the same procedure as Appendix B.4 but evaluate the biases at 10 different checkpoints during training (spread log-equidistantly between the first and last epoch). At each checkpoint, we draw 5 mini-
38

---

Published as a conference paper at ICLR 2025
Eigenvectors computed mini-bate
Bg
B
13
Directional slope
0.2
0.0
-022
Permuted cigenvector index
Permuted cigenvector index
Permuted cigenvector index
Directional curvature
10a
10'-
10
10
30
75
100
25
75
100
50
75
100
Eigenvector index
Eigenvector index
Eigenvector index
Figure 19: Directional slopes and curvatures are biased. The experimental setting is similar to Figure 2, but we use test problem (C) (WIDE RESNET on CIFAR- 100) with the GGN curvature proxy HB - GB + 3I at batch size 256. For the top panel, we switch the order and sign of the eigenvectors such that the orange dots are all above zero and descending order. There is strong, systematic bias, particularly in the curvature: Computing the eigenvectors and directional curvatures on same data results over-estimation that very pronounced along the first few eigenvectors but then decays more quickly than in test problem (A).
Eigenvectors computed on By with batch size
Bd
|al]
[Bp]
[Bp] 512
Dircctional curvature
10a
10a_
10
50
100
50
100
50
100
50
100
Eigenvector index
Eigenvector index
Eigenvector index
Eigenvector index
Figure 20: Directional curvatures with K-FAC. We use the CIFAR-100 100 dataset with the fully trained ALL-CNN-C model. For each mini-batch size {8, 32, 128, 512), we draw one mini-batch and compute the top 100 eigenvectors U1 U100 of the K-FAC matrix KB,- We show the directional curvatures u, KB.. up p {1, ...... 100} on the same mini-batch (m as all other mini-batches in the training set (m (1, )) as and their average as Similar Figure we observe a systematic bias in the curvature. The bias decreases with increasing mini-batch size.
batches, compute the top 100 cigenvectors of the corresponding GGN- CGN-based quadratics, and finally evaluate the relative errors the directional slopes and curvatures (where the ground truth the full-batch quadratic's slope/curvature). For each mini-batch, this distribution of 100 relative curvature and slope biases is represented as a dot (at the mean relative bias) and a vertical line ranging from 25% to the 75% percentile in Figure 21. While the biases in the slope remain relatively stable over the course of training, the biases in the curvature increase over more than orders magnitude up to relative error of order 10 in epoch 349 (which is consistent with the results in Figure 2). This suggests that the eigenspaces different curvature matrices become more and more misaligned as training progresses steadily increasing the need for effective debiasing strategies.
39

---

Published as a conference paper at ICLR 2025
Directional slope
Directional curvature
10
10
Relative emoo
Relative error
10
100
10
10
Epoch
Epoch
Figure 21: Curvature biases increase over the course of training. We evaluate the relative slope and curvature biases at different checkpoints during training for the ALL-CNN-C model CIFAR-100 100. 5 mini-batches are drawn per checkpoint. For each mini-batch, the relative biases are represented dot (at the mean relative bias) and vertical line ranging from the 25% the 75% percentile. While the biases in the slope remain relatively stable over the course of training, the biases in the curvature increase (over more than orders of magnitude).
B.10 ADDITIONAL EXPERIMENT: THE CURVATURE BIAS INCREASES WITH P
In high-dimensional spaces, it becomes increasingly unlikely that random vectors are aligned. While the eigenspaces of the curvature matrices is thus conceivable that their overlap decreases as the number of parameters P' increases. If this hypothesis is true, the curvature biases should become more pronounced large models.
completely random, they are subject noise. It
Experimental procedure & results. To test this hypothesis, we use test problem (D) that implements simple convolutional neural network with variable width and depth (for details, see Appendix B.1). For each fully trained model, we evaluate the relative error between C (0,; B) and _ g(.; for the Gs's top 100 eigenvectors at batch size 128. This procedure is repeated for 5 different mini-batches for each of the 9 models, resulting in a total of 45 error distributions (see Figure 22). The results confirm our hypothesis: The relative errors tend to increase with the number of parameters P. In the massively overparameterized regime, the biases might thus become even more relevant and effective debiasing strategies are needed.
Figure 22:
Increasing curvature biases with P. We train 9 convolutional neural networks with dif- ferent widths and depths for 100 epochs the CIFAR-10 10 dataset using ADAM with standard hy- perparameters. For the fully trained model, we eval- uate the relative error 12 ((.;;
2.0
Relative error
1.5
; B) - 2 ((.;; DD] 162, (0.: D) 1 for Gs's top 100 eigenvectors batch size [ 128. different tini-batches are used per model resulting in a total of 45 error dis- tributions (each consisting of 100 numbers). These distributions are represented by their median (as dot and the 25% and 75% percentiles (as a line segment -) P. The 5 distributions for each model are slightly spread along the z-axis for better visibil- ity. The experiment confirms hypothesis: The relative errors tend to increase with the number of parameters P.
0.5
10a
Number of parameters P
model 1.32
model 7:32
model 7-64
model 4.32
model 4.64
model 4.128
model 1.64
model 128
model 7-128
40


Reviews:
[{"review_id": "7CmV0i5Rwd", "reviewer": "Reviewer_SWR8", "summary": "This paper investigates biases in mini-batch quadratic approximations, for second-order optimization in deep learning. Authors show that when computing quadratic on mini-batches rather than the full dataset, it introduces curvature distortions, leading to overly narrow approximations. The bias, assigned to “regression to the mean”, affects Newton steps in optimizers and Laplace approx for uncertainty quantification. They provide a “two-batch” strategy to decouple direction selection from magnitude estimation, to align mini batch approx more closely with full-batch behaviour with minimal computational overhead. They demonstrate the improve stability of their strategy on CG and K-FAC experiments.", "questions": "— How do you anticipate the method to scale or adapt with different architecture and larger dataset, which have different curvature profiles and optimisation challenges?\n— Could you clarify the criteria for selecting mini-batches in the two-batch debasing process? Is there any benefit to use overlapping vs. non-overlapping mini-batches, or would a mini-batch with a specific loss profile provide better results? \n— Can you provide more details on runtime/memory costs associated with the two-batch approach? While authors claim it roughly doubles, it would be helpful to understand how this really scales with batch size, and whether costs remain manageable as model size and dataset complexity grow", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Authors study a well-known but often overlooked issue in mini-batch approx within second-order optimisations and uncertainty quantification for deep learning: the bias in curvature estimates due to mini-batch sampling. While the issue is not new, the approach of a two-batch debasing strategy is an effective and novel (to my knowledge) way to address it and improve the accuracy of the approximates without suffering too much of a computational overhead. This strategy could generalise well across various settings in ML where the full-batch quadratics are intractable. The paper is supported by nice empirical experiments on CG and K-FAC demonstrating the effectiveness of their solution. The authors provide a theoretical intuition for the bias arising from that “regression to the mean” adding an interesting insight.  The experiments are systematically presented and show clear performance difference between single-batch and debasing method . While very technical, authors make an effort to make it as digestible as possible, with the inclusion of figures to visualise the bias and the impact of their debasing strategy.", "weaknesses": "The experimental scope is pretty limited, with only ablations on CNN architecture and CIFAR-10/100 datasets. Testing on a wider variety of architectures, such as transformers or larger datasets, would strengthen a lot the paper’s claims on general applicability. Another limitations lies in the lack of ablations to investigate the sensitivity of the two-batch approach to mini-batch size, as well as the composition of mini-batches used for debasing. Without these, it’s unclear how the two-batch batch method would perform under different resource constraints or mini-batch sampling strategies. Last, the paper discusses the computational trade-offs, but does not provide sufficient detail to me on actual runtime and memory costs for different debasing configuration, leaving with an incomplete understanding of the true method’s scalability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Rom1V7yxSp", "reviewer": "Reviewer_3MNZ", "summary": "This paper addresses the systematic bias in mini-batch quadratic approximations used for second-order optimization and uncertainty quantification in deep learning. Mini-batch methods distort the quadratic shape, making uncertainty estimates unreliable. The authors provide a theoretical explanation, discuss the impact on optimization, and propose debiasing strategies to improve accuracy.", "questions": "How well do the proposed debiasing strategies scale for large foundation models with millions or billions of parameters?\n\nCan debiased mini-batch quadratics improve uncertainty quantification in foundation models, especially in tasks requiring precise uncertainty, like language modeling or image synthesis?\n\nCan the theoretical insights and debiasing techniques be effectively adapted for transformer-based foundation models?\n\nCould reduced bias in mini-batch quadratics enhance transfer learning by improving local loss landscape approximations during fine-tuning?\n\nHow might the debiasing methods apply to multi-modal foundation models, where different data modalities add complexity to curvature estimation?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This is an interesting theoretical investigation collaborated with numerical justifications. It teaches me something that I would expect with a nice and convincing narrative. The overall clarity and language of this paper is good. The insights are well-delivered. The theoretical analyses are mostly backed up by numerical experiments.", "weaknesses": "I did not identify any obvious weaknesses.\n\nI did not check the proofs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aKjRO8htft", "reviewer": "Reviewer_ddR4", "summary": "This paper demonstrates a systematic bias in stochastic 2nd order approximation of empirical risk for neural networks. The study start from an observation showing that a particular curvature structure on minibatch Hessian is biased towards underestimation. Then, authors proposes a debiasing method.", "questions": "- I recommend to replace $u_1$ and $u_2$ in Figure 1 by leading eigenvectors of full-batch Hessian to better grasp my comment in weaknesses. I expect to see that increasing batch size will lead to a better approximation of the curvature. \n- What is exactly the quantity that you want to estimate and why it is important to estimate? Directional curvature on which random directions?\n- What do you exactly mean by bias? To prove an estimate has bias, we need to take average. But, experiments in Figure 1 are not computing an average. How can I conclude from these experiments that the estimate is biased?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- A computationally-cheap estimate of 2nd order information is an important topic for optimization \n- Ample experimental observations \n- Designing computational-cheap experiments to study 2nd order derivative of training loss of neural networks", "weaknesses": "I believe the notion of unbiased estimate is not well defined and studied. Define the Hessian associate with full-batch training loss as  $$H(x) = \\frac{1}{n} \\sum_{i=1}^n \\nabla^2 f_i(x)$$. Similarly, we define the unbiased estimate of this matrix as $$H'(x) = \\frac{1}{m} \\sum_{k=1}^m \\nabla^2 f_{i_k}(x)$$ where $i_k$ are uniformly drawn from $\\{1,\\dots, n\\}$. \nFor each fixed vector $d$, we have $$E d^\\top H' d^\\top = d^\\top H d$$. Thus, the **directional curvature**, defined in the paper, is unbiased. But, why this paper observe a bias in the estimate? The main issue is that they consider random directions $d$ depending on minibatches. I do not know why we need to choose random $d$s? \n\nI believe the quantity of interest is not well defined and motivated, here. It is easy to provide an unbiased estimation of  $q$, directional curvature and directional slop. But, the authors want to estimating another quality that they do not exactly defined and it is not clear how its estimation connect with Newton's method or Laplace approximation. \n\nIf you want to estimate the maximum/minimum eigenvalue of the Hessian or even its condition number, the empirical Hessian provides an asymptotically unbiased estimate (see for example https://arxiv.org/pdf/1912.10754).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iu6vcoA2vF", "reviewer": "Reviewer_9wSH", "summary": "Often one would like to obtain a second-order Taylor approximation to a deep learning objective based on a minibatch.  This paper shows that the use of a minibatch results in a strong overestimation of the top Hessian eigenvalues.  The basic problem is that the directions of highest curvature for a minibatch tend to slightly 'overfit' (my words) that minibatch, and do not generalize to other minibatches or to the full-batch objective.  This paper shows that this bias can be corrected with a scheme that uses two separate minibatches - one to compute the top eigenvectors, and another to compute the directional curvature along these eigenvectors.  The paper discusses applications in both optimization (conjugate gradient method for computing a newton step), and uncertainty quantification (laplace's method)", "questions": "Regarding the motivation in section 3.1 -- just because the gradient lives in the top subspaces does _not_ mean that all the optimization progress occurs in that space.  To the contrary, the evidence suggests that the vast majority of the motion along the top Hessian eigenvectors is oscillatory, and 'cancels out.'  Long-term loss reduction is due to the components of the gradient that lie along the low-curvature Hessian eigenvectors.  For example, see \"Does SGD really happen in tiny subspaces?\" by Minhak Song, Kwangjun Ahn, Chulhee Yun: https://arxiv.org/abs/2405.16002.  \n\nI don't think this is fatal for the paper, since you are basically only using your improved estimate of the the top Hessian eigenvalues in order to properly compute the conjugate gradient _step size_, which will depend on the top Hessian eigenvalues.  But if you actually tried to optimize only along those eigenvectors, it wouldn't work (as the linked paper shows).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper offers a simple solution to a problem of interest to multiple areas of deep learning.", "weaknesses": "In the application to optimization, the paper does not actually carry out an end-to-end optimization run with a second-order method using the proposed conjugate gradient scheme; instead, the paper just uses the proposed scheme to compute a single newton step.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)