{'fraction_non_alphanumeric': 0.051552264312161776, 'fraction_numerical': 0.01597835374537169, 'mean_word_length': 4.459648577204168, 'pattern_counts': {'":': 0, '<': 2, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Isotropic SGD: a Practical Approach to Bayesian Posterior Sampling

In this work we define a unified mathematical framework to deepen our understanding of the role of stochastic gradient (SG) noise on the behavior of Markov chain Monte Carlo sampling (SGMCMC) algorithms. Our formulation unlocks the design of a novel, practical approach to posterior sampling, which makes the SG noise isotropic using a fixed learning rate that we determine analytically, and that requires weaker assumptions than existing algorithms. In contrast, the common traits of existing \sgmcmc algorithms is to approximate the isotropy condition either by drowning the gradients in additive noise (annealing the learning rate) or by making restrictive assumptions on the \sg noise covariance and the geometry of the loss landscape. Extensive experimental validations indicate that our proposal is competitive with the state-of-the-art on \sgmcmc, while being much more practical to use.


Introduction
Despite mathematical elegance and some promising results restricted to simple models, standard stochastic gradient (SG) methods for MCMC-based Bayesian posterior sampling Welling and Teh [2011], Ahn et al. [2012], Patterson and Teh [2013], Chen et al. [2014], Ma et al. [2015] fall short in dealing with the complexity of the loss landscape of deep models Draxler et al. [2018], Garipov et al. [2018], for which stochastic optimization poses serious challenges Chaudhari and Soatto [2018], Maddox et al. [2019]. Moreover, existing methods are often unpractical, as they require ad-hoc, sophisticated vanishing learning rate schedules, and hyper-parameter tuning.
In general, SG-MCMC algorithms inject random noise to SG descent algorithms: the covariance of such noise and the learning rate are tightly related to the assumptions on the loss landscape, which together with the SG noise, determine their sampling properties Ma et al. [2015]. However, current SG-MCMC algorithms applied to popular complex models such as Deep Nets, cannot satisfy the simplifying assumptions on loss landscapes and on the behavior of the SG noise covariance, while operating with practical requirements, such as non-vanishing learning rates and ease of use. A recent work Mandt et al. [2017] argues for fixed step sizes, but settles for variational approximations of simple quadratic losses. In a generic Bayesian deep learning setting, none of the existing implementations of SG-MCMC methods converge to the true posterior without learning rate annealing.
as it can clarify how various methods compare. Most importantly, this endeavor can suggest novel and more practical algorithms relying on fewer tunable parameters and less restrictive assumptions.
In this work, we present a mathematical framework that emphasizes the role of SG noise covariance and learning rate on the behavior of SG-MCMC algorithms (Sec. 2). As a result, the equivalence between learning rate annealing and the injection of noise with extremely large variance becomes clear, and this allows us to propose a novel, practical SG-MCMC algorithm (Sec. 3) that produces (approximate) posterior samples with a fixed, easy to derive, learning rate. Furthermore, our approach can be readily applied to pre-trained models: after a "warm-up" phase to compute SG noise estimates, it can efficiently perform Bayesian posterior sampling. The proposed SG-MCMC method is a valid, theoretically sound, and simple alternative to popular techniques, that have shortcomings when it comes to the assumptions they rely on Gal and Ghahramani [2016].
We evaluate SG-MCMC algorithms (Sec. 4) through an extensive experimental campaign, where we compare our approach to a number of alternatives, including Monte Carlo Dropout (MCD) [Gal and Ghahramani, 2016] and Stochastic Weighted Averaging Gaussians (SWAG, [Maddox et al., 2019]), which have been successfully applied to the Bayesian deep learning setting. Our results indicate that our approach is simpler to use than other families of SG-MCMC methods, and offer performance that are competitive to the state-of-the-art, in terms of accuracy and uncertainty.

Preliminaries and Related Work
Consider a data-set of m−dimensional observations D = {U i } N i=1 and a statistical model defined through a likelihood function p(D|θ) parameterized through a d-dimensional set of parameters θ. Given a prior p(θ) the posterior over the parameters is obtained by means of Bayes theorem as: where p(D) is also known as the model evidence, defined as the integral p(D) = p(D|θ) p(θ)dθ.
The posterior distribution is necessary in order to obtain predictive distributions for new test observations U * : p(U * |D) = p(U * |θ)p(θ|D)dθ. We focus on Monte Carlo methods to obtain an estimate of this predictive distribution, by averaging over N MC samples obtained from the posterior over θ, that is θ (i) ∼ p(θ|D): Since Eq. (1) is analytically intractable Bishop [2006], unless the prior is conjugate to the likelihood, we work with an unnormalized version of the logarithm of the posterior density, and express the negative logarithm of the joint distribution of the data-set D and parameters θ as: For computational efficiency, we use a mini-batch stochastic gradient g(θ), which guarantees that the estimated gradient is an unbiased estimate of the true gradient ∇f (θ), and we assume that the randomness due to the mini-batch introduces a Gaussian noise: 1 g(θ) ∼ N (∇f (θ), 2B(θ)), where the matrix B(θ) denotes the SG noise covariance, which depends on the parametric model, the data distribution and the mini-batch size.
A survey of algorithms to sample from the posterior using SG methods can be found in Ma et al. [2015]. In Appendix A we complement this section with a full derivation, well-known facts and definitions, for completeness. As shown in the literature , Mandt et al. [2017], there are structural similarities between SG-MCMC algorithms and stochastic optimization methods, and both can be used to draw samples from posterior distributions. In what follows, we define a unified notation to compare many existing algorithms in light of the role played by their noise components.
Stochastic gradient descent (SGD) can be studied through the following stochastic differential equation (SDE) Gardiner [2004], Kushner and Yin [2003], Ljung et al. [1992], when the learning rate η is small enough: (4) Here we use a generic form of the SDE, with variable z instead of θ, that accommodates SGD variants, with and without momentum. It is typical SG-MCMC practice Welling and Teh [2011], Ahn et al. [2012], Patterson and Teh [2013], Mandt et al. [2017] to allow the stochastic process induced by Eq. (4) to go through an initial adaptation phase where the learning rate is annealed, followed by fixing the learning rate to a small value, to ensure the process reaches and maintains a stationary distribution thereafter.
Definition 1 A distribution ρ(z) ∝ exp(−φ(z)) is said to be a stationary distribution for the SDE of the form (4), if and only if it satisfies the following Fokker-Planck equation (FPE): Note that, the operator ∇ applied to matrix D(z) produces a row vector whose elements are the divergences of the D(z) columns Chen et al. [2014].
In general, the stationary distribution does not converge to the desired posterior distribution, i.e., φ(z) = f (z), as shown by Chaudhari and Soatto [2018]. Additionally, given an initial condition for z t , its distribution is going to converge to ρ(z) only for t → ∞.
Next, we revisit known approaches to Bayesian posterior sampling, and interpret them as variants of an SGD process, using the FPE formalism. In what follows, we use n to indicate discrete time, and t for continuous time.
Gradient methods without momentum. The generalized update rule of SGD, described as a discrete-time stochastic process, writes as: δθ n = −ηP (θ n−1 )(g(θ n−1 ) + w n ), where P (θ n−1 ) is a user-defined preconditioning matrix, and w n is a noise term, distributed as w n ∼ N (0, 2C(θ n )), with a user-defined covariance matrix C(θ n ). Then, the corresponding continuous-time SDE is Gardiner [2004]: We denote by C(θ) the covariance of the injected noise and by Σ(θ t ) = B(θ t ) + C(θ t ) the composite noise covariance, which combines the SG and the injected noise.
We define the stationary distribution of the SDE in Equation (6) as ρ(θ) ∝ exp(−φ(θ)), noting that when C = 0, the potential φ(θ) differs from the desired posterior f (θ) Chaudhari and Soatto [2018]. The following theorem, which is an adaptation of known results in light of our formalism, states the conditions for which the noisy SGD converges to the true posterior distribution (proof in appendix A.2).
An extension to SGLD is Stochastic Gradient Fisher Scoring (SGFS) Ahn et al. [2012], which can be tuned to switch between sampling from an approximate posterior, using a non-vanishing learning rate, and the true posterior, by annealing the learning rate to zero. SGFS uses preconditioning, P (θ) ∝ B(θ) −1 . Generally, however, B(θ) is ill conditioned: many of its eigenvalues are almost zero Chaudhari and Soatto [2018], and computing B(θ) −1 is problematic. Moreover, using a nonvanishing learning rate, the conditions of Theorem 1 are met only if, at convergence, ∇ (B(θ) −1 ) = 0 , which would be trivially true if B(θ) was constant. However, recent work Draxler et al. [2018], Garipov et al. [2018] suggests that this condition is difficult to justify.
The Stochastic Gradient Riemannian Langevin Dynamics (SGRLD) algorithm Patterson and Teh [2013] extends SGFS to the setting in which ∇ (B(θ) −1 ) = 0 . The process dynamics are adjusted by adding the term ∇ (B(θ) −1 ) which, however, cannot be easily estimated, restricting SGRLD to cases where it can be computed analytically.
The work in  considers a regularized diagonal preconditioning matrix P (θ), derived from the SG noise. However, in the sampling phase, the method neglects the term ∇ P (θ) and the regularization term is a user-defined parameter, which is not trivial to tune.
The approach in Mandt et al. [2017] investigates constant-rate SGD (with no injected noise), and determines analytically the learning rate and preconditioning that minimize the Kullback-Leibler (KL) divergence between an approximation and the true posterior. Moreover, it shows that the preconditioning used in SGFS is optimal, in the sense that it converges to the true posterior, when B(θ) is constant and the true posterior has a quadratic form.
In summary, to claim convergence to the true posterior distribution, existing approaches require either vanishing learning rates or assumptions on the SG noise covariance that are difficult to verify in practice, especially when considering deep models. We instead propose a novel practical method, that induces isotropic SG noise and thus satisfies Theorem 1. We determine analytically a fixed learning rate and we require weaker assumptions on the loss geometry.
Gradient methods with momentum. Momentum-corrected methods emerge as a natural extension of SGD approaches. The general set of update equations for (discrete-time) momentum-based algorithms is: where P (θ n−1 ) is a preconditioning matrix, M is the mass matrix and A(θ n−1 ) is the friction matrix Chen et al. [2014], Neal et al. [2011]. Similarly to the first order counterpart, the noise term is distributed as w n ∼ N (0, 2C(θ n ))). Then, the SDE to describe continuous-time system dynamics is: where P (θ t ) 2 = P (θ t )P (θ t ), and P (θ t ) is symmetric. The theorem hereafter (proof in Appendix A.3) describes the conditions for which the SDE converges to the true posterior distribution.
Theorem 2 Consider dynamics of the form (7) and define the stationary distribution for θ t as In the naive case, where P (θ) = I, A(θ) = 0, C(θ) = 0, the conditions in Theorem 2 are not satisfied and the stationary distribution does not correspond to the true posterior Chen et al. [2014]. To generate samples from the true posterior, it is sufficient to set P (θ) = I, A(θ) = ηB(θ), C(θ) = 0 (as in Eq. (9) in Chen et al. [2014]).
Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) Chen et al. [2014] suggests that estimating B(θ) can be costly. Hence, the injected noise C(θ) is chosen such that C(θ) = η −1 A(θ), where A(θ) is user-defined. When η → 0, the following approximation holds: Σ(θ) C(θ). It is then trivial to check that conditions in Theorem 2 hold without the need for explicitly estimating B(θ). A further practical reason to avoid setting A(θ) = ηB(θ) is that the computational cost for the operation . This however, severely slows down the sampling process.
Stochastic Gradient Riemannian Hamiltonian Monte Carlo (SGRHMC) is an extension to SGHMC Ma et al. [2015]), which considers a generic, space-varying preconditioning matrix P (θ) derived from information geometric arguments Girolami and Calderhead [2011]. SGRHMC suggests to set P (θ) = G(θ) − 1 2 , where G(θ) is the Fisher Information matrix. To meet the requirement ∇ P (θ) = 0 , it includes a correction term, −∇ P (θ). The injected noise is set to C(θ) = η −1 I − B(θ), consequently Σ = η −1 I, and the friction matrix is set to A(θ) = P (θ) 2 . With all these choices, Theorem 2 is satisfied. While appealing, the main drawbacks of this method are the need for an analytical expression of ∇ P (θ), and the assumption for B(θ) to be known.
Recently, the work in Zhang et al. [2020] suggests to use a cyclical learning rate schedule to better explore the loss landscape and sample more efficiently. While the idea is appealing, it introduces further hyper-parameters to tune, which is opposite to our quest for a simple, easy to use method.
From a practical standpoint, momentum-based methods suffer from the complexity and fragility of hyper-parameter tuning, including the learning rate schedule and those that govern the simulation of a second-order Langevin dynamics. The method we propose in this work can be applied to momentumbased algorithms as well; then, it can be viewed as an extension to the work in Springenberg et al. [2016], albeit addressing more complex loss landscapes. However, we leave this avenue of research for future work.

