S PARSE M O E S MEET E FFICIENT E NSEMBLES
Anonymous authors
Paper under double-blind review
A BSTRACT
Machine learning models based on the aggregated outputs of submodels, either at the activation or prediction levels, lead to strong performance.
We study the interplay of two popular classes of such models: ensembles of neural networks and sparse mixture of experts (sparse MoEs).
First, we show that the two approaches have complementary features whose combination is beneﬁcial.
Then, we present partitioned batch ensembles , an efﬁcient ensemble of sparse MoEs that takes the best of both classes of models.
Extensive experiments on ﬁne-tuned vision Transformers demonstrate the accuracy, log-likelihood, few-shot learning, robustness, and uncertainty improvements of our approach over several challenging baselines.
Partitioned batch ensembles not only scale to models with up to 2.7B parameters, but also provide larger performance gains for larger models.
I NTRODUCTION
Neural networks typically use all their parameters to process an input.
Sustaining the growth of such models—reaching today up to 100B parameters (Brown et al., 2020)—is challenging, e.g., due to their high computational and environmental costs (Strubell et al., 2019; Patterson et al., 2021).
In this context, sparse mixtures of experts (sparse MoEs) employ conditional computation (Bengio et al., 2013) to combine multiple submodels (experts) and route examples to certain experts (Shazeer et al., 2017;
Lepikhin et al., 2021; Fedus et al., 2021; Riquelme et al., 2021;
Yang et al., 2021).
Conditional computation can decouple the growth of the number of parameters from the training and inference costs, by only activating a subset of the overall model in an input-dependent fashion.
Paralleling this trend, the deployment of ML systems in safety-critical ﬁelds, e.g., medical diagnosis (Dusenberry et al., 2020b) and self-driving cars (Levinson et al., 2011), has motivated the development of reliable deep learning, e.g., for calibrated and robust predictions (Ovadia et al., 2019).
Among the approaches, ensembles of neural networks have remarkable performance for calibration and accuracy under dataset shifts (Ovadia et al., 2019).
These methods improve reliability by aggregating the predictions of individual submodels (ensemble members).
While sharing conceptual similarities, these two classes of models—MoEs and ensembles—have different properties.
Sparse MoEs adaptively combine their experts depending on the inputs, and the combination generally happens at internal activation levels.
Ensembles typically combine several models in a static way and at the prediction level.
Moreover, these two classes of models tend to be benchmarked on different tasks: few-shot classiﬁcation for MoEs (Riquelme et al., 2021) and uncertainty-related evaluation for ensembles (Ovadia et al., 2019; Gustafsson et al., 2020).
C ONTRIBUTIONS :
In this paper, we study the interplay between sparse MoEs and ensembles.
This results in two sets of contributions.
Contribution 1: Complementarity of MoEs and ensembles.
We show that sparse MoEs and ensembles have complementary features and beneﬁt from each other.
Speciﬁcally:
• The adaptive computation in sparse MoEs and the static combination in ensembles are orthogonal, with additive beneﬁts when associated together.
Their association results in insightful performance versus FLOPs trade-offs while varying the ensemble size and sparsity.
• In sparse MoEs, combining models at the prediction level leads to improved uncertainty estimates.
• Over tasks where either sparse MoEs or ensembles are known to perform well, naive—and computationally expensive—ensembles of MoEs provide the best predictive performance.
Our benchmarking effort includes the ﬁrst evaluation of sparse MoEs on uncertainty-related vision tasks, which builds upon the empirical work of Riquelme et al.
hhhh 3 ,h 2 ,h 1 ,h 1 , 2 h 1 ,h 2 , 2 h 2 ,h 3 , 2 h 3 ,
T ILE gate K ( W 1 h i, 1 ) gate K ( W 2 h i, 2 )
MLP 2 ( h i, 1 )
MLP 1 ( h i, 1 )
MLP 3 ( h i, 1 )
MLP 5 ( h i, 2 )
MLP 4 ( h i, 2 )
MLP 6 ( h i, 2 )
E 1 = { 1 , 2 , 3 }
D ISPATCH ( K =2)
E 2 = { 4 , 5 , 6 } ˆ h 2 ,ˆ h 1 ,ˆ h 3 ,ˆ h 2 , 2 ˆ h 2 ,ˆ h 1 , 2 ˆ h 1 ,ˆ h 3 , 2 ˆ h 3 ,
C OMBINE ¯ h 1 ,¯ h 1 , 2 ¯ h 1 ,¯ h 2 ,¯ h 2 , 2 ¯ h 2 ,¯ h 3 ,¯ h 3 , 2 ¯ h 3 ,
S OFTMAX
E NSEMBLE (M EAN )
V I T B LOCK n − 4 ×
N ORM
MSA +
N ORM
P BE-M O E
+
P -M O E B LOCK
V I T B LOCK
P -M O E B LOCK
V I T B LOCK
C LASSIFIER
[Figure]
P -M O E
[Figure]
Contribution 2: Partitioned batch ensembles.
We propose partitioned batch ensembles (pBE), see Figure 1, an efﬁcient ensemble approach tailored to sparse MoEs.
Speciﬁcally:
• pBE improves over sparse MoEs across metrics including few-shot performance, likelihood and calibration error
pBE matches the performance of deep ensembles for 30%-43% fewer FLOPs.
• pBE gracefully scales up to vision Transformers with up to 2.7B parameters.
• pBE is both simple (requiring only minor implementation changes) and convenient because standard sparse-MoE checkpoints can be used directly to initialize pBEs for ﬁne-tuning.
P RELIMINARIES
We focus on classiﬁcation tasks where we learn classiﬁers of the form f ( x ; θ ) based on some training data D = { ( x n , y n ) } Nn =1 .
A pair ( x n , y n ) corresponds to an input x n ∈ R P together with its label y n ∈ { 1 , . .
, C } belonging to one of the C classes.
The model f ( · ; θ ) is parametrized by θ and outputs a C -dimensional probability vector.
We use ◦ to refer to matrix element-wise product.
2. V ISION T RANSFORMERS AND S PARSE M O E S
Vision Transformers.
Throughout the paper, we choose the model f to be a vision Transformer (ViT)
(Dosovitskiy et al., 2021).
ViT is growing in popularity for vision, especially in transferlearning settings where it was shown to outperform convolutional networks while requiring fewer pre-training resources.
ViT operates at the level of patches.
An input image is split into equal-sized patches (e.g., 32 × 32 , 16 × 16 , or 14 × 14 pixels) whose resulting sequence is (linearly) embedded and processed by a Transformer (Vaswani et al., 2017).
The operations in the Transformer then mostly consist of a succession of multiheaded self-attention (MSA) and MLP layers.
ViT is deﬁned at different scales (Dosovitskiy et al., 2021): S(mall), B(ase), L(arge) and H(uge); see speciﬁcations in Appendix A.
For example, ViT-L/16 stands for a large ViT with patch size 16 × 16 .
Sparse MoEs and V-MoEs.
The main feature of sparsely-gated mixture-of-experts models (sparse
MoEs) lies in the joint use of sparsity and conditional computation (Bengio et al., 2013).
In those models, we only activate a small subset of the network parameters for a given input , which allows the total number of parameters θ to grow while keeping the overall computational cost constant.
The subparts of the network that are activated on a per-input fashion are known as experts .
Central to our study, Riquelme et al.
(2021) recently extended ViT to sparse MoEs.
Their extension, referred to as V-MoE, follows the successful applications of sparse models in NLP (Shazeer et al.,
Riquelme et al.
(2021) show that V-MoEs dominate their “dense” ViT counterparts on a variety of tasks for the same computational cost.
In the speciﬁc case of V-MoEs, the experts are placed in the MLP layers of the Transformer, a design choice reminiscent of Lepikhin et al.
(2021) in NLP.
Given the input h ∈ R D of such a layer, the output of a single MLP ( h ) is replaced by
[Equation]
where the routing weights { g e ( h ) } Ee =1 combine the outputs of the E different experts { MLP e } Ee =1 .
To sparsely select the experts, top K sets all but the K largest weights to zero.
The router parameters W ∈ R E × D are trained together with the rest of the network parameters.
We call the layer deﬁned by (1) an MoE layer.
In practice, the weights { g e ( h ) } Ee =1 are obtained by a noisy version of the routing function top K ( softmax ( W h + σε )) with ε ∼ N ( 0 , I ) , which mitigates the nondifferentiability of top K when combined with auxiliary losses (Shazeer et al., 2017).
We use the shorthand gate K ( z )
=
top K ( softmax ( z + σε )) and take σ = 1 /E (Riquelme et al., 2021).
In this paper, we take the “last- n ” setting of Riquelme et al.
(2021) wherein only a few MoE layers are placed at the end of the Transformer ( n = 2 for the { S, B, L } scale and n = 5 for H).
This setting retains most of the performance gains of V-MoEs while greatly reducing the training cost.
2. E NSEMBLES OF N EURAL N ETWORKS
Ensembles.
We build on the idea of ensembles, which is a known scheme to improve the performance of individual models (Hansen & Salamon, 1990;
Geman et al., 1992; Krogh & Vedelsby, 1995; Opitz & Maclin, 1999; Dietterich, 2000; Lakshminarayanan et al., 2017).
Formally, we assume a set of M model parameters Θ = { θ m } Mm =1 .
We refer to M as the ensemble size .
Prediction proceeds by computing 1 M (cid:80) θ
∈ Θ f ( x ; θ ) , i.e., the average probability vector over the M models.
To assess the diversity of the predictions in the ensemble, we will use the KL divergence
D KL ( f ( x t ; θ m ) (cid:107) f
( x t ; θ m (cid:48) )) between the predictive distributions f ( x t ; θ m ) and f ( x t ; θ m (cid:48) ) , averaged over the test input x t and all pairs ( m, m (cid:48) ) of ensemble members.
Batch ensembles.
Ensembles differ in the way Θ is deﬁned.
Central to our study, batch ensembles (BE) (Wen et al., 2019) build the ensemble as a collection of submodels, with the parameters θ m ∈
Θ sharing components.
This mitigates the computational and memory cost of ensembling, enabling one to improve the performance of the original model at little extra cost.
We focus on the example of a single dense layer in f with parameters U ∈ R D × L , assuming no bias.
BE deﬁnes
M copies of parameters { U m } Mm =1 so that U m =
U ◦ ( r m s (cid:62) m ) , where U are parameters shared across ensemble members, and r m and s m are separate D - and L -dimensional vectors for ensemble member m .
Given an input, the BE produces M outputs, and the M outputs are averaged after applying all layers.
Despite the simple rank-1 parametrization, BE leads to remarkable predictive performance and robustness (Wen et al., 2019).
Notably, the efﬁciency of BE relies on tiling the inputs to simultaneously predict with the M ensemble members, an insight that we also exploit
. 2. U PSTREAM P RE - TRAINING AND D OWNSTREAM F INE - TUNING
Large-scale Transformers pre-trained on upstream tasks were shown to have strong performance when ﬁne-tuned on smaller downstream tasks, across a variety of domains (Devlin et al., 2018;
Dosovitskiy et al., 2021; Radford et al., 2021).
We follow this paradigm and focus on the ﬁnetuning of models pre-trained on JFT-300M (Sun et al., 2017), similar to Riquelme et al.
We will thus assume the availability of already pre-trained ViT and V-MoE model checkpoints.
Our assumption relies on the growing popularity of transfer learning, e.g. Kolesnikov et al.
(2020), and the increasing accessibility of pre-trained models in repositories such as www.tensorflow. org/hub or www.pytorch.org/hub .
The ﬁne-tuning of all the approaches we study here, including extensions of ViT and V-MoE, will be either directly compatible with those checkpoints or require only mild adjustments, e.g., reshaping or introducing new downstream-speciﬁc parameters (see Appendix C).
Also, unless otherwise mentioned, the performance we report will always be downstream, e.g., for ImageNet (Deng et al., 2009) or Cifar10/100 (Krizhevsky, 2009).
In all our comparisons, we will use the downstream training ﬂoating point operations per second (FLOPs), or GFLOPs (i.e., 10 9 × FLOPs), to quantify the computational cost of the different methods.
[Table]
S PARSE M O E S MEET E NSEMBLES
As illustrated in Table 1, sparse MoEs and ensembles have different properties.
For instance, ensembles typically do not use conditional computation and just statically combine members at the prediction level.
This contrasts with sparse MoEs where the different experts are combined at internal activation levels while enjoying per-input adaptivity through the routing logic; see (1).
In terms of cost, sparse MoEs are usually designed to match the cost of their dense counterparts whereas ensembles, in their simplest forms, will typically lead to a substantial overhead.
In this section, we study the extent to which these properties are complementary and may beneﬁt from each other.
In
Section 5, we further evaluate this complementarity on tasks where either sparse MoEs or ensembles are known to perform well, e.g., few-shot and out-of-distribution (OOD) evaluations, respectively.
More details about the experiments in this section can be found in Appendix B.6.
3. S TATIC VERSUS A DAPTIVE C OMBINATION
[Figure]
[Figure]
We ﬁrst focus on the interplay between the (a) static combination in ensembles and (b) the adaptive combination of experts in sparse MoEs.
To this end, we study the performance of downstream deep ensembles (i e., with all ensemble members having the same upstream checkpoint) formed by M independent VMoEs with E experts per MoE layer and a sparsity level K (the larger K , the more selected experts).
The parameter M controls the static combination, while K and E impact the adaptive combination of experts in each sparse MoE model.
We report in Figure 2 the ImageNet performance and compute cost for ensembles with varying choices of K and M , while keeping E = 32 ﬁxed.
We focus on K rather than E as the axis to explore adaptive computation, as we ﬁnd that the performance changes for E plateau relatively quickly (see Figure 8 in the Appendix).
Also, by ﬁxing E = 32 , we match more closely the experimental setup of Riquelme et al.
The architecture of the V-MoE is ViT-S/32; see details in Appendix B.6.1.
We make the following observations:
Cumulative effect.
In the absence of ensembles ( M = 1 ), and given a ﬁxed number of experts, Riquelme et al.
(2021) already reported an increase in performance as K gets larger.
Interestingly, we observe that for each value
[Table]
of K , it is also beneﬁcial to increase the ensemble size M .
In other words, the static combination of ensembles is beneﬁcial when applied to sparse MoEs.
This observation is perhaps surprising since adaptive combination may already encapsulate the effect of static combination.
Figure 3, and
Appendix H.1, show that the combination of static and adaptive ensembling is beneﬁcial to NLL for a range of ViT families.
We also see that the beneﬁts of static ensembling are similar for V-MoE and ViT (which does not have any adaptive ensembling).
Taking FLOPs into account.
Without any computational constraints, the previous observation would favor approaches with the largest values of K and M .
However, different values of ( K, M ) lead to different computational costs, as measured here by FLOPs, with ( K, M ) =
(1 , 1) being the cheapest.
Figure 2b shows, as expected, that the number of FLOPs grows more quickly along the
M axis than along the K axis.
To capture the various trade-offs at play, in Figure 2c we report the logarithm of the normalized gains in log likelihood
LL ( K,M ) − LL (1 , 1) GFLOPs ( K,M ) − GFLOPs (1 , 1) when going from ( K, M ) =
(1 , 1) to other choices of ( K, M ) .
Interestingly, it appears more advantageous to ﬁrst grow K , i.e., the adaptive combination, before growing M .
3. F EATURE - LEVEL VERSUS P REDICTION - LEVEL E NSEMBLING
As highlighted in Table 1, an ensemble of size M outputs M predictions for a given input (thereafter, averaged) while sparse MoEs only produce a single prediction.
We study the impact of this differentiating property.
To this end, we propose a simple variant of sparse MoEs wherein the last
MoE layer of the form (1) is replaced by
[Equation]
where we have assumed MLP e ( h ) ∈ R Q .
Instead of summing the expert outputs like in (1), we stack the K selected expert contributions (as a reminder, gate K zeroes out the E − K smallest weights).
Keeping track of those K contributions makes it possible to generate K predictions per input as in the classiﬁer of Figure 1, thus capturing model uncertainty around the true prediction.
Table 2 compares the ImageNet performance—negative log likelihood (NLL), classiﬁcation error and expected calibration error (ECE) (Guo et al., 2017)—of
this naive multi-head method with the standard V-MoE. For K = 2 , the multi-head method provides small but statistically signiﬁcant gains in NLL and ECE.
However, its classiﬁcation error is worse.
On the other hand, for K = 4 , both the NLL and classiﬁcation error for multi-head are worse than V-MoE, despite an even larger improvement in ECE.
In fact, the multi-head for K = 4 performs worse in terms of NLL, classiﬁcation error, and diversity than for K = 2 .
Note that the KL diversity metric indicates that the multi-head variant is unable to provide diverse predictions (e.g., a downstream ensemble of two VMoEs with K = 1 provides a much higher diversity of 0.073, see Table 10).
Following Havasi et al.
; Soﬂaei et al.
(2020), a possible ﬁx to this problem would be to consider a multi-head and multi-input approach, but we show in Appendix F that this strategy is not effective in our context.
Thus, while prediction level ensembling can be beneﬁcial for uncertainty calibration of MoEs, a different strategy is required such that the error is not worse.
We propose a better approach next.
P ARTITIONED B ATCH E NSEMBLE
Equipped with the insights from Section 3, we describe partitioned batch ensemble (pBE), with the goal of keeping the strengths of both sparse MoEs and ensembles.
Conceptually, we can view pBE as jointly learning an ensemble of smaller sparse MoEs, where all the layers that do not contain experts are shared across the members, e.g., the self-attention layers.
As its name indicates, pBE is inspired by batch ensemble in that (a) ensemble members have shared parameters—here, the sharing is tailored to sparse MoEs—and (b) we reuse the idea of tiled representations.
4. T HE A RCHITECTURE
There are two main components in PBE:
Disjoint subsets of experts as ensemble members.
We change the structure of (1) by partitioning the set of E experts into M sets of E/M experts (we assume that E is a multiple of M ).
We denote this partition by ∪ Mm =1 E m , for example E 1 = { 1 , 2 , 3 } and E 2 = { 4 , 5 , 6 } for E = 6 and M = 2 .
The M sets of E/M experts play the role of the M ensemble members.
Intuitively, the ensemble members have separate parameters for independent predictions, while efﬁciently sharing parameters among all non-expert layers.
Instead of having a single routing function gate K ( W · ) like in (1), we apply separate routing functions { gate K ( W m · ) } Mm =1 to each member of the partition.
Note that this does not affect the total number of parameters since W has E rows while each W m has E/M rows.
A similar partitioning of the experts was proposed in Yang et al.
but not exploited to create different ensemble members, in particular not in conjunction with tiled representations, which we show to be required to get performance gains (see comparison in Section 4.2.1).
Tiled representation.
To jointly handle the predictions of the M ensemble members, we tile the inputs by a factor M , as proposed in Wen et al.
Tiling naturally ﬁts into the formalism of sparse MoEs, as illustrated by the connection we draw between BE and sparse MoEs (Appendix I).
This enables a simple implementation of pBE on top of an existing one of sparse MoEs.
Because of the tiling, a given image patch has M different representations that, when entering an
MoE layer, are each routed within their respective parts of the partition ∪ Mm =1 E m .
Formally, consider some tiled inputs H ∈ R B × M × D where B refers to the batch size and h i,m ∈ R D is the representation of the i -th input for the m -th member.
The routing logic in pBE can be written as pBE-MoE ( h i,m ) =
(cid:88)
[Equation]
where the routing weights are now parametrized by W m ∈ R ( E/M ) × D ; see Figure 1.
To echo the observations from Section 3, we can ﬁrst see that pBE brings together the static and adaptive combination of ensembles and sparse MoEs, which we found to be complementary.
However, we have seen that static ensembling comes at the cost of a large increase in FLOPs, thus we opt for an efﬁcient ensembling approach.
Second, we “split” the MoE layers along the axis of the experts, i.e., from E experts to M times E/M experts.
We do so since we observed that the performance of sparse MoEs tends to plateau quickly as the number of experts grows.
Finally, pBE retains the important property of ensembles to output multiple predictions per input, which we also saw to be beneﬁcial for uncertainty calibration.
In a generic implementation, we tile a batch of B inputs X ∈ R B × P by a factor M to obtain the tiled inputs X tiled =
[ X ; . . .
; X ] ∈ R ( M · B ) × P and the model processes f
( X tiled ; θ ) .
Since tiling in pBE has an effect only from the ﬁrst MoE layer onwards, we postpone the tiling operation to that stage, thus saving all prior computations in non MoE-layers that would have been redundant otherwise.
For example, for L/16 and K = M = 2 , we can save about 47% of the FLOPs.
We apply the same optimization for all the efﬁcient ensembles methods we compare to in Appendix F.
We provide further implementation details of pBE in Appendix D.
4. A BLATION S TUDIES : P ARTITIONING AND T ILING
Our method introduces two changes to V-MoEs: (a) the partitioning of the experts and (b) the tiling of the representations.
In this section, we assess the separate impact of each of those changes and show that it is indeed their combination that explains the performance gains.
We summarize the results of the ablation in Table 3 where we show the ImageNet performance of the different variants of V-MoE. All models have a ViT-B/32 base architecture and K
= M = 2 .
[Table]
[Table]
[Table]
P ARTITIONING WITHOUT T ILING
We ﬁrst compare pBE with a variant of V-MoE where we only partition the set of experts (“Only partitioning”).
In that variant, each input h i ∈
R D (note the dropping of the index m due to the absence of tiling) can select K experts in each part of the partition ∪ Mm =1 E m , resulting in a total of
K × M selected experts per input.
Formally, (3) becomes (cid:80) Mm =1 (cid:80)
e ∈E m g e,m ( h i ) · MLP e ( h i ) .
The expert prototyping of Yang et al.
(2021) leads to a similar formulation.
As shown in Table 3, “Only partitioning” is not competitive with pBE across all metrics.
We do not report the KL since without tiling, “Only partitioning” does not output multiple predictions per input.
4.2. T ILING WITHOUT P ARTITIONING
We now compare pBE with the variant where only the tiling is enabled (“Only tiling”).
In that case, we have tiled inputs H ∈ R B × M × D applied to the standard formulation of (1).
Compared with (3), there is no mechanism to enforce the M representations of the i -th input across the ensemble members, i.e., { MoE ( h i,m ) } Mm =1 , to be different.
Indeed, without partitioning, each h i,m could select
K identical experts.
As a result, we expect “Only tiling” to output M similar predictions across ensemble members.
We capture this intuition in Table 3 where we observe that the KL for “Only tiling” is orders of magnitude smaller than for pBE.
To mitigate this effect, we also tried to increase the level of noise σ in gate K (by a factor { 2 , 4 } ), to cause the expert assignments to differ across { h i,m } Mm =1 .
While we do see an increase in KL, “Only tiling” still performs worse than pBE across all metrics.
Interestingly, we can interpret “Only tiling” as an approximation, via M samples, of the marginalization E ε 1 ,...,ε (cid:96)
[ f ( x ; θ )] with respect to the noise { ε l } (
cid:96)l =1 in the (cid:96) MoE layers of f ( · ; θ ) (further assuming the capacity constraints of the experts, as described in Riquelme et al.
(2021), does not bias the M samples).
E VALUATION
We now benchmark pBE against V-MoE. As a baseline we also include results for downstream ensembles of V-MoE and ViT. These ensembles offer a natural baseline against pBE as they also use a single upstream checkpoint, are easy to implement, and provide consistent improvements upon V-MoE. In Appendix G, we compare with upstream ensembles that require multiple upstream checkpoints (Mustafa et al., 2020).
In Appendix F, we compare with other efﬁcient ensembling approaches: MIMO (Havasi et al., 2020), BE (Wen et al., 2019), and MC Dropout (Gal & Ghahramani, 2016).
All results correspond to the average over 8 (for { S, B, L } single models) or 5 (for H single models and all up/downstream ensembles) replications.
In Appendix H we provide standard errors as well as results for additional datasets and metrics.
Following Riquelme et al.
(2021), we compare the predictive-performance vs. compute cost trade-offs for each method across a range of ViT families.
In the results below, pBE uses ( K, M ) =
(1 , 2) , and V-MoE uses K = 1 .
Experimental details, including about our upstream training, downstream ﬁne-tuning, hyperparameter sweeps and our (linear) few-shot evaluation can be found in Appendix B.
Our main ﬁndings are as follows:
(a) V-MoE versus ViT .
Predictive performance and robustness .
• Ensembles help V-MoE just as much as ViT. Ensembles were expected to beneﬁt ViT models.
However, Figure 3 and Figure 4 suggest that ensembling provides similar gains for V-MoE models in terms of both few-shot performance and NLL.
We believe this has not been observed before.
Moreover, a downstream ensemble with four H/14 V-MoEs leads to a 88.8% accuracy on ImageNet
[Figure] [Figure]
impressive
% for
• ViT consistently provides better ECE than V-MoE. Surprisingly, despite V-MoE tending to have better NLL than ViT (Figure 3), their ECE is worse (Figure 5).
• ECE is not consistent for different ViT/V-MoE families.
We see the ECE, unlike other metrics presented in this work, provides no consistent trends as we increase the ViT family size (Figure 5).
• V-MoE outperforms ViT in OOD detection.
With L/32 being the only exception, V-MoE outperforms ViT on a range of OOD detection tasks (Figure 5).
• For smaller ViT families, V-MoE outperforms ViT in the presence of distribution shift.
In contrast to the OOD detection results, Figure 6 shows that for smaller ViT families V-MoE improves on the performance of ViT, however, as the ViT family becomes larger, this trend reverses.
(b) Partitioned Batch Ensembles .
Predictive performance and robustness .
• pBE improves classiﬁcation performance.
As shown in Figure 4, pBE is either on or very near to the Pareto frontiers for NLL and 10-shot classiﬁcation error, despite the fact that these are metrics for which ensembles and V-MoE, respectively, are known to perform well.
Furthermore, Figuresand 15 show that even starker conclusions hold on Cifar10/100.
• pBE performs best at the largest scale.
The difference in predictive performance between pBE and V-MoE—or ensembles thereof—increases as the ViT family becomes larger (Figures 4, 5 and 6).
See Appendix E for further motivation of this point.
• pBE tends to be Pareto efﬁcient in the presence of distribution shift .
Figure 6 shows that pBE is more robust to distribution shift for larger ViT families, despite the opposite being true for V-MoE. • pBE improves ECE over ViT and V-MoE. Despite V-MoE providing poor ECE, pBE does not suffer from this limitation (Figure 5).
Furthermore, for most ViT families, pBE also provides better ECE than V-MoE ensembles.
• pBE does not provide consistent OOD detection performance.
Firstly, Figure 5 shows that for small ViT families, pBE performs worse than V-MoE and (even ViT in some cases).
Nevertheless, as above, the relative performance improves for larger ViT families such that pBE becomes Pareto efﬁcient for two dataset pairs.
Secondly, pBE seems to perform better on the more difﬁcult near
N LL ( l o w e
r
i s b e tt e r)
ImageNet-C (average) 3.
ImageNet-C (average)
ImageNet-A   0.
ImageNet-V downstream log(GFLOPs) (lower is better) pBE
V-MoE
ViT
S/
B/
L/ L/
H/ 1 Model 2 Models 4 Models
OOD detection task (Cifar10 vs. Cifar100).
These results, although sometimes subtle, are consistent across OOD detection metrics and dataset pairs, as shown in Appendix H.
R ELATED W ORK
Mixture of Experts.
MoEs (Jacobs et al., 1991; Jordan & Jacobs, 1994; Chen et al., 1999; Yuksel et al., 2012; Eigen et al., 2014) combine the outputs of different submodels, or experts , in an inputdependent way.
Sparse MoEs only select a few experts per input, enabling to greatly scale models while keeping the prediction time constant.
Sparse MoEs have been used to build large language models (Shazeer et al., 2017;
Lepikhin et al., 2021; Fedus et al., 2021).
Recently, sparse MoEs have been also successfully applied to vision problems (Riquelme et al., 2021; Yang et al., 2021; Lou et al., 2021; Xue et al., 2021).
Our work builds on the V-MoE architecture proposed by Riquelme et al.
(2021), which is based on the vision Transformer (ViT)
(Dosovitskiy et al., 2021).
While previous work studied ViT’s calibration and robustness (Minderer et al., 2021; Fort et al., 2021; Paul & Chen, 2021; Mao et al., 2021), we are the ﬁrst to study the robustness of V-MoE models.
Ensembles.
Ensemble methods combine several different models to improve generalization and uncertainty estimation.
In their simplest form, they can be inefﬁcient because they consist of multiple models themselves potentially expensive.
To reduce test time, Xie et al.
and Hinton et al.
(2015) respectively use compression and distillation mechanisms.
To reduce training time, ensembles can be constructed with cyclical learning-rate schedules to snapshot models along the training trajectory (Huang et al., 2017; Zhang et al., 2019).
Our work builds on batch ensemble (Wen et al., 2019) where a single model encapsulates an ensemble of networks, a strategy also explored by Lee et al.
; Havasi et al.
; Antor ´ an et al.
Dusenberry et al.
(2020a);
Rame et al.
Wenzel et al.
(2020) extended BE to combine models with different hyperparameters.
C ONCLUSIONS AND F UTURE W ORK
Our study of the interplay between sparse MoEs and ensembles has shown that these two classes of models are symbiotic.
Partitioned batch ensemble exempliﬁes those mutual beneﬁts—as illustrated by its accuracy, log-likelihood, few-shot learning, robustness, and uncertainty calibration improvements over several challenging baselines in a range of benchmarks.
While our study has focused on downstream ﬁne-tuned models, we believe that an extension to the upstream case would also result in a fruitful investigation.
Similarly, although we have focused on computer vision, our approach should be readily applicable to the modelling of text, where sparse MoEs have been shown to be remarkably effective.
With the growing prevalence of sparse MoEs in NLP (Patterson et al., 2021), the questions of understanding and improving the robustness and reliability of such models become increasingly important.
Furthermore, the computational scale at which those models operate make those questions even more challenging to tackle.
We believe that our study, and approaches such as pBE, make steps in those directions.
Ethics Statement.
Our research lies at the intersection of two topics where we hope that our work can make positive contributions.
First, following the conclusions of Patterson et al.
(2021), we develop an approach based on sparse
MoEs that were shown to reduce the environmental footprint of standard “dense” models.
Second, for any decision-making process, it is critical to be able to reliably trust the uncertainty output by
ML systems.
In particular, this desirable property has a growing importance within a context where those systems are being widely deployed in safety-critical ﬁelds such as self-driving cars and medical diagnosis.
We think that approaches such as pBE can help make progress in this area.
Reproducibility Statement.
We are fully aware that (i) relying on the proprietary JFT-300M dataset for our upstream models, as well as (ii) not having already open-sourced code are two obstacles for reproducibility.
We are focusing on a lightweight open-sourced version of V-MoE and pBE together with the release of checkpoints pre-trained on ImageNet-21k.
We are actively working to release those with the camera-ready version of the paper.
We would like to stress the fact that we have provided as many details as possible about the experimental settings (see Appendix B) and the hyperparameter sweeps.
Furthermore, we aimed to provide well-considered and fair experimental setups for all of the baselines used in this work.
For instance, to illustrate the fairness of our approach, our experimental design choices tend to make the baselines we compare to more competitive; see Table 6.
Moreover, we performed the evaluation of our models using the open-sourced library robustness metrics (Djolonga et al., 2020).
Finally, we have also reported results that could be regarded as “negative” with the hope to inform other researchers about those ﬁndings.
For instance, in Appendix F.3, we systematically described in detail the failures we encountered while trying to apply MIMO to ViT and V-MoE. Similarly, we have transparently reported—in the core paper—results of pBE for OOD detection where our approach seems to perform worse than the baselines in several settings.
A V I T M ODEL S PECIFICATIONS
Following Dosovitskiy et al.
(2021), we recall the speciﬁcations of the ViT models of different scales in Table 4.
B E XPERIMENT S ETTINGS B. U PSTREAM S ETTING
For all our upstream experiments, we scrupulously follow the setting described in Riquelme et al.
(2021), see their Section B.2 in their appendix.
For completeness, we just recall that S/32 models are trained for 5 epochs while B/ { 16, 32 } and L/32 models are trained for 7 epochs.
For L/16 models, both 7 and 14 epochs can be considered (Dosovitskiy et al., 2021;
Riquelme et al., 2021)
; we opted for 7 epochs given the breadth of our experiments.
Finally, the H/14 model are trained for 14 epochs.
In particular, the models are all trained on JFT-300M (Sun et al., 2017).
This dataset contains about 305M training and 50 000 validation images.
The labels have a hierarchical structure, with a total of 18 291 classes, leading on average to 1.89 labels per image.
B. D OWNSTREAM S ETTING
During ﬁne-tuning, there is a number of common design choices we apply.
In particular:
• Image resolution: 384.
• Clipping gradient norm at: 10.0. • Optimizer: SGD with momentum (using half-precision, β = 0 . 9 ).
• Batch size: 512.
• For V-MoE models, we ﬁnetune with capacity ratio C = 1 . 5 and evaluate with C = 8 .
We use the following train/validation split partitions depending on the dataset:
[Equation]
[Table]
All those design choices follow from Riquelme et al.
and Dosovitskiy et al.
B. H YPERPARAMETER S WEEP FOR F INE - TUNING
In all our ﬁne-tuning experiments, we use the sweep of hyperparameters described in Table 5.
We use the recommendations from Dosovitskiy et al.
(2021) and Riquelme et al.
(2021), further considering several factors { 0.5, 1.0, 1.5, 2.0 } to sweep over different numbers of steps.
Riquelme et al.
(2021) use a half schedule (with the factor 0.5) while Dosovitskiy et al.
(2021) take the factor 1.0.
[Table] [Table]
We show in Table 6 the impact of this enlarged sweep of hyperparameters in the light of the results reported in Riquelme et al.
We notably tend to improve the performance of ViT and V-MoE (especially for smaller models), which thus makes the baselines we compare to more competitive.
B. D ETAILS ABOUT THE (L INEAR ) F EW - SHOT E VALUATION
We follow the evaluation methodology proposed by Dosovitskiy et al.
; Riquelme et al.
(2021) which we recall for completeness.
Let us rewrite our model f with parameters θ = { Q , θ (cid:48) } as f ( x ; θ ) =
softmax ( Q φ ( x ; θ (cid:48) )) where Q ∈ R C × S corresponds to the parameters of the last layer of f with the S -dimensional representation φ ( x ; θ (cid:48) )
∈ R S .
In linear few-shot evaluation
[Table]
we construct a linear classiﬁer to predict the target labels (encoded as one-hot vectors) from the S -dimensional feature vectors induced by φ ( · ; θ (cid:48) ) ; see Chapter 5 in Hastie et al.
(2017) for more background about this type of linear classiﬁers.
This evaluation protocol makes it possible to evaluate the quality of the representations φ learned by f .
While Dosovitskiy et al.
; Riquelme et al.
(2021) essentially focus on the quality of the representations learned upstream on JFT by computing the (linear) few-shot accuracy on ImageNet, we are interested in the representations after ﬁne-tuning on ImageNet.
As a result, we consider a collection of 8 few-shot datasets (that does not contain ImageNet):
• Caltech-UCSD Birds 200 (Wah et al., 2011) with 200 classes, • Caltech 101 (Bansal et al., 2021) with 101 classes, • Cars196 (Krause et al., 2013) with 196 classes,
• Cifar100 (Krizhevsky, 2009) with 100 classes, • Colorectal histology (Kather et al., 2016) with 8 classes, • Describable Textures Dataset (Cimpoi et al., 2014) with 47 classes, • Oxford-IIIT pet
(Parkhi et al., 2012) with 37 classes and • UC Merced (Yang & Newsam, 2010) with 21 classes.
In the experiments, we compute the few-shot accuracy for each of the above datasets and we report the averaged accuracy over the datasets, for various number of shots in { 1 , 5 , 10 , 25 } .
As commonly deﬁned in few-shot learning, we understand by s shots a setting wherein we have access to s training images per class label in each of the dataset.
To account for the different scales of accuracy across the 8 datasets, we also tested to compute a weighted average, normalizing by the accuracy of a reference model (ViT-B/32).
This is reminiscent of the normalization carried out in Hendrycks & Dietterich (2019) according to the score of AlexNet.
We found the conclusions with the standard average and weighted average to be similar.
B.4. S PECIFIC C ONSIDERATIONS IN THE E NSEMBLE C ASE
For an ensemble with M members, we have access to M representations { φ ( x ; θ (cid:48) m ) }
Mm =1 for a given input x .
We have explored two ways to use those representations:
Joint : We concatenate the M representations { φ ( x ; θ (cid:48) m ) } M m =1 into a single “joint” feature vector in R M × S , remembering that each φ ( x ; θ (cid:48) m ) ∈ R S .
We then train a single a linear classiﬁer to predict the target labels from the “joint” feature vectors.
• Disjoint :
For each of the M representations { φ ( x ; θ (cid:48) m ) } Mm =1 , we separately train a linear classiﬁer to predict the target labels from the feature vectors induced by φ ( x ; θ (cid:48) m ) .
We then average
the predictions
of
the
M
linear
classiﬁers trained in this fashion.
In Table 7, we report a comparison of those approaches.
We aggregate the results over all ensemble models (namely, pBE and upstream ViT/V-MoE ensembles of size 2 and 4) and over 8 replications, for the ViT families S/32, B/32 and L/32.
The results indicate that “joint” and “disjoint” perform similarly.
Throughout our experiments, we use the “joint” approach because it eased some implementation considerations.
B. L IST OF D ATASETS
For completeness, in addition to the few-shot datasets listed in Appendix B.4, we list the datasets used for downstream training and evaluation in this work.
• ImageNet (ILSVRC2012)
(Deng et al., 2009) with 1000 classes and 1281167 training examples.
• ImageNet-C
(Hendrycks & Dietterich, 2019), an ImageNet test set constructed by applying 15 different corruptions at 5 levels of intensity to the original ImageNet test set.
(We report the mean performance over the different corruptions and intensities.)
• ImageNet-A
(Hendrycks et al., 2019), an ImageNet test set constructed by collecting new data and keeping only those images which a ResNet-50 classiﬁed incorrectly.
• ImageNet-V2 (Recht et al., 2019), an ImageNet test set independently collected using the same methodology as the original ImageNet dataset.
• Cifar
(Krizhevsky, 2009) with 10 classes and 50000 training examples.
• Cifar10-C
(Hendrycks & Dietterich, 2019), a Cifar10 test set constructed by applyingdifferent corruptions at 5 levels of intensity to the original Cifar10 test set.
(We report the mean performance over the different corruptions and intensities.)
• Cifar
(Krizhevsky, 2009) with 100 classes and training 50000 examples.
• Oxford Flowers 102 (Nilsback & Zisserman, 2008) with 102 classes and 1020 training examples.
[Table]
• Oxford-IIIT pet
(Parkhi et al., 2012) with 37 classes and [Table] training examples.
• SVHN (Netzer et al., 2011) with 10 classes.
Places
(Zhou et al., 2017) with 365 classes.
• Describable Textures Dataset (DTD) (Cimpoi et al., 2014) with 47 classes.
B. A BLATION D ETAILS B.6. S TATIC VERSUS A DAPTIVE A BLATION D ETAILS
The setup for the experiments in Figures 2 and 8 differs slightly the other experiments in this paper.
Speciﬁcally, while for all other experiments we used upstream V-MoE checkpoints with ( K, E ) =
(2 , 32) , for these experiments we matched the upstream and downstream checkpoints.
We did this to avoid a checkpoint mismatch as a potential confounder in our results.
B.6. F EATURE - LEVEL VERSUS P REDICTION - LEVEL E NSEMBLING A BLATION D ETAILS
The “Naive Multi-Head” method presented in Section 3.2 was trained in almost the same manner as the vanilla V-MoE, the only difference being the handling of multiple predictions.
This was accomplished by using the average ensemble member cross entropy as described for pBE in Appendix D.
On the other hand, in order to compute the evaluation metrics presented in Table 2, we ﬁrst averaged predictions of the model and then used the average prediction when calculating each metric.
C C OMPATIBILITY AND A DAPTATION OF THE U PSTREAM C HECKPOINTS
Throughout the paper, we make the assumption that we can start from existing checkpoints of ViT and V-MoE models (trained on JFT-300M; see Appendix B.1).
We next describe how we can use those checkpoints for the ﬁne-tuning of the extensions of ViT and V-MoE that we consider in this paper.
In all our experiments that involve V-MoEs, we consider checkpoints with K = 2 and E = 32 , which is the canonical setting advocated by Riquelme et al.
C. P ARTITIONED B ATCH E NSEMBLES ( P BE)
In the case of pBE, the set of parameters is identical to that of a V-MoE model.
In particular, neither the tiled representation nor the partitioning of the experts transforms the set of parameters.
To deal with the fact that the single routing function gate K ( W · ) of a V-MoE becomes separate routing functions { gate K ( W m · ) } M m =1 , one for each part of the partition, we simply slice row-wise
W ∈ R E × D into the M matrices
W m ∈ R ( E/M ) × D .
C. B ATCH E NSEMBLES (BE)
We train BE starting from ViT checkpoints, which requires to introduce downstream-speciﬁc parameters.
Following the design of V-MoEs, we place the batch-ensemble layers in the MLP layers of the Transformer.
Let us consider a dense layer in one of those MLPs, with parameters U ∈ R D × L , in absence of bias term.
In BE, the parametrization of each ensemble member has the following structure
U m =
U ◦ ( r m s (cid:62) m ) where { r m } Mm =1 and { s m } Mm =1 are respectively D - and L -dimensional vectors.
A standard ViT checkpoint provides pre-trained parameters for U .
We then introduce { r m } Mm =1 and { s m } Mm =1 at ﬁne-tuning time, following the random initialization schemes proposed in Wen et al.
(2019); see details in the hyperparameter sweep for BE in Appendix F.1.
C. MIMO
We train MIMO models from V-MoE checkpoints.
The only required modiﬁcations are to the input and output parameters of the checkpoints.
The linear input embedding must be modiﬁed to be compatible with input images containing M times as more channels, as required by the multipleinput structure of MIMO.
Similarly, the ﬁnal dense layer in the classiﬁcation head must be modiﬁed to have M times more output units, following the multiple-output structure in MIMO.
Concretely, the embedding weight W in ∈ R
H
× W × 3 × D is replicated in the third (channel) dimension, resulting in W MIMO , in ∈ R
H × W × 3 · M × D , where H and W are the height and width of the convolution and D is the hidden dimension of the ViT family (speciﬁed in Table 4).
The output layer weight W out ∈ R D
× C is replicated in the second (output) dimension, resulting in
W MIMO , out ∈ R H × C · M , where C is the number of classes.
The output layer bias b out ∈ R C is replicated resulting in b MIMO , out ∈ R C × M .
Finally, in order to preserve the magnitude of the activation for these layers, W MIMO , in and W MIMO , out are scaled by 1 /M .
D I MPLEMENTATION D ETAILS OF P BE
We provide details about the training loss and the regularizer used by pBE.
D. T RAINING L OSS
Since pBE outputs M predictions { f ( x ; θ m ) } Mm =1 for a given input x , we need to adapt the choice of the training loss L accordingly.
Following the literature on efﬁcient ensembles (Wen et al., 2019;
Dusenberry et al., 2020a; Wenzel et al., 2020), we choose the average ensemble-member cross entropy
[Equation]
instead of other alternatives such as the ensemble cross-entropy
[Equation]
that was observed to generalize worse (Dusenberry et al., 2020a).
D. A UXILIARY L OSSES
Inspired by previous applications of sparse MoEs in NLP (Shazeer et al., 2017), Riquelme et al.
(2021) employ regularizers, also referred to as auxiliary losses , to guarantee a balanced usage of the E experts.
Two auxiliary losses—the importance and load losses, see Appendix A in Riquelme et al.
(2021) for their formal deﬁnitions—are averaged together to form the ﬁnal regularization term that we denote by Ω .
As a reminder, let us recall the notation of the routing function
[Equation]
with W ∈ R E × D and ε ∼ N ( 0 , I ) .
Consider a batch of B inputs { h i } Bi =1 that we represent by
H ∈ R B × D .
Finally, let us deﬁne
[Equation]
where we emphasise that ε B × E is a matrix of Gaussian noise entries in R B × E .
The regularization term Ω used by Riquelme et al.
(2021) can be seen as a function that depends on A and HW (cid:62) .
In the context of partitioned batch ensemble, the set of E experts is partitioned into M groups of
E/M experts, whose partition is denoted by ∪ Mm =1 E m ; see Section 4.1.
With the introduction of the
M routing functions { gate K ( W m · ) } Mm =1 with each W m ∈ R ( E/M ) × D , the matrix A becomes accordingly partitioned into { A m } Mm =1 where each A m ∈ R B × ( E/M ) .
Since we want to enforce a balanced usage of the E/M experts in each part E m of the partition, we thus redeﬁne the regularization as the average regularization separately applied to each part of the partition
[Equation]
We found this option to work better in practice.
To guarantee a fair comparison, we also applied Ω partition to the “Only partitioning” model in the ablation study of Section 4.2.1.
Following Riquelme et al.
(2021), the regularization parameter controlling the strength of Ω partition was set to 0.1 throughout the experiments.
E P BE AND V-M O E R ELATIVE I MPROVEMENTS PER V I T F AMILY
In Section 5 we claim that pBE performs best at the largest scale.
In this section we motivate that claim in more detail.
Speciﬁcally, we consider two metrics of improvement in performance.
Firstly, we consider the percentage improvement in NLL for both pBE and V-MoE versus vanilla ViT.
Secondly, we consider a normalised version of this improvement.
We consider this second metric to take into account the “difﬁculty” in further improving the NLL of larger ViT family models.
Intuitively, the larger the ViT family, the better the corresponding NLL will be, and the more difﬁcult it will be to improve on that NLL.
The normalisation we apply is based on the gradient of the NLL with respect to FLOPs.
Indeed, this gradient captures the typical variation of NLL at a particular amount of FLOPs.
The ratio of this gradient at the FLOPs values (i.e., the instantaneous change in NLL at those FLOPs values) for two
ViT families is a measure of the relative difﬁculty in increasing the NLL.
Thus, we can use this ratio to normalise our results.
To be more concrete, let us deﬁne the mapping
[Equation]
We estimate ϕ and its gradient by ﬁtting a linear model to the (NLL , FLOPs) pairs for each
ViT family, using the data of the standard ViT models we trained.
We use feature expansion [1 . 0 , log(FLOPs) , log(FLOPs) 2 , log(FLOPs) 3 ] and solve for the parameters of the linear model via ordinary least squares.
We determine the gradient of this function at each FLOPs value using automatic differentiation in JAX (Bradbury et al., 2018).
See Figure 7 for the resulting ﬁt and an indication of the gradients.
The normalised values are calculated as:
[Equation]
where v is one of the ViT families, i.e., S/32, B/32, L/32, L/16, or H/14.
Note that this normalisation leaves the improvement for H/14 the same.
We tried to normalize with respect to other choices of
ViT family, different from H/14.
Our conclusions are robust, in the sense that both the ordering and the monotonic behavior with respect to scale are preserved.
Using the ratio for normalisation also has the advantage that the normalisation is less sensitive to the particular parameterisation of ϕ .
Table 8 shows both the difﬁculty-normalised and original improvements (without normalisation).
Looking ﬁrst at the original improvements, we can see that while both pBE and V-MoE have smaller (downstream) GFLOPs 0. N LL ϕ Tangent ViT
[Table]
improvements over ViT for larger families, pBE’s improvements decrease more slowly.
Furthermore [Table] by comparing the normalised improvements we see that pBE’s improvements actually grow monotonically when taking difﬁculty into account.
This is not the case for V-MoE.
F E FFICIENT E NSEMBLE C OMPARISONS
In this section [Table] we compare partitioned batch ensembles (pBE) to several popular efﬁcient ensemble approaches, namely MIMO (Havasi et al., 2020), batch ensemble (BE) (Wen et al., 2019) [Table]
and MC
Dropout (Gal & Ghahramani, 2016).
Table 9 reports the ImageNet performance of those different techniques, when all models are based on a ViT-B/32 architecture.
We start by highlighting the most salient conclusions of the experiment and defer to the next subsections the descriptions of the different competing techniques.
We make the following observations:
• BE built upon ViT improves the performance of ViT in terms of NLL, classiﬁcation error and ECE.
However, the resulting increase in FLOPs makes BE a less viable option compared to pBE.
• MC dropout V-MoE is on par with standard V-MoE in terms of NLL and classiﬁcation error, while it improves the ECE.
For all values of K , we observe that the performance tends to improve as the number of samples,
i
e., M , increases.
However, already for M in { 2 , 4 } , the resulting increase in FLOPs makes MC dropout V-MoE a less favorable option compared to pBE.
• Perhaps surprisingly (see detailed investigations in Appendix F.3), MIMO V-MoE does not lead to improvements compared with V-MoE.
In fact, for higher ensembles sizes, MIMO V-MoE results in worse performance than standard V-MoE.
Moreover, increasing the batch [Table] [Table] repetition parameter of MIMO (“BR” in Table 9) further worsens the results.
Interestingly, we can see that MIMO does not manage to produce diverse predictions, as illustrated by the small values of KL.
• pBE offers the best performance vs. FLOPs trade-offs, e.g., when looking at ( K, M ) =
(1 , 2) and ( K, M ) =
We notably observe that the diversity of the predictions in pBE is orders of magnitude larger than that of the other ensemble approaches.
We brieﬂy recall the optimization explained in Section 4.1 to save redundant computations:
In the “last- n ” setting of Riquelme et al.
(2021), it is sufﬁcient to tile the representations only when entering the ﬁrst MoE layer/dropout layer/batch-ensemble layer for respectively pBE/MC dropout
V-MoE/BE.
We apply this optimization to all the efﬁcient ensemble methods.
F. B ATCH E NSEMBLES
Following the design of V-MoEs, we place the batch-ensemble layers in the MLP layers of the
Transformer, following the “last- n ” setting of Riquelme et al.
; see Section 2.1.
The vectors of the rank-1 parametrization introduced at ﬁne-tuning time (see Appendix C) need to be initialized and optimized.
Following the recommendation from Wen et al.
(2019), we consider the following hyperparameters in addition to the common sweep described in Table 5:
• Initialization: Either a random sign vector with entries in {− 1 , 1 } independently drawn with probability 12 or a random Gaussian vector with entries independently drawn from N (1 , 0 . 5) .
• Learning-rate scale factor: The vectors of the rank-1 parametrization are updated with a learning rate scaled by a factor in { 0 . 5 , 1 , 2 } .
F. MC D ROPOUT V-M O E S
For MC dropout V-MoE, we take the available ﬁne-tuned V-MoEs and enable dropout at prediction time.
Indeed, as described in Table 5, all V-MoE models already have a 0.1 dropout rate in the experts.
F. MIMO V-M O E S
Following Havasi et al.
we consider two MIMO-speciﬁc hyperparameters, in addition to the hyperparameters listed in Table 5:
• Input replication probability: { 0 . 5 , 0 .
625 , 0 . 75 } • Batch repetitions: { 1 , 2 }
Our preliminary investigations also considered lower input repetition probabilities and higher batch repetitions.
However, lower input repetition probabilities tended to result in poorer performance.
While higher batch repetitions did help to some extent, the additional computational cost made it impractical.
Given the surprising result that an ensemble size of M = 2 provides no performance improvement over the standard V-MoE and that increasing M further provides worse performance, there seems to be some incompatibility between MIMO and V-MoE.
In fact, our investigations revealed that ViT is the source of the problems since applying MIMO to vanilla ViT without experts resulted in the same trends as for V-MoE.
Thus we hypothesise that the differences between ViT and ResNet—the architecture to which MIMO was originally applied by Havasi et al.
(2020)—are responsible for
MIMO’s poor performance when applied to ViT.
Difference 1: Class token.
One of the differences between ViT and ResNet is that ViT makes use of a special learnable class token to classify an image (see Dosovitskiy et al.
(2021) for details).
ResNet on the other hand makes use of the representation from an entire image for classiﬁcation.
We tried two strategies to mitigate this difference:
We applied the global average pooling (GAP) and multi-head attention pooling (MAP) classiﬁcation strategies introduced in Dosovitskiy et al.
(2021) and Zhai et al.
(2021), respectively.
In short, both of these methods make use of all the tokens from an image for classiﬁcation.
However, neither of these strategies made a signiﬁcant difference to the rel- ative performance of MIMO and ViT.
In fact, the choice of classiﬁcation method was the least impactful hyperparameter in our sweep.
Rather than learning a single class token, we learnt M class tokens.
This strategy resulted in MIMO with M = 2 outperforming ViT.
However, for M > 2 the improvement was small enough that ViT still outperformed MIMO.
Difference 2: Attention.
The other major difference between ViT and ResNet is the building block for each model.
While ResNets are primarily composed of convolution operations, ViT makes heavy used of attention.
We hypothesised that attention is less suited to separating the information for M input images stored in the channel dimension of a single image.
We tried two strategies to mitigate this potential issue:
We applied the hybrid architecture, described in Dosovitskiy et al.
(2021), in which the input sequence to ViT is formed by CNN feature maps.
We used ResNet14 and ResNet50.
In both cases, we found that strategy boosted the performance of ViT and MIMO equally.
Rather than concatenating images in the channel dimension, we concatenated them in the width dimension, resulting in 3 times as many patches for ViT to process.
This strategy was successful in the sense that the MIMO performance for M > 2 improved signiﬁcantly.
However, the signiﬁcant additional computational cost made it an infeasible solution.
Our ﬁndings suggest that MIMO and ViT are indeed somewhat incompatible.
Unfortunately, none of our proposed solutions to this problem provided high enough predictive performance increases (or indeed low enough computational cost increases in some cases) to warrant immediate further investigation.
G U PSTREAM & D OWNSTREAM VERSUS D OWNSTREAM - ONLY E NSEMBLES
In Section 5, and Appendix H include downstream deep ensembles (down-DE) of V-MoE, and in some cases ViT, as a baseline.
This choice was motivated by the fact that like ViT, V-MoE, and pBE, down-DE requires a only a single upstream checkpoint, which all of the methods more comparable.
However, it is clear that using different upstream checkpoints and then further ﬁne-tuning each of these with different random seeds to construct an upstream deep ensemble (up-DE) would result in more varied ensemble members and as a result, a better performing ensemble.
This idea has recently been explored by Mustafa et al.
Thus, for completeness, we also investigate the effects of upstream ensembling on V-MoE. Tablecompares the performance of upstream and downstream V-MoE ( K = 1 ) ensembles of sizes M =and M = 4 .
Across the range of metrics, for both ImageNet and ImageNet-C, for all ViT families, and for both values of M , we see that up-DE outperforms down-DE.
In fact, up-DE with M =is very often better than or equal to down-DE with M = 4 .
This is especially true for the diversity metrics, which indicates that diversity is indeed the driver for improved performance in up-DE.
Not shown in the table is the very large computational cost associated with training upstream ensembles.
U nd e r r e
v
i e w a
[Table]
a c on f e r e n ce p a p e r a t
I C L
R
[Table] [Table]
H A DDITIONAL E XPERIMENTAL R ESULTS
In this section we expand on various experiments presented in sections 3 and 5.
In experiments considering multiple ViT families we also include B/16 which was excluded from the main text for clarity.
H. S TATIC VERSUS A DAPTIVE C OMBINATION
Here we continue the investigation into static versus adaptive combination from Section 3.1.
Individual gains with respect to E, K and M .
Figure 8 shows the effect of increasing the the various ‘ensemble size’ parameters for a deep ensemble of V-MoEs.
In particular, we investigate the static combination axis M (the number of ensemble members), as well as the two adaptive axes— K (the number of experts chosen per patch) and E (the total number of experts).
When investigating the effect of K , we ﬁx E = 32 and average over M ∈ { 1 , .., 8 } .
Similarly, when investigating M , we ﬁx E = 32 and average over K ∈ { 1 , .., 8 } .
When investigating the effect of
E we ﬁx K = 2 and average over M ∈ { 1 , .., 8 } .
As a result of this procedure the exact values of the curves are not directly comparable.
However, we can still examine the relative trends.
Speciﬁcally, we note that while the variation in K and M curves is roughly of the same size, the variation in the E curve is smaller.
We also note that there is very little variation beyond E =(note the difference in the scales of the axes for the curves).
These observations motivate the design of pBE, where we split the sub-models along the E axis, in order to better take advantage of the experts.
[Figure]
Extended Results for the Cumulative Effects of Static and Adaptive
Combination.
In Figure 9 we extend the ImageNet NLL results, presented in Figure 3, to a range of other datasets and performance metrics.
We see that in most cases, the trends are the same as before.
That is, (static) ensembling tends to improve the performance of ViT and V-MoE equally.
The two exceptions are
ECE and OOD detection for ViT-S/32 where we see that larger ensemble sizes can result in decreased performance.
These results indicate that the for small ViT models, larger ensembles can have slightly lower quality uncertainty estimates.
The trend for ImageNet-C performance is also not as consistent with ensembling sometimes helping ViT or V-MoE less (as indicated by the changes in ordering on the y-axis).
[Figure]
H
[Figure]
H. E XTENDED R ESULTS FOR F EW - SHOT L EARNING
In Figure
[Figure]
we extend the few-shot learning results of Figure 4 to also include 1, 5, and 25-shot.
Additionally
[Figure]
we show results for the weighted aggregation strategy mentioned in Appendix B.4.
We conﬁrm the result that few-shot performance for pBE gets better, relative to the other baselines, with larger ViT families.
Additionally, we see that pBE performance seems to get better, again relative to the other baselines, with more shots.
This phenomenon can most easily be noticed by comparing the results for S/32 across the different numbers of shots.
Finally, we see that the trends with and without the weighted mean are the same
[Figure]
H. E XTENDED R ESULTS FOR OOD D ETECTION
Here we extended the OOD results of Figure 5.
Speciﬁcally, we add Cifar100 as an in-distribution dataset and Describable Textures Dataset (DTD) (Cimpoi et al., 2014) as an OOD dataset.
We also add area under the receiver operating characteristic (AUC (ROC)) and area under the precision-recall curve (AUC (PR)) as metrics.
Figures 11 and 12 contain the results with Cifar10 and Cifar100 as the in-distribution datasets, respectively.
As in Figure 5, we see that pBE performs better (relative to the other baselines) for larger ViT families.
Furthermore, pBE seems to perform better in near OOD detection (i.e. Cifar10 versus
Cifar100, and vice versa) than far OOD detection.
Finally, we see that these trends are consistent across ODD metrics.
[Figure]
[Figure] [Figure]
C
I F
A
Rv s .
C
I F A
[Figure]
FPR@
[Figure]
(lower is better)
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
.96 AUC (ROC) (higher is better) 0.
AUC (PR) (higher is better)
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
C
I F
A
Rv s .
D
T D
[Figure]
[Figure] [Figure]
[Figure] [Figure]
C
I F
A
Rv s .
P l ace s
[Figure] [Figure]
C
I F
A
Rv s .
S
VHN
[Figure] [Figure]
[Figure]
H. E XTENDED R ESULTS FOR I MAGE N ET
In this section we extend the results for ImageNet and the corrupted variants presented in Figures 4, 5 and 6.
In addition to NLL (and ECE for standard ImageNet), Figure 13 provides classiﬁcation error and ECE for all ImageNet variants.
Most of the trends observed in Section 5 remain true:
• pBE tends to be Pareto efﬁcient in the presence of distribution shift.
• For smaller ViT families, V-MoE outperforms ViT in the presence of distribution shift.
• pBE improves ECE over ViT and V-MoE. • pBE improves classiﬁcation performance.
• ViT consistently provides better ECE than V-MoE.
However, there are some exceptions:
• ImageNet-A classiﬁcation error.
All models (including pBE) under-perform relative to ViT-S/32 and ViT-H/14.
• ECE for ImageNet-C, ImageNet-A, and ImageNet-V2.
Interestingly for the non
- standard ImageNet variants, and in particular for ImageNet-A, there is a strong correlation between lower ECE and larger ViT families.
H. A DDITIONAL C IFAR 10, C IFAR 100, F LOWERS , AND P ETS R ESULTS
Here we extend the results for ImageNet and the corrupted variants presented in Figures 4, 5 andto 4 additional datasets.
Figures 14, 15, 16 and 17 provide results for Cifar10, Cifar100, Oxford
Flowers 102, and Oxford IIIT Pet, respectively.
As in Appendix H.4, we ﬁnd that the results are similar to those in Section 5.
Compared to ImageNet, for Cifar10, Cifar10-C, and Cifar100, pBE seems to perform even better relative to the other baselines.
Note, for example, that pBE is Pareto efﬁcient (even for S/32) in the cases of Cifar10-C and Cifar100 NLL.
As in Appendix H.4, we see that the ECE has a stronger downward trend with respect to increased ViT family size for shifted test data.
For Flowers and Pets, where we only have results for smaller ViT families, pBE seems to under perform.
However, the performance for L/32 is better than for S/32 and B/32 which suggests that the results for these datasets are consistent with the other datasets presented in this work and, therefore, that we should expect pBE’s predictive performance to keep improving with larger models.
H. P BE AND V-M O E WITH LARGER VALUES OF K AND M
Figure 18 and Figure 19 show the effect of varying K on pBE and V-MoE, and the effect of varying
M on pBE, respectively.
We make the following observations:
• In almost all cases, increasing K or M does not result in Pareto efﬁcient models.
• For V-MoE, increasing K seems to help in most cases, except for ECE performance where it usually hurts.
• For pBE, going from K = 1 to K = 2 seems to help in most cases but going from K =to K = 4 usually hurts.
Going from K = 1 to K = 4 still helps but to a lesser extent than from K = 1 to K = 2 .
• For pBE, increasing M either doesn’t make a consistent and signiﬁcant difference or hurts (e.g. in OOD detection).
These conclusions should, however, be considered with caution.
Recall that the upstream checkpoints used for ﬁne-tuning all V-MoE and pBE models in this work are V-MoE models with K = 2 .
Thus, the results in this experiment are confounded by upstream and downstream checkpoint mismatch for all pBE models and all V-MoE models with K (cid:54) = 2 .
For example, we hypothesise that it is more difﬁcult to train downstream pBE models with larger values of M from upstream V-MoE models because in each partition some common expert specialisations will need to be duplicated.
I m a g e N e t
NLL (lower is better)
Error (lower is better) 0.
ECE (lower is better) 0.
I m a g e N e tC       0. I m a g e N e tA      0.
I m a g e N e tV          0.
downstream log(GFLOPs) (lower is better)
S/ pBE
B/
V-MoE
L/
ViT
B/ 1 Model
L/ 2 Models
H/
[Figure] [Figure]
[Figure] [Figure]
NLL (lower is better) 0.
Error (lower is better) 0.
ECE (lower is better)  0.
[Figure]
downstream log(GFLOPs) (lower is better)
S/ pBE
B/
V-MoE
L/
ViT
B/ 1 Model
L/ 2 Models
[Figure]
NLL (lower is better)
Error (lower is better)  0.
ECE (lower is better) downstream log(GFLOPs) (lower is better)
S/ pBE
B/
V-MoE
L/
ViT
B/ 1 Model
L/ 2 Models
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
I m a g e N e t
NLL (
[Figure]
Error (lower is better) 0.
ECE (lower is better) 1.
[Figure] [Figure]
[Figure]
I m a g e N e tC
[Figure] [Figure]
is better)
Error (
[Figure]
is better) 0.
[Figure] [Figure]
[Figure]
ECE (
[Figure]
is better)
[Figure]
ss
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
is
[Figure] [Figure]
ViT
S/
B/
L/ 1 Model 2 Models
[Table]
Table 11 provides the downstream training FLOPs for various pBE, V-MoE, and ViT conﬁgurations.
These numbers correspond to the x-values of the points in the ﬁgures presented in Sectionand Appendix H. Table 12 provides the percentage difference in FLOPs between the pBE, V-MoE and down-DE models most commonly used in this work.
Note that the percentage differences for
H/14 do not follow the trend of the other sizes, e.g. that the percentage difference between pBE and
V-MoE gets smaller for larger sizes, due to the fact that for H/15 we use a last-5 conﬁguration rather than the last-2 conﬁguration used for the other ViT families.
[Table]
S/
B/ B/
L/ L/
H/ pBE vs. V-MoE
[Table]
[Table] [Table]
[Table]
pBE vs. down-DE -29.
[Table]
ABLES WITH
[Table]
TANDARD
[Table]
The tables in this section provide the mean values and corresponding standard errors for many of the results depicted in ﬁgures throughout Section 5 and Appendix H.
U nd e r r e
v
i e w a s a c on f e r e n ce p a p e r a t
I C L
R
[Table] [Table]
B/ pBE
[Table] [Table]
[Table]
U nd e r r e
v
i e w a s a c on f e r e n ce p a p e r a t
I C L
R
[Table]
L/
[Table]
-MoE
[Table] [Table]
[Table] [Table]
[Table]
M
NLL ↓
E RROR ↓
ECE
[Table] [Table]
[Table] [Table]
U nd e r r e
v
i e w a s a c on f e r e n ce p a p e r a t
I C L
R
[Table] [Table]
U nd e r r e
v
i e w a s a c on f e r e n ce p a p e r a t
I C L
R
[Table] [Table]
[Table]
n
[Table]
C L
[Table] [Table]
[Table] [Table]
[Table]
I F ROM B ATCH E NSEMBLES TO S PARSE M O E S
Wen et al.
(2019) have shown that, given a batch of B inputs X ∈ R B × P , a single forward pass can efﬁciently compute the predictions of all the ensemble members { f ( X ; θ m ) }
Mm =1 .
By appropriately tiling the inputs of the network X tiled ∈ R ( M · B ) × P by a factor M , each internal operation per ensemble member can then be vectorized.
We take the previous example of a dense layer with parameters U ∈ R D
×
L
and we assume the layer receives the tiled inputs { H m } Mm =1 where H m ∈ R B × D .
We need to compute for each ensemble member H m U m
=
H m
[ U ◦ ( r m s (cid:62) m )] .
Denoting by h i,m ∈ R D the i -th input in H m , we have h (cid:62) i,m U m =
E (cid:88) e =
g e ( h i,m ) · expert e ( h i,m ) with M = E, (cid:26) g e ( h i,m )
= 1 if e = m, g e ( h i,m ) =
otherwise (5) and expert e ( z ) =
z (cid:62)
[ U ◦ ( r e s (cid:62) e )] .
Although (5) may appear as a convoluted way of writing the operations in batch ensembles, it unveils a connection with (1).
Indeed, operations in batch ensembles can be seen as a speciﬁc sparse MoE, e.g., with binary routing weights depending only on the position in the tiled inputs.
While Wen et al.
(2019) primarily tiled the inputs for the sake of efﬁciency, it also induces some form of conditional computation, an insight that we exploit here.
J N EW R ESULTS FOR THE R EBUTTAL J. N EW R ESULTS FOR THE F EATURE - LEVEL VERSUS P REDICTION - LEVEL A BLATION
Table 25 shows our updated and expanded results for the feature-level versus prediction-level ablation, described in Section 3.2.
The original conclusion—i.e., ensembling at the prediction level is helpful for calibration but that the naive multi-head approach provides worse error, NLL, and diversity—seems to hold for the K = 8 case.
Note that the results for K = 2 and K = 4 differ slightly from the original version in Table 2 due to a change in experimental setup.
In the original version, we matched K for the up- and downstream models in the multi-head case (i.e., the pretrained and ﬁne-tuned models both use either K =or K = 4 ).
We did this to be as fair as possible to the multi-head model.
While it is known, from Riquelme et al.
(2021), that V-MoE models are fairly robust to the choice of upstream model (e.g., K UPSTREAM can be set to 2 and we can just change K DOWNSTREAM , which is the approach we take throughout our paper–see Appendix C), this was not known for the multi-head model.
Unfortunately, due to time constraints we are unable to train an upstream model with K = 8 .
Thus, for the updated results we have opted for a single upstream model with K = 2 in all cases.
As a side observation, we can notably observe that the multi-head variant is more sensitive than V-MoE to the choice of the upstream model, e.g., the performance worsens from ( K UPSTREAM , K DOWNSTREAM )
= (4 , 4) to ( K UPSTREAM , K DOWNSTREAM ) =
[Table]
J. N EW R ESULTS FOR THE S TATIC VERSUS A DAPTIVE A BLATION
Figure 20 repeats the experiment of Figure 2a with two new random seeds.
The new results show a smoother improvement in LL as K increases.
In particular, we see that there is nothing special about K = 5 or K = 6 .
This indicates that the anomalous drop at K = 6 in the original results was simply due to noise in the training of the upstream model.
We note that, like Riquelme et al.
(2021), we have observed that training noise is more pronounced in the smallest (i.e., S/32) models.
[Figure]
J. N EW L OAD B ALANCING A BLATION
Figure 21 shows a new experiment exploring the effect of the load balancing loss on diversity and predictive performance.
We see that for both V-MoE and pBE the predictive performance, as measured by NLL, Error, and ECE are largely insensitive to the strength of the load balancing loss.
Similarly, the diversity of the predictions of pBE—as measured by the KL—mildly depends on that regularization.
Only when the load balancing loss strength becomes excessively large (4 orders of magnitude larger than standard values) do all the performance metrics plummet.
This hyperparameter does not allow us to increase the diversity and thereby the predictive performance of our models.
J. N EW E XPERT I NITIALIZATION A BLATION
Figure 22 shows a new experiment exploring the inﬂuence of the initialization of the experts.
In particular, we add Gaussian noise with varying standard deviations to the initial weights of the expert MLPs.
We notably show that more diverse initializations (larger standard deviations) do not translate to any clear performance gain.
Note that this new experiment takes place upstream , since downstream the experts are already initialized (we just ﬁne-tune them from the upstream checkpoint).
J.
[Equation]
[Table]
Table 26 compares the parameter counts for ViT and V-MoE/pBE models in each ViT family.
[Table]
[Figure]
e
[Figure]
i
[Figure]
[Figure]