T RAINING INDEPENDENT SUBNETWORKS FOR ROBUST PREDICTION
Brain Team balajiln@
Google Health adai@google.com Dustin Tran Google Research Brain Team trandustin@google.com
A BSTRACT
Recent approaches to efﬁciently ensemble neural networks have shown that strong robustness and uncertainty performance can be achieved with a negligible gain in parameters over the original network.
However, these methods still require multiple forward passes for prediction, leading to a signiﬁcant computational cost.
In this work, we show a surprising result: the beneﬁts of using multiple predictions can be achieved ‘for free’ under a single model’s forward pass.
In particular, we show that, using a multi-input multi-output (MIMO) conﬁguration, one can utilize a single model’s capacity to train multiple subnetworks that independently learn the task at hand.
By ensembling the predictions made by the subnetworks, we improve model robustness without increasing compute.
We observe a signiﬁcant improvement in negative log-likelihood, accuracy, and calibration error on CIFAR10, CIFAR100,
ImageNet, and their out-of-distribution variants compared to previous methods.
I NTRODUCTION
Uncertainty estimation and out-of-distribution robustness are critical problems in machine learning.
In medical applications, a conﬁdent misprediction may be a misdiagnosis that is not referred to a physician as during decision-making with a “human-in-the-loop.”
This can have disastrous consequences, and the problem is particularly challenging as patient data deviates signiﬁcantly from the training set such as in demographics (Dusenberry et al., 2020b; Filos et al., 2019).
Using a distribution over neural networks is a popular solution stemming from classic Bayesian and ensemble learning literature (Hansen & Salamon, 1990; Neal, 1996), and recent advances such as BatchEnsemble and extensions thereof achieve strong uncertainty and robustness performance (Wen et al., 2020;
Dusenberry et al., 2020a; Wenzel et al., 2020).
These methods demonstrate that signiﬁcant gains can be had with negligible additional parameters compared to the original model.
However, these methods still require multiple (typically, 4-10) forward passes for prediction, leading to a signiﬁcant runtime cost.
In this work, we show a surprising result: the beneﬁts of using multiple predictions can be achieved “for free” under a single model’s forward pass.
The insight we build on comes from sparsity.
Neural networks are heavily overparameterized models.
The lottery ticket hypothesis (Frankle & Carbin, 2018) and other works on model pruning (Molchanov y 1 : p(Dog)=91% y 2 : p(Cat)=95% y 3 : p(Firetruck)=71% x 1 : x 2 : x 3 : (a) Training y 1 : p(Cat)=97% y 2 : p(Cat)=5% y 3 : p(Cat)=81% x 1 : x 2 : x 3 : Avg.y: p(Cat)=61% (b) Testing et al., 2016; Zhu & Gupta, 2017) show that one can prune away 70-80% of the connections in a neural network without adversely affecting performance.
The remaining sparse subnetwork, called the winning ticket, retains its predictive accuracy.
This suggests that a neural network has sufﬁcient capacity to ﬁt 3-4 independent subnetworks simultaneously.
We show that, using a multi-input multioutput (MIMO) conﬁguration, we can concurrently train multiple independent subnetworks within one network .
These subnetworks co-habit the network without explicit separation.
The advantage of doing this is that at test time, we can evaluate all of the subnetworks at the same time, leveraging the beneﬁts of ensembles in a single forward pass .
Our proposed MIMO conﬁguration only requires two changes to a neural network architecture.
First, replace the input layer: instead of taking a single datapoint as input, take M datapoints as inputs, where M is the desired number of ensemble members.
Second, replace the output layer: instead of a single head, use M heads that make M predictions based on the last hidden layer.
During training, the inputs are sampled independently from the training set and each of the M heads is trained to predict its matching input (Figure 1a).
Since, the features derived from the other inputs are not useful for predicting the matching input, the heads learn to ignore the other inputs and make their predictions independently.
At test time, the same input is repeated M times.
That is, the heads make
M independent predictions on the same input, forming an ensemble for a single robust prediction that can be computed in a single forward pass (Figure 1b).
The core component of an ensemble’s robustness such as in Deep Ensembles is the diversity of its ensemble members (Fort et al., 2019).
While it is possible that a single network makes a conﬁdent misprediction, it is less likely that multiple independently trained networks make the same mistake.
Our model operates on the same principle.
By realizing multiple independent winning lottery tickets, we are reducing the impact of one of them making a conﬁdent misprediction.
For this method to be effective, it is essential that the subnetworks make independent predictions.
We empirically show that the subnetworks use disjoint parts of the network and that the functions they represent have the same diversity as the diversity between independently trained neural networks.
Summary of contributions.
We propose a multi-input multi-output (MIMO) conﬁguration to network architectures, enabling multiple independent predictions in a single forward pass “for free.”
Ensembling these predictions signiﬁcantly improves uncertainty estimation and robustness with minor changes to the number of parameters and compute cost.
We analyze the diversity of the individual members and show that they are as diverse as independently trained neural networks.
We demonstrate that when adjusting for wall-clock time, MIMO networks achieve new state-ofthe-art on CIFAR10, CIFAR100, ImageNet, and their out-of-distribution variants.
M ULTI - INPUT MULTI - OUTPUT NETWORKS
The MIMO model is applicable in a supervised classiﬁcation or regression setting.
Denote the set of training examples
X = { ( x ( n ) , y ( n ) ) }
Nn =1 where x ( n ) is the n th datapoint with the corresponding label y ( n ) and N is the size of the training set.
In the usual setting, for an input x , the output [Figure] of the neural network ˆ y is a probability distribution p θ (ˆ y | x ) , 1 which captures the uncertainty in the predictions of the network.
The network parameters θ are trained using stochastic gradient descent (SGD) to minimize the loss L ( θ ) on the training set, where the loss usually includes the negative log-likelihood and a regularization term R (such as the L2 regularization):
L ( θ ) =
E ( x , y ) ∈
X [ − log p θ ( y | x )]
+ R ( θ ) .
In the MIMO conﬁguration, the network takes M inputs and returns M outputs (Figure 1), where each output is a prediction for the corresponding input.
This requires two small changes to the architecture.
At the input layer, the M inputs { x 1 , . . .
, x M } are concatenated before the ﬁrst hidden layer is applied and at the output layer, the network gives M predictive distributions { p θ ( y 1 | x 1 , . . .
, x M ) , . .
, p θ ( y M | x 1 , . . .
, x M ) } correspondingly.
Having M input and M outputs require additional model parameters.
The additional weights used in the MIMO conﬁguration account for just a 0.03 % increase in the total number of parameters and 0.01 % increase in ﬂoating-point operations (FLOPs).
The network is trained similarly to a traditional neural network, with a few key modiﬁcations to account for the M inputs and M outputs (Figure 1a).
During training, the inputs x 1 , . . .
, x M are sampled independently from the training set.
The loss is the sum of the negative log-likelihoods of the predictions and the regularization term:
[Equation]
which is optimized using stochastic gradient descent.
Note that the sum of the log-likelihoods is equal to the log-likelihood of the joint distribution (cid:80)
Mm =1 log p θ ( y m | x 1 , . . .
, x M )
= log p θ ( y 1 , . . .
, y M | x 1 , . . .
, x M ) since the input-output pairs are independent.
Hence a second interpretation of MIMO is that it is simply training a traditional neural network over M -tuples of independently sampled datapoints.
At evaluation time, the network is used to make a prediction on a previously unseen input x (cid:48) .
The input x (cid:48) is tiled M times, so x 1 = . . .
=
x M = x (cid:48)
(Figure 1b).
Since all of the inputs are x (cid:48) , each of the outputs independently approximate the predictive distribution p θ ( y m | x (cid:48) , . . .
, x (cid:48) )
≈ p ( y (cid:48)
| x (cid:48) ) (for m = 1 . . .
M ).
As an ensemble, averaging these predictions improves the predictive performance, leading to the combined output p θ ( y (cid:48) | x (cid:48) )
= 1 M (cid:80) Mm =1 p θ ( y m | x (cid:48) , . . .
, x (cid:48) ) .
Unlike Bayesian methods requiring multiple weight samples, or even parameter-efﬁcient methods like BatchEnsemble, MIMO’s advantage is that all of the ensemble members can be calculated in a single forward pass.
As a result, MIMO’s wall-clock time is almost equivalent to a standard neural network. [Figure] [Figure] [Figure]
2. I LLUSTRATION OF MIMO ON A SYNTHETIC REGRESSION EXAMPLE
Before applying MIMO to large-scale vision models, we ﬁrst illustrate its behavior on a simple one-dimensional regression problem.
We consider the noisy function from Blundell et al.
(2015) (see
Figure 2, left), with a training and test set of N = 64 and 3000 observations respectively.
We train a multilayer perceptron with two hidden-layers, composed of (32, 128) units and ReLU activations.
For different ensemble sizes M ∈ { 1 , . . . , 5 } , we evaluate the resulting models in terms of expected mean squared error E M .
If we denote by ˆ f M the regressor with M ensemble members learned over X , we recall that E M = E ( x (cid:48) ,y (cid:48) )
∈
X test [ E X
[( ˆ f M ( x (cid:48) , . . .
, x (cid:48) ) − y (cid:48) ) 2 ]] , where E X [ · ] denotes the expectation over training sets of size N .
We make two main observations.
First, in the example of M = 3 in Figure 2 (left), we can see that MIMO can learn a diverse set of predictors.
Second, the diverse predictors obtained by MIMO translate into improved performance, as seen in Figure 2 (right).
Moreover, in the regression setting, we can decompose E M into its (squared) bias and variance components (Sec. 2.5 in Hastie et al.
More formally, we have
E M
[Equation]
where ¯ f M = E X [ ˆ f M ] .
The bias-variance decomposition nicely captures the strength of MIMO.
While learning a neural network over M -tuples induces a slight bias compared to a standard model with M = 1 , i.e. the individual members perform slightly worse (Figure 2, middle-left), this is compensated by the diverse predictions of the ensemble members that lead to lower variance (Figure 2, middle-right).
MIMO yields an improvement when the model has sufﬁcient capacity to ﬁt M >diverse, well-performing ensemble members.
10.561G FLOPs to evaluate.
We provide an interactive notebook that reproduces these experiments.
[Figure]
[Figure]
U NDERSTANDING THE SUBNETWORKS
The mapping of each input-output pair in a MIMO conﬁguration is referred to as a subnetwork.
In this section, we show that the subnetworks converge to distinct local optima and they functionally behave as independently trained neural networks.
3. L OSS - LANDSCAPE ANALYSIS
Using multiple inputs is the key to training diverse subnetworks.
The subnetworks learn independently, since features derived from each input are only useful for the corresponding output.
In contrast, in a naive multiheaded architecture, where the input is shared, but the model has separate M outputs, the outputs rely on the same features for prediction, which leads to very low diversity.
To showcase this, we replicate a study from Fort et al.
We look at the SmallCNN model ( convolutional layers with 16, 32, 32 ﬁlters respectively) trained on CIFAR-10, and linearly interpolate between the three subnetworks in weight space.
In the case of MIMO, we interpolate the input and output layers, since the body of the network is shared, and for the naive multiheaded model, we only interpolate the output layers, since the input and the body of the network is shared.
Analogously to
Deep Ensembles, the subnetworks trained using MIMO converge to disconnected modes in weight space due to differences in initialization, while in the case of the naive multiheaded model, the subnetworks end up in the same mode (Figure 3, left).
Figure 3 (right) shows the disagreement i.e. the probability that the subnetworks disagree on the predicted class.
MIMO’s disconnected modes yield diverse predictions, while the predictions of the naive multiheaded model are highly correlated.
3. F UNCTION SPACE ANALYSIS
We visualize the training trajectories of the subnetworks in MIMO in function space (similarly to Fort et al.
As we train the SmallCNN architecture ( M = 3 ), we periodically save the predictions on the test set.
Once training is ﬁnished, we plot the t-SNE projection (Maaten & Hinton, 2008) of the predictions.
We observe that the trajectories converge to distinct local optima (Figure 4).
For a quantitative measurement of diversity in large scale networks, we look at the average pairwise similarity of the subnetwork’s predictions at test time, and compare against other efﬁcient ensemble methods.
The average pairwise similarity is
D = E [ D ( p θ ( y 1 | x , x . . . x ) , p θ ( y 2 | x , x . . . . . .
x ))] , where D is a distance metric between predictive distributions and ( x , y ) ∈ X .
We consider two distance metrics.
Disagreement: whether the predicted classes agree, D disagreement ( P 1 , P 2 ) =
I(arg max ˆ y P 1 (ˆ y ) =
arg max ˆ y P 2 (ˆ y )) and Kullback–Leibler divergence: D KL ( P 1 , P 2 ) =
E P 1 [log P 1 ( y ) − log P 2 ( y )] .
When the ensemble members give the same prediction at all test points, both their disagreement and KL divergence are 0.
The ﬁrst efﬁcient ensemble approach we compare against is the aforementioned naive multiheaded model, where the input and the body of the neural network is shared by the ensemble members, but each member has its own output layer.
Next, TreeNet (Lee et al., 2015), where the input and the ﬁrst two residual groups are shared, but the ﬁnal residual group and the output layer are trained separately for each member.
Finally, BatchEnsemble (Wen et al., 2020), where the members share network parameters up to a rank-1 perturbation, which changes information ﬂow through the full network.
[Figure] We ﬁnd that the naive multiheaded model fails to induce diversity: the predictions of its subnetworks are nearly identical on all test points as shown in Figure 4 (right).
TreeNet and BatchEnsemble have more diversity, although there is still signiﬁcant correlation in the predictions.
MIMO has better diversity than prior efﬁcient ensemble approaches and it matches the diversity of independently trained neural networks.
These results allow us to pinpoint the source of robustness: The robustness of MIMO comes from ensembling the diverse predictions made by the subnetworks, thus MIMO faithfully replicates the behaviour of a Deep Ensemble within one network.
3. S EPARATION OF THE SUBNETWORKS
To show that the subnetworks utilize separate parts of the network, we look at the activations and measure how they react to changes in each of the M inputs.
Namely, we calculate the conditional variance of each pre-activation in the network with respect to each individual input.
For input x 1 :
Var( a i | x 2 ) =
E x 2 [Var x 1 ( a i ( x 1 , x 2 )]
( M = 2) , and Var( a i | x 2 , x 3 ) =
E x 2 , x 3 [Var x 1 ( a i ( x 1 , x 2 , x 3 )] ( M = 3) where a i is the value of i -th pre-activation in the function of the M inputs.
For reference, there are 8190 pre-activations in a ResNet28-10.
We can estimate the conditional variance by ﬁxing xand calculating the variance of a i w.r.t.
x 2 ∈ X (when M = 2 , x 2 , x 3 ∈ X when M = 3 ) and ﬁnally averaging over the possible values x 1 ∈ X .
The conditional variance is analogously deﬁned for x 2 , . . .
, x M .
If the conditional variance of an activation is non-zero w.r.t.
an input, that means that the activation changes as the input changes and therefore we consider it part of the subnetwork corresponding to the input.
If the subnetworks are independent within the network, we expect that the conditional variance of each activation is non-zero w.r.t.
one of the inputs, the subnetwork to which it belongs, and close-to zero w.r.t.
all the other inputs.
When we plot the conditional variances, this is exactly what we see.
In Figure 4 (left), we see the histogram of the pre-activations in the network.
Each point has two corresponding values: the conditional variance w.r.t.
the two inputs.
As we can see, all activations have non-zero conditional variance w.r.t.
one of the inputs and close-to zero w.r.t.
the other.
Figure 4 (middle-left) shows the scatterplot of the activations for M = 3 .
We see that, similarly to M = 2 , almost all activations have non-zero conditional variance w.r.t.
one of the inputs and close-to zero conditional variance w.r.t.
the others.
Since almost all activations are part of exactly one of the subnetworks, which we can identify by calculating the conditional variances, we conclude that the subnetworks separate within the network.
This implies an extension to Frankle & Carbin (2018): the subnetworks realize separate winning lottery tickets within a single network instance.
3. T HE OPTIMAL NUMBER OF SUBNETWORKS
A natural question that arises is the ideal number of subnetworks M to ﬁt in a network.
Too few subnetworks do not fully leverage the beneﬁts of ensembling, while having too many quickly reaches the network capacity, hurting their individual performances.
Ideally, we are looking to ﬁt as many subnetworks as possible without signiﬁcantly impacting their individual performances.
Figure 5 shows the performance of both the individual subnetworks and the ensemble as M varies.
M = 1 is equivalent to a traditionally trained network i.e. the performance of the subnetwork matches the performance of the ensemble since there is only one subnetwork.
As M grows, we can see that the performance of the subnetworks slowly declines as they utilize more and more of the network Test Accuracy (%) ResNet50Wide ResNet 0.0 0.2 0.4 0.6 0.8 1. Test Log-likelihood
ResNet50Wide ResNet (
[Figure] Input repetition   6 # batch repetitions 94. Test Accuracy (%)    6 # batch repetitions 0.
Test Log-likelihood
[Figure]
capacity.
The performance of the ensemble, however, peaks between M = 2 and M = 4 , where the beneﬁts of ensembling outweigh the slight decrease in performance of the individual subnetworks.
Interestingly, the accuracy peaks earlier than the log-likelihood, which suggests that ensembling is more beneﬁcial for the latter.
In Appendix C, we further illustrate how MIMO exploits the capacity of the network.
In particular, we study the performance of MIMO when the regularization increases (both in terms of L1 and L regularization), i.e. when the capacity of the network is increasingly constrained.
In agreement with our hypothesis that MIMO better utilizes capacity, we observe that its performance degrades more quickly as the regularization intensiﬁes.
Moreover, the larger M , the stronger the effect.
Interestingly, for the L1 regularization, we can relate the performance of MIMO with the sparsity of the network, strengthening the connection to Frankle & Carbin (2018).
3. I NPUT AND BATCH REPETITION
MIMO works well by simply adding the multi-input and multi-output conﬁguration to an existing baseline, and varying only one additional hyperparameter (Section 3.4’s number of subnetworks
M ).
We found two additional hyperparameters can further improve performance, and they can be important when the network capacity is limited.
Input repetition Selecting independent examples for the multi-input conﬁguration during training forces the subnetworks not to share any features.
This is beneﬁcial when the network has sufﬁcient capacity, but when the network has limited excess capacity, we found that relaxing independence is beneﬁcial.
For example, ResNet50 on ImageNet (He et al., 2016) does not have sufﬁcient capacity to support two independent subnetworks ( M = 2 ) in MIMO conﬁguration.
Our proposed solution is to relax independence between the inputs.
Instead of independently sampling x 1 and x 2 from the training set during training, they share the same value with probability ρ .
That is, x 1 is sampled from the training set and x 2 is set to be equal to x 1 with probability ρ or sampled from the training set with probability 1 − ρ .
Note that this does not affect the marginal distributions of xand x 2 , it merely introduces a correlation in their joint distribution.
Figure 6a shows the performance as ρ varies.
At ρ = 0 , the subnetworks are independent but their performance is limited by the network capacity.
As ρ grows, the subnetworks share increasingly more features, which improves their performance.
However, as ρ approaches 1 , the subnetworks become highly correlated and the beneﬁt of ensembling is lost.
Unlike ResNet50, Wide ResNet50 has more capacity and beneﬁts less from input repetition (roughly 78-79% vs 74-77%).
Batch repetition For stochastic models, most notably MC dropout and variational Bayesian neural nets, drawing multiple approximate posterior samples for each example during training can improve performance as it reduces gradient noise w.r.t.
the network’s model uncertainty, e.g., Dusenberry et al.
(2020a).
We achieve a similar effect by repeating examples in the minibatch: this forms a new minibatch size of, e.g., 512 · 5 (batch size and number of batch repetitions respectively).
Like the choice of batch size which determines the number of unique examples in the SGD step (Shallue et al., 2018), varying the number of repetitions has an implicit regularization effect.
Figure 6b shows performance over the number of batch repetitions, where each batch repetition setting indicates a box plot over a sweep of 12 settings of batch size, learning rate, and ensemble size.
Higher repetitions typically yield a slight boost.
B ENCHMARKS
We described and analyzed MIMO.
In this section, we compare MIMO on benchmarks building on
Uncertainty Baselines.
This framework allows us to benchmark the performance and to compare against high-quality, well-optimized implementations of baseline methods (see framework for further baselines than ones highlighted here).
We looked at three model/dataset combinations: ResNet2810/CIFAR10, ResNet28-10/CIFAR100, and ResNet50/ImageNet.
MIMO’s code is open-sourced.
Baselines Our baselines include the reference implementations of a deterministic deep neural network (trained with SGD), MC-Dropout, BatchEnsemble, and ensemble models, as well as two related models, Naive multihead and TreeNet.
Thin networks use half the number of convolutional ﬁlters while wide models use double.
See Appendix B for the details on the hyperparameters.
Metrics To measure robustness, we look at accuracy, negative log-likelihood (NLL), and expected calibration error (ECE) on the IID test set as well as a corrupted test set where the test images are perturbed (e.g. added blur, compression artifacts, frost effects) (Hendrycks & Dietterich, 2019).
Appendix D includes ImageNet results for 5 additional out-of-distribution datasets.
To measure computational cost, we look at how long it takes to evaluate the model on a TPUv2 core, measured in ms per example.
[Table] [Table]
[Table]
The metrics show that MIMO signiﬁcantly outperforms other single forward pass methods on all three benchmarks.
It approaches the robustness of a Deep Ensemble, without increasing the computational costs.
R ELATED WORK
Multi-headed networks have been previously studied by Lee et al.
; Osband et al.
Tran et al.
In this approach to efﬁcient ensembles, the input and part of the network are shared by the members while the ﬁnal few layers and the outputs are separate.
The advantage of the approach is that the computational cost is reduced compared to typical ensembles, since many layers are shared, but the ensemble diversity (and resulting performance) is quite lacking.
MIMO’s multi-input conﬁguration makes a signiﬁcant impact as each ensemble member may take different paths throughout the full network.
Further, MIMO has lower computational cost than multi-headed approaches, since all of the layers except the ﬁrst and last are shared.
Related efﬁcient ensemble approaches include BatchEnsemble, Rank-1 BNNs, and hyper batch ensembles (Wen et al., 2020; Dusenberry et al., 2020a;
Wenzel et al., 2020).
In these methods, most of the model parameters are shared among the members, which reduces the memory requirement of the model, but the evaluation cost still requires multiple forward passes.
Interestingly, like MIMO, these methods also apply a multi-input multi-output conﬁguration, treating an ensemble of networks as a single bigger network; however, MIMO still outperforms BatchEnsemble.
We believe important insights such as Section 3.5’s input independence may also improve these methods.
Finally, there are simple heuristics which also retain efﬁcient compute such as data augmentation, temperature scaling, label smoothing, contrastive training.
These methods are orthogonal to MIMO and they can provide a performance boost, without increasing the computational cost.
C ONCLUSIONS
We propose MIMO, a novel approach for training multiple independent subnetworks within a network.
We show that the subnetworks separate within the model and that they behave as independently trained neural networks.
The key beneﬁts of MIMO are its simplicity, since it does not require signiﬁcant modiﬁcations to the network architecture and it has few hyperparameters, and also its computational efﬁciency, since it can be evaluated in a single forward pass.
Our empirical results conﬁrm that MIMO improves performance and robustness with minor changes to the number of parameters and the compute cost. A CKNOWLEDGEMENTS
Marton Havasi is funded by EPSRC.
We thank Ellen Jiang for helping with writing.
A P SEUDOCODE
Algorithm
1: for t = 1 . . .
N iter do 2: ( x 1: M , y 1: M ) ∼ U ( X ) 3: p θ ( y 1 | x 1: M ) . . .
p θ ( y M | x 1: M )
← MIMO( x 1: M ) 4:
L M ( θ ) ← (cid:80) Mm =1 − log p θ ( y m | x 1: M )
+ R ( θ ) 5: θ ← θ − (cid:15) ∇ L M ( θ ) (cid:46) (cid:15) is the learning rate.
6: end for
Algorithm 2 Evaluate( x (cid:48) ) 1: p θ ( y 1 | x 1: M = x (cid:48) ) . . .
p θ ( y M | x 1: M = x (cid:48) )
← MIMO( x 1: M = x (cid:48) ) 2: return 1 M (cid:80) Mm =1 p θ ( y m | x 1: M = x (cid:48) )
B H YPERPARAMETERS
For the ResNet28-10/CIFAR models, we use a batch-size of 512, a decaying learning rate of 0.
(decay rate 0.1) and L2 regularization 2e-4.
The Deterministic, Dropout and Ensemble models are trained for 200 epochs while BatchEnsemble, Naive multihead and TreeNet are trained forepochs.
For the ResNet50/ImageNet models, we use a batch-size of 4096 and a decaying learning rate of 0.
(decay rate 0.1) and L2 regularization 1e-4.
The Deterministic, Dropout and Ensemble models are trained for 90 epochs, the BatchEnsemble model is trained for 135 epochs and Naive multihead and
TreeNet are trained for 150 epochs.
Regarding model speciﬁc hyperparameters, Dropout uses a 10% dropout rate and a single forward pass at evaluation time.
Both Ensemble and BatchEnsemble models use M = 4 members, since this provides most of the beneﬁts of ensembling without signiﬁcantly increasing the computational costs.
The TreeNet architecture
MIMO For MIMO, we use the hyperparameters of the baseline implementations wherever possible.
For the ResNet28-10/CIFAR models, we use a batch-size of 512 with decaying learning rate of 0.1 (decay rate 0.1), L2 regularization 3e-4, 250 training epochs, and a batch repetition of 4.
For the ResNet50/ImageNet models, we use a batch-size of 4096 with decaying learning rate of 0.
(decay rate 0.1), L2 regularization 1e-4, 150 training epochs, and batch repetition of 2.
This makes the training cost of MIMO comparable to that of BatchEnsemble and Ensemble models.
For the
ResNet28-10/CIFAR experiments, we use M = 3 subnetworks because it performs well in both accuracy and log-likelihood.
For ResNet50/
ImageNet the model has lower capacity so we used
M = 2 with ρ = 0 . 6 .
C MIMO BETTER EXPLOITS THE NETWORK CAPACITY : PERFORMANCE VERSUS REGULARIZATION STRENGTH
In this section, we further illustrate how MIMO better exploits the capacity of the network, through the lens of its sensitivity to regularization.
Our experimental protocol is guided by the following rationale
• Regularization controls the capacity of the network: the higher the regularization, the more constrained the capacity.
• MIMO makes better use of the capacity of the network: the more ensemble members, the more capacity is exploited.
• As a result, MIMO should be more sensitive to the constraining of the capacity of the network.
And the more ensemble members (i
e., the larger M ), the stronger the effect should be.
We consider the ResNet28-10/CIFAR10 and ResNet28-10/CIFAR100 settings used in the main paper where we additionally vary the L1 (respectively L2) regularization while keeping the other L
[Figure]
(respectively L1) term equal to zero.
We display in Figures 7-8 the accuracy and log-likelihood over those regularization paths (averaged over three repetitions of the experiments)
[Figure]
As previously hypothesised, we observe that MIMO is indeed more sensitive to the constraining of the capacity of the network as the regularization increases.
Moreover, the larger the ensemble size, the stronger the effect.
In the case of the L1 regularization, we can show how the accuracy and log-likelihood evolve with respect to the sparsity of the network 6 .
We report those results in Figurewhere we can observe the same phenomenon.
D A DDITIONAL I MAGE N ET OOD R ESULTS
In the following table, we evaluate trained ResNet-50 models on 7 datasets.
ImageNet, ImageNet-C,
ImageNet-A, and ImageNetV2 each display three metrics: negative log-likelihood, accuracy, and expected calibration error respectively.
ImageNet-C further includes mCE (mean corruption error) in parentheses.
ImageNet-Vid-Robust, YTTB-Robust, and ObjectNet use their own pre-deﬁned stability metrics.
These experiments were expensive to run, so we were only able to obtain them for a smaller set of methods.
We ﬁnd these results are consistent with the main text’s benchmarks, showing MIMO consistently outperforms methods not only on corrupted images, but also across distribution shifts.
[Table]
[Figure] [Figure]