PaperlayerView exact revision on arXiv

Source-linked AI summary

BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning

Yeming Wen, Dustin Tran, Jimmy Ba

arXiv:2002.06715v2cs.LGstat.ML

TL;DR

Neural ensembles improve accuracy and uncertainty but incur costs that grow linearly with the number of networks. BatchEnsemble shares weights and uses rank-one member-specific factors to reduce these costs while enabling efficient parallelization. Across benchmark, uncertainty, and lifelong-learning tasks, it achieves competitive performance and scales to 100 sequential ImageNet tasks.

  • Problem

    Neural ensembles provide useful accuracy and uncertainty estimates, but their training, testing, and memory costs grow linearly with ensemble size, limiting practical and lifelong-learning applications.

  • Method

    BatchEnsemble forms each member’s weights from a shared weight and a member-specific rank-one fast weight using a Hadamard product, with parallel updates within a mini-batch.

  • Results

    BatchEnsemble delivers competitive accuracy and uncertainty across classification, translation, and out-of-distribution tasks, comparable Split-CIFAR-100 performance to progressive neural networks, and scaling to 100 sequential Split-ImageNet tasks.

  • Takeaways & Limitations

    BatchEnsemble enables efficient ensembles and lifelong learning by reducing the computational and memory bottlenecks of typical ensembles.

  • Takeaways & Limitations

    Rank-one task perturbations may limit expressiveness for substantially varied tasks, and the shared weight is trained only on the first task, restricting subsequent transfer.

Abstract

from arXiv · show

Ensembles, where multiple neural networks are trained individually and their predictions are averaged, have been shown to be widely successful for improving both the accuracy and predictive uncertainty of single neural networks. However, an ensemble's cost for both training and testing increases linearly with the number of networks, which quickly becomes untenable. In this paper, we propose BatchEnsemble, an ensemble method whose computational and memory costs are significantly lower than typical ensembles. BatchEnsemble achieves this by defining each weight matrix to be the Hadamard product of a shared weight among all ensemble members and a rank-one matrix per member. Unlike ensembles, BatchEnsemble is not only parallelizable across devices, where one device trains one member, but also parallelizable within a device, where multiple ensemble members are updated simultaneously for a given mini-batch. Across CIFAR-10, CIFAR-100, WMT14 EN-DE/EN-FR translation, and out-of-distribution tasks, BatchEnsemble yields competitive accuracy and uncertainties as typical ensembles; the speedup at test time is 3X and memory reduction is 3X at an ensemble of size 4. We also apply BatchEnsemble to lifelong learning, where on Split-CIFAR-100, BatchEnsemble yields comparable performance to progressive neural networks while having a much lower computational and memory costs. We further show that BatchEnsemble can easily scale up to lifelong learning on Split-ImageNet which involves 100 sequential learning tasks.

1 INTRODUCTION

Neural ensembles improve predictive performance and uncertainty but become impractical as training, testing, and memory costs grow with ensemble size. BatchEnsemble addresses this bottleneck through shared weights and rank-one member-specific weights while retaining competitive empirical performance across tasks.

  • Ensembles improve neural-network performance and predictive uncertainty by combining outputs from multiple models.
  • Ensemble training and testing costs increase linearly with ensemble size because each member requires separate computation and weights.
  • BatchEnsemble generates each member’s weights by combining one shared weight matrix with a member-specific rank-one fast weight through a Hadamard product.
  • BatchEnsemble supports parallelization across devices and within devices, allowing multiple ensemble members to update simultaneously for a mini-batch.
  • BatchEnsemble achieves a favorable accuracy, running-time, and memory trade-off across image classification, machine translation, calibration, and uncertainty tasks.

2 BACKGROUND

The background reviews ensemble methods, uncertainty evaluation, and lifelong learning. It highlights calibration and out-of-distribution uncertainty assessment, while identifying catastrophic forgetting and memory growth as central lifelong-learning challenges.

  • Bagging combines different models through strategies such as averaging or majority voting to improve generalization performance.
  • Research on ensembles includes reducing test-time and training-time costs through compression, distillation, and checkpoint-based combinations.
  • Implicit ensembles such as dropout use weight-sharing subnetworks and are generally inexpensive during training and testing.
  • Uncertainty quality is assessed using calibration, out-of-distribution generalization, and contextual-bandit performance.
  • Lifelong learning presents catastrophic forgetting because models tend to lose knowledge from earlier tasks after training on subsequent tasks.
  • Capacity-expansion methods such as progressive neural networks prevent forgetting by preserving prior-task networks but incur significant memory consumption as tasks accumulate.

3 METHODS

