Deep Ensembles: A Loss Landscape Perspective

1 INTRODUCTION . Consider a typical classification problem , where xn ∈ RD denotes the D-dimensional features and yn ∈ [ 1 , . . . , K ] denotes the class label . Assume we have a parametric model p ( y|x , θ ) for the conditional distribution where θ denotes weights and biases of a neural network , and p ( θ ) is a prior distribution over parameters . The Bayesian posterior over parameters is given by p ( θ| { xn , yn } Nn=1 ) ∝ p ( θ ) N∏ n=1 p ( yn|xn , θ ) . ( 1 ) Computing the exact posterior distribution over θ is computationally expensive ( if not impossible ) when p ( yn|xn , θ ) is a deep neural network . A variety of approximations have been developed for Bayesian neural networks , including Laplace approximation ( MacKay , 1992 ) , Markov chain Monte Carlo methods ( Neal , 1996 ; Welling & Teh , 2011 ; Springenberg et al. , 2016 ) , variational Bayesian methods ( Graves , 2011 ; Blundell et al. , 2015 ; Louizos & Welling , 2017 ; Wen et al. , 2018 ) and Monte-Carlo dropout ( Gal & Ghahramani , 2016 ; Srivastava et al. , 2014 ) . While computing the posterior is challenging , it is usually easy to perform maximum-a-posteriori ( MAP ) estimation , which corresponds to a mode of the posterior . The MAP solution can be written as the minimizer of the following loss ( negative log likelihood + negative log prior ) : θ̂MAP = argmin θ L ( θ , { xn , yn } Nn=1 ) = argmin θ − log p ( θ ) − N∑ n=1 log p ( yn|xn , θ ) . ( 2 ) The MAP solution is computationally efficient , but only gives a point estimate and not a distribution over parameters . Deep ensembles , proposed by Lakshminarayanan et al . ( 2017 ) , train an ensemble of neural networks by initializing at M different values and repeating the minimization multiple times which could lead to M different solutions , if the loss is non-convex . ( Lakshminarayanan et al . ( 2017 ) found adversarial training provides additional benefits in some of their experiments , but we will ignore adversarial training and focus only on ensembles with random initialization in this paper . ) Given finite training data , many parameter values could equally well explain the observations , and capturing these diverse solutions is crucial for quantifying epistemic uncertainty ( Kendall & Gal , 2017 ) . Bayesian neural networks learn a distribution over weights , and a good posterior approximation should be able to learn multi-modal posterior distributions in theory . Deep ensembles were inspired by the bootstrap ( Breiman , 1996 ) , which has nice theoretical properties . However , it has been empirically observed by Lakshminarayanan et al . ( 2017 ) ; Lee et al . ( 2015 ) that training individual networks with just random initialization is sufficient in practice and using the bootstrap even hurts performance in some cases ( e.g . for small ensemble sizes ) . Furthermore , Ovadia et al . ( 2019 ) and Gustafsson et al . ( 2019 ) independently benchmarked existing methods for uncertainty quantification on a variety of datasets and architectures , and observed that ensembles tend to outperform approximate Bayesian neural networks in terms of both accuracy and uncertainty , particularly under dataset shift . These empirical observations raise an important question : Why do ensembles trained with just random initialization work so well in practice ? One possible hypothesis is that ensembles tend to sample from different modes1 in function space , whereas variational Bayesian methods ( which minimize DKL ( q ( θ ) |p ( θ| { xn , yn } Nn=1 ) ) might fail to explore multiple modes even though they are effective at capturing uncertainty within a single mode . See Figure 1 for a cartoon illustration . Note that while the MAP solution is a local minima for the training loss by definition , it may not necessarily be a local minima for the validation loss . Recent work on understanding loss landscapes ( Fort & Jastrzebski , 2019 ; Draxler et al. , 2018 ; Garipov et al. , 2018 ) allows us to investigate this hypothesis . Note that prior work on loss landscapes has focused on mode-connectivity and low-loss tunnels , but has not explicitly focused on how diverse the functions from different modes are , beyond an initial exploration in Fort & Jastrzebski ( 2019 ) . Our findings show that : • The functions sampled along a single training trajectory or subspace thereof ( e.g . diagonal Gaussian , low-rank Gaussian and Dropout subspaces ) tend to be very similar in predictions ( while potential far away in the weight space ) , whereas functions sampled from different randomly initialized trajectories tend to be very diverse . • Solution modes are connected in the loss landscape but they are distinct in the space of predictions . Low-loss tunnels create functions with near-identical low values of loss along the path , however these functions tend to be very different in function space , changing significantly in the middle of the tunnel . 2 BACKGROUND . The loss landscape of neural networks ( also called the objective landscape ) – the space of weights and biases that the network navigates – is typically a very high dimensional function and therefore could potentially be very complicated . However , many empirical results show interesting properties of the loss surface . Goodfellow & Vinyals ( 2014 ) observed that the loss along a linear path from an initialization to the corresponding optimum is monotonically decreasing , encountering no significant 1We use the term mode to refer to unique functions fθ ( x ) . Due to weight space symmetries , different parameters could correspond to the same function , i.e . fθ1 ( x ) = fθ2 ( x ) even though θ1 6= θ2 , but we ignore this aspect and leave it to future work . obstacles along the way . Li et al . ( 2018 ) demonstrated that constraining optimization to a random , low-dimensional hyperplane in the weight space leads to results comparable to full-space optimization , provided that the dimension exceeds a modest threshold . This was geometrically understood and extended in ( Fort & Scherlis , 2019 ) . Garipov et al . ( 2018 ) ; Draxler et al . ( 2018 ) demonstrate that while a linear path between two independent optima hits a high loss area in the middle , there in fact exist continuous , low-loss paths connecting any pair of optima . These observations are unified into a single phenomenological model in ( Fort & Jastrzebski , 2019 ) . While independent , low-loss optima in the loss landscape are connected , Fort & Jastrzebski ( 2019 ) provide an early indication that in fact they represent very different functions in terms of their predictions . Therefore the connectivity can not be due to trivial symmetries of the network which would keep the input–output mapping intact . 3 VISUALIZING FUNCTION SIMILARITY ACROSS INITIALIZATIONS . We train convolutional neural networks on the CIFAR-10 ( Krizhevsky , 2009 ) dataset : • SmallCNN : channels [ 16,32,32 ] for 10 epochs which achieves 64 % test accuracy . • MediumCNN : channels [ 64,128,256,256 ] for 20 epochs which achieves 70 % test accuracy . • ResNet20v1 : for 200 epochs which achieves 90 % test accuracy . We use the Adam optimizer ( Kingma & Ba , 2015 ) for training and to make sure the effects we observe are general , we validate that our results hold for vanilla stochastic gradient descent ( SGD ) as well , which we do not show in this paper . We use batch size 128 and dropout 0.03 for training SmallCNN and MediumCNN . To generate weight space and prediction space similarity results , we use a constant learning rate of 1.6× 10−3 , unless specified otherwise . We do not use any data augmentation with those two architectures . For ResNet20v1 , we use the data augmentation and learning rate schedule used in Keras examples2 . The overall trends are consistent across all architectures , datasets , and other hyperparameter and non-linearity choices we explored . 3.1 SIMILARITY OF FUNCTIONS WITHIN AND ACROSS TRAJECTORIES . First , we compute the similarity between different checkpoints along a single trajectory . We plot the cosine similarity in weight space in Figure 2 ( a ) and the disagreement in function space , defined as the fraction of points the checkpoints disagree on , in Figure 2 ( b ) . We observe that the checkpoints along a trajectory are largely similar both in the weight space and the function space . Next , we evaluate how diverse the final solutions from different random initializations are . The functions from different initialization are different , as demonstrated by the similarity plots in Figure 3 . Comparing this with Figures 2 ( a ) and 2 ( b ) , we see that functions within a single trajectory exhibit higher similarity and functions across different trajectories exhibit much lower similarity . Next , we take the predictions from different checkpoints along the individual training trajectories from multiple initializations and compute a t-SNE plot ( Maaten & Hinton , 2008 ) to visualize their similarity in function space . More precisely , we take the softmax output for a set of points , flatten the vector and use it as the input to the t-SNE plot . Figure 2 ( c ) shows that the functions explored by different trajectories ( denoted by circles with different colors ) are far away , while functions explored within a single trajectory ( circles with the same color ) tend to be much more similar . 3.2 SIMILARITY OF FUNCTIONS ACROSS SUBSPACES FROM EACH TRAJECTORY . In addition to the checkpoints along a trajectory , we also construct subspaces based on each individual trajectory . Scalable Bayesian methods typically compute statistics based on the weights along a trajectory , hence visualizing the diversity of functions between the subspace helps understand the difference between Bayesian neural networks and ensembles . We use a representative set of four subspace sampling methods : a random subspace , a Monte Carlo dropout , a diagonal Gaussian approximation , and a low-rank covariance matrix Gaussian approximation . In the descriptions of the methods , let ~w0 be the current weight-space position ( the weights and biases of our trained neural net ) around which we will construct the subspace . 2https : //keras.io/examples/cifar10_resnet/ • Random subspace sampling : We start at an optimized solution ~w0 and choose a random direction v̂ in the weight space . We step in that direction by choosing different values of t and looking at predictions at configurations ~w0 + tv̂ . We do this for many random directions v̂ . • Monte Carlo dropout subspace : We start at an optimized solution ~w0 and apply dropout with a randomly chosen pkeep to it . We do this many times , each time choosing a random pkeep , and look at predictions at dropoutpkeep ( ~w0 ) . • Diagonal Gaussian subspace : We start at an optimized solution ~w0 and look at the most recent iterations of training proceeding it . For each trainable parameter wi , we calculate its mean meani and standard deviation std ( wi ) . To sample solutions from the subspace , we draw each parameter independently as wi ∼ N ( meani , stdi ) . We repeat this many times and obtain predictions in each . This corresponds to sampling from a normal distribution with a diagonal covariance matrix . • Low-rank Gaussian subspace : We start at an optimized solution ~w0 and look at the most recent iterations of training proceeding it . For each trainable parameter wi , we calculate its mean meani . For a rank-k approximation , we calculate top k principal components of the weight vectors in the most recent iterations of training { ~pi ∈ Rparams } k. We sample from a k-dimensional normal distribution and obtain the weight configurations as ~w ∼ ~mean + ∑ iN ( 0k , 1k ) ~pi . Figure 4 shows that functions sampled from a subspace ( denoted by colored squares ) corresponding to a particular initialization , are much more similar to each other . While some subspaces are more diverse , they still do not overlap with functions from another randomly initialized trajectory . Diversity versus Accuracy plots To illustrate the difference in another fashion , we sample functions from a single subspace and plot diversity ( as measured by disagreement between predictions ) versus accuracy in Figure 5 . Comparing these subspace points ( colored dots ) to the baseline optima ( green star ) and the optima from different random initializations ( denoted by red stars ) , we observe that random initializations are much more effective at sampling diverse and accurate solutions , than subspace based methods constructed out of a single trajectory . The diversity score used above quantifies the difference of two functions , by measuring fraction of points on which their predictions differ . We chose this approach due to its simplicity ; one could also compute the KL-divergence or other distances between the output probability distributions Let ddiff denote the fraction of predictions on which the two functions differ . It is 0 when the two functions make identical class predictions , and 1 when they differ on every single example . To account for the fact that the lower the accuracy of a function , the higher its potential ddiff due to the possibility of the wrong answers being random and uncorrelated between the two functions , we normalize this by ( 1− a ) , where a is the accuracy . For a reference function f∗ of accuracy a∗ and a function f of accuracy a whose predictions are obtained by randomly perturbing the predictions of f∗ , the expected fractional difference is ddiff = ( C − 1 ) ( a∗ − a ) / ( a∗C − 1 ) , where C is the number of classes . If the function f of accuracy a were entirely independent of f∗ , then the expected fractional difference would be ddiff = ( 1− a∗ ) a+ ( 1− a ) a∗ + ( 1− a∗ ) ( 1− a ) ( C − 2 ) / ( C − 1 ) . Those two limiting behaviours – the function f being derived from f∗ by a perturbation , and f and f∗ being completely independent – form the two dashed lines in Figure 5 . We refer to Appendix D for further details on the limiting curves . The diversity reached is not as high as the theoretical optimum even for the independently initialized and optimized solutions , which provides scope for future work .

Summary: This paper analyzes ensembling methods in deep learning from the perspective of the loss landscapes. The authors empirically show that popular methods for learning Bayesian neural networks produce samples with limited diversity in the function space compared to modes of the loss found using different random initializations. The paper also considers the low-loss paths connecting independent local optima in the weight-space. The analysis shows that while the values of the loss and accuracy are nearly constant along the paths, the models corresponding to different points on a path define different functions with diverse predictions. The paper also demonstrates the complementary benefits of using subspace sampling/weight averaging in combination with deep ensembles and shows that relative benefits of deep ensembles are higher.