Source-linked AI summary

Sparse MoEs meet Efficient Ensembles

James Urquhart Allingham, Florian Wenzel, Zelda E Mariet, Basil Mustafa, Joan Puigcerver, Neil Houlsby, Ghassen Jerfel, Vincent Fortuin, Balaji Lakshminarayanan, Jasper Snoek, Dustin Tran, Carlos Riquelme Ruiz, Rodolphe Jenatton

arXiv:2110.03360v2cs.LGcs.CVstat.ML

TL;DR

Naive neural-network ensembles improve reliability but become increasingly infeasible as models grow, motivating efficient combinations of ensembles and sparse MoEs. E^3 combines static ensembling with sparse-MoE adaptive routing by tiling representations and partitioning experts into disjoint member-specific subsets. E^3 improves few-shot error, likelihood, and calibration error over sparse MoEs, matches deep ensembles, and uses 30% to 45% fewer FLOPs.

  • Problem

    Naive neural-network ensembles improve reliability but become increasingly infeasible as models grow, motivating efficient combinations of ensembles and sparse MoEs.

  • Method

    E^3 combines static ensembling with sparse-MoE adaptive routing by tiling representations and partitioning experts into disjoint member-specific subsets.

  • Results

    E^3 improves few-shot error, likelihood, and calibration error over sparse MoEs, matches deep ensembles, and uses 30% to 45% fewer FLOPs.

  • Takeaways & Limitations

    E^3 yields accuracy, log-likelihood, few-shot learning, robustness, and uncertainty calibration improvements over several challenging baselines across a range of benchmarks.

  • Takeaways & Limitations

    The study focuses on downstream fine-tuned models and assumes access to pre-trained ViT and V-MoE checkpoints.

Abstract

from arXiv · show

Machine learning models based on the aggregated outputs of submodels, either at the activation or prediction levels, often exhibit strong performance compared to individual models. 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 beneficial. This includes a comprehensive evaluation of sparse MoEs in uncertainty related benchmarks. Then, we present Efficient Ensemble of Experts (E$^3$), a scalable and simple ensemble of sparse MoEs that takes the best of both classes of models, while using up to 45% fewer FLOPs than a deep ensemble. Extensive experiments demonstrate the accuracy, log-likelihood, few-shot learning, robustness, and uncertainty improvements of E$^3$ over several challenging vision Transformer-based baselines. E$^3$ not only preserves its efficiency while scaling to models with up to 2.7B parameters, but also provides better predictive performance and uncertainty estimates for larger models.

1 Introduction

Sparse MoEs reduce computation through input-dependent expert routing, while ensembles improve reliability by aggregating predictions but incur substantial cost. The paper studies their complementarity and introduces E^3, which improves sparse-MoE performance while using fewer FLOPs than deep ensembles.

  • The work targets more reliable large-scale fine-tuned models, motivated by the calibration and robustness benefits of ensembles under dataset shifts.These benefits are relevant to safety-critical applications such as medical diagnosis and self-driving cars.
  • Sparse MoEs and ensembles provide additive benefits because adaptive computation and static prediction combination are orthogonal.The performance–compute frontier can be explored by varying ensemble size and MoE sparsity.
  • Naive ensembles of sparse MoEs achieve the best predictive performance across tasks favoring either sparse MoEs or ensembles, but are computationally expensive.The benchmarking includes an evaluation of sparse MoEs on uncertainty-related vision tasks.
  • E^3 is a scalable, simple ensemble approach tailored to sparse MoEs and can scale to models with up to 2.7B parameters.The method requires minor implementation changes and can be fine-tuned from standard sparse-MoE checkpoints.
  • E^3 improves few-shot error, likelihood, and calibration error over sparse MoEs while matching deep-ensemble performance with 30% to 45% fewer FLOPs.

2 Preliminaries

