ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition

TL;DR

ALINE tackles the challenge of simultaneously selecting informative data and performing rapid Bayesian inference under budgets and data constraints. It introduces a transformer-based framework trained with reinforcement learning that uses a reward derived from self-estimated information gain to jointly amortize both inference and data acquisition. The method supports flexible, runtime-targeted acquisition goals, enabling selective querying of parameter subsets or predictive tasks while delivering instant posterior and predictive updates. Empirical results across regression active learning, Bayesian experimental design benchmarks, and psychometric modeling demonstrate fast, accurate inference and efficient data point selection, with notable runtime advantages over non-amortized approaches.

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.

Figures (13)

• Figure 1: Conceptual workflow of Aline, demonstrating its capability to sequentially query informative data points and perform rapid posterior or predictive inference based on the gathered data.
• Figure 2: The Aline architecture. The model takes historical observations (context set), candidate points (query set), and the current inference goal (target set) as inputs. These are transformed by embedding layers and subsequently by transformer layers. Finally, an acquisition head determines the next data point to query, while an inference head performs the approximate Bayesian inference.
• Figure 3: Predictive performance on active learning benchmark functions (RMSE $\downarrow$). Results show the mean and 95% confidence interval (CI) across 100 runs.
• Figure 4: Hyperparameter inference performance on synthetic GP functions.
• Figure 5: Results on psychometric model. RMSE (mean $\pm$ 95% CI) when targeting (a) threshold & slope and (c) guess & lapse rates, with Aline's corresponding query strategies shown in (b) and (d).

Theorems & Definitions (6)

• Proposition 1
• Proposition 2
• Proposition : Proposition \ref{['prop:sEPIG_definition']}
• proof
• Proposition : Proposition \ref{['proposition:combined']}
• proof