Sampling by layer-wise Isotropization
We present a simple and practical approach to inject noise to SGD iterates to perform Bayesian posterior sampling. Our goal is to sample from the true posterior distribution (or approximations thereof) using a constant learning rate, and to rely on more lenient assumptions about the geometry of the loss landscape that characterize deep models, compared to previous works.
From Theorem 1, observe that P (θ), Σ(θ) are instrumental to determine the convergence properties of SG methods to the true posterior. We consider the constructive approach of designing ηP (θ) to be a constant, diagonal matrix, constrained to be layer-wise uniform: By properly setting parameters λ (p) , we achieve the simultaneous result of a non-vanishing learning rate and a well-conditioned preconditioning matrix. This implies a layer-wise learning rate η (p) = 1 λ (p) for the p-th layer, without further preconditioning. We can now state (see proof in Appendix B.1), as a corollary to Theorem 1, that our method guarantees convergence to the true posterior distribution.
An ideal method. Now, let's consider an ideal case, in which we assume the SG noise covariance B(θ) and Λ to be known in advance. The procedure described in Algorithm 1 illustrates a naive SG method that uses the injected noise covariance C(θ) to sample from the true posterior.
Algorithm 1 Idealized posterior sampling This deceivingly simple procedure is guaranteed to generate samples from the true posterior, with a non-vanishing learning rate 2 . However, it cannot be used in practice as B(θ) and Λ are unknown. Also, the algorithm is costly, as it requires computing (Σ−B(θ)) 1 2 , which requires O(d 3 ) operations, and C(θ) 1 2 , which costs O(d 2 ) multiplications. Next, we describe a practical approach, where we 2 Note that instead of computing the gradient of f (θ), we compute the (mini-batch) gradient of f (θ) N , similarly to the notation used in Mandt et al. [2017], that we indicate with ∇f (θ t ) N . use approximations at the expense of generating samples from the true posterior distribution. Note that Mandt et al. [2017] suggests to explore a related preconditioning, but does not develop the idea.
A practical method: Isotropic SGD. To make the idealized sampling method practical, we require additional assumptions which are milder than what is required by current approaches in the literature, as we explain at the end of this section.
Assumption 2 The signal to noise ratio (SNR) of a gradient is small enough such that, in the stationary regime, the un-centered variance of the gradient is a good estimate of the true variance Springenberg et al. [2016], Saxe et al. [2019]. Hence, combining with assumption 1, b(θ) Assumption 3 In the stationary regime, the maximum of the variances of noise components, layer by layer, are fixed constants (similarly to Zhu et al. [2019]): β (p) = max j∈Ip b j (θ), where I p is the set of indexes of parameters belonging to p th layer.
Note that Assumption 2 and Assumption 3 must hold only in the stationary regime when the process reaches the bottom valley of the loss landscape. Given our assumptions, and our design choices, it is then possible to show (see appendix B.2) that the optimal (i.e., the smallest possible) Λ = λ (1) , . . . , λ (1) , . . . , λ (N l ) , . . . λ (N l ) satisfying Corollary 1 can be obtained as λ (p) = β (p) .
Since we cannot assume B(θ) to be known, in what follows we discuss two approaches to estimate its components. A simple method to estimate B(θ) is as follows (see Appendix B.3): we compute is the portion of stochastic gradient corresponding to the p-th layer. Our estimates can be extended to use a moving average approach. Our empirical validation, however, indicates that this simple method does not produce stable estimates.
Indeed, a shared assumption of SG-MCMC methods is that SG noise is Gaussian. While this assumption can be justified with the C.L.T. for relatively simple models (linear models or simple feed-forward networks), its validity has been challenged in the deep learning domainŞimşekli et al. [2019a,b] (see Appendix B.3 for a detailed discussion), suggesting that, for complex architectures, the noise distribution is heavy tailed. Then, the hypothesis is that the various components of SG noise follow an α-stable distribution: w ∼ p(w), where +∞ −∞ exp (j2πwt) p(w)dw = exp(−|ct| α ), with α ∈ (0, 2], where α, c can vary across different parameters. When α = 2, p(w) becomes Gaussian, but for α < 2, its variance goes to infinity, thus estimating the SG noise covariance is problematic.
Prior worksŞimŠekli [2017], suggest to define a stochastic process governed by an SDE that uses Lévy Noise instead of a Brownian motion. However, this approach comes with its own challenges, such as the approximation of a fractional derivative, and the use of full gradients. Instead, we propose the following approximate method: we consider that the SG noise follows a Gaussian distribution, with parameters set to minimize the l 2 -distance between p(w) = .3 for details). We estimate the parameters of the α-stable distribution by extending Levy Vehel et al. [2018] to space varying settings, and derive the equivalent variancesσ 2 that minimize the l 2 distance. Then, the λ (p) are computed as 1 2 max j (σ j (θ) (p) ) 2 . In Sec. 4 we use this method, as it is more stable than the simple approach discussed above.
Ultimately, the composite noise matrix Σ = Λ is a layer-wise isotropic covariance matrix, which inspires the name of our proposed method as Isotropic SGD (I-SGD). Once all parameters λ (p) have been estimated, the layer-wise learning rate is determined automatically: for the p-th layer, the learning rate is η (p) = 1 N λ (p) . The practical implementation of I-SGD is shown in Algorithm 2. The computational cost of I-SGD is as follows. Similarly to Chen et al. [2014], we define the cost of computing a gradient mini-batch as C g (N b , d). Then (see Appendix B.3), the computational cost for estimating the noise covariance scales as O(d) logarithm computations. The computational cost of generating random samples with  In appendix C we report a simple numerical example where the true posterior is analytically available: visual inspection indicates that our method produces a sensible predictive posterior distribution.
Assumptions and convergence to the true posterior. Our theory shows that the ideal version of I-SGD (Corollary 1 holds, and B(θ) is known) converges to the true posterior with a constant learning rate. This is not the case for existing work. Even when B(θ) is assumed to be known, SGFS requires the correction term ∇ B(θ) −1 = 0. Also, both SGRLD and SGRHMC require computing ∇ B(θ) −1 , for which an estimation procedure is elusive. The method in Springenberg et al. [2016] needs a constant, diagonal B(θ), a condition that does not necessarily hold for deep models.