The preliminaries define ViT-based classifiers, sparse MoEs, V-MoEs, and neural-network ensembles. These models differ in how they allocate computation, combine submodels, and trade predictive performance against inference cost.

  • Vision Transformers: ViT processes images as sequences of linearly embedded patches and is used throughout the paper as the underlying vision model.The paper focuses on transfer-learning settings in which ViT models are fine-tuned downstream.
  • Sparse MoEs: Sparse MoEs activate only a small, input-dependent subset of parameters, allowing parameter growth while keeping computational cost constant.Their experts are subparts of the network activated on a per-input basis.
  • V-MoEs: V-MoE places sparse experts in Transformer MLP layers and can outperform dense ViT counterparts at the same computational cost.The paper adopts a “last-n” setting that places only a few MoE layers near the Transformer’s end to reduce training cost.
  • Sparse MoEs: Sparse MoE routing uses a noisy topK(softmax(W h + σε)) function that retains only the K largest expert weights.The router parameters W are trained jointly with the rest of the network, and σ = 1/E is used.
  • Ensembles: Ensembles average the probability vectors of M models, typically combining predictions statically rather than using conditional computation.Batch ensembles reduce cost by sharing parameters while producing and averaging M member outputs.

3 Sparse MoEs meet Ensembles

Sparse MoEs and neural-network ensembles provide complementary adaptive and static aggregation, with additive benefits when combined. Increasing adaptive combination before static ensembling offers a favorable compute–performance trade-off.

  • Increasing ensemble size M improves sparse-MoE performance for every tested sparsity level K, showing static aggregation adds benefits beyond adaptive routing.
  • Larger ensembles of smaller MoEs outperform fewer larger experts: (M = 4, E = 8) is better than (M = 2, E = 16), which is better than (M = 1, E = 32) on NLL.
  • Ensembling benefits ViT and V-MoE similarly across model scales, indicating the gain from static prediction-level combination is not specific to adaptive MoEs.
  • The compute cost grows faster with M than K, so increasing adaptive combination K before ensemble size M is more advantageous under a FLOPs trade-off.
  • Simple sparse-MoE ensembling achieves strong predictive performance but loses computational efficiency, motivating an efficient combination of both approaches.

4 Efficient Ensemble of Experts

Efficient Ensemble of Experts (e3) jointly learns smaller sparse MoEs by combining disjoint expert subsets with tiled representations. The ablations show that both mechanisms, especially prediction diversity from partitioning, are needed for strong performance.

  • 4.1 The Architecture: e3 combines smaller sparse MoEs with shared non-expert layers, using disjoint expert subsets and tiled representations to create efficient ensemble members.Tiling gives each input M representations, while separate routers assign them to their respective expert subsets.
  • 4.2 Ablation Studies: Partitioning and Tiling: Partitioning alone is not competitive with e3, while tiling alone produces similar predictions and performs worse than e3 across all metrics.Without tiling, partitioning does not output multiple predictions; without partitioning, tiling can route representations to identical experts.
  • 4.2.3 Tiling with Increasing Parameter Sharing: Increasing expert overlap decreases prediction diversity and correspondingly worsens NLL, classification error, and ECE, supporting diversity as a key ingredient of e3’s performance.The same trends hold for K = 1 and K = 2.
  • 4.2.3 Tiling with Increasing Parameter Sharing: A simple Multi-pred variant underperforms e3: it cannot produce diverse predictions, has worse classification error, and only slightly improves NLL over V-MoE for K = 2.For K = 4, Multi-pred is worse than K = 2 on NLL, classification error, and diversity.
  • 4.3 Comparison with Alternative Efficient Ensembling Strategies: Common efficient ensembling strategies such as batch ensembles, MC Dropout, and MIMO provide alternative comparisons, but the supplied passage does not report their complete Table 5 outcomes.

5 Evaluation

