Bayesian Inference for Large Scale Image Classification
Authors: Jonathan Heek
Nal Kalchbrenner
Venue:  (2020)
Keywords: image classification
bayesian inference
mcmc
imagenet

Abstract: Bayesian inference promises to ground and improve the performance of deep neural networks. It promises to be robust to overfitting, to simplify the training procedure and the space of hyperparameters, and to provide a calibrated measure of uncertainty that can enhance decision making, agent exploration and prediction fairness.
Markov Chain Monte Carlo (MCMC) methods enable Bayesian inference by generating samples from the posterior distribution over model parameters.
Despite the theoretical advantages of Bayesian inference and the similarity between MCMC and optimization methods, the performance of sampling methods has so far lagged behind  optimization methods for large scale deep learning tasks.
We aim to fill this gap and introduce ATMC, an adaptive noise MCMC algorithm that estimates and is able to sample from the posterior of a neural network.
ATMC dynamically adjusts the amount of momentum and noise applied to each parameter update in order to compensate for the use of stochastic gradients.
We use a ResNet architecture without batch normalization to test ATMC on the Cifar10 benchmark and the large scale ImageNet benchmark and show that, despite the  absence of batch normalization, ATMC outperforms a strong optimization baseline in terms of both classification accuracy and test log-likelihood. We show that ATMC is intrinsically robust to overfitting on the training data and that ATMC provides a better calibrated measure of uncertainty compared to the optimization baseline.