Experiments
We first study I-SGD using standard UCI data-sets Dua and Graff . We compare I-SGD (with the Gaussian approximation to the estimated α-stable distribution) to SGHMC Chen et al. [2014], and to alternative approaches to approximate Bayesian inference, including MCD Gal and Ghahramani [2016], and SWAG Maddox et al. [2019], the variational SGD approach (V-SGD Mandt et al. [2017]) and CSGHMC Zhang et al. [2020]. appendix B.3 includes additional implementation details on I-SGD. appendix D presents detailed configurations of all methods we compare and additional results.   (2), for I-SGD and SGHMC, and 10 000 samples for MCD. All our experiments use 10-splits. We note that the task of tuning our competitors is not trivial. 3 In this set of experiments we omit results for SWAG, which we keep for more involved scenarios.
Tab. 1 is an overview of our results: MCD and I-SGD are roughly similar, in terms of RMSE. Instead, results for MNLL are more nuanced: all methods are essentially on par, with I-SGD outperforming alternatives in some cases. Additionally, our experiments confirm the practical viability of I-SGD, which does not require tedious and creative parameter tuning during the sampling phase.
Classification tasks, with deeper models. With the exception of MCD, to compare methods on a fair ground, we use pre-trained models, such that all approaches produce samples from the same conditions. As an alternative, we can use a standard optimizer to train a model from scratch: then, I-SGD sampling can be triggered when stationarity conditions are met.
Next, we compare I-SGD, MCD, SWAG and V-SGD using first the MNIST dataset, and a simple LENET-5 CNN LeCun et al. [1998]. We use a pre-trained LENET-5, and proceed with a "warm-up" phase to estimate the SG noise covariance (this is needed both in SWAG, V-SGD and I-SGD). Then, the posterior sampling phase can begin. At test time we use 30 samples for all methods. Results are obtained averaging over 5 independent seeds. As commonly done in the literature, we also use temperature scaling, which helps stabilizing algorithm dynamics. All methods are compared based on the test accuracy (ACC) and the MNLL. Additionally, we compute out-of-distribution entropy diagrams by training on MNIST , and testing on both MNIST and NOT-MNIST. This is common practice to check whether the entropy of the predictions on NOT-MNIST is higher than the entropy of the predictions on MNIST . Tab. 2 indicates that all methods are essentially equivalent in terms of accuracy. Instead, I-SGD outperforms others in terms of MNLL. Fig. 1 indicates that I-SGD achieves a better separation of in-distribution (orange) and out-of-distribution (blue) predictions.
We now use RESNET-18, RESNET-56 and VGG-16 networks on CIFAR10. For this set of experiments we compare I-SGD, MCD, SWAG and V-SGD, reporting test accuracy and MNLL. We also report a comparison between I-SGD and CSGHMC using the results in Zhang et al. [2020] for RESNET-18, albeit the MNLL was not available. These results are summarized in Tab. 3, Tab. 8 and Tab. 9. At test time we use 30 samples for all methods. Results are obtained averaging over five independent seeds. I-SGD is competitive with other methods. Moreover, I-SGD is the the easiest to tune, as its learning rate is fixed and computed automatically.