The evaluation finds that ensembling benefits V-MoE across predictive and uncertainty metrics, while E^3 is generally Pareto efficient, especially at larger scales, with inconsistent OOD detection.

  • 5.1 V-MoE vs. ViT: Ensembling provides V-MoE gains comparable to ViT across few-shot performance, NLL, ECE, OOD detection, and distribution-shift robustness.A downstream ensemble of four H/14 V-MoEs reaches 88.8% ImageNet accuracy.
  • 5.2 Efficient Ensemble of Experts: E^3 is on or near the Pareto frontiers for NLL and 10-shot classification error, with similar conclusions for CIFAR10/100 NLL.These results compare E^3 against V-MoE and ensemble baselines across classification and likelihood metrics.
  • 5.2 Efficient Ensemble of Experts: E^3 improves relative to V-MoE as ViT families grow, reaching up to 8.33% better shifted-dataset NLL at the largest scale.It improves over V-MoE on averaged ImageNet-C, ImageNet-A, and ImageNet-V2 NLL for all families except S/32.
  • 5.2 Efficient Ensemble of Experts: E^3 improves ECE over ViT and V-MoE, although ECE trends are less consistent across model families than other metrics.ViT consistently provides better ECE than V-MoE, while E^3 avoids V-MoE’s poor ECE behavior.
  • 5.2 Efficient Ensemble of Experts: E^3 does not provide consistent OOD detection performance, performing worse for small ViT families but becoming Pareto efficient for larger families on two dataset pairs.The evaluation measures OOD detection using false positive rate at 95% precision on CIFAR10 fine-tuned models.
  • Summary: Across evaluation settings, no single model is best everywhere; E^3 performs well and is usually Pareto efficient, particularly for larger ViT families.V-MoE outperforms ViT in accuracy, while its uncertainty estimates vary by downstream application.

6 Related Work

Related work frames sparse MoEs and ensembles as complementary approaches to combining submodels, while emphasizing ensembles’ uncertainty benefits and computational cost.

  • Mixture of Experts: Sparse MoEs route each input to only a few experts, enabling model scaling while keeping prediction time constant.MoEs combine expert outputs in an input-dependent way; sparse variants select only a subset of experts.
  • Ensembles: Ensembles combine multiple models to improve generalization and uncertainty estimation, especially when members make complementary errors.Standard ensembles are inefficient because each member can itself be computationally expensive.
  • Bayesian Neural Networks: Bayesian neural networks provide an alternative route to uncertainty quantification by treating weights as random variables and marginalizing predictions.Because neural-network weight posteriors are intractable, BNNs require approximate inference methods.

7 Conclusions and Future Work

The paper concludes that sparse MoEs and ensembles are symbiotic, with E^3 improving multiple performance and uncertainty measures while retaining efficiency, especially for large models.

  • 7 Conclusions and Future Work: E^3 improves accuracy, log-likelihood, few-shot learning, robustness, and uncertainty calibration over challenging baselines across benchmarks.The authors describe E^3 as an example of mutual benefits between sparse MoEs and ensembles.
  • 7 Conclusions and Future Work: The study’s experiments focus on downstream fine-tuned computer-vision models, leaving upstream and from-scratch extensions as future work.Preliminary from-scratch results are provided, while broader text-model applications are suggested as a future direction.
  • 7 Conclusions and Future Work: E^3 provides larger relative improvements over V-MoE for larger models.This scaling result is documented in Appendix D.
  • 7 Conclusions and Future Work: The appendices extend evaluation to additional datasets, metrics, model sizes, efficient ensemble approaches, and sensitivity analyses.They also report implementation details, experimental settings, parameter counts, FLOPs, and standard errors.
  • 7 Conclusions and Future Work: The experimental setup follows established ViT and V-MoE specifications, including JFT-300M pretraining and specified fine-tuning choices.The appendix records model scales, training schedules, image resolution, batch size, gradient clipping, and V-MoE capacity settings.

A.5 Details about the (Linear) Few-shot Evaluation