Under review as a conference paper at ICLR 2020
BAYESIAN INFERENCE
FOR LARGE SCALE IMAGE CLASSIFICATION
A Anonymous authors Paper under double-blind review
A ABSTRACT
Bayesian inference promises to ground and improve the performance of deep neural networks. It promises to be robust to overfitting, to simplify the training procodure and the space of hyperparameters, and provide a calibrated measure of uncertainty that can enhance decision making, agent exploration and prediction fairness. Markov Chain Monte Carlo (MCMC) methods enable Bayesian inference by generating samples from the poste- rior distribution over model parameters. Despite the theoretical advantages of Bayesian inference and the similarity between MCMC and optimization methods, the performance of sampling methods has so far lagged behind optimization methods for large scale deep learning tasks. We aim fill this gap and introduce ATMC, an adaptive noise MCMC algorithm that estimates and is able sample from the posterior of neural network. ATMC dynamncalll adjusts the amount of momentum and noise applied each parameter update in order to compensate for the use of stochastic gradients. We use ResNet architecture without batch normalization to test ATMC on the Cifar10 benchmark and the large scale ImageNet bench- mark and show that, despile the absence of batch normalization, ATMC out perporrs a strong optimization baseline in terms of both classification accuracy and test log- ligliikelhoood. We show that ATMC is intrinsically ro- bust to overfitting on the training data and that ATMC provides a better calibrated mcasure of uncertainty compared to the optimization bascline.
INTRODUCTION
In contrast to optimization approaches in machine learning that derive a single estimate for the weights of neural network, Bayesian inference aims at deriving posterior distribution over the weights of the network. This makes it possible to sample model instances from the distribution over the weights and offers unique advantages. Multiple model instances can be aggregated to obtain robust uncertainty estimates over the network's predictions; uncertainty estimates crucial in domains such medical diagnosis and autonomous driving where following a model's incorrect predictions can result in catastrophe (Kendall & Gal (011)). Sampling a distribution, as opposed to optimizing a loss, less prone to overlitting and more training doesn't decrease test performance. Bayesian inference can also be applied differential privacy, where each individual sample has increased privacy guarantees (Wang et al. 2015), and to reinforcement learning, where one can leverage model uncertainty to balance between exploration and exploitation (Osband & Van Roy 2017).
Traditional Markov Chain Monte Carlo (MCMC) methods like HMC (Neal et al. 2011 are a standard class of methods for generating samples from the posterior distribution over model parameters. These methods are seldom applied in deep learning because they have traditionally failed to scale well with large datasets and many parameters (Rajaratnam & Sparks 2015). Stochastic Gradient MCMC (SG-MCMC) methods have fared somewhat better in scaling to large datasets due to their close relationship to stochastic optimization methods. For example the SGLD sampler (Welling & Teh 2011) amounts to performing stochastic gradient descent while adding Gaussian noise to each parameter update. Despite these improvements, samplers like SGLD are only guaranteed to converge to the correct.

---

Under review as a conference paper at ICLR 2020
Algorithm 1 The ATMC sampler. The algorithm accepls the initialized model parameters 0o: step size h, pre-conditioner m, and momentum noise De-
1:
procedure ATMC TRAINING h, m, m, D)
2:
Po
3:
En 1-
4:
while t < < T do
5:
G, - minibatch.gradien
0t)
6:
TI - random_ )
7:
ai f max(D- - 0e
8:
Be + ai + E.
9:
Pi+h cahh pt
exp (8,
B.
G: +
exp 2 h
Cx2 77
10:
Oath 0: +
Pith
m
Eith + E:-
Pith
m
12: ( +
distribution when the step size is annealed to zero; additional control variates have been developed to mitigate this some extent (Ahn et al. 2012 Ding et al. 2014).
The objective of this work is to make Bayesian inference practical for deep learning by mak- ing SG-MCMC methods scale to large models and datasets. The contributions described in this work fall in three categories. We first propose the Adaptive Thermostat Monte Carlo (ATMC) sampler that offers improved convergence and stability. ATMC dynamically adjusts the amount of momentum and noise applied to each model parameter. Secondly, we improve an existing second order numerical integration method that needed for the ATMC sampler. Third, since ATMC, like other SG-MCMC samplers, is not directly compat- ible with stochastic regularization methods such as batch normalization (BatchNorm) and Dropout (see Sect. [4. we construct the ResNet++ network by taking the original ResNet. architecture (He et al. 2016). removing Batch and introducing SELUs (Klambaner et al. 2017). Fixup initialization (Zhang al.) 2019a) and weight normalization (Salimans & Kingmaa 2016). We design ResNeet++ so that its parameters are easy to sample from and the gradients are well-behaved even in the absence of BatchNorm.
We show that the ATMC sampler is able to outperform optimization methods in terms of accuracy, log-likelihood and uncertainty calibration in the following settings. First, when using the + architecture for both the ATMC sampler and the optimization baseline, the ATMC sampler significantly outperforms the optimization baseline on both Cifar-10 10 and ImageNet. Secondly, when using the standard Res.et for the optimization baseline and the ResNet++ for the ATMC sampler, multiple samples the ATMC that approximate the predictive posterior of the model are still able to outperform the optimization baseline ImageNet. Using the + architecture, the ATMC sampler reduces the need for hyper- pareme-ar tuning since it does not require carly stopping, does not use stochastic regularization, is not prone to fittingt on the training data and avoids a carefully tuned learning rale decay schedule.
2 ATMC SAMPLER
In this section we define the Stochastic Differential Equation (SDE) that gives rise to the ATMC sampler described in Algorithm [ A detailed background and framework for con- structing SDEs that converge to target distribution can be found in (Ma et al. 2015).
2.1 GENERAL FORM OF THE SDE
Our starting point for constructing the ATMC sampler is the framework of Stochastic Dif- ferential Equations. We are interested in SDEs that converge to a distribution p(2) over the vector E Rdd for which we evaluate V llog((2). Because only the gradient of p(z()
2

---

Under review as a conference paper at ICLR 2020
is required, it is sufficient to define energy function H22) log C up stant C. As a consequence, we can sample from the posterior distribution p(0 x) by only evaluating the energy function gradient (H(0) v log p(z, 0). The general form of SDEs converging to p(z) for which only the gradient of p(z) is required is as follows Ma et al., 2015
((()d (2D(s) r₂(=) Σ (ii((=) Q1)(=)]
de
dzj
where (=) a positive-defnitte matrix that determines the amount of noise, Q(=) is skew-symmetric matrix that mixes energy between variables, W, is a Wiener process, and [(z) is a correction factor that compensates for dynamics that depend on the current state 2. 'he ATMC sampler that we propose is an instance of ( for specific definitions of (z), D(2), and Q(2).
2.2 ENERGY FUNCTION
We start by defining the energy function (2).. The energy function for the model posterior p(0|z) is defined by the loss function (()) log p(x, 0). Because the dataset x is generally large, we would like only evaluate mini-batch loss Σ(0). However, naively using stochastic gradient in (1) will result in significant bias (Chen et al. 2014). Motivated by the Central Limit Theorem, the stochastic gradient assumed follow Gaussian distribution Vā(0) N(VL(()) B) where the covariance is additionally assumed be diagonal and constant w.r.t. 0. The energy function for the ATMC sampler is defined as:
1 H(O, ((() + k(p) + [ diag( B)
(2)
2m
where p is the momentum, (pp) defines the momentum distribution, and & is control variate referred as the temperature. Both p and & have the same dimensionality as 0. The hyper-parameter m controls the strength of the coupling between & and p. The distribution of the control variate p(E) depends on the amount of noise B the stochastic gradient estimate ((0).
2.3 NOISE ROBUST DYNAMICS
Next we define the dynamics Q(z) and D(2) such that the SDE that results from be simulated without the need to evaluate B:
can
O
am) + J 9) Q(0.p.E)
0 -1
((),p.C)=
mVK(p)
(3)
-mVK(p)
where a(E) a non- neg-tieae function that determines how the temperature & affects the amount of noise added to the momentum update.
We first illustrate the resulting SDE by using a simpler Gaussian momentum distribution K(p) = |p| /2m). Note that the variance of the momentum Var(p) = m is reused in ( and 3 to control the strength of the coupling between & and p. This will result in
temperature control with a momentum friction term proportional to &, unlike previously
reported thermostat MCMC methods (Ding et. al.] 2014 Lu et al. 2016) where the friction

---

Under review as a conference paper at ICLR 2020
term is proportional to g/m. We substitute the dynamics Q(2) and D(2) defined in into
3 and
energy function (2) defined in
d0
p/m v [(0) 3(E)p dt + p²/m-
0
))av. dw,, B(E) a(E) + (4)
dp
v (2(())
de
where we use V((0)dtt dt+ +CC)) replace the gradient of the loss with the mini- batch estimate. The momentum is dampened by a friction term ((5) that depends on the
choice of (g)). The stochastic gradient noise does not show up (4) due to the particular
choice of energy function (()) and dynamics (z)), D(2). Note however this analysis relies on the assumption that the covariance of the stochastic gradient noise is constant in and single temperature variable per parameter can only correct for diagonal covariance B. We do not expect that this assumption will hold in practice and the approximation will therefore lead to bias in the samples. However, annealing the step size h will reduce the error due to bani-batching together with other sources of discretization error (Welling & Teh 2011
2.4 ADAPTIVE NOISE THERMOSTAT
Finally, we must choose a function a(E) which controls the amount of noise and momentum damping 3(E) Previous work uses the Nosé-Hoover thermostat that is defined by a(g) De where D. is a constant determining the amount of noise added to the momentum update (Ding et al., 2014). Although the Nosé- Hoover thermostat is able correct the stochastic gradient noise B, the correction comes at the cost of slower convergence because additional friction ((() is applicd as increascs. Another drawback of the Nosé-Hoover thermostat is that it causes negative friction when < DD. In the negative friction phase B(5) < 0, previous gradient terms are amplisied rather than dampened. Although this behavior mathematically sound find that it can cause exploding momentum variables.
Our choice of a<{) based on the idea that negative friction should not occur and conver- gence speed should not be reduced by the stochastic gradient noise. Based on this intuition, we define the ATMC sampler by a(E) = max D - E, 0). The ATMC sampler is best charac- terized by the various temperature stages. For 0 < D. the total amount of noise added to the momentum is De and the friction coefficient B(E) = De- At this stage, the stochastic gradient noise is compensated for by adding less noise to the momentum update. If B > De the dominant stage will be & > D resulting in 3(5) < D. and zero noise being added to the momentum. Finally, when & < 0 the friction coefficient B(E) = D. and the noise added to the momentum is proportional to D. c. Thus, the momentum always experiences minimum amount of friction B({) 2 De determined by the syper-parameter De and the noise added to the momentum update is automatically adjusted based on the amount of noise present in the stochastic gradients.
2.5 MOMENTUM ENERGY FUNCTION
Following Lu et al. 2016 we generalize the momentum energy function K(p) to the symmetric hyperbolic distribution which is defined as follows (Lu et al. 2016
K(p) = E me²
D m²c²
+1-
(5)
where m and e are hyper-parameters. The Gaussian kinetic energy K(p) |Ipll /(2m) is special case obtained by taking the limit - xo. The magnitude of parameter updates | is determined by the gradient of the momentum:
| p M(p)
p²
| = |(
M(p) = m/
+
(6)
m

---

Under review as a conference paper at ICLR 2020
Hence, the hyperbolic distribution results in relativistic momentum dynamics where the parameter updates are upper bounded by and the pre-conditioner (()) depends on p. The [I]
average update magnitude (p) = 1/ /m for > m. Consequently, the parameters m and e are interpretable hyper-parameters controlling the average and maximum parameter update per step together with the step size h.
The SDE we derive in (4) and integrate in Sec. 8 uses a Gaussian momentum enorgy function for clarity. Deriving ATMC with a different momentum distribution like the hyperbolic distribution amounts to substituting (2), (3 and the alternative momentum distribution into (T. For the hyperbolic distribution, the dynamic friction coefficient }({) will also
depend on p. For the numerical integration of with a hyperbolic momentum distribution
we assume ((E) to be constant in p.
IMPROVED NUMERICAL INTEGRATOR FOR MCMC SAMPLERS
In this section we construct the numerical integrator required to numerically approximale the ATMC sampler defined in 4. An efficient numerical integrator be constructed by splitting the SDE into two terms:
p/ p/m m O
0
dp de
dt +
-v 3(p, E)p
dt
v 2
0
dW,
(7)
p² m
0
B
Hence, we obtain a linear ODE in part (A) that updates the parameters and the ther- mostats & and a linear SDE in part (B) that updates the momentum p. The operators that simulate these dynamics exactly for a time step h are denoted of and d4. respectively. Using the Strang splitting scheme yields a second order method (Chen et al. 2015):
(8)
The first operator d is given by
(b(())) 02 h
&
+
-
(9)
The second operator ΦB is an instance the Ornstein Uhlenbeck process which can also be computed analytically follows (Nelson
1967
dz((2)) (0. (1E,) Pt (EEEV (2((E))((E) n E.) E,)",
(10)
7a(E)) expla 3(E,))] -
3(E,)
where 7). isotropic Gaussian noise. Previous work Chen al. 2015 on higher order integrators for samplers splits the SDE into three parts where the third term is obtained from separating the friction term from the other terms in the momentum update Óp. By integrating (10) exactly the gradient step and the noise and gradient term directly affected by the friction. An exact momentum update provides additional robustncss to large gradients because the temperature will increase in order to compensate for momentum updates that would lead to excessively large sleps. Another advantage of a two-way split. integrator is that the first. and last steps in (8) can be fused together such that only momentum update is performed per iteration. Algorithm I shows the pseudocode for the ATMC sampler with the split integrator defined in 9 and 10
4
THE RESNET++ ARCHITECTURE
The generalization performance of large neural nets trained using optimization depend on stochastic regularization methods like Dropout. (Srivastava et. al. 2014) and BatchNorm
5

---

Under review as a conference paper at ICLR 2020
Conv Ixl
BatchNorm
ReI.U
Conv 3x.3
BatchNorm
ReI.U
Conv Ixl
BatchNorm
ReL.U
Conv Ixl
ScLU
Cony 3x3
ScLU
Conv Ixl
ScLU
Figure 1: Residual blocks in respectively the ResNet and ResNet++ architectures.
X19.903%
SGD
SG0 BlachNorm
SGD Batchtomm
SGNHT predictive
predictive
posterior predictive
ATMC postenir predictive
ideal
ideal
200933
90.3%
50%
02%
Figure Calibration plot for Cifar10
Figure 3: Calibration plot for ImageNet
loffe Szegedy 2015 These methods implicitly add noise into the model parameters
Kingma al. 2015 Teye et al. 2018) and significantly boost training performance and generalization for image classifiers. These methods can be interpreted as a coarse approx- imation of Bayesian Inference (Kingma
al. 2015 Teye al., 2018). But stochastic gradient sampler like ATMC already adds the necessary amount of noise and combined with BatchNorm or Dropout leads to underfitting. We thus define BatchNorm free version of ResNet called ResNet+ that includes SELUs (Klambauer et al. (017]. Fixup initialization (Zhang et al., (019a)) and weight normalization (Salimans & Kingma 2016) (see Fig. []. We use ATMC to fill the significant gap in performance due to the absence of BatchNorm in ResNet++..
4.1 SELU
We find the SELU activation to work well in BatchNorm free networks. SELU forces the statistics of the activations towards zero mean and unit variance (Klambauer et al.) (2017). The SELU activation function additionally has a non-zero gradient everywhere which could improve the mixing of the sampler by providing more informative gradient.
4.2 FIXUP INITIALIZATION
ResNets are known to scale well with depth (He et al. 2016). However, the additive effect of the residual branch causes the magnitudes of the activations increase with the number of residual connections. Fixup is recently proposed initialization method that mitigates the exploding residual branch problem without using BatchNorm (Zhang al. 2019a). We use a simplified version of Fixup by initializing the scales of the final layer in each residual branch a small constant.
4.3 WEIGHTT NORMALIZATION
We use weight normalization Salimans & Kingma 2016 to separate the direction and scale of each linear feature vector
ii)
(1) (6)
(11)

---

Under review as a conference paper at ICLR 2020
Table 1: Performance on Cifar10 with ResNet-56 model. The posterior predictive estimate
using sample of the posterior parameters at the end of cach learning rate cycle.
Setup
Top acc. [%]
NLL [Nats]
SGD
91.5
0.370
SGD + BatchNorm
94.4
0.243
ATMC
(single sample)
92.4
0.303
ATMC
(Posterior predictive)
93.9
0.194
SGNHT
(single sample)
91.7
0.343
SGNHT (Postcrior predictive)
93.5
0.211
where of is the direction vector and @s is the magnitude of a feature vector ((0,, Weight normalization does not depend on batch statistics and is compatible with MCMC methods.
The scale of the direction vector does not affect the outputs of the model. does however affect the effective step size (Wu et al., 2018). Therefore the prior the direction vector (0 chosen such that forced to unit length
[ (10001--11
-11 )
p(ol') x exp
(12)
The prior the scales p(o.) is problem-speeific and can for example be chosen to encode a preference for structurally sparse models.
EXPERIMENTS
The experiments presented here aim demonstrate that the ATMC sampler is competitive with a well-tuned optimization baseline for large-scale datasets and models. We use the TensorFlow official implementation of ResNet-56 and ResNet-50 on Cifar10 and ImageNet, respectively. We compare our ATMC sampler to an optimization bascline with and with- out BatchNorm. For the optimization bascline without BatchNorm we use the + architecture as described in Sec. 4 For the baseline with BatchNorm found standard ResNet with Xavier initialization and the ReLU non-linearity to work better.
For the ATMC sampler report both the performance of single sample and the estimated posterior predictive based on a finite number of samples. Similar to earlier work (Zhang et al. 2019b) found that many fewer samples needed when cyclic step size he + cOS(TT mod[t,n]) with cycle length n is used. The final sample in each cycle is used to estimate the posterior predictive.
hp
For ResNet++ we further use a group Laplace prior p(0,) x exp(- ||.11| /b) with regularize the scales of each linear feature in ResNet++. The momentum noise is chosen as D. - -log(0.9)/la such that the friction applied to the momentum is at least 0.9.
to
5.1 CIFAR 10
For Cifar10 we choose the step size hg = 0.001 and the cycle length is set to 50 epochs. The momentum hyper-parameters are m (0.0003/ho)) and c 0.001 /hg such that the average speed and maximum speed per step are 0.0003 and 0.001, respectively. The number of convolution filters doubled to 32 compared to the original ResNet-56 implementation. We use single V100 GPU with batch size of 128. The sampler runs for 1000 epochs and we start collecting samples for the posterior predictive after 150 epochs. The optimization baseline converges in 180 epochs. We also report the results of sampling with sampler based Nosé Hoover thermostats (SGNHT) (Ding et al. 2014 Lu et al. 2016) applied to the Res architecture.

---

Under review as a conference paper at ICLR 2020
Table 2:
Performance on ImageNet with ResNet-50 model. The posterior predictive is
estimated using sample of the posterior paramcters at the end of cach learning rate cycle.
Setup
Top 1 acc. [%]]
NLL [Nats]
SGD
70.9
1.24
SGD + BatchNorm
76.2
0.947
ATMC
(single sample)
74.2
1.08
ATMC
(Posterior predictive)
77.5
0.883
SGNHT
(single sample)
73.1
1.15
SGNHT
(Posterior predictive)
76.4
0.941
Table lists the test set performance for Cifar10. A single sample from the posterior already outperforms the baseline without BatchNorm by a significant margin in both test accuracy and lig-liielihood. Using BatchNorm significantly improves the generalization of the optimization bascline. It outperforms the estimate of the posterior predictive in accuracy yet it does not have better test log- ligliiklihhod.
To further analyze the quality of the uncertainty estimates, we group each model's prediction in 8 equally sized bins based the confidence ([;)( where Wy is the maximum probability class for example T4- If the probabilities are well-calibrated, the average confidence should be close the average accuracy. Figure shows the calibration of the uncertainty estimates for the posterior predictive and optimization baselines. The posterior predictive is calibrated for the least confident predictions px;; < 0.9 and shows less bias towards overconfidence compared to the models trained with SGD.
5.2 IMAGENET
For the ImageNet experiments we use an initial step size bg = 0.0005 and a cycle length of 20 epochs. The other hyper-parameters for the sampler the same as for the Cifar10 experiments. We use a single Google Cloud TPUv3 with batch size of 1024. We did not observe a significant difference in wall clock time per training step between SGD and ATMC on the same model. Each training step using ResNet- BatchNorm model takes 20% longer in wall clock time compared to single train step using ResNet++. Samples for the posterior predictive are collected after 150 epochs and the sampler runs for 1000 epochs. The optimization baseline converges in 90 epochs.
Table 2 lists the results for ImageNet classification. A single sample from the posterior outperforms the optimization baseline without BatchNorm. The posterior predictive based on ATMC outperforms the optimizer with BatchNorm by wide margin in both accuracy and test log-likelihood. We note that the sampler runs significantly longer (10x) compared to the optimization baseline because it takes long time for the posterior predictive esti- mate to converge. However, the posterior predictive of ATMC matches the accuracy of the optimization baseline with NarchNorm (76.2%) after 240 epochs.
Figure shows the quality of the uncertainty for various levels of confidence. Again, the ATMC based posterior predictive produces much better calibrated predictions and is almost perfectly calibrated for low confidence predictions p(:](r;) < 0.9 and shows less bias towards overconfidence compared to the optimization baseline.
DISCUSSION
The empirical results show possible to sample the posterior distribution of neural net- works on large scale image classification problems like ImageNet. A major obstacle for sampling the posterior of ResNets in particular is the lack compatibility with Batch- Norm. Using recent advances in initialization and the SELU activation function to stabilize and speed up training of ResNets without resorting to BatchNorm. Nonetheless, observe that BatchNorm still offers unique advantage in terms of generalization per-
able

---

Under review as a conference paper at ICLR 2020
formance. We hope that future work will allow the implicit inductive bias that BatchNorm has to be transferred into an explicit prior that is compatible with sampling methods.
Multiple posterior samples provide much more accurate estimate of the posterior pre- dictive, and consequently much better accuracy and uncertainty estimates. For inference, making predictions using large ensemble of models sampled from the posterior can be costly. Variational Inference methods can be used to quickly charactcrize a local mode of the posterior (Blundell al., 2015)). More recent work shows that rumning estimate of the mean and variance of the parameters during training can also be used to approximate a mode of the posterior (Maddox et al. 2019). Methods like distillation could potentially be used to compress high-quality ensemble into single network with a limited computational budget. (Balan et al. 2015).
Although the form in is very general, alternative methods for dealing with stochastic gradients have been proposed in the literature. One approach is estimate the covariance of the stochastic gradient noise B explicitly and use it correct and pre-condition the sampling dynamics (Ahn et al. 2012 Li et al. 2016).
Other sampling methods are not based on SDE that converges to the larget distribution. Under some conditions stochastic optimization methods can be interpreted as such a biased sampling method (Mandt et al.) 2017). Predictions based multiple samples from the trajectory of SGD have been used successfully for obtaining uncertainty estimates in large scale Deep Learning (Maddox et al. 2019). However, these methods rely tuning hyper- paramoters such a way that just the right amount of noise inserted.
CONCLUSION
This work introduces the ATMC sampler, a robust posterior sampling method that scales to large deep learning problems. To the best of our knowledge, we are the first to successfully train neural networks using MCMC on ImageNct. In a BatchNorm free setting, single sample from the posterior generated by ATMC outperforms the optimization baseline. A posterior predictive estimate outperforms the optimization baseline with BatchNorm on ImageNet. Based on these empirical results we hope the ATMC sampler will enable new applications of Bayesian inference in deep learning.
REFERENCES
Sungjin Ahn, Anoop Korattikara, and Max Welling. Bayesian posterior sampling via stochastic gradient fisher scoring, arXiv preprint arxiv: 6380, 2012.
Anoop Korattikara Balan, Vivek Rathod, Kevin P Murphy, and Max Welling. Bayesian dark knowledge. Advances in Newral Information Processing Systems, pp. 3438 1138-3446, 2015.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight un- certainty in neural network. In International Conference on Machine Learning, pp. 1613- 1622, 2015.
Changyou Chen, Nan Ding, and Lawrence Carin. On the convergence of stochastic gra- dient meme algorithms with high-order integrators. In Advances in Neural Information Processing Systems, 2015.
Tianqi Chen, Emily Fox, and Carlos Guestrin. Stochastic gradient hamiltonian monte carlo. In International conference on machine learning, pp. 1683-1691, 2014.
Nan Ding, Youhan Fang, Ryan Babbush, Changyou Chen, Robert D Skeel, and Hartmut Neven. Bayesian sampling using stochastic gradient thermostats. In Z. Ghahramani, M. Welling, C. Cortes, N. D. Lawrence, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems 27, pp. 3203 3211. Curran Associates, Inc., 2014.

---

Under review as a conference paper at ICLR 2020
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pottern recognation, pp. 770-778, 2016.
Sergey Toffc and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pp. 448 456, 2015.
Alex Kendall and Yarin Gal. What uncertainties do we need in bayesian deep learning for computer vision? In Advances in neural information processing systems, pp. 5574 5584, 2017.
Durk P Kingma, Tim Salimans, and Max Welling. Variational dropout and the local repa- rameterization trick. In Advances in Neural Information Processing Systems, pp. 2575- 2583, 2015.
Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self- normalizing neural networks. In Advances in neurol information processing systems, pp. 971 980, 2017.
Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin. Preconditioned stochas- tic gradient langevin dynamics for deep neural networks. In Thirticth AAAI Conference on Artificial Intelligence, 2016.
Xiaoyu Lu, Valerio Perrone, Leonard Hasenclever, Yee Whye Teh, and Sebastian J Vollmer. Relativistic monte carlo. orXiv preprint arXiv:1609.04330 2016.
Yi-An Ma, Tianqi Chen, and Emily Fox. A complete recipe for stochastic gradient meme. In C. Cortes, N. D. Lawrence, D. D. Lee, M. Sugiyarna, and R. Garnett (eds.), Advances in Neural Information Processing Systems 28, pp. 2917-2925. Curran Associates, Inc., 2015.
Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncortainty in deep lcarning. arXiv preprint arXiv:1902 1et.1024226666 2019.
Stephan Mandt, Matthew D Hoffman, and David M Blei. Stochastic gradient descent as approximate bayesian inference. The Journal of Machine Learning Research, 18(1):4873 4907, 2017.
Radford M Neal et al. Meme using hamiltonian dynamics. Hondbook of markov chain monte carlo, 2(11):2, 2011.
Edward Nelson. Dynamical theories of Brownian molion, volume 3. Princeton university press, 1967.
lan Osband and Benjamin Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 2701-2710. JMLR. org, 2017.
Bala Rajaratnam and Doug Sparks. Meme-based inference in the era of big data: A funda- mental analysis of the convergence complexity of high-dimensional chains. arXiv preprint arXi::1508. 00947, 2015.
Tim Salimans and Durk Kingma. Weight normalization: A simple reparameterization to accelerate training of deep neural networks. In Advances in Neural Information Processing Systems, pp. 901 909, 2016.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, llya Sutskever, and Ruslan Salakhut- dinov. Dropout: simple way to prevent neural networks from overfitting, The Journal of Machine Learning Research, 15(1)=1999-1958, 2014.
Mattias Teye, Hossein Azizpour, and Kevin Smith. Bayesian uncertainty estimation for batch normalized deep networks. Xiv preprint 180 06455, 2018.
10

---

Under review as a conference paper at ICLR 2020
Yu-Xiang Wang, Stephen Fienberg, and Alex Smola. Privacy for free: Posterior sampling and stochastic gradient monte carlo. In International Conference on Mochine Learning. pp. 2493 2502, 2015.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp- 681-688, 2011.
Xiaoxia Wu, Rachel Ward, and Léon Bottou. Wngrad: learn the learning rate in gradient descent. ar preprint ar Xiv: 02865, 2018.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint 1rXi::101.009922 2019a.
Ruqi Zhang, Chunyuan Li, Jianyi Zhang, Changyou Chen, and Andrew Gordon Wil- son. Cyclical stochastic gradient meme for bayesian deep learning, arXiv preprint arXi::1902.0800 03932, 2019b.
11


Reviews:
[{"review_id": "rJlQBJSP9H", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a novel MCMC algorithm (ATMC) that estimates and samples from the posterior distribution of neural network weights. The motivation for this approach is that applying Bayesian inference to deep learning should lead to less overfitting and better uncertainty-calibrated models. Unlike previous work, the proposed method scales to large models (ResNet) and data sets (ImageNet).\n\nIn addition to the main contribution, this work improves an existing numerical integrator necessary for the ATMC sampler. Moreover, in order to apply this method to ResNet, the authors use a modified version of ResNet without stochastic regularization (batch normalization and dropout).\n\nEmpirical results show that the proposed method outperforms baselines on accuracy, log-likelihood, and uncertainty calibration. The single-sample posterior already yields decent results, beating a batchnorm-free version of ResNet.\n\nOverall, the paper is well-written, fluently readable, and provides a clear presentation of motivations and methods. Experiments, although not very extensive, are described in sufficient detail and corroborate the claims. Related previous work is cited throughout the paper, although there is no explicit section for it.\n\n\nWeaknesses:\n\nWhile the authors claim that the need for hyperparameter tuning is reduced, they use a cyclic step size with parameter n, a Laplace prior with parameter b, a momentum noise with parameter 0.9, and dataset-specific parameters h0, m, and c. This weakens (or even contradicts) the claim, and raises the question of how much the choice of these hyperparameters affects performance.\n\nOn ImageNet, ATMC doesn't seem to be well calibrated, and the authors do not discuss the fact that the calibration is basically on par with SGNHT.\n\n\nMinor:\n- page 4: practise -> practice\n- page 7: extra space after \"compatible with MCMC methods\"\n- page 7: problem specific -> problem-specific", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I did not assess the derivations or theory."}, "extra_text": {}}, {"review_id": "BJlHL35H9B", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose the adaptive thermostat Monte Carlo sampler for feedforward neural networks. The proposed approach dynamically adjust the amount of momentum and noisy applied to each model parameter during updates. ResNet++ (ResNet without batchnorm/dropout but adding SELU, fixup and weight normalization) is introduced. Further, the authors claim that the need for hyperparameter setup is reduced provided that early stopping, stochastic regularization and carefully tuned learning rate schedules are not required.\n\nThe authors highlight some practical issues with the Nose-hoover thermostat, however, recognize its mathematical soundness. When ATMC is described (temperature stages), a motivation is provided but not justified theoretically.\n\nIn (10) and (11), \\gamma_1(), \\gamma_2(), \\eta_t and a are introduced without definition or further explanation.\n\nThe authors claim that the need for hyperparameter setup is reduced, however, in the experiments they use a cyclic step size with length n=50 (20 for ImageNet), a Laplace prior with parameter b=5, momentum noise with parameter 0.9, pre-conditioner parameter 0.0003 and c parameter 0.001. The impact of these choices on performance is not described. Further, the number of filters is doubled relative to ResNet-56 without explanation.\n\nThe calibration curves in Figure 3 are underwhelming. ATMC is better than SGD but not necessarily well calibrated. Also, note that x and y scales are heavily biased toward 1.\n\nIn summary, the proposed approach needs to be described in more detail and the experiments are not very satisfying given the claims made by the authors in the Introduction.\n\nMinor:\n- In (1) W is not defined.\n- In (1) the dimensionality of D, Q and \\Gamma is not defined but their elements are used.\n- In (2) m is only defined after (3), in fact, only called by its name, pre-conditioner, in Algorithm 1.\n- In Section 2.3 there is a reference to the step size, though not introduced until discretization later in Section 3.\n- In (7) \\beta() is a function of p, but not in other instances, e.g., (4), (10) and (11).\n- Move Algorithm 1 closer to definition.\n- In (13), d is not defined.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory."}, "extra_text": {}}, {"review_id": "BJl3XMGhtS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThe main idea of the paper is the introduction of ATMC, an adaptive noise MCMC algorithm that dynamically adjusts the momentum and noise applied to each parameter update while sampling from the posterior of a neural network. A modified version of ResNet architecture called ResNet++ has been introduced in the paper and later used to train it on ImageNet using MCMC which is great work considering the complexities associated. Furthermore, the authors claim ATMC to be robust to overfitting and it provides a measure of uncertainty with predictions. The paper is well written with clear flow and good mathematical explanation. \n\nQuestions from authors:\n\n1. The approach of Stochastic Gradient MCMC (SG-MCMC) has been there for quite some time and in this paper, there is no clear explanation of the advantages of ATMC over SG-MCMC methods. The authors can argue that the introduction of dynamically adjustable momentum and noise can make the paper novel and unique. However, in Stochastic gradient Markov chain Monte Carlo (SG-MCMC) [1], the authors used a meta-learning algorithm to learn Hamiltonian dynamics with state-dependent drift and diffusion which makes the system scalable to large datasets, very similar to this work. Also, changing the momentum variable in Hamiltonian dynamics with thermostat variable in stochastic gradient thermostats, the meta learner algorithm of [1] can be used to make the stochastic gradient thermostats approach adaptive. What are the advantages of the approach proposed in this paper over [1]?\n\n2.  The paper follows the approach of stochastic gradient thermostats as introduced in the paper  [2] and builds upon it to make it adaptive to improve stability and for better convergence.  The advantages are clearly mentioned in the paper except for the case when 0 < ξ < D_c the total amount of noise added to the momentum is D_c and the friction coefficient β(ξ) = D_c. At this stage, the authors claim that 'the stochastic gradient noise is compensated for by adding less noise to the momentum update.' How is this done? Can authors please explain it a bit more.\n\n3.  For the experiments, the authors introduced ResNet++ architecture which is based on ResNet architecture but is novel in its design with the use of SELU and removal of BatchNorm and with different initialization schemes. The design idea is well explained by the authors. The experiments were performed on CIFAR10 and ImageNet dataset to show the large scale scaling of the approach. A comparison with SGNHT is provided in the paper but a fair comparison with other approaches like SGHMC [3], PSGLD [4] and MCMC [1] is missing which might have proved the effectiveness of the approach when compared to other methods. Any comments on why the experiments were restricted to the ones mentioned in the paper?\n\n4. Lastly, the ResNet++ network is trained on ImageNet and CIFAR10 datasets. There is no clear mention of the time duration it took for training and evaluating the network. The authors mentioned ‘BatchNorm did result in an overhead of roughly 20% compared to the ResNet++ model’, which is not clear. I would urge the authors to explain the above line and provide an estimate of the training time and testing time. \n\nA closing remark: The mention of the acknowledgement section in a double-blind review is not advisable and in future please refrain from doing so.\n\nReferences:\n[1] https://openreview.net/pdf?id=HkeoOo09YX#page=11&zoom=100,0,754\n[2] http://people.ee.duke.edu/~lcarin/sgnht-4.pdf\n[3] https://arxiv.org/pdf/1402.4102.pdf\n[4] https://arxiv.org/abs/1512.07666\n\n\n\n\n\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a variant of Hamiltonian Monte Carlo for Bayesian inference in deep learning.

Although the reviewers acknowledge the ambition, scope and novelty of the paper they still have a number of reservations regarding experimental results and claims (regarding need for hyperparameter tuning). The overall score consequently falls below acceptance.

Rejection is recommended. These reservations made by the referees should definitely be addressable before next conference deadline so looking forward to see the paper published asap.
Decision: Reject