ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition

TL;DR

ALINE is a unified framework that combines Bayesian inference and active data acquisition using a transformer trained with reinforcement learning, enabling instant, accurate inference and strategic data collection.

Contribution

It introduces ALINE, a novel transformer-based approach that jointly optimizes Bayesian inference and active learning through reinforcement learning, adaptable to various tasks.

Findings

ALINE achieves accurate instant inference in regression and Bayesian design tasks.

It efficiently selects informative data points, improving data acquisition strategies.

ALINE can target specific model parameters or predictive tasks effectively.

Abstract

Many critical applications, from autonomous scientific discovery to personalized medicine, demand systems that can both strategically acquire the most informative data and instantaneously perform inference based upon it. While amortized methods for Bayesian inference and experimental design offer part of the solution, neither approach is optimal in the most general and challenging task, where new data needs to be collected for instant inference. To tackle this issue, we introduce the Amortized Active Learning and Inference Engine (ALINE), a unified framework for amortized Bayesian inference and active data acquisition. ALINE leverages a transformer architecture trained via reinforcement learning with a reward based on self-estimated information gain provided by its own integrated inference component. This allows it to strategically query informative data points while simultaneously refining its predictions. Moreover, ALINE can selectively direct its querying strategy towards specific subsets of model parameters or designated predictive tasks, optimizing for posterior estimation, data prediction, or a mixture thereof. Empirical results on regression-based active learning, classical Bayesian experimental design benchmarks, and a psychometric model with selectively targeted parameters demonstrate that ALINE delivers both instant and accurate inference along with efficient selection of informative points.