The linear few-shot evaluation probes representations learned after fine-tuning across eight datasets, using either joint or disjoint ensemble features, which perform similarly.

  • A.5 Details about the (Linear) Few-shot Evaluation: Linear few-shot evaluation trains a classifier from S-dimensional features produced by the representation ϕ(·; θ′), assessing representation quality.The protocol predicts one-hot target labels from learned feature vectors.
  • A.5 Details about the (Linear) Few-shot Evaluation: The study evaluates fine-tuned representations on eight few-shot datasets rather than ImageNet, using 1, 5, 10, and 25 shots per class.Reported accuracy is averaged across datasets; weighted averaging gives similar conclusions to the standard average.
  • A.5 Details about the (Linear) Few-shot Evaluation: For ensemble representations, the joint approach concatenates member features into one vector and trains one classifier, whereas the disjoint approach averages separate classifiers’ predictions.The two approaches perform similarly, and the joint method is used because it simplifies implementation.

C.1 Training Loss

The paper trains efficient ensembles with an average member-wise cross-entropy and balances expert routing separately within each partition. Its efficiency relies on delayed tiling, while difficulty-normalized comparisons show e3’s gains improve with model scale.

  • C.1 Training Loss: e3 uses the average ensemble-member cross-entropy rather than ensemble cross-entropy, which was observed to generalize worse.
  • C.2 Auxiliary Losses: Expert balancing is implemented by averaging the auxiliary regularization separately across each partition of E/M experts.The regularizer combines importance and load losses, and the partitioned version was also applied to the comparison ablation.
  • C.3 Memory Requirements versus V-MoE: Delayed tiling keeps e3’s memory increase mild because only the final MoE layers are replicated; for ViT-L with M = 2, activation memory rises 12.5% while parameter memory is unchanged.The last-n configuration has 21 layers before tiling and 3 after it.
  • C.4 Scaling Analysis: e3’s difficulty-normalized NLL improvement grows monotonically with model scale, unlike V-MoE’s, although both methods’ raw improvements decrease for larger families.The comparison is reported in Table 10 for e3 with (K, M) = (1, 2) and V-MoE with K = 1 versus ViT.
  • From Batch Ensembles to Sparse MoEs: Batch ensembles can be viewed as sparse MoEs with binary routing induced by tiled-input positions, motivating e3’s combination of conditional computation and ensemble structure.Unlike batch ensembles, e3 assigns substantially larger member-specific parameter sets through expert subsets and routers.

G Efficient Ensemble Comparisons

Efficient ensemble comparisons show that e3 provides the strongest performance–FLOPs trade-off among the evaluated efficient ensembles. Its advantage is associated with much greater predictive diversity, whereas MIMO V-MoE often fails to improve over standard V-MoE.

  • Batch Ensemble: BE improves ViT’s NLL, classification error, and ECE, but its FLOPs increase makes it less viable than e3.
  • MC Dropout: MC Dropout V-MoE improves ECE while remaining comparable to V-MoE in NLL and classification error, but sampling with M ∈ {2, 4} makes it less favorable than e3 computationally.
  • MIMO V-MoE: MIMO V-MoE does not improve over V-MoE and becomes worse at larger ensemble sizes, while its small KL values indicate insufficient prediction diversity.Applying MIMO to vanilla ViT produces the same trend, suggesting an incompatibility associated with ViT rather than only its expert layers.
  • Overall comparison: e3 offers the best performance-versus-FLOPs trade-offs at (K, M) = (1, 2) and (2, 2), with prediction diversity orders of magnitude larger than competing efficient ensembles.The comparison uses ViT-B/32 models and reports means ± standard errors over 8 replications.
  • Ablations: More diverse expert initialization does not produce a clear performance gain, and excessively strong load-balancing regularization is needed before predictive metrics or diversity change substantially.
  • Prediction diversity: e3’s large gap between individual members and their ensemble resembles deep ensembles because partitioned experts provide many non-shared parameters and high induced diversity.The paper describes e3 as approximating a deep ensemble of smaller V-MoE models.

J Upstream & Downstream versus Downstream-only Ensembles