Conclusion
SG methods allowed Bayesian posterior sampling algorithms, such as MCMC, to regain relevance in an age when data-sets have reached extremely large sizes. However, despite mathematical elegance and promising results, standard approaches from the literature are restricted to simple models. The sampling properties of these algorithms are determined by simplifying assumptions on the loss landscape, which do not hold for deep networks. SG-MCMC algorithms require vanishing learning rates, which force practitioners to develop creative annealing schedules that are often model specific and difficult to justify.
We have attempted to target these weaknesses by suggesting a simpler algorithm that relies on fewer parameters and mild assumptions compared to the literature. We introduced a unified mathematical notation to deepen our understanding of the role of the SG noise and learning rate on the behavior of SG-MCMC algorithms. We presented a practical variant of the SGD algorithm, which uses a constant learning rate, and an additional noise to perform Bayesian posterior sampling. Our proposal is derived from an ideal method, which guarantees that samples are generated from the true posterior. When the noise terms are empirically estimated, our method automatically determines the learning rate, and it offers a very good approximation to the posterior, as demonstrated by our experimental campaign.

A.1 The minibatch gradient approximation
Starting from the gradient of the logarithm of the posterior density: it is possible to define its minibatch version by computing the gradient on a random subset I N b with cardinality N b of all the indexes. The minibatch gradient g(θ) is computed as By simple calculations it is possible to show that the estimation is unbiased (E(g(θ)) = ∇f (θ)).
If the minibatch size is large enough, invoking the central limit theorem, we can state that the minibatch gradient is normally distributed: g(θ) ∼ N (∇f (θ), 2B(θ)).

