Title: ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition
Venue: NeurIPS 2025
Accepted: yes
Keywords: Amortized inference

Bayesian experimental design

Active learning

Bayesian inference

Neural processes
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper develops the Amortized Active Learning and INference Engine (ALINE) which employs a sim2real training setup to train a transformer-based backbone to simultaneously perform amortised Bayesian inference and amortised active data acquisition. Though a flexible approach to conditioning the model can be deployed on a range of different acquisition tasks. ALINE is compared to standard GP approaches on a GP regression / OOD task, it is evaluated on two classical Bayesian experimental design tasks, and a psychometric experimental design task.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nThe paper is tightly and clearly written and well-polished. I like table 1 for situating the current work in the context of past work. The method was neat with a single transformer-based backbone to perform three distinct tasks (forming the predictive, inferring parameters, and implementing the policy). The experiments were also well done and I especially liked the final experiment on experimental design for efficiently eliciting psychometric functions.\n\nWeaknesses\n\nWhereas neural processes can be trained directly on real data without the need for sim2real transfer, the current approach to supporting inference for parameters generally relies on a sim2real setup as it requires access to ground truth parameters at training time. This is not necessarily a huge limitation (we have to specify models when performing Bayesian inference and the TabPFN line have work has shown the sim2real transfer in NPs can work well for the predictive), but it does nevertheless require the sim2real gap to be bridged (which regular Bayesian inference does not have to address). I’d point this out more centrally in the development of the paper. At the moment this is mentioned a little cryptically in the “Limitations section” whereas I think it should be stated up front as a key aspect of the method.\n\nI terms of novelty, I thought that although the method was a combination of well-known components, the combination was quite sophisticated and sensibly put together.\n\nMinor point: In the related work I’d mention \"Environmental sensor placement with convolutional Gaussian neural processes” which uses Gaussian neural processes for active learning to support sensor placement and performs well compared to standard baselines. \n\nhttps://www.cambridge.org/core/journals/environmental-data-science/article/environmental-sensor-placement-with-convolutional-gaussian-neural-processes/F466DBCE3FA1088E04335D98225FA572\n\nFigure 2 / Section 5.1: It was good to see the comparison to the more standard approach of using a GP with an acquisition function, but I was a bit surprised that there was not a comparison to an existing neural process based active learning approach. This might enable you to decouple improvements that arise because of the use of the neural process like architecture and improvements that arise because of the more sophisticated policy used for acquisition.\n\ne.g. see section 5 in \"Meta-Learning Stationary Stochastic Process\nPrediction with Convolutional Neural Processes” or the reference mentioned above. https://papers.nips.cc/paper/2020/file/5df0385cba256a135be596dbe28fa7aa-Paper.pdf\n\nMore generally, although I liked the experiments that were present in the paper, I thought the experiments were slightly on the light side.\n\nQUESTIONS:\nThe RL component of the objective depends on the model predictions e.g. q(\\theta | D_c), but are gradients of the RL loss wrt q used in training or are these stopped? It seems like you only want policy gradients flowing through this objective?\n\nHow well does the approach generalise to scenarios (e.g. combinations of targets) at deployment time?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAmortization is a powerful tool in Bayesian inference and more recently in Bayesian experimental design. The paper proposes a unified approach for amortized Bayesian inference and data acquisition and is designed to handle both predictive and parameter targets. Specifically, the proposed method Aline, given an accumulated history and a target (predictive or parameter) can pick the next data point to acquire and estimate the posterior (or posterior predictive). To achieve this, Aline uses a modification of the Transformer Neural Processes architecture, with separate heads for the acquisition policy and inference. The method uses the sEIG bound from DAD, and a variant of EPIG called sEPIG for the predictive case as the objectives for data acquisition policy and a factorized likelihood for the inference head. The training of the acquisition policy is framed as a RL problem to maximize the variational lower bounds (which define the per step rewards). The method is evaluated on some synthetic benchmarks for active learning and standard tasks used in the Bayesian experimental design literature. Aline outperforms the task specific baselines in most experiments. The authors finally present results on a psychometric model with domain specific baselines and show that Aline matches the performance of these baseline while being computationally more efficient.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n* To the best of my knowledge, this work is the first to propose joint training of an amortized acquisition and inference network. I really like this idea since it unifies these two very intimately linked procedures and demonstrates the benefits of exploiting the structure between the two. \n* Another strength of the approach is that in addition to allowing specification of predictive or parameter targets, it also allows specification of subsets of parameters as the target which can provide a lot of flexibility in problems with known parameter structure.\n* The paper was extremely well written. The problem is clearly set up and each component of the approach is described clearly in sufficient detail. \n* On the tasks considered, the performance is similar to or slightly better than task specific methods while being faster in terms of the training time than some of the baselines and competitive in terms of the deployment time. \n* The authors also include the code which is helpful for reproducibility, in. addition to very detailed description in the paper itself.\n\nWeaknesses:\n* One of the main shortcomings of the paper is limited the scale of the empirical evaluations. All the experiments are limited to fairly low dimensional problems. I do understand these amortized networks can be expensive to train but I believe at demonstrating at least some scalability can be very useful. \n* Another weakness which is already mentioned in the paper is the choice of assuming factorized likelihood over the targets. While a potential solution in the form of an autoregressive extenions is discussed, it is left for future work. \n* While the training procedures for the individual parts is sound and well justified, there is not discussion of the stability of the joint training procedure. Based on my understanding there is no guarantee for the joint training procedure to converge (though I am a bit unsure about this and happy to hear the authors' thoughts).\n\nQUESTIONS:\n* The experiments only consider problems with continuous $x$ and $y$ and it is unclear how the method applies to cases where $x$ is discrete for instance?\n* Is it possible to switch between predictive and parameter targets in the middle of a rollout?\n* Could you elaborate on potential ways to make the method support inputs with arbitrary dimensionality? \n* What are the most important bottlenecks for scaling this to more complicated and higher dimensional problems?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes ALINE, amortized active learning and inference engine, a unified framework that jointly performing amortized Bayesian inference and optimal experimental design. ALINE employs a transformer-based architecture with two specialized heads, an inference head for variational posterior inference and an acquisition head for data selection policy. The authors leverage a query-target cross-attention mechanism that enables flexible targeting of specific parameter subsets or prediction goals, rather than optimize the expected information gain (EIG) over all the parameters. ALINE is trained using reinforcement learning (RL) where the reward is to tighten the variational posterior lower bound.\n\nSTRENGTHS AND WEAKNESSES:\n__Strengths__:\n\n1. ALINE presents a unified framework to jointly amortize both posterior inference and Bayesian optimal experimental design (BOED) within a single neural architecture. The dual-head architecture (inference head for posterior estimation and acquisition head for design policy) with joint training is well-motivated and more align with the practical setting, where the ultimate goal of performing BOED is to get an informative posterior for inference under a limited budget.  \n\n2. The query-target cross-attention mechanism enables ALINE to dynamically adapt its acquisition strategy at runtime to focus on specific parameter subsets or predictive goals..\n\n__Weaknesses__:\n\n1. This paper has significant overlap with the already-published vsOED work [1]. Both papers: (a) tackle sequential BOED with RL-based amortization, (b) handle flexible design criteria and prediction goals, (c) use similar reward formulations (ALINE's incremental reward is essentially equivalent to vsOED's variational-one-point-IIG), (d) evaluate on identical benchmarks (Location Finding) against the same baselines (DAD, RL-BOED), and (e) claim similar efficiency advantages. From my perspective, the contribution of ALINE is the unified architecture for both posterior inference and BOED and the query-target cross-attention mechanism applied in the framework. However, the core problem motivation and methodological approach are no longer novel given [1]'s prior publication. To properly claim the contribution of this paper, the authors should add a thorough discussion comparing the ALINE and vsOED, and also include vsOED in the baselines for a direct empirical comparison.\n\n2. ALINE requires explicit likelihood for training (due to log-probability computations in the reward signal and training loss), while vsOED [1] handles both explicit and implicit likelihood cases through its actor-critic framework. This makes ALINE more limited than existing work.\n\n3. The paper lacks formal theoretical analysis regarding convergence properties, optimality guarantees, or theoretical relationships to information-theoretic objectives. \n\n4. (typo) Line 210: Should the \"locations\" here be deleted as not specifically discussing the Location Finding task here?\n\n\n[1] Shen, W., Dong, J., & Huan, X. (2025). Variational sequential optimal experimental design using reinforcement learning. Computer Methods in Applied Mechanics and Engineering, 444, 118068.\n\nQUESTIONS:\nPlease see the comments in the Weaknesses part.\n\nLIMITATIONS:\nThe authors discussed the limitations in the Section 6. They did not mention the Broader impacts but provided justification in the Checklist.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the **Amortized Active Learning and Inference Engine (ALINE)**, a unified framework for amortized Bayesian inference and active data acquisition that relies on a transformer-based backbone. At run time, ALINE ingests the current data history along with a user-specified target specifier (any subset of parameters or predictive outputs). In a single forward pass it\n\n1. produces an approximate posterior or predictive distribution, and\n2. proposes a candidate query point expected to deliver the greatest information gain for the specified target.\n\nTraining combines a maximum-likelihood objective for the inference head with a reinforcement-learning objective for the query head, where each step’s reward is the self-estimated improvement in the log-probability of the chosen targets.\n\nEmpirically, ALINE matches or exceeds strong baselines—such as Gaussian processes with uncertainty sampling or expected information gain—on standard active-learning benchmarks and outperforms them on out-of-distribution tasks. It also achieves competitive performance on classical Bayesian experimental-design problems (e.g., Location-Finding, CES, and psychometric-model tasks) while remaining orders of magnitude faster at deployment than non-amortized alternatives.\n\nSTRENGTHS AND WEAKNESSES:\n**STRENGTHS**\n\n1. Provides a unified amortization framework for both inference and acquisition. The target-specifier mechanism lets a single trained model focus on any subset of parameters or predictions at run time.\n\n2. The paper is clearly written and well organized, making it easy to follow.\n\n3. Delivers strong empirical performance while enabling much faster deployment.\n\n**WEAKNESSES**\n\n1. The empirical study is confined to Gaussian-process benchmarks; broader experiments would strengthen the evidence.\n\n2. The rationale for combining inference and active data acquisition is not fully developed. What, exactly, would be lost by handling them separately, and what concrete computational or practical gains come from unifying them?\n\n3. Correct joint inference requires autoregressive prediction, yet this is deferred to future work.\n\nQUESTIONS:\n1. **Equation 6 appears to be incorrect.** It should be expressed in an autoregressive form, because the terms are not conditionally independent given the dataset.\n\n2. **Why are the parameters and their inference not modeled autoregressively?** An autoregressive treatment seems the principled approach.\n\n3. **Consider relocating the Related-Work section to follow the Introduction** to improve the paper’s narrative flow.\n\nLIMITATIONS:\nYes, the authors have discussed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-08"}