Upstream ensembles of V-MoE outperform downstream-only ensembles across the reported datasets, metrics, model families, and ensemble sizes. Upstream diversity is especially beneficial for diversity measures and can make a smaller ensemble competitive with a larger downstream ensemble.

  • Cost and setup: Upstream ensembling requires different upstream checkpoints and additional fine-tuning, whereas the downstream comparison uses a single upstream checkpoint.The paper notes that upstream ensembles therefore have more varied members and better performance, but their training cost is very large.
  • Comparison: Upstream deep ensembles outperform downstream-only V-MoE ensembles across ImageNet, ImageNet-C, reported metrics, ViT families, and M ∈ {2, 4}.The comparison concerns V-MoE ensembles with K = 1.
  • Comparison: An upstream ensemble with M = 2 is often better than or equal to a downstream-only ensemble with M = 4.
  • Interpretation: The advantage is especially pronounced on diversity metrics, supporting diversity as the driver associated with improved ensemble performance.

K Preliminary ImageNet Results without Pre-training

Without pre-training, large sparse MoEs are difficult to train because their parameter count causes severe overfitting, but regularized e3 nevertheless outperforms V-MoE on ImageNet NLL and classification error. Additional analyses show that ensemble size, few-shot settings, and OOD behavior influence the relative gains.

  • Training challenge: Large sparse MoEs trained from scratch on ImageNet severely overfit because of their massive parameter count, requiring multiple regularization techniques.
  • Pre-training-free results: e3 outperforms V-MoE in both NLL and classification error for an S/32 backbone trained on ImageNet without pre-training.The comparison is reported in Table 16 over three seeds.
  • Expert allocation: The variation from E is smaller than from K or M, with little change beyond E = 8; this motivates splitting experts across ensemble members.The exact curves are not directly comparable because each parameter sweep averages over different values of the other variables.
  • Robustness and uncertainty: Across extended datasets and metrics, e3’s relative advantage increases for larger ViT families, while larger ensembles can reduce ECE and OOD detection for ViT-S/32.
  • Cumulative effects: For fixed total experts, ensembles with more members and fewer experts per member outperform single larger V-MoEs on NLL, Error, and ECE.The reported ordering is (M = 4, E = 8) better than (M = 2, E = 16), better than (M = 1, E = 32).
  • Few-shot learning: e3’s few-shot performance improves relative to baselines with larger ViT families and more shots, with the same trends under weighted and unweighted aggregation.

L.6 Extended Results for ImageNet

Extended evaluations largely preserve the paper’s main trends: e3 improves calibration and classification, while its distribution-shift advantages are strongest for larger ViT families. Results for larger K and M are mixed and confounded by checkpoint mismatch.

  • e3 tends to be Pareto efficient under distribution shift, whereas V-MoE outperforms ViT for smaller ViT families.
  • e3 improves ECE and classification performance across the extended ImageNet evaluations, while ViT consistently provides better ECE than V-MoE.
  • Additional datasets: For CIFAR10, CIFAR10-C, and CIFAR100, e3 performs relatively better than competing baselines and is Pareto efficient for CIFAR10-C and CIFAR100 NLL, including S/32.
  • Additional datasets: For Flowers and Pets, e3 underperforms on the available smaller ViT families, although L/32 improves over S/32 and B/32.
  • Varying K and M: Increasing K helps V-MoE in most cases but usually hurts ECE; for e3, K = 2 generally helps while larger K or M often hurts or lacks consistent benefit.
  • Varying K and M: The K and M ablations are confounded because models were fine-tuned from upstream V-MoE checkpoints with K = 2, while increasing M also reduces experts per ensemble member.

M.3 Summary for NLL under Distribution Shift

Under distribution shift, e3 generally outperforms V-MoE in NLL across ViT families, with the clearest gains at larger scales, while implementation and checkpoint choices constrain interpretation.

  • e3 outperforms V-MoE in shifted-data NLL for all ViT families except S/32, with improvements reported across ImageNet-C, ImageNet-A, and ImageNet-V2.
  • Ablations: Multi-pred underperforms e3: its classification error is always worse, and its predictions lack the diversity indicated by KL measurements.
Loading 2110.03360v2…