A.2 Gradient methods without momentum
The SDE from discrete time We start from the generalized updated rule of SGD: δθ n = −ηP (θ n−1 )(g(θ n−1 ) + w n ).
As the learning rate goes to zero (η → dt), similarly to the previous case, we can interpret the above difference equation as a discretization of the following FPE Proof of Theorem 2. As before we assume that the stationary distribution has form ρ(z) ∝ exp(−φ(z)). The corresponding FPE is: Notice that since ∇ D (z) = 0 we can rewrite: If ∇ P (θ) = 0, in fact: since ∇ 0 −P (θ) P (θ) 0 = 0 and the second term is zero due to the fact that 0 −P (θ) P (θ) 0 is anti-symmetric while ∇ ∇ (φ(z))ρ(z) is symmetric.

B.2 Proof of optimality of Λ
Our design choice is to select λ (p) = β (p) . By the assumptions the matrix B(θ) is diagonal, and consequently C(θ) = Λ − B(θ) is diagonal as well. The preconditioner Λ must be chosen to satisfy the positive semi-definite constraint, i.e. C(θ) ii ≥ 0 ∀i, ∀θ. Equivalently, we must satisfy λ (p) − b j (θ) ≥ 0 ∀j ∈ I p , ∀p, ∀θ, where I p is the set of indexes of parameters belonging to p th layer. By assumption 3, i.e. β (p) = max k∈Ip b k (θ), to satisfy the positive semi-definite requirement in all cases the minimum valid set of λ (p) is determined as λ (p) = β (p) .

