[Paper Review] BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning

Andreas Kirsch, Joost van Amersfoort+1 more|arXiv (Cornell University)|Jun 19, 2019

Machine Learning and Algorithms34 references192 citations

TL;DR

BatchBALD provides a tractable, greedy 1-1/e approximate acquisition function that jointly selects diverse batches by modeling the mutual information between a batch of labels and model parameters, improving data efficiency over BALD in deep Bayesian active learning.

ABSTRACT

We develop BatchBALD, a tractable approximation to the mutual information between a batch of points and model parameters, which we use as an acquisition function to select multiple informative points jointly for the task of deep Bayesian active learning. BatchBALD is a greedy linear-time $1 - \frac{1}{e}$-approximate algorithm amenable to dynamic programming and efficient caching. We compare BatchBALD to the commonly used approach for batch data acquisition and find that the current approach acquires similar and redundant points, sometimes performing worse than randomly acquiring data. We finish by showing that, using BatchBALD to consider dependencies within an acquisition batch, we achieve new state of the art performance on standard benchmarks, providing substantial data efficiency improvements in batch acquisition.

Motivation & Objective

• Motivate data efficiency in deep learning via active learning with Bayesian neural networks.
• Extend BALD to jointly select batches of data points by accounting for dependencies and diversity within the batch.
• Provide a tractable algorithm with guarantees that yields diverse, informative batches.
• Demonstrate improved data efficiency and training speed in standard image datasets.

Proposed method

• Define BatchBALD as the mutual information between a batch of labels and model parameters given a data batch.
• Propose a greedy, 1-1/e approximate algorithm to select batches efficiently (Algorithm 1).
• Show BatchBALD reduces overlap by using the union of batch information via μ*(⋃ yi ∩ ω) instead of summing independent BALD scores.
• Derive a Monte Carlo estimator to compute BatchBALD scores with k samples of ω and MC estimates for joint entropies.
• Implement efficient computation by factorizing p(y1:n|ω) and using caching/matrix operations to handle larger batches.
• Provide complexity bounds showing BatchBALD is tractable (O(b c min{c^b, m} |D_pool| k)).

Experimental results

Research questions

• RQ1 Does BatchBALD improve data efficiency by accounting for joint information and diversity within acquisition batches compared to naive BALD?
• RQ2 Can a greedy 1-1/e approximate BatchBALD batch selection achieve near-optimal performance with tractable computation?
• RQ3 How does BatchBALD perform on standard AL benchmarks (MNIST, EMNIST) and transfer learning (CINIC-10) relative to BALD and random acquisition?
• RQ4 What is the impact of batch size on diversity and training efficiency in deep Bayesian active learning?
• RQ5 What are practical implementation considerations (MC dropout consistency, reinitialization, and computation) that affect BatchBALD performance?

Key findings

• BatchBALD outperforms BALD when acquisition size increases, maintaining strong performance with larger batches.
• BatchBALD yields more diverse acquired data, and incurs higher diversity in predicted class distributions (EMNIST results).
• On MNIST, BatchBALD with acquisition size 10 performs close to the ideal BALD with acquisition size 1 and is faster overall than BALD with size 1.
• BatchBALD improves data efficiency across MNIST, EMNIST, and CINIC-10 transfer-learning experiments, compared to BALD and random acquisition.
• A naive BALD batch can perform poorly in replicated data settings, whereas BatchBALD remains robust to duplicates and replication in data.