ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition

Key points

• Problem: Fields with limited budgets and data, such as clinical trials and materials science, need systems that strategically acquire informative data and perform inference on it right away. Amortized Bayesian inference is fast but typically assumes passive data collection, while amortized data acquisition selects informative points well, but the inference update on new data often sits outside the amortization and may need separate, costly procedures such as MCMC (Markov chain Monte Carlo). Existing amortized acquisition methods usually optimize information gain about all parameters or one fixed predictive target, which makes it hard to switch goals when only some parameters matter.
• Input and output: ALINE takes the observations gathered so far (context set), the candidate points that could be queried next (query set), and the current inference goal (target set). In a single forward pass it returns the posterior over the targeted parameters or the predictive distribution and selects the next point to query from the candidates. This is a method paper that unifies amortized Bayesian inference and active data acquisition in one model, with experiments on simulated tasks and on the HPO-B benchmark of real hyperparameter performance data.
• Method: One transformer processes the three inputs together, and an added query-target cross-attention lets candidate points attend directly to the target set so that selection fits the current goal. The policy is trained with policy-gradient RL using a per-step reward equal to the increase in the log-probability of the targets under its own inference head, and the authors show this objective is a variational lower bound on the total information gain. Training starts with a warm-up phase in which only the inference part is trained on randomly chosen points, after which both parts are trained jointly, and gradients from the policy loss update only the policy parameters.
• Implementation: The two heads share 3 transformer layers (4 attention heads each, internal embedding dimension 32), the inference head represents posteriors and predictive distributions with a 10-component Gaussian mixture, and the policy head is a 2-layer MLP that outputs a probability for each candidate point. The active learning model is trained only on synthetic functions drawn from Gaussian Process (GP) priors, the whole study used about 5000 GPU hours on AMD MI250X GPUs, and training one model takes about 20 hours. The official repository README lists the training script train_aline.py and folders for hydra configuration (config), model, tasks, loss, distributions, utils, and evaluation notebooks, and gives training commands for GP active learning in 1D and 2D, Location Finding, CES, and the psychometric model.
• Results: On Location Finding, ALINE reached a total EIG lower bound of 8.91, above every baseline including the amortized BED methods RL-BOED (7.70), DAD (Deep Adaptive Design, 7.33), and vsOED (7.30), and on CES it scored 13.50, second to RL-BOED (14.60). In active learning it performed comparably to the best GP-based methods on in-distribution GP functions and had lower RMSE (root mean squared error) than the baselines on 3 of 4 benchmark functions unseen in training. On the psychometric model, when the guess and lapse rates were the targets, ALINE and Psi-marginal, an existing method that can focus on a parameter subset, inferred them significantly better than QUEST+, which targets all parameters together, and ALINE took 0.002 s per step to propose the next stimulus, against 0.07 s for QUEST+ and 0.02 s for Psi-marginal.
• Limits: ALINE is trained with predefined, fixed priors, so a different prior requires retraining, and it estimates marginal posteriors per parameter, leaving joint posterior estimation to future work. Performance can degrade when deployment observations differ substantially from the training data, and on the Three Hump Camel function, whose output range extends beyond that of the pre-training GP functions, ALINE and ACE-US (the amortized inference model ACE paired with uncertainty sampling) both showed reduced accuracy. The architecture is tailored to fixed input dimensionalities and discrete design spaces, and the authors list dimension-agnostic models and continuous experimental designs as future work.

Paper links

External research summaries. These are not HDATF publications or measured product results.