B.3 Estimation of λ (p)
The case of Gaussian SG noise. We here give additional details on the estimation of λ (p) . The simple and naive estimation described in the paper is the following: λ (p) = max j∈Ip (g j (θ) (p) ) 2 . For the Gaussian SG noise case we found however the following (safe) looser estimation of the maximum noise covariance to be more stable: . From a practical point of view, we found the following filtering procedure to be useful and robust: where an exponential moving average is performed with estimation momentum determined by µ.
Notice that during sampling, the same smoothing can be applied to the tracking of B(θ). We refer to the variant of I-SGD implemented using this estimator as I-SGD-G. In this supplement we also considered the case of having a unique, and not layerwise, learning rate, that we indicate by justapposing the (SLR) acronym to the right of the methods. In this case, the unique equivalent λ is computed as p λ (p) .
The case of Heavy Tailed Noise. A shared assumption of SG-MCMC methods is the Gaussianity of SG noise. While this can be justified with the C.L.T. for relatively simple models (linear models or simple feed-forward networks), this assumption has been challenged in the deep learning domaiņ Simşekli et al. [2019a,b] suggesting that from complex architectures the noise distribution is heavy tailed. In particular, the hypothesis is that the noise follows and α-stable distribution, i.e.
where α ∈ [0, 2]. Notice that except for particular cases, p(w) can not be expressed in closed form.
In general, when α < 2 the variance of the distribution goes to infinity and thus dealing with all where γ = 0.5772156649015329 . . . is the Euler-Macheroni constant. Notice that the computational cost for estimation of the two quantities is dominated by the calculation of logarithms, in fact for a full sequence of N independent samples the cost is for the estimation of α N + N 2 absolute values, N + N 2 logarithms, 2N sums, with a per sample cost roughly equal to the cost of 1 logarithm evaluation, and for the estimation of c the cost is N logarithms, sums and absolute values (and thus similarly the cost is dominated by the log evaluation). When considering vectors of independent d−dimensional samples, the computational cost scales as O(d) logarithms.
Notice that for the I-SGD-α version we treated biases and weights of the layers as unique groups of parameters. Additional details on estimation. Having chosen one of the two variants I-SGD-α or I-SGD-G for the estimation of λ (p) , that we generically indicate as I-SGD-X, we still need to clarify what are the possibilites for the estimation of the parameters λ p before the sampling. We considered three schemes: • I-SGD-X-A: the starting point is a freshly initialized model. The estimation is performed while training, similarly to Mandt et al. [2017] and Ahn et al. [2012], and a filtered version of the instantaneous estimation is stored; • I-SGD-X-B: we start from a pre-trained model, and a warm-up phase is necessary. We continue the training during the warm-up phase and collect a filtered version of estimates, as for the previous case.
• I-SGD-X-C: we start from a pre-trained model, and therefore a warm-up phase is necessary. We freeze the network and estimate λ (p) using an adequeate number of mini-batches.
Summarizing, all the possible combinations are I-SGD-α/G-A/B/C. While not always the best performing, we found the I-SGD-α version the more stable across a large range of hyperparameters, and in the spirit of practicality, in the main paper we report only results obtained with this version.