BatchEnsemble generates ensemble members from shared weights and rank-one member-specific factors, then computes them in parallel within a device. The design reduces overhead and extends to lifelong learning by assigning task-specific fast weights while preserving prior task models.

  • 3.1 BATCHENSEMBLE: BatchEnsemble generates each member’s weight by elementwise multiplying shared weights with a member-specific rank-one fast-weight matrix.Each member has trainable vectors r_i and s_i, whose outer product forms the rank-one factor.
  • 3.1 BATCHENSEMBLE: Vectorization computes multiple ensemble members’ forward passes in parallel within a device using mini-batch matrix operations.Inputs can be divided into M sub-batches so each receives its corresponding ensemble weight.
  • 3.1 BATCHENSEMBLE: At testing, repeating the input batch lets all members process the same examples in one forward pass before predictions are averaged.The effective batch size becomes B · M for B inputs and M members.
  • 3.2 COMPUTATIONAL COST: 10% more parameters are required for a size-4 BatchEnsemble ResNet-32, compared with 3X more for a naive ensemble.The additional parameters are member-specific vectors rather than full weight matrices.
  • 3.3 BATCHENSEMBLE AS AN APPROACH TO LIFELONG LEARNING: BatchEnsemble applies lifelong learning by assigning one ensemble member to each task and training only the relevant fast weights on subsequent tasks.Previously learned task models remain unchanged, preventing catastrophic forgetting without storing previous-task data.
  • 3.3 BATCHENSEMBLE AS AN APPROACH TO LIFELONG LEARNING: Rank-1 task perturbations may lack expressiveness for substantially varied tasks, while shared weights trained only on the first task restrict transfer between later tasks.The paper leaves lateral connections between previously learned task features for future work.

4 EXPERIMENTS

BatchEnsemble is evaluated across lifelong learning, machine translation, classification, and corrupted-data calibration tasks. It generally matches or improves relevant baselines while reducing computational or memory costs.

  • 4.1 LIFELONG LEARNING: BatchEnsemble achieves comparable accuracy to PNN on Split-CIFAR100 with 4X speed-up and 50X less memory consumption.The comparison evaluates average accuracy, forgetting, and cost over 20 tasks.
  • 4.1 LIFELONG LEARNING: 100 sequential Split-ImageNet tasks require 30M BatchEnsemble parameters versus 25M for a vanilla ResNet-50.PNN cannot learn all 100 tasks because of its memory consumption.
  • 4.2 MACHINE TRANSLATION: 1.5X faster convergence is achieved by the big BatchEnsemble Transformer than by a single big Transformer on WMT16 English-German.BatchEnsemble also gives lower validation perplexity, although it does not achieve a better BLEU score than the single model.
  • 4.3 CLASSIFICATION: BatchEnsemble reaches better accuracy than single model and MC-dropout on CIFAR-10/100 while remaining faster than sequentially training individual models.The BatchEnsemble and MC-dropout experiments use 375 epochs, 50% more iterations than single-model training.
  • 4.4 CALIBRATION ON CORRUPTED DATASET: BatchEnsemble and dropout ensemble achieve comparable accuracy across corruption skew intensities, while BatchEnsemble has better average calibration as skew increases.Combining BatchEnsemble with dropout improves accuracy and calibration while retaining efficient memory use.

5 DIVERSITY ANALYSIS

The diversity of BatchEnsemble members depends on available training data. With limited data, BatchEnsemble can approach naive-ensemble diversity, helping explain stronger gains on CIFAR-100 than CIFAR-10.

  • 5 DIVERSITY ANALYSIS: BatchEnsemble has intermediate diversity between naive and dropout ensembles, matching the reported ordering of their accuracies.The diversity metric measures disagreement among ensemble members on the test set.
  • 5 DIVERSITY ANALYSIS: Limited training data leaves parameters near initialization, so initialization diversity substantially determines ensemble-member diversity.The paper uses this intuition to explain why rank-1 perturbations can produce diverse members with fewer data.
  • 5 DIVERSITY ANALYSIS: BatchEnsemble diversity is comparable to naive-ensemble diversity when training data is limited.CIFAR-100 has 500 training points per class, compared with 5000 for CIFAR-10.

6 CONCLUSION

BatchEnsemble is presented as an efficient method for ensembling and lifelong learning. It removes typical ensemble computation and memory bottlenecks while supporting lifelong learning across up to 100 tasks.

  • BatchEnsemble is an efficient method for ensembling and lifelong learning.
  • BatchEnsemble removes the computation and memory bottlenecks of typical ensemble methods.
  • BatchEnsemble enables lifelong learning on up to 100 tasks.

A DATASET DETAILS

The experiments use CIFAR-10/100, WMT translation datasets, and sequential Split-CIFAR100 and Split-ImageNet tasks. The datasets are partitioned and preprocessed using task-specific procedures described for each benchmark.

  • CIFAR-10 and CIFAR-100 each contain 50K training and 10K test images with 32x32 pixels.
  • WMT experiments use English-German and English-French translation datasets with standard preprocessing and 37K shared BPE sub-word units.
  • Split-CIFAR100 randomly divides CIFAR-100 into T tasks, each containing 100/T classes and a task-specific final classifier.
  • Split-ImageNet randomly divides ImageNet into T tasks, each containing 1000/T classes and its own final linear classifier.

