Speeding Up BatchBALD: A k-BALD Family of Approximations for Active Learning

TL;DR

This paper introduces k-BALD, an efficient approximation of BatchBALD for active learning with Bayesian neural networks, significantly reducing computation time while maintaining performance, and includes a dynamic k selection method for larger datasets.

Contribution

The paper presents k-BALD, a novel approximation method for BatchBALD that is computationally more efficient and adaptable through a dynamic k selection strategy.

Findings

k-BALD is significantly faster than BatchBALD on MNIST.

k-BALD maintains similar performance to BatchBALD.

Dynamic k selection improves efficiency for larger datasets.

Abstract

Active learning is a powerful method for training machine learning models with limited labeled data. One commonly used technique for active learning is BatchBALD, which uses Bayesian neural networks to find the most informative points to label in a pool set. However, BatchBALD can be very slow to compute, especially for larger datasets. In this paper, we propose a new approximation, k-BALD, which uses k-wise mutual information terms to approximate BatchBALD, making it much less expensive to compute. Results on the MNIST dataset show that k-BALD is significantly faster than BatchBALD while maintaining similar performance. Additionally, we also propose a dynamic approach for choosing k based on the quality of the approximation, making it more efficient for larger datasets.