C Toy Model
Next, we consider a simple numerical example whereby it is possible to analytically compute the true posterior distribution. We define a simple 1-D regression problem, in which we have D trigonometric basis functions: f (x) = w cos(ωx − π/4), where w ∈ R D×1 contains the weights of D features and ω ∈ R D×1 is a vector of fixed frequencies. We consider a Gaussian likelihood with variance 0.1 and prior p(w) = N (0, I D ); the true posterior over w is known to be Gaussian and it can be calculated analytically.
To assess the quality of the samples from the posterior obtained by I-SGD, in Fig. 2 we show the predictive posterior distribution (estimated using Eq. (2)) of I-SGD, in comparison to the "ground truth" posterior. Visual inspection indicates that there is a good agreement between predictive posterior distributions, especially in terms of uncertainty quantification for test points far from the input training distribution.

D Experimental Methodology
We hereafter present additional implementation details and experiments. We variants of I-SGD described in Section B. We report results for two preconditioner design choices: either tailored learning rates for each layer, as described in the main paper, or a unique learning rate where the equivalent λ is the sum of all λ (p) . We refer to this last version as single learning rate (SLR).

D.1 Regression tasks, with simple models
At test time we use 100 samples to estimate the predictive posterior distribution, using Eq. (2), for I-SGD and SGHMC, and 10 000 samples for MCD. All our experiments use 10-splits. The considered batch size is 128 for all methods. In this set of experiments we use I-SGD-α-A with µ = 0.5 during both warm-up and sampling. For both I-SGD-α-A and SGHMC the warm-up has been set to 2000, and we do store a sample every 2000 iterations (the keepevery value is 2000).