B IMPLEMENTATION DETAILS

Implementation details cover weight regularization, diversity initialization, translation training, and classification settings. BatchEnsemble uses shared and fast weights while training ensemble members with task-specific configurations.

  • Weight Decay: BatchEnsemble regularizes the shared weight, optionally using the mean weight over ensemble members, while leaving fast weights unregularized.
  • Diversity Encouragement: Random-sign fast-weight initialization and separate input sub-batches provide sufficient diversity without requiring an additional diversity loss.
  • Machine Translation: Translation models are trained for 100K steps for the base Transformer and 180K steps for the big Transformer on 4 NVIDIA P100 GPUs.
  • Classification: Classification uses mini-batches of 128, ensemble size 4, and 32 training examples per ensemble member.

C PREDICTIVE UNCERTAINTY

The uncertainty evaluation examines out-of-distribution predictive entropy and calibration using ECE. BatchEnsemble produces higher uncertainty on unseen classes than a single model and more calibrated predictions than single neural networks.

  • BatchEnsemble is evaluated on out-of-distribution examples and Expected Calibration Error.
  • Out-of-distribution uncertainty: On unseen classes, BatchEnsemble and naive ensembles produce higher predictive uncertainty than a single model.
  • Out-of-distribution uncertainty: Predictive entropy is compared between known CIFAR-10 classes and unknown CIFAR-100 classes.
  • Calibration: Expected Calibration Error groups predictions into M confidence bins and measures the difference between averaged accuracy and confidence.
  • Calibration: BatchEnsemble makes more calibrated predictions than single neural networks.

D UNCERTAINTY ON BANDITS

BatchEnsemble is evaluated as an uncertainty model in contextual bandits, where reliable uncertainty supports exploration and exploitation. It achieves strong average performance across bandit problems, including the best mean value at ensemble size 8.

  • Reliable uncertainty estimation is central to contextual bandits because decisions must balance exploration and exploitation.
  • Ensemble members support Thompson sampling by providing alternative policies from which the algorithm samples.
  • BatchEnsemble with ensemble size 8 achieves the best mean value on the contextual bandits task.
  • BatchEnsemble with ensemble sizes 4 and 8 both outperform Dropout in average bandit performance.
  • On CIFAR-10-C, BatchEnsemble offers a promising trade-off among accuracy, uncertainty, and cost, while combining it with dropout improves uncertainty prediction.

E.1 DIVERSITY METRIC

The diversity analysis compares BatchEnsemble with naive and dropout ensembles across accuracy and prediction-disagreement measures. BatchEnsemble retains substantial diversity, particularly when training data is limited, where it is comparable to naive ensembles and exceeds dropout ensembles.

  • The diversity metric measures the fraction of test examples on which two models disagree, ranging from 0 for identical predictions to 1 for complete disagreement.
  • BatchEnsemble’s diversity falls between naive and dropout ensembles when trained on the full CIFAR-10 dataset, matching the reported accuracy ordering.
  • Figure 8 plots validation accuracy against diversity for BatchEnsemble, naive ensemble, and dropout ensemble models trained with 100%, 50%, 20%, or 10% of the data.
  • E.2 DIVERSITY METRIC ON PARTIAL TRAINING SET: With 50% of training data, BatchEnsemble achieves diversity comparable to naive ensemble and significantly higher than dropout ensemble.
  • E.2 DIVERSITY METRIC ON PARTIAL TRAINING SET: With 20% and 10% of training data, BatchEnsemble remains comparable to naive ensemble and more diverse than dropout ensemble.
  • E.2 DIVERSITY METRIC ON PARTIAL TRAINING SET: Because diversity reflects epistemic uncertainty, the results indicate that BatchEnsemble members become more diverse than dropout members under limited training data.

F COMPARISON TO NAIVE ENSEMBLE OF SMALL MODELS

Under a matched memory budget, BatchEnsemble is compared with a naive ensemble of smaller models and with member-level prediction visualizations. BatchEnsemble achieves higher accuracy than the smaller-model naive ensemble while preserving ensemble-like prediction diversity.

  • Under the same memory budget, a naive ensemble of four ResNet14x4 models achieves lower accuracy than BatchEnsemble.The smaller-model naive ensemble still has roughly 10% memory overhead relative to BatchEnsemble.
  • BatchEnsemble is therefore the better choice than the small-model naive ensemble under the reported matched-memory setup.
  • The visualized examples focus on cases where a single model is confidently wrong but the ensemble predicts correctly, despite some members also making mistakes.
  • On selected CIFAR-10 examples, BatchEnsemble preserves prediction diversity among members similarly to a naive ensemble.
Loading 2002.06715v2…