[Paper Review] BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning

Yeming Wen, Dustin Tran+1 more|arXiv (Cornell University)|Feb 16, 2020

Domain Adaptation and Few-Shot Learning54 references127 citations

TL;DR

BatchEnsemble proposes a parameter-efficient ensemble method where each member’s weights are a Hadamard product of shared weights and a per-member rank-1 perturbation, enabling fast, memory-efficient ensembles and scalable lifelong learning.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for effective ensembling with reduced computational and memory costs.
• Introduce BatchEnsemble as a parameter-efficient alternative to traditional ensembles.
• Demonstrate BatchEnsemble’s performance on classification, translation, and lifelong learning benchmarks.
• Show that BatchEnsemble provides calibrated predictions and competitive uncertainty estimates.

Proposed method

• Define each ensemble member’s weight as the Hadamard product W_i = W ∘ (r_i s_i^T) where W is shared and r_i, s_i are per-member vectors.
• Vectorize computations so that multiple ensemble members are updated in parallel within a single mini-batch, enabling device-level and intra-device parallelism (Y = φ(((X ∘ R) W) ∘ S)).
• Use a testing strategy that averages predictions across ensemble members by expanding the mini-batch to B·M so all members process the same inputs in one forward pass.
• Apply BatchEnsemble to lifelong learning by training a shared W and a single pair of fast weights for the first task and only training new fast weights for subsequent tasks.
• Evaluate uncertainty calibration and out-of-distribution performance, comparing to MC-dropout and naive ensembles.

Experimental results

Research questions

• RQ1 Can BatchEnsemble achieve competitive accuracy and uncertainty estimates with substantially lower memory and computation than traditional ensembles?
• RQ2 How well does BatchEnsemble scale to lifelong learning with many sequential tasks?
• RQ3 What is the impact of BatchEnsemble on calibration and out-of-distribution robustness?
• RQ4 How does BatchEnsemble perform across vision, language, and translation tasks compared to standard baselines?

Key findings

• BatchEnsemble achieves similar accuracy and uncertainty to traditional ensembles while dramatically reducing cost:test-time speedup and memory reduction of about 3x at an ensemble size of 4.
• In lifelong learning, BatchEnsemble attains competitive accuracy to progressive neural networks with far lower memory and computation, scalable up to 100 sequential tasks.
• BatchEnsemble provides well-calibrated predictions on corrupted/corrupted-like data, with competitive calibration compared to dropout ensembles and potential gains when combined with dropout.
• Across CIFAR-10/100, WMT14 EN-DE/EN-FR, and out-of-distribution tasks, BatchEnsemble demonstrates strong performance and faster convergence in Transformer-based settings (encoder self-attention layers).
• Diversity analysis shows BatchEnsemble can achieve diversification close to naive ensembles with limited training data, while benefiting from larger networks.