D.2 Classification task, simple CONVNET
For the LENET-5 on MNIST experiment, we consider I-SGD variants, MCD, and SWAG. We moreover consider samples obtained by SGD trajectories where the learning rate is derived as in Mandt et al. [2017] (V-SGD). At test time we use 30 samples for all methods. The batch size is 128, the temperature is 10 −5 and the keepevery and warm-up periods are 100 and 100 respectively. For SWAG we continued the training of the network for 5 epochs with learning rate 0.01.
We report results in Table 6, in general all methods perform similarly. 99.41 ± 0.03 0.0224 ± 0.0012 Table 6: Performance comparison of LENET-5 on MNIST

D.3 Classification task, deeper models
In the main paper we report results of RESNET-56 on CIFAR10 , using I-SGD-α-C,SWAG, MCD. Here we add results for I-SGD-G-B,I-SGD-G-C and V-SGD. At test time we use 30 samples for all methods. For I-SGD-α-C the batch size is 64, temperature is 10 −5 , warm-up and keepevery are 800 and 4000 respectively. For the Gaussian noise implementations the batch size is 64, the estimation momentum µ is 0.9 and the keepevery and warm-up periods are 100 and 1000 respectively. For SWAG we used the default parameters described in Maddox et al. [2019]. Notice that for the I-SGD-G-C version we treated biases and weights of the layers as unique groups of parameters. We report results in Table 7. We notice that the various I-SGD versions perform competitively and the (SLR) versions are worse in terms of performance.
We hereafter report additional results for CIFAR10 classification using VGG-16. We used the same parameters as for RESNET-56. We do omit results for I-SGD-G-B because we encountered numerical problems. Results are reported in Table 8. For results in Table 9, we use the same configuration as for the previous experiments: batch size is 64, temperature is 10 −5 , warm-up and keepevery are 800 and 4000 respectively.