ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition
Authors: Daolang Huang
Xinyi Wen
Ayush Bharti
Samuel Kaski
Luigi Acerbi
Venue: NeurIPS 2025 spotlight ()
Keywords: Amortized inference
Bayesian experimental design
Active learning
Bayesian inference
Neural processes
Primary area: probabilistic_methods
TL;DR: We introduce a novel framework that seamlessly integrates amortized Bayesian inference and active data acquisition, featuring adaptive strategies that can optimize for diverse, user-specified learning objectives at deployment.

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.

ALINE: Joint Amortization for Bayesian Inference and Active Data Acquisition
Daolang Huang 1.2
Xinyi Wen 1.2.3
Ayush Bharti Samuel 1,s
Luigi Acerbi
ELLIS Institute Finland
2 Department of Computer Science, Aalto University, Finland Department of Computer Science, University of Helsinki, Finland
4 Department of Computer Science, University of Manchester, UK
{daolang- huang, xinyi wen, ayush. bharti, ayarh.b samuel kamuel.kaski) . .ii
Abstract
Many critical applications, from autonomous scientific discovery to personalized medicine, demand systems that both strategically acquire the most informa- tive data and instantancously 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 be collected for instant inference. To tackle this issue, we intro- duce 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 des- ignated predictive tasks, optimizing for posterior estimation, data prediction, or mixture thereof. Empirical results on regression-based active learning, classi- cal Bayesian experimental design benchmarks, and psychometric model with selectively targeted parameters demonstrate that ALINE delivers both instant and accurate inference along with efficient selection of informative points.
Introduction
Bayesian inference [28] and Bayesian experimental design [61] offer principled mathematical means for reasoning under uncertainty and for strategically gathering data, respectively. While both are foundational, they introduce notorious computational challenges. For example, in scenarios with continuous data streams, repeatedly applying gold-stan gold-standard inference methods such as Markov Chain Monte Carlo (MCMC) [13] to update posterior distributions can be computationally demanding, leading to various approximate sequential inference techniques [10, 18], yet challenges in achieving both speed and accuracy persist. Similarly, in Bayesian experimental design (BED) or Bayesian active learning (BAL), the iterative estimation and optimization of design objectives can become costly, especially in sequential learning tasks requiring rapid design decisions [21, 30, 55], such as in psychophysical experiments where the goal is to quickly infer the subject's perceptual or cognitive parameters [70, 57]. Moreover, common approach in BED involves greedily optimizing for single-
Equal contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Table 1: Comparison of different methods for amortized Bayesian inference and data acquisition.
Amortized Inference
Amortized Data Acquisition
Method
Posterior
Predictive
Posterior
Predictive
Flexible
Neural Processes [27, 26, 53]
X
Neural Posterior Estimation [49, 54, 34, 60]
X
DAD [23], RL.-BOED [9]
X
X
/
RL-sCEE [8]. vsOED [65]
J
x
/
X
X
AAL [46]
X
x
X
/
X
JANA [59], Simformer [31], ACB [15]
/
/
X
X
X
ALINE (this work)
V
V
V
V
step objectives, such as the Expected Information Gain (EIG), which measures the anticipated reduction in uncertainty. However, this leads to myopic designs that may be suboptimal, which do not fully consider how current choices influence future learning opportunities [23].
To address these issues, a promising avenue has been the development of amortized methods, leveraging recent advances in deep learning [75]. The core principle of amortization involves pre- training a neural network on a wide range of simulated or existing problem instances, allowing the network to "meta-leamy (such as a posterior distribution or an optimal experimental design). Consequently, at deployment, this trained network can provide solutions for new, related tasks with remarkable efficiency, often in a single forward pass. Amortized Bayesian inference (ABI) methods such as neural posterior estimation [49, 54, 34, 60] that target the posterior, or neural processes [27, 26, 52, 53] that target the posterior predictive diistibutionn-iee or other approximate inference methods. Similarly, methods for amortized data acquisition, including amortized BED [23, 41, 40, 9] and BAL [46], instantaneously propose the next design that targets the learning of the posterior or the posterior predictive distribution, respectively, using a deep policy byroosinnnn the iterative optimization of complex information thhooretic objectives.
direct mapping from problem's 's features (like observed data) to its solution
almost instantaneous results for new data, bypassing MCMC
While amortized approaches have significantly advanced Bayesian inference and data acquisition, progress has largely occurred in parallel. ABI offers rapid inference but typically assumes passive data collection, not addressing strategic data acquisition under limited budgets and data constraints common in fields such as clinical trials [7] and material sciences [48, 44]. Conversely, amortized data acquisition excels at selecting informative data points, but often the subsequent inference update based on new data is not part of the amortization, potentially requiring separate, costly procedures like MCMC. This separation means the cycle of efficiently deciding what data to gather and then instantaneously updating beliefs has not yet been seamlessly integrated. Furthermore, existing amortized data acquisition methods often optimize for information gain across all model parameters or fixed predictive target, lacking the flexibility to selectively target specific subsets of parameters or adapt to varying inference goals. This is significant drawback in scenarios with nuisance parameters [58] or when the primary interest lies in particular aspects of the model or predictions-whie might not be fully known in advance. A unified framework that jointly amortizes both active data acquisition and inference, while also offering flexible acquisition goals, would therefore be highly beneficial.
In this paper, we introduce Amortized Active Learning and INference Engine (ALINE), novel framework designed to overcome these limitations by unifying amortized Bayesian inference and active data acquisition within a single, cohesive system (Table 1). ALINE utilizes a transformer-bant based architecture [69] that, in a single forward pass, concurrently performs posterior estimation, generates posterior predictive distributions, and decides which data point to query next. Critically, and in contrast to existing methods, ALINE offers flexible, targeted acquisition: it can dynamically adjust runtime its Bata-gathering strategy to focus on any specified combination of model parameters or predictive tasks. This is enabled by an attention mechanism allowing the policy to condition on specific inference goals, making it particularly effective in the presence of nuisance variables [58] or for focused investigations. ALINE is trained using a self- self-guided reinforeement learning objective; the reward is the improvement in the log-probability of its own approximate posterior over the selected targets, a principle derived from variational bounds on the expected information gain [24]. Extensive experiments on diverse tasks demonstrate ALINE's ability to simultaneously deliver fast, accurate inference and rapidly propose informative data points.

---

2 Background
Consider a parametric conditional model defined on some space y C Rdy of output variables y given inputs (or covariates) Z E X C IRxx, and parameterized by 0 E θ S R Let Dr = {(x,, ys ) be a collection of T data points (or context) and p(Dr 0) = p(yi:r Z1.T. f) denote the likelihood function associated with the model, which we assume to be well-specified in this paper (i.e., it matches the true data generation process), Given a prior distribution p(0), the classical Bayesian inference or prediction problem involves estimating either the posterior distribution p(0 (D)) x p(Dr 10)((), or Dr)
the posterior predictive distribution P(yī:a< T Dr) = Ep(o| Pr) [p(3f:M : 1T 8, Dr over tar- get outputs Vi.M == (Vi,... Viz) corresponding to a given set of target inputs ai.as = (#1, ..., Tis)-
Estimating these quantities repeatedly via approximate inference methods such as MCMC can be computationally costly [28], motivating the need for amortized inference methods.
Amortized Bayesian inference (ABI). ABI methods involve training a conditional density network qs- parameterized by learnable weights o. to approximate either the posterior predictive distribu- tion (v [x] Dr) se ply L ( M. Dr) [26, 42, 33, 11, 38, 12, 53, 52], the joint posterior qo(0 PT) s p(0 Dr) [49, 54, 34, 60], or both [59, 31, 15]. These networks are usually trained by minimizing the negative ooglliieldhood (NLL) objective with respect to Φi
{ (p())p(Pr ()p(at (0) [log g (Vi | Dr)), (()) = { (p(0)p(Dr 0) [log yo (0) Dr)],
(predictive tasks)
(posterior estimation)
where the expectation is over datasets simulated from the generative process p O)p(O). Once trained, (s can then perform instantaneous approximate inference on new contexts and unseen data points with a single forward pass. However, these ABI methods do not have the ability to strategically collect the most informative data points to be included in Df in order to improve inference outcomes.
Amortized data acquisition. BED [47, 14, 63, 61] methods aim sequentially select the next input (or design parameter) x to query in order to maximize the Expected Information Gain (EIG), that is, the information gained about parameters upon observing y:
EIG(x) := F plv x) [M(( - H [p(0 [r, [.,))],,
where H is the Shannon entropy H[p(-)] -Ep(.) [log p(-]( Directly computing and optimizing EIG sequentially at each step of an experiment is computationally expensive due to the nested expectations, and leads myopic designs. Amortized BED methods address these limitations by offline learning design policy network Te X' x y + X. parameterized by eb [23, 40, 9], such that at any step t the policy Tे proposes query 2₁ TT P 0) to acquire data point y/2 ~ p(y xr, 0). forming D, = D U [(2,, }t)}. To propose non-myopic designs, To is trained by maximizing tractable lower bounds of the total EIG over T-step sequential trajectories generated by the policy Te.
(FIG Eppcr (.) [Mp(()] H |
By pre-compiling the design strategy into the policy network, amortized BED methods allow for near-instantaneous design proposals during the deployment phase via a fast forward pass. Typically, these amortized BED methods are designed to maximize information gain about the full set of model parameters 8. Separately, for applications where the primary interest lies in reducing predictive uncertainty rather than parameter uncertainty, objectives like the Expected Predictive Information Gain (EPIG) [67] have been proposed, so far in non-amortized settings:
(x) = Ept [a) [p[["" xx [x*)] - H [p((" |x",,,") x)]]
This measures the EIG about predictions y" at target inputs xx drawn from a target input distribution p. (x*). Notably, current amortized data acquisition methods are inflexible: they are generally trained to learn about all parameters 0 (via objectives like sEIG) and lack the capability to dynamically target specific subsets of parameters or adapt their acquisition strategy to varying inference goals at runtime.
Related work. A major family of ABI methods is Neural Processes (NPs) [27, 26], that learn a mapping from the observed context data points to predictive distribution for new target points. Farly NPs often employed bl.e-based encoders [27, 26, 38, 1], while recent works utilize more advanced attention and transformer architectures [42, 53, 52, 19, 20, 5, 4]. Complementary to
3

---

Next query
Approximate inference
instant
instant
Experimentation
Policy Head
Inference Head
Posterior
3
Tt
Selective Mask
Step
Step
Transformer
Posterior Predictive
Data Embedder
1/1
3/t
ALINE
Step
Step
Updated history
Posterior incrovemem
Conceptual workflow of ALINE, demonstrating its capability to sequentially query informa- tive data points and perform rapid posterior or predictive inference based on the gathered data.
Figure 1:
NPs, methods within simulation-base inference [17] focus amortizing the posterior distribution [54, 49, 34, 60, 51, 75J. More recently, methods for amortizing both the posterior and posterior predictive distributions have been proposed [31, 59, 15]. Specifically, ACE [15] shows how to flexibly condition on diverse user-specified inference targets, method ALINE incorporates for its own flexible inference capabilities. Building this principle of goal-directed adaptabilityy ALINE advances by integrating a learned policy that dynamically tailors the data acquisition strategy to the specified objectives. Existing amortized BED BAL methods [23, 40, 9, 46] that learn offline design policy do not provide real-time estimates of the posterior, unlike ALINE. Recent exceptions include methods like RL-SCEE [8] and vsOED [65], that use variational posterior bound of EIG to provide amortized posterior inference via a separate proposal network. Compared to these methods, ALINE uses single, unified architecture where the same model performs amortized inference for both posterior and posterior predictive distributions, and learns the flexible acquisition policy.
3 Amortized active learning and inference engine
Problem setup. We aim develop a system that intelligently acquires sequence of T' informative data points, DD = {(.,, y: }=)} to enable accurate and rapid Bayesian inference. This system must be flexible: capable of targeting different quantities of interest, such subsets of model parameters or future predictions. To formalize this flexibility, we introduce a targer specifier, denoted by & E E, which defines the specific inference goal. We consider two primary types of targets: (1) Parameter targets (E)) with the goal to infer specific subset of model parameters 0g. where S C (1,..., L} is an index set of parameters of interest. For example, 4i1.2; would target the joint posterior of 0₁ and 02 while (!......) targets all parameters, aligning with standard BED. We define S (S1,..., Sis) as the collection of all predefined parameter index subsets the system can target. (2) Predictive targets (5%.) ), where the objective is to improve the posterior predictive distribution ply" |x: Dr) for inputs x drawn from a specified target input distribution p. (x*). For simplicity, and following Smith et al. [67], we consider single target distribution p. (x*) in this work. The set of all target specifiers that ALINE is trained to handle is thus E = (s&s U (cy.' }. We assume a discrete distribution p(g) over these possible targets, reflecting the likelihood or importance each specific goal.
To achieve both instant, informative querying and accurate inference, we propose to jointly learn an amortized inference model qa and an acquisition policy T within a single, integrated architecture. Given the accumulated data history D,-1 and a specific target & E E, the policy T. selects the next query Ir, designed be most informative for that target. Subsequently, the new data point (z,, Ve is observed, and the inference model qu updates its estimate of the corresponding posterior or posterior predictive distribution. A conceptual workflow of ALINE is illustrated Figure 1. the remainder of this section, we detail the objectives for training the inference network ys (Section 3.1) and the acquisition policy म. (Section 3.2), discuss their practical implementation (Section 3.3), and describe the unified model architecture (Section 3.4).

---

3.1 Amortized inference
We use the inference network io to provide accurate approximations the true Bayesian posterior p(0 Dr) or posterior predictive distribution p(y" a Dr), given the acquired data Dr. We train to via maximum-likeliee (Eq. 1). Specifically, for parameter targets Ç
S.. our objective is:
(o) _E(())(DD 0) loggo Pr)] ~ -pp())(D 0) [Eis log go (0 Pr)],
(5)
where we adopt a diagonal or mean field approximation, where the joint distribution is obtained as a product of marginals (0 Dr) Iles ( Dr). Analogously, for predictive targets = cp.
we assume a factorized likelihood over targets sampled from the target input distribution p. (a*):
C (6) (P 10,, M ii.. .0) [log qo qo(Vi.s 1 Dr-)
(()p)) flp, (1 drlii or (i.
Σ m
log qo (um am Dr)
(6)
The factorized form of these training objectives is a common scalable choice in the neural process literature [26, 52, 53, 15] and is more flexible than it might seem, as conditional marginal distributions can be extended to represent full joints autoregressively [53, 1, 15J. However, a full autoregressive model would require multiple forward passes compute the reward signal for our policy each training step. making the learning process computationally intractable. Therefore, for simplicity and tractability, within the scope of this paper we focus on the marginals, leaving the autoregressive exten- sion to future work. Eqs. and form the basis for training the inference component qg. Optimizing them minimizes the Kullback-Leibe or predictive) defined by the underlying generative process and the model's approximations de [52]. Learning an accurate ds crucial as not only determines the quality of the final inference output but also serves as the basis for guiding the data acquisition policy, as we see next.
(KL) divergence between the true target distributions (posterior
3.2 Amortized data acquisition
The quality of inference from le depends critically on the informativeness of the acquired dataset Dr. The acquisition policy T. is thus responsible for actively selecting sequence of query-data pairs (Iz, yr) to maximize the information gained about specific target &.
When targeting parameters 0g (i.e., { = <g), the objective is the total Expected Information Gain (sEIGeg) about Og over the T'step trajectory generated by Te (see Eq. 3):
sEIGgs (v) 1x(0)p())1=.00)) [log p(0s Dr)] + H[p(Os)).
(7)
For completeness, include a derivation of sEIG,; in Appendix A.1 which is analogous to that of sEIG in [23]. Directly optimizing sEIG is generally intractable due to its reliance on the unknown true posterior p(0s] Dr). We circumvent this by substituting p(As) (Dr) with its approximation (0s0s] Dr) (from Eq. 5), yielding the tractable objective Jg for training Tw:
Ig ( = (p0())(PP m..9) Eies log da (O: Dr)] + H(p(0s))
(8)
The inference objective (Bq. 5) and this policy objective are thus coupled: (g)()) depends on data Dr acquired through Ty. and Jg (1) depends on the inference network ds-
Similarly, when targeting predictions for & < we aim to maximize information about p(y" xx,, Dr) for x² (="). We extend the Expected Predictive Information Gain (EPIG) framework [67] to the amortized sequential setting, defining the total sEPIG:
Proposition 1. The total expected predictive information gain for design policy z over a data trajectory of length T is:
sEPIG( == Bp. mm) ( [p((" [x*)]
H Hp((" 1,", Dr))]] (p (p. =*)0) [log xx", (r)] + Ep.(xr) | [p[p")" [=*)]]-
This result adapts Theorem [23] the predictive case (see Appendix A.2 for proof) and, unlike single-step EPIG (Eq. 4), considers the entire trajectory Dr given the policy me.
5

---

Now, similar to Eq. 8, we use the inference network go(y" |=* Dr) to replace the true posterior predictive distribution p(y" xx" Dr) in Proposition obtain our active learning objective:
Jz. (+b := (Pr mp. (x*)p(y" =r .n) [log goly" x², Dr)] + (a. [H[p(y" x |=)]]| (9)
Finally, the following proposition proves that our acquisition objectives, Jh and JH.- are variational lower bounds on the true total information gains (sEIGog and sEPIG, respectively), making them principled tractable objectives for our goal. The proof is given in Appendix A.3.
Proposition 2.
Let the policy zy generate the trajectory DDr With (0 (D)) approximating we have (bb - sEIGes (1)) and
p(Os| Dr), and q0(y" |z", ,D)) approximating ply" Jg. (v) I sEPIG( b) Moreover,
SEIGUs s (1) Jz((4) Ep(Dr ( KL(p(0s (T) Dr)llaa(Os [ and
SEPIG( - Jy." (1)) = Ep. (x* (p(r m. [KL(p)) j=", Dr))l [v"l"o xx,D)))]-
This principle of using approximate posterior (or predictive) distributions to bound information gain is foundational in Bayesian experimental design (e.g.. [24]) and has been extended to sequential amortized settings [8, 65]. Maximizing these J objectives thus encourages policies that increase information about the targets. The tightness of these bounds is governed by the expected KL divergence between the true quantities and their approximation, with a more accurate go leading to tighter bounds and more effective training signal for the policy. Additionally, our objectives solely rely on ALINE'S variational posterior, which does not require an explicit likelihood, making naturally applicable to problems with implicit likelihoods where only forward sampling is possible.
Data acquisition objective for ALINE. To handle any target E E from a user-specified set, we unify the previously defined acquisition objectives. We define J(w: {) based on the type of target &:
[1](,,, J(h.E) - TH. (v), (v),
(
if c if ey.
The final objective for learning the policy network Te. denoted J=, is the expectation of [1b,)) taken over the distribution of possible target specifiers p(E): J=(() (Eg-p(s)(((v)
3.3 Training methodology for ALINE
The policy and inference networks. Te and 9a. in ALINE are trained jointly. Training T. for sequential data acquisition over -step horizon is naturally framed as a reinforcement learning (RL) problem.
Policy network training (Tt). To guide the policy, we employ a dense, per-step reward signal Rs rather than relying on a sparse reward at the end of the trajectory. This approach, common in amortized experimental design [9, 8, 37], helps stabilize and accelerate learning. The reward R. quantifies the immediate improvement in the inference quality provided by Je upon observing a new data point (z,, 3a). specifically concerning the current target Ç. defined based on the one-step change in the log- probabiiisssss from our acquisition objectives (Eqs. 8 and 9):
Eles (log go (Or (D)) log (0) DD-1)))
if [ = c
R&(S)=
(10)
1
E (log go (yth | D.) log go (yim ath D.-1)), if Ex.
M
As per common practice, for gradient stabilization take averages (not sums) over predictions, which amounts to a constant relative rescaling of our objectives. The policy ny. is then trained using policy gradient (PG) algorithm with per-episode loss:
Lro(w) = -T- 7² (E) log T (xr(DDr-1.))
(11)
which maximizes the expected cumulative 7-discounted reward over trajectories [68]. Gradients from this policy loss only update the policy parameters vh. They are not propagated back the inference network go to ensure each component has clear and distinct objective.

---

2441 ~ TTp(1D.)
Acquisition Head
Inference Head
Transformer Layers
1
(4())
fa(et)
Fa(r)) Edlnn)
Embedding Layers
16
8
41
9/1
Fec
3/+
1
&
Or
45
Context Set
Query Set
Target Set
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, acquisition head determines the next data point to query, while an inference head performs the approximate Bayesian inference.
Inference network training (ga). For the per-step rewards Rg to be meaningful, the inference net- work 96 must provide accurate estimates of posteriors predictive distributions at each intermediate step t of the acquisition sequence, not just at the final step T. Consequently, the practical training objective for ya, denoted EN))(,), encourages this step-wise accuracy. In practice, training proceeds in episodes. For each episode: (1) A ground truth parameter set 0 is sampled from the prior p(0). (2) target specifier& is sampled from P(E). (3) If the target is predictive (g = 53," ). M target inputs {xt, 1m=1, are sampled from p= (z*), and their corresponding true outcomes {v ) m=1 are simulated using 8, similarly as [67]. The negative log- log-likelihood loss ( for do in an episode then computed by averaging over the T' acquisition steps and predictions, using the Monte Carlo estimates
of the objectives defined in Eqs. 5 and
Di= Eies log q(0) Dz), Lt=1 ΣM m= log q(v, I D.),
if cB
(o)
(12)
if Ey
Joint training. To ensure do provides reasonable reward signal early in training, we employ initial warm-up phase where only g, is trained, with data acquisition (Iz, Ur) guided by random actions instead of Ty. After the warm-up, Q and Th are trained jointly. A detailed step-by-step training algorithm is provided in Appendix B.
3.4 Architecture
ALINE employs a single, integrated neural architecture based on Transformer Neural Processes (TNPs) [53, 15], a network architecture that has been successfully applied to various amortized sequential decision-making settings [37, 1, 50, 39, 76]. ALINE leverages TNPs to concurrently manage historical observations, propose future queries, and condition predictions on specific, potentially varying, inference objectives. An overview of ALINE'S architecture is provided in Figure 2.
The model inputs are structured into three sets. Following standard TNP-based architecture, the context set D, {(T,, b=) comprises the history of observations, and the target set ' contains specific target specifier ç. To facilitate active data acquisition, we incorporate a query set Q = {~% J. of candidate points. In this paper, we focus on discrete pool-based setting for consistency, though can be straightforwardly extended to continuous design spaces (e.g.. [64]). Details regarding the embeddings of these inputs are provided in Appendix B.2.
7

---

Synthetic 1D
Synthetic 2D
Higdon 1D
0.6
0.6
RMSE
0.4
0.2
0.2
0.2
0.0
Goldstein-Pree 2D
Ackley 2D
Gramacy 2D
1.00
0.0
0.100
RMSE
0.6
0.78
0.075
0.1
0.50
0.050
0.25
0.025
0.00
40
Namber of Steps
Number Steps
Number of Steps
CP RS
CP-US
GPVR
CP-EPIC
ACE-US
ALINE (ours)
Figure 3: Predictive performance on active learning benchmark functions (RMSE 4). Results show the mean and 95% confidence interval (CI) across 100 runs.
Standard transformer attention mechanisms process these embedded representations. Self-attention operates within the context set, capturing dependencies within D, Both the query and target set then employ cross-attention to attend to the processed context set representations. To enable the policy T to dynamically adapt its acquisition strategy based on the specific inference goal ç, we introduce an additional query-target cross-attention mechanism to allow the query candidates to directly attend to the target set. This allows the evaluation of each candidate r² to be informed by its relevance the different potential targets ç. Examples of the attention mask are shown in Figure A1.
Finally, two specialized output heads operate these processed representations. The inference head (1a), following ACE [15], uses a Gaussian mixture to parameterize the approximate posteriors and posterior predictives. The acquisition head (TT) generates policy over the query set (x l'a.). drawing on principles from policy-based design methods [37, 50]. This unified design, which leverages shared transformer backbone with specialized heads, significantly improves parameter efficiency by avoiding the need for separate encoders for the two tasks, leading to faster training and more efficient deployment.
4 Experiments
We now empirically evaluate ALINE'S performance in different active data acquisition and amortized inference tasks. We begin with the active learning task in Section 4. where we want to efficiently minimize the uncertainty over an unknown function by querying T data points. Then, we test ALINE'S policy on standard BED benchmarks in Section 4.2. In Section 4.3, we demonstrate the benefit of ALINE'S flexible targeting feature in psychometric modeling task [72]. Finally, to demonstrate the scalability of ALINE, we test it on high-dimensional task of actively exploring hyperparameter performance landscapes; the results are presented in Appendix D.I. The code to reproduce our experiments is available at: bttts::/ggthii com/ www.nnnollaggaalllaa.aa
4.1 Active learning for regression and hyperparameter inference
For the active learning task, ALINE is trained on a diverse collection of fully functionn drawn from Gaussian Process (GP) [62] priors (see Appendix C. for details). We evaluate ALINE'S performance under both in-distribution and out-of-distributtin settings. For the in-distribution setting, ALINE evaluated on synthetic functions sampled from the same GP prior that used during training. In the out-of-die distribbbtiin setting, we evaluate ALINE on benchmark functions (Higdon, Goldstein-Price, Ackley, and Gramacy) unseen during training, to assess generalization beyond the training regime. We compare ALINE against non-amortized GP models equipped with standard acquisition functions such as Uncertainty Sampling (GP-US), Variance Reduction (GP-VR) [74], EPIG (GP-EPIG) [67], and Random Sampling (GP-RS). Additionally, we include an amortized neural

---

Table 2: Results on BED benchmarks. For the EIG lower bound, we report the mean + 95% CI across 2000 runs (200 for VPCE). For deployment time, we use the mean + 95% CI from 20 runs.
Location Finding
Constant Flasticity of Substitution
EIG lower
Training time (h)
Deployment time (s)
EIG lower bound (1)
Training time (h)
Deployment time
bound (1)
Random
5.17:10.05
N/A
N/A
9.05:10026
N/A
N/A
VPCE [25]
5.25-10222
N/A
146.59-10.09
9.40:10.27
N/A
788.90 1.03
DAD [23]
7.33
7.24
0.0001:0.00
10.77=10.15
13.70
0.0001
vsOED [65]
7.30:10.06
4,31
0.0002:10.00
12.12.10.18
0.49
0.0003-10.00
RL-BOED [9]
7.70-10.06
63.29
0.0003
14.60:10.10
67.28
0.00041:0.00
ALINE (ours) 8.91-10.04
21.20
0.03-10.00
13,50:10.15
13.29
0.04=00.0
process baseline, the Amortized Conditioning Engine (ACE) [15], paired with Uncertainty Sampling (ACE-US), specifically evaluate the advantage of ALINE'S learned acquisition policy over using standard acquisition function with an amortized inference model. The performance metric is the root mean squared error (RMSE) of the predictions on a held-out test set.
Results for the active learning task in Figure 3 show that ALINE performs comparably the best-performing GP-based methods for the in-distribution setting. Importantly, for the out-of-distributin setting, ALINE outperforms the baselines in 3 out of the 4 bench- mark functions. These results highlight the advantage of ALINE'S end-to-end learning strategy, which obviates the need for kernel spec- ification using GPs or explicit acquisition function selection. Further evaluations on additional benchmark functions (Gramacy ID, Branin, Three Hump Camel), visualizations of ALINE'S 's sequential querying strategy with corresponding predictive updates, and a comparison of average inference times are provided in Appendix D.2.
Log Probabbliggg 1
ACE US
ALINE (ours)
103
0.4
10
15
20
25
Number of Steps
Figure 4: Hyperparameter in- ference performance on syn- thetic GP functions.
Additionally, for the in-distribution setting, we test ALINE'S ' capability to infer the underlying GP's hyperparameters- wwtthut retraining, by leveraging ALINE'S flexible target specification at runtime. For baselines, we use ACE-US and ACE-RS, since ACE [15] is also capable of posterior estimation. Figure 4 shows that ALINE yields higher log probabilities of the true parameter value under the estimated posterior each step compared the baselines. This is due to the ability to flexibly switch ALINE'S acquisition strategy to parameter inference, unlike other active learning methods. We also visualize the obtained posteriors in Appendix D.2.
4.2 Benchmarking Bayesian experimental design tasks
We test ALINE on two classical BED tasks: Location Finding [66] and Constant Elasticity of Substitution (CES) [3], with two- and -ix-dimensional design space, respectively. As baselines, we include a random design policy, gradient-basee method with variational Prior Contrastive Estimation (VPCE) [25], and three amortized BED methods: Deep Adaptive Design (DAD) [23], vsOED [65], and RL-BOED [9]. Details of the tasks and the baselines are provided in Appendix C.2.
To evaluate performance, we compute lower bound of the total EIG, namely the sequential Prior Contrastive Estimation lower bound [23]. As shown in Table 2, ALINE surpasses all the baselines in the Location Finding task and achieves competitive performance on the CES task, outperforming most other methods, except RL-BOED. Notably, ALINE'S training time is reduced as its reward is based on the internal posterior improvement and does not require large number of contrastive samples to estimate sEIG. While ALINE'S 's deployment time is slightly higher than MLP-based amortized methods due to the computational cost of its transformer architecture, it remains orders magnitude faster than non-amortized approaches like VPCF. Visualizations of ALINE'S inferred posterior distributions are provided in Appendix D.3.
4.3 Psychometric model
Our final experiment involves the psychometric modeling task [72]-a fundamental scenario in behavioral sciences, from neuroscience clinical settings [29, 57, 73], where the goal is to infer parameters governing observer's responses to varying stimulus intensities. The psychometric

---

Threshold & Slope
Guess Rate & Lapse Rate
1.5
0.30
QURST+ Pri-ntaginal
Stimuly Values
threshold
Stanuli Values
Query points
RMSE
ALINE (ours)
RMSE.
1.0
0.20
0.5
10
20
30
10
20
30
10
20
30
10
20
30
Number Steps
Number Steps
Number of Steps
Number Steps
(b)
(d)
Figure 5: Results on psychometric model. RMSE (mean + 95% CI) when targeting (a) threshold & slope and (c) guess & lapse rates, with ALINE'S corresponding query strategies shown in (b) and (d).
function used here is characterized by four parameters: threshold, slope, guess rate, and lapse rate; see Appendix C.3.1 for details. Different research questions in psychophysics necessitate focusing on different parameter subsets. For instance, studies on perceptual sensitivity primarily target precise estimation of threshold and slope, while investigations into response biases or attentional phenomena might focus on the guess and lapse rates. This where ALINE'S 's unique flexible querying strategy can be used to target specific parameter subsets of interest.
We compare ALINE with two established adaptive psychophysical methods: QUEST+ [70], which targets all parameters simultaneously, and Psi-marginal [58], which can marginalize over nuisance parameters focus specified subset, non-amortized gold-standard method for flexible acquisi- tion. We evaluate scenarios targeting either the threshold and slope parameters or the guess and lapse rates. Details of the baselines and the experimental setup are in Appendix C.3.2.
Figure 5 shows the results. When targeting threshold and slope (Figure Sa), which are generally easier to infer, ALINE achieves results comparable to baselines. When targeting guess and lapse rates (Figure Sc). QUEST+ performs sub-optimally as its experimental design strategy is dominated by the more readily estimable threshold and slope parameters. In contrast, both Psi-marginal and ALINE lead to significantly better inference than QUEST+ when explicitly targeting guess and lapse rates. Moreover, ALINE offers 10x speedup over these non-amortized methods (See Appendix D). We also visualize the query strategies adopted by ALINE in the two cases: when targeting threshold and slope (Figure 5b), stimuli are concentrated near the estimated threshold. Conversely, when targeting guess and lapse rates (Figure 5d), ALINE appropriately selects 'easy" stimuli at extreme values where mistakes can be more readily attributed to random behavior (governed by lapses and guesses) rather than the discriminative ability of the subject (governed by threshold and slope). To further demonstrate ALINE'S runtime flexibility, we conduct two additional investigations detailed in Appendix D.4. First, we show that a single pre-trained ALINE model can dynamically switch its acquisition target mid- experiment. Second, we validate its ability to generalize to novel combinations of targets that were not seen during training, showing the effectiveness of the query-targe cross- cross-attention mechanism.
5 Conclusion
We introduced ALINE, unified amortized framework that seamlessly integrates active data acqui- sition with Bayesian inference. ALINE dynamically adapts its strategy to target selected inference goals, offering flexible and efficient solution for Bayesian inference and active data acquisition.
Limitations & future work. Currently, ALINE operates with pre-defined, fixed priors, necessitating re-training for different prior specifications. Future work could explore prior amortization [15, 71], to allow for dynamic prior conditioning. As ALINE estimates marginal posteriors, extending this joint posterior estimation, potentially via autoregressive modeling [11, 351, is a promising direction. Note that, at deployment, we may encounter observations that differ substantially from the training data, leading to degradation in performance. This issue can potentially be tackled by combining ALINE with robust approaches such as [22, 36). Lastly, ALINE'S current architecture is tailored to fixed input dimensionalities and discrete design spaces, common practice with TNPs [53, 50, 76, 37]. Generalizing ALINE to cemensionnagnostic [45] and to support continuous experimental designs [64] are valuable avenues for future research.
10

---

Acknowledgements
DH, LA and SK were supported by the Research Council Finland (Flagship programme: Finnish Center for Artificial Intelligence FCAI, 359207). The authors acknowledge the research environment provided by ELLIS Institute Finland. LA was also supported by Research Council of Finland grants 358980 and 356498. SK was also supported by the UKRI Tuuring Al World-Leading Researcher Fellowship, [FP/W002973/1). AB was supported by the Research Council of Finland grant no. 362534. The authors wish to thank Aalto Science-ITT project, and CSC-IT Center for Science, Finland, for the computational and data storage resources provided.
References
[1] Andersson, T. R., Bruinsma, W. P., Markou, S., Requeima, J., Coca-Castro, A., Vaughan, A., Ellis, A.-L.,, Lazzara, M. A., Jones, D., Hosking, S., et al. (2023). Environmental sensor placement with convolutional gaussian neural processes. Environmental Data Science, 2:e32. 7
[2] Arango, S. P., Jomaa, H. S., Wistuba, M., and Grabocka, J. (2021). Hpo-b: A large-scale reproducible benchmark for black-box hpo based on openml. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track. 24
[3] Arrow, K. J., Chenery, H. B., Minhas, B. S., and Solow, R. M. (1961). Capital-labor substitution and economic efficiency. The Review of Economics and Statistics, 43(3):225 220-200.222220 9, 21
[4] Ashman, M., Diaconu, C., Kim, J., Sivaraya, L., Markou, S., Requeima, J., Bruinsma, W. P., and Turner, R. E. (2024a). Translation equivariant transformer neural processes. In International Conference on Machine Learning, pages 1 1924-1944 PMLR. 3
[5] Ashman, M., Diaconu, C., Weller, A., and Turner, R. F. (2024b). In-context in-context learning with transformer neural processes. In Symposium Advances in Approximate Bayesian Inference, pages 1-29, PMLR. 3
[6] Barlas, Y. Z. and Salako, K. (2025). Performance comparisons of reinforcement learning algorithms for sequential experimental design. arXiv preprint arXiv:2503.05900 27
[7] Berry, S. M., Carlin, B. P., Lee, J. J., and Muller, P. (2010). Bayesian adaptive methods for clinical trials. CRC press. 2
[8] Blau, T., Bonilla, E., Chades, and Dezfouli, A. (2023). Cross-entropy estimators for sequential experiment design with reinforcement learning. arXiv preprint arXiv:2305. 18435. 2, 4,
[9] Blau, T., Bonilla, E. V., Chades, I., and Dezfouli, A. (2022). Optimizing sequential experimental design with deep reinforcement learning. In International conference on machine learning, pages 2107-2128. PMLR. 2, 3, 4, 6, 9, 21, 22, 23
[10] Broderick, T., Boyd, N., Wibisono, A., Wilson, A. C., and Jordan, M. I. (2013). Streaming variational bayes. Advances in neural information processing systems, 26.
[11] Bruinsma, W., Markou, S., Requeima, J., Foong, A. Y., Andersson, T., Vaughan, A., Buonomo, A., Hosking, S., and Turner, R. F. (2023). Autoregressive conditional neural processes. In The Eleventh International Conference on Learning Representations. 3, 5, 10
[12] Bruinsma, W., Requeima, J., Foong, A. Y., Gordon, J., and Turner, R. E. (2020). The gaussian neural process. In Third Symposium on Advances in Approximate Bayesian Inference.
[13] Carpenter, B., Gelman, A., Hoffmann, M. D., Lee, D., Goodrich, B., Betancourt, M., Brubaker, M., Guo, J., Li, P., and Riddell, A. (2017). Stan: A probabilistic programming language. Journal of statistical software, 76:1-32. 1
[14] Chaloner, K. and Verdinelli, I. (1995). Bayesian experimental design: A review. Statistical science, pages 273-304.
[15] Chang, P. F., Loka, N., Huang, D., Remes, U., Kaski, S., and Acerbi, L. (2025). Amortized probabilistic conditioning for optimization, simulation and inference. In International Conference on Artificial Intelligence and Statistics. PMLR. 2, 3, 4, 5, 7, 8, 9, 10, 21
11

---

[16] Chen, X., Wang, C., Zhou, Z., and Ross, K. W. (2021). Randomized ensembled double q learning: Learning fast without a model. In International Conference on Learning Representations. 23
[17] Cranmer, K., Brehmer, J., and Louppe, G. (2020). The frontier of simulation- based inference. Proceedings of the National Academy of Sciences, 117(48):3055533000 4
[18] Doucet, A., De Freitas, N., Gordon, N. J., al. (2001). Sequential Monte Carlo methods in practice, volume 1. Springer.
[19] Feng, L., Hajimirsadeghi, H., Bengio, Y., and Ahmed, M. O. (2023). Latent bottlenecked attentive neural processes. In The Eleventh International Conference on Learning Representations.
3
[20] Feng, L., Tung, F., Hajimirsadeghi, H., Bengio, Y., and Ahmed, M. O. (2024). Memory efficient neural processes via constant memory attention block. In International Conference Machine Learning, pages 13365-13386. PMLR.
[21] Filstroff, L., Sundin, I., Mikkola, P., Tiulpin, A., Kylmäoja, J., and Kaski, S. (2024). Targeted active learning for bayesian decision-making Transactions on Machine Learning Research.
[22] Forster, A., Ivanova, D. R., and Rainforth, T. (2025). Improving robustness to model misspecifi- cation in bayesian experimental design. In 7th Symposium on Advances in Approximate Bayesian Inference Workshop Track. 10
[23] Foster, A., Ivanova, D. R., Malik, I., and Rainforth, T. (2021). Deep adaptive design: Amortizing sequential bayesian experimental design. In International Conference on Machine Learning, pages 3384-3395. PMLR. 2, 3, 4, 5, 9, 16, 22, 27
[24] Foster, A., Jankowiak, M., Bingham, E.,, Horsfall, P., Teh, Y. W., Rainforth, T., and Goodman, N. (2019). Variational bayesian optimal experimental design. Advances in Neural Information Processing Systems, 32. 2, 6, 21
[25] Foster, A., Jankowiak, M., O'Meara, M., Teh, Y. W., and Rainforth, T. (2020). A unified stochas- tic gradient approach to designing bayesian-optimal experiments. In International Conference on Artificial Intelligence and Statistics, pages 2959-29999 2969. PMLR. 9, 22
[26] Garnelo, M., Rosenbaum, D., Maddison, C., Ramalho, T., Saxton, D., Shanahan, M., Tch, Y. W., Rezende, D., and Eslami, S. A. (2018a). Conditional neural processes. In International conference on machine learning, pages 1704-1713. PMLR. 2, 3,5
[27] Garnelo, M., Schwarz, J., Rosenbaum, D., Viola, F., Rezende, D. J., Eslami, S., and Teh, Y. W. (2018b). Neural processes. arXiv preprint arXiv: prXiiv1807.000 180.00162.. 2,3
[28] Gelman, A., Carlin, J. B., Stern, H. S., Dunson, D. B., Vehtari, A., and Rubin, D. B. (2013). Bayesian Data Analysis. CRC Press. 1,3
[29] Gitaie-Dotan, S., Kanai, R., Bahrami, B., Rees, G., and Saygin, A. P. (2013). Neuroanatomical correlates of biological motion detection. Neuropsychologia, 51(3):457-463. 9
[30] Giovagnoli, A. (2021). The bayesian design of adaptive clinical trials. International journal of environmental research and public health, 18(2):530.
[31] Gloeckler, M., Deistler, M., Weilbach, C. D., Wood, F., and Macke, J. H. (2024). All-in-one simulation-based inference. In International Conference on Machine Learning, pages 15735- 15766. PMLR. 2, 3,4
[32] Glorot, X. and Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. In Teh, Y. W. and Titterington, M., editors, Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 249-256, Chia Laguna Resort, Sardinia, Italy. PMLR. 22
[33] Gordon, J., Bruinsma, W. P., Foong, A. Y., Requeima, J., Dubois, Y., and Turner, R. E. (2020). Convolutional conditional neural processes. In International Conference on Learning Representations.
12

---

[34] Greenberg, D., Nonnenmacher, M., and Macke, J. (2019). Automatic posterior transformation for likelihood-fFi inference. International conference on machine learning, pages 2404-2414.
PMLR. 2,4
[35] Hassan, C., Loka, N., Li, C.-Y,, Huang, D., Chang, P. E., Yang, Y., Silvestrin, F., Kaski, S., and Acerbi, L. (2025). Efficient autoregressive inference for transformer probabilistic models, arXiv preprint urXiv:2510.09477 10
[36] Huang, D., Bharti, A., Souza, A., Acerbi, L.,, and Kaski, S. (2023a). Learning robust statistics for iimulation-basee inference under model misspecification. Advances in Neural Information Processing Systems, 36:7289- -310. 10
[37] Huang, D., Guo, Y., Acerbi, L.,, and Kaski, S. (2025). Amortized bayesian experimental design for decision-making. Advances in Neural Information Processing Systems, 37:109460-109488 6,
7, 8, 10
[38] Huang, D., Haussmann, M., Remes, U., John, S., Clarté, G., Luck, K., Kaski, S., and Acerbi, L. (2023b). Practical equivariances via relational conditional neural processes. Advances in Neural Information Processing Systems, 36:29201-29238.
[39] Hung, Y. H., Lin, K.-J., Lin, Y.-H.. Wang, C.-Y., Sun, C., and Hsich, P.-C. (2025). Boformer: Learning to solve multi-ohjective bayesian optimization via non-markovian rl. In The Thirteenth International Conference on Learning Representations.
[40] Ivanova, D. R., Foster, A., Kleinegesse, S., Gutmann, M. U., and Rainforth, T. (2021). Implicit deep adaptive design: Policy-based experimental design without likelihoods. Advances in Neural Information Processing Systems, 34. 2, 3, 2,
[41] Ivanova, D. R., Hedman, M., Guan, C., and Rainforth, T. (2024). Step-DAD: Semi-Amortized Policy Based Bayesian Experimental Design. ICLR 2024 Workshop on Data-centric Machine Learning Research (DMLR). 2, 21
[42] Kim, H., Mnih, A., Schwarz, J., Gamelo,, M., Eslami, A., Rosenbaum, D., Vinyals, O., and Teh,
Y. W. (2019). Attentive neural processes. In International Conference on Learning Representations.
3
[43] Kontsevich, L. L. and Tyler, C. W. (1999). Bayesian adaptive estimation of psychometric slope and threshold. Vision research, 39(16):7729-2737. 24
[44] Krause, A., Guestrin, C., Gupta, A., and Kleinberg, (2006). Near-optimal sensor placements: Maximizing information while minimizing communication cost. In Proceedings of the 5th international conference on Information processing in sensor networks, pages 2-10.
[45] Lee, H., Jang, C., Lee, D. B., and Lee, J. (2025). Dimension agnostic neural processes. The Thirteenth International Conference on Learning Representations. 10
[46] Li, C.-Y,, Toussaint, M., Rakitsch, B., and Zimmer, C. (2025). Amortized safe active learning for real-time decision decisi:n-making: Pretrained neural policies from simulated nonparametric functions. arXiv preprint ar 2511v:2501.15458. 2.
[47] Lindley, D. V. (1956). On measure of the information provided by experiment. The Annals of Mathematical Statistics, 2764):986 3
[48] Lookman, T., Balachandran, P. V., Xue, D., and Yuan, R. (2019). Active learning in mate- rials science with emphasis on adaptive sampling using uncertainties for targeted design. npj Computational Materials, 5(1):21. 2
[49] Lueckmann, J.-M., Goncalves, P. J., Bassetto, G., Öcal, K., Nomnenmacher, M.. and Macke, J. H. (2017). Flexible statistical inference for mechanistic models of neural dynamics. Advances in neural information processing systems, 30. 2, 3,
[50] Maraval, A., Zimmer, M., Grosnit, A., and Bou Ammar, H. (2024). End-to-end meta-bayesian optimisation with transformer neural processes. Advances Neural Information Processing Systems, 36. 7, 8, 10
13

---

[51] Mittal, S., Bracher, N. L., Lajoie. G., Jaini, P., and Brubaker, M. (2025). Amortized in-context bayesian posterior estimation. arXiv preprint arXiv:2502.0660/
[52] Müller, S., Hollmann, N., Arango, P., Grabocka, J., and Hutter, (2022). Transformers can do bayesian inference. In International Conference Learning Representations, 2, 3,5
[53] Nguyen, T. and Grover, A. (2022). Transformer neural processes: Uncertainty-awan meta learning via sequence modeling. In International Conference on Machine Learning, pages 16569- 16594. PMLR. 3, 5, 10
[54] Papamakarios, G. and Murray, (2016). Fast E-free inference of simulation models with bayesian conditional density estimation. Advances in neural information processing systems, 29.
2, 3,
[55] Pasek, J. and Krosnick, J. (2010). Optimizing survey questionnaire design in political science: Insights from psychology.
[56] Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V. Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., et al. (2011). Scikit-leam: Machine learning in python. the Journal of machine Learming research, 12:2825-2836 21,
[57] Powers, A. R., Mathys, C., and Corlett, R. (2017). Pavlovian conditioningginnuuee hal- lucinations result from overweighting of perceptual priors. Science, 357(6351):5596666 1,
[58] Prins, N. (2013). The psi-marginal adaptive method: How to give nuisance parameters the attention they deserve (no more, no less). Journal of vision, 13(7):3-3, 2, 10, 24
[59] Radev, S. T., Schmitt, M., Pratz, V., Picchini, U., Köthe, U., and Bürkner, P.-C. (2023a). Jana: Jointly amortized neural approximation of complex bayesian models. Uncertainty in Artificial Intelligence, pages 1695-1706. PMLR. 2, 3,4
[60] Radev, S. T., Schmitt, M., Schumacher, L.., Elsemüller, L.,, Pratz, V., Schälte, Y., Köthe, U., and Bürkner, P.-C. (2023b). Bayesflow: Amortized bayesian workflows with neural networks. Journal of Open Source Software, 8889):5702. 2, 3,
[61] Rainforth, T., Foster, A., Ivanova, D. R., and Bickford Smith, F. (2024). Modern bayesian experimental design. Statistical Science, 39(1):100-114, 1,3
[62] Rasmussen, C. E. and Williams, C. K. (2006). Gaussian Processes for Machine Learning. MIT Press.
[63] Ryan, E. G., Drovandi, C. C., McGree, J. M., and Pettitt, A. (2016). A review of modern computational algorithms for bayesian optimal design. International Statistical Review, 84(1):128- 154.
[64] Schulman, J., Levine, S., Abbeel, P., Jordan, M., and Moritz, P. (2015). Trust region policy optimization. In International conference on machine learning, pages 1889-1897. PMLR. 7, 10
[65] Shen, W., Dong, J., and Huan, X. (2025). Variational sequential optimal experimental de- sign using reinforcement learning. Computer Methods in Applied Mechanics and Engineering,
444:1 18068. 2, 4, 6, 9, 22, 27
[66] Sheng, X. and Hu, Y.H.. (2004). Maximum likelihood multiple-source localization using acoustic energy measurements with wireless sensor networks. IEEE transactions signal processing, 33(1):44-53, 91
[67] Smith, F. B., Kirsch, A., Farguhar, S., Gal, Y., Foster, A., and Rainforth, T. (2023). Prediction- oriented bayesian active learning. In International Conference on Artificial Intelligence and Statistics, pages 7331-7348 PMLR. 3, 4, 5, 17, 20
[68] Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. (1999). Policy gradient methods for reinforcement learning with function approximation. Advances neural information processing systems, 12. 6
14

---

[69] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.. Gomez, A. N.. Kaiser, L., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30. 2
[70] Watson, A. B. (2017). Quest+: A general multidimennional bayesian adaptive psychometric method. Journal of Vision, 17())10-11.. 1, 10, 24
[71] Whittle, G., Ziomek, J., Rawling, J., and Osborne, M. A. (2025). Distribution transformers: Fast approximate bayesian inference with on-the-fly prior adaptation. arXiv preprint arXiv:2502.02463.
10
[72] Wichmann, F. A. and Hill, N. J. (2001). The psychometric function: I. Fitting, sampling, and goodness of fit. Perception & Psychophysics, 62933:1293-1313. 8,
[73] Xu, C.,, Hülsmeier, D., Buhl, M., and Kollmeier, B. (2024). How does inattention influence the robustness and efficiency of adaptive procedures in the context of psychoacoustic assessments via smartphone? Trends Hearing, 28:2331216524128880
[74] Yu, K., Bi, J., and Tresp, V. (2006). Active learning via transductive experimental design. In Proceedings of the 23rd international conference Machine learning, pages 1081-1088. 8, 20
[75] Zammit-Mangion, A., Sainsbury-Dale, M., and Huser, R. (2024). Neural methods for amortised parameter inference. arXiv e-prints, pages arXiv-2404. 2,
[76] Zhang, X., Huang, D., Kaski, S., and Martinelli, J. (2025). Pabbo: Preferential amortized black-box optimization, In The Thirteenth International Conference on Learning Representations.
7, 10
15

---

Appendix
The appendix is organized follows:
In Appendix A, we provide detailed derivations and proof's for the theoretical claims made regarding information gain and variational bounds.
In Appendix B, we present the complete training algorithm and the specifics of the ALINE model.
In Appendix C, we provide comprchensive details for each experimental setup, including task descriptions and bascline implementations.
In Appendix D, we present additional experimental results, including further visualizations, performance more benchmarks, and analyses of inference times.
In Appendix E, we provide an overview of the computational resources and software dependencies for this work.
A Proofs of theoretical results
A.1 Derivation of total EIG for 05
Following Eq. 3, we can write the expression for the total expected information gain sEIGgg about parameter subset 0g 0 given data Dy generated under policy Te as:
sEIGgg (w) H[p(0s)] + Eptes, Dr (x Dog p(0s Dr)),
(A1)
E:
where p(0s) is the marginal prior for Os. and p(0s, Dr ).) is the joint distribution of 0g and Dr under Ty. Now, let OR Os be the remaining component of not included in 0g. Then, we can express E1 from Eq. Al as
E
1(0s(0s| ( (r)pp(ss, Dr kx)d0s
/ logp(0s (DT) / p(0s, OR. Dr d00 d0s
logr(0s DT) p(0, Dr mx)d0
Ep( (T) Po [T)]].
(A2)
Plugging the above expression in Eq. A1 and noting that p(0, Dr (x)) = p(0)p(Dr |0, T.)) we arrive at the expression for sEIGa, in Eq. 7.
A.2 Proof of Proposition 1
Proposition (Proposition 1). The total expected predictive information gain for a design policy T over a data trajectory of length T is:
sEPIG()) = E. ( (H[(() |=*)] - Hlply" ,²,Dr)]]
(Pp(t)p(Dr mo.0).. (x" (ply" =*.0) [logply" x=", Dr)] + Fp. (H(("" [=*)]]]
Proof. Let p. (xx) ) be the target distribution over inputs x* for which we want to improve predictive performance. Let gd be the corresponding target output. The single-step EPIG for acquiring data (x, y) measures the expected reduction in uncertainty (entropy) about y* for random target x* ~ (x*):
EPIG(x)) Ep. (x* Ip(u (w) Hpp("" |x=)] H [p(y" | v)]]-
Following Theorem in [23], the total EPIG, is the total expected reduction predictive entropy
from the initial prediction p(y" x=*) to the final prediction based on the full history p(y" | Dr):
sE =
Ep, (ar )p(Dr (x) ) [M [p(y" |x*)] - H [p(y" x= ,D)]]]
(A3)
Ep, (ar )p(, (w) [ppp lr .r, [lyg |=", Dr)ll + Ep. (xr) [H |x*)]]
(A4)
= Ep, far "y | [ly" |0= Dr)] + (p. (H [a*)]]-
(A5)
16

---

Here, Eq. A3 follows from conditioning EPIG on the entire trajectory Dr instead of single data point y, Eq. A4 follows from the definition of entropy H and Eq. A5 follows from noting that
p( (DT,,y +T)**) p(Dr kw)p(y" xx' Dr). Next, we combine the expectations and express the
joint distribution p(Dr, y² m. = p(0)p((T) Ty. 0)p(y* 1 0)df, where, following [67],
we assume conditional independence between Dr and y" given O. This yields:
( By(0)p(Dr | ke.0), (xr)p(r 1 1=,,0) [log(y" ,0" Dr)] (x, (H(')""(((( |x<[)=>>))]]]
which completes our proof.
A.3 Proof of Proposition 2
Proposition (Proposition 2). Let the policy Th generate the trajectory Dr. With (0(0s
approximating p(0s (D),, and golv" (r) approximating p(y" |x", ,Dr), we have (b - aEIGo,(1) and Jy. (1b) S sEPIG(()). Moreover,
sEIGos (1b) (b) By(Dr (m)) (L(p((s (ra (0r)))))) Dr))], SEPIG( Jp. (1b) = (x mm (KL(py) a Dr)ll x², Dr))]-
and
(A6)
(A7)
Proof. Using the expressions for sEIGG and Th from Eq. 7 and Eq. 8, respectively. and noting that log (0s Dr) = Eies log gs Dr), we can write the expression for sEIGa, (+i) - -T (1) as:
sEIGas (1b) - IZ(+)= = Ep(0)p(Pr mw. flog [logg(0s DT)]
Ep(0)((Dr 1) [log 4s (0s Dr)
log p(0s| Dr)
(p(Dr,00xo)
(A8)
L (0s Dr)
log p Dr)
(p(DD,0sITT)
(A9)
La (0s Dr)
E Ep(as Dr) log p(0s) ( Dr)
= Ep(Dr mm)
(A10)
(0s Dr)
= EppDr |m,)
[KI(p(0s] (Dr) 1r)) (D)))]-
(A11)
Here, Eq. A8 follows from the fact ((0)pp(p)))) Tw. 0) p((DDr Tw). Eq. A9 follows from Eq. A2, Eq. A10 follows from the fact that pl (rr, 0s TT.)) pDDr p+)s(0s) Dr), and Eq. All follows from the definition of KL. divergence.
Since the KL divergence always non non-negative (KL(PIIQ) 2 0), its expectation over trajectories p(Dr [T) must also be non-negative. Therefore:
(&h) S sEIGg, (v)
(A 12)
Now, we consider the difference between sEPIG and Jy. (v²);
log ply" xx,, (T)
SEPIG( - Jy. (bb Ep(O)p(Dr mg.0)). (p()" |
4oh (yy |x*, Dr)
(A13)
ply ;²+
Ep, (m (p(P) I m Bpls" | (D))
log
go (y xxt, (,,DDr)
Similar to the previous case, the inner expectation is the definition of the KL divergence between the true posterior predictive ply" xx* Dr) and the variational approximation qo("" { Dr):
SEPIG (+)) Jy. (v) Ep, (xr PPr (T) [LL(" |x Pr))llo()) |2" (r)rr
(A14)
Since the KL. divergence is always non- non-negative, therefore:
TH. (+) < sEPIG
(A15)
which completes the proof.
17

---

B Further details on ALINE
B.1 Training algorithm
Algorithm ALINE Training Procedure
1: Input: Prior p(0), likelihood p(y zz, 0), target distribution p(E), query horizon T, total training episodes Emaax warm-up episodes Ewarr-
2: Output: Trained ALINE model (4d. TT).
3: for epoch = to Emar do
4: Sample parameters 0 p(().
5: Sample target specifier set Ç ~ p(5) and corresponding targets Os or {rm Dm JM
6: Initialize candidate query set Q.
7:
for / 1 to T do
8:
if epoch & Ewarm then
9:
Select next query Ia uniformly at random from Q.
else
11:
Select next query Ia (- (-1 1,) from Q.
end if
13:
Sample outcome yt p(yl Iz. xt. 0). Update history D, D,-1 ((r1)
15:
Update query set Q 4 @ {
epoch S Ewarr then
17:
C = LNLL (Fq. 12)
18:
else
19:
Calculate reward R2 (Eq. 10) L LNLL (Eq. 12) + LPG (Eq. 11)
20:
21:
end if
22: Update ALINE using& L.
23:
end for
24:
end for
B.2 Architecture and training details
In ALINE, the data first processed by different embedding layers. Inputs (context r, query candidates z4, target locations zf}) are passed through shared nonlinear embedder fa. Observed outcomes Y/ are embedded using a separate embedder fy. For discrete parameters, we assign unique indicator & to each parameter 0, which is then associated with unique, learnable embedding vector, denoted as f(().. We compute the final context embedding by summing the outputs of the respective embedders: E ((fa(z,) + (y()))) Query and target sets embedded as and ET (either [[f= } hEs)))) Both fa and fy are MLPs
g [(f (zf))) consisting of an initial linear layer, followed by a ReLU activation function, and final linear layer. For all our experiments, the embedders embedding dimension of 32.
feedforward dimension of 128 and project inputs
The core of our architecture is transformer network. We employ configuration with 3 transformer layers, each equipped with attention heads. The feedforward networks within each transformer layer have dimension of 128. The model's internal embedding dimension, consistent across the transformer layers and the output of the initial embedding layers, is 32. These transformer layers process the embedded representations of the context, query. and target sets. The interactions between these sets are governed by specific attention masks, visually detailed Figure AI, where a shaded element indicates that the token corresponding to its row is permitted to attend to the token corresponding to its column.
ALINE has two specialized output heads. The inference head, responsible for approximating posteriors and posterior predictive distributions, parameterizes a Gaussian Mixture Model (GMM) with 10 components. The embeddings corresponding to the inference targets are processed by 10 separate MIPs, one for each GMM component. Fach ML. outputs parameters for component; a mixture weight, mean, and standard deviation. The standard deviations passed through Softplus
18

---

(#1.
(a., sk)
(#1.
zf
(.,
>
4
(syrth)
(suth)
(Ztr
(a)
(b)
Figure A1: Example attention masks in ALINE's transformer architecture. (a) Mask for a predictive
target& { = Sp. (b) Mask for parameter target <Qn) Shaded squares indicate allowed attention.
activation function ensure positivity, and the mixture weights are normalized using Softmax function. The policy head, which generates a probability distribution over the candidate query points, 2-layer MLP with a feedforward dimension of 128. Its output passed through Softmax function to ensure that the probabilities of all actions sum to unity. The architecture of ALINE is shown in Figure 2.
ALINE is trained using the Adam optimizer with a weight decay of 0.01. The initial learning rate is set to 0.001 and decays according to cosine annealing schedule.
C Experimental details
This section provides details for the experimental setups. Appendix C.1 outlines the specifics for the active learning experiments Section 4.1, including the synthetic function sampling procedures (Appendix C.1.1), implementation details for baseline methods (Appendix C.1.2), and training and evaluation details for these tasks (Appendix C.1.3). Next, Appendix C.2 we describe the details of BED tasks, including the task descriptions (Appendix C.2.1), implementation of the baselines (Appendix C.2.2), and the training and evaluation details (Appendix C.2.3). Lastly, Appendix C.3 contains the specifics of the psychometric modeling experiments, detailing the psychometric function we use (Appendix I. and the setup for the experimental comparisons (Appendix C.3.2).
C.1 Active learning for regression and hyperparameter inference
C.1.1 Synthetic functions sampling procedure
For active learning tasks, ALINE is trained exclusively on synthetically generated Gaussian Process (GP) functions. The procedure for generating these functions is as follows. First, the hyperparameters of the GP kernels, namely the output scale and lengthscale(s), distributions. For multi-dimennsonal input spaces (d > 1), there is Piso 0.5 probability that an isotropic kernel is used, meaning that all input dimensions share common lengthscale. Otherwise, an anisotropic kernel is employed, with a distinct lengthscale sampled for each input dimension. Subsequently, kernel function is chosen randomly from a pre-defined set, with each kernel having uniform probability of selection. In our experiments, we utilize the Radial Basis Function (RBF), Matérn 3/2, and Matérn 5/2 kernels.
sampled from their respective prior
The kernel's 's output scale is sampled uniformly from the interval (0.1,, The lengthscale(s) are sampled from (0.1, 2) x ddg. Input data points x are sampled uniformly within the range [-5, for each dimension. Finally, Gaussian noise with fixed standard deviation of 0.01 is added to
19

---

Sample 1: 22, sode 0.33
Sample la 1.5,, scale 0.29
Sample 3: lam 192, scale
Sample la 0.33 scale 0.1
2
Sample lu sode
Sample 6: Is 1.01, soae
Sanple T: 0 saal 0.99
Semple w565 scale 0.44
-2.5
0.0
-25
Sample
Somple
Sumple
Sumple
-5.0
225
0.0
5.0
-3.0
-2.5
0.0
2.5
5.0
-5.0
-2.5
0.0
25
5.0
-5.0
-2.5
0.0
2.5
5.0
Sample 13: la 061, scalem 0:34
Sample 16:h m 0.54 wale
Sumple 15: 0.88 sode 0.41
Sumple 16: b 0.24 sodl 0
-5.0
-2.5
Sample 17: 102 stalemOU1 0.17
Sample 18: 1.19 wale 0.54
Sumple 10: la 0.16 sall 0.72
Sumple 20: bom 182, sode
255
0.0
0.0
Figure A2: Examples of randomly sampled ID synthetic GP functions used to train ALINE.
the true function output y for each sampled data point. Figure A2 illustrates some examples of the synthetic GP functions generated using this procedure.
C.1.2 Details of acquisition functions
We compare ALINE with four commonly used AL acquisition functions. For Random Sampling (RS), we randomly select one point from the candidate pool as the next query point.
Uncertainty Sampling (US) is simple and widely used AL acquisition strategy that prioritizes points where the model is most uncertain about its prediction:
US(x) = VVly x. D),
(A16)
where Vly x, DJ is the predictive variance given the current training data D.
Variance Reduction (VR) [74] aims select a candidate point that expected maximally reduce the predictive variance over pre-defined test set {x), LM
Jm=1, which is defined as:
.M 2m=1 (Cov post (x*, ,a))²
VR(x)
(A17)
V/y 2 D
Covpost (x*, xx) is the posterior covariance between the latent function values at x and x, given the history D {(Xuain- }r where Xuain comprises all currently observed inputs with Strain being their corresponding outputs. It is computed as:
Covpost posa (a* x)) k<xx x²) k((= Xmmin
((Xurainnn))))) X)).
(A (A18)
Here, k(.))) ) is the GP kernel function, K Kiain 6(Xirain Xirim), and a is the noise variance.
Expected Predictive Information Gain (EPIG) [67] measures the expected reduction in predictive uncertainty on a target input distribution p. (a*). Following Smith et al. [67], for Gaussian predictive
distribution, the EPIG for a candidate point can be expressed as:
V
xr DJV("" x*
EPIG Fpp.(x log
(A19)
Vyy
x.
D
D
x D - Covpou x2)
20

---

In practice, we approximate it by averaging over m sampled test points:
M
Vly x, DJV Wuth
Im D
(PI 2
log
(A20)
2M
m2= V/y x. DJV Win Fim DJ Covpost(177) x)2
C.1.3 Training and evaluation details
For both ID and 2D imput scenarios, ALINE is trained for 2 . 10 epochs using batch size of 200. The discount factor y for the policy gradient loss is set to 1. For the GP-based baselines, we utilized Gaussian Process Regressors implemented via the scikit-leans library [56]. The hyperparameters of the GP models are optimized at each step. For the ACE baseline [15], we use a transformer architecture and inference head design consistent with our ALINE model.
All active learning experiments are evaluated with a candidate query pool consisting of 500 points. Each experimental run commenced with an initial context set consisting of single data point. The target set size for predictive tasks is set to 100.
C.2 Benchmarking Bayesian experimental design tasks
C.2.1 Task descriptions
Location Finding [66] is a benchmark problem commonly used in BED literature [24, 40, 9, 41]. The objective is to infer the unknown positions of K hidden sources, {0p R Tk=1 by strategically selecting sequence of observation locations, x Rd Each source emits signal whose intensity attenuates with distance following an inverse-square law. The total signal intensity at an observation location is given by the superposition of signals from all sources:
p( x) x) -
(A21)
04 x=
where are known source strength constants, and 6, m > 0 are constants controlling the background level and maximum signal intensity, respectively. In this experiment, we use K 1, d 2, 0g
b = 0.1 and m = 10-1, and the prior distribution over each component of a source's 's location
0g (fp,1, ...., 0R,a) is uniform over the interval [0, 1].
The observation is modeled as the log- log-transformed total intensity corrupted by Gaussian noise: log y 0, ((og ((0, (0),x)), o
(A22)
where we use a = 0.5 in our experiments.
Constant Elasticity of Substitution (CES) [3] considers behavioral economics problem in which a participant compares two baskets of goods and rates the subjective difference utility between the baskets on sliding scale from 0 to 1. The utility of a basket z, consisting of K goods with different values, is characterized by latent parameters f (p, ax u). The design problem is select pairs of baskets, (z,z') E [0. 100]2K to infer the participant's latent utility parameters.
The utility of basket is defined using the constant elasticity of substitution function, as:
U(2)
Σ
2
(A23)
The prior of the latent parameters specified as: Beta(1, 1) a Dirichlet(1x)) log u ~ (1,
(A24)
The subjective utility difference between two baskets is modeled as follows: n ~ / ( (U(x) U (z')) w² 72 (1+ || ='11²²)
(A25)
y clip(sigmoid()) e).
In this experiment, choose K = 3, 0.005 and 2 22
21

---

C.2.2 Implementation details of baselines
We compare ALINE with four baseline methods. For Random Design policy, we randomly sample design from the design space using a uniform distribution.
VPCE [25] iteratively infers the posterior through variational inference and maximizes the myopic Prior Contrastive Estimation (PCE) lower bound by gradient descent with respect to the experimental design. The hyperparameters used in the experiments are given in Table A1.
Table Al: Additional hyperparameters used in VPCE [25].
Parameter
Location Finding
CES
VI gradient steps
1000
1000
VI learning rate
10-3 3
10 10-3
Design gradient steps
2500
2500
Design learning rate
10-3
10-3
Contrastive samples L
500
10
Expectation samples
500
10
Deep Adaptive Design (DAD) [23] learns an amortized design policy guided by sPCE lower bound. For design policy т, and L 2 0 contrastive samples, sPCE over sequence of T experiments is defined as:
p (Dr 0. π)
LT(TT L.) = Rpt Dr(s)(ra))) log
(A26)
L+I Et= (Dr
02,
where the contrastive samples 0I.,. are drawn independently from the prior p(0). The bound becomes tight as L, + oo, with a convergence rate of O(L-1).
The design network comprises an MLP encoder that encodes historical data into a exedddimensional representation, and an MLP emitter that proposes the next design point. The encoder processes the concatenated designnobservation pairs from history and aggregates their representations through pooling operation.
Following the work of Foster et al. [23], the encoder network consists of two fully connected layers with 128 and 16 units with ReLU activation applied to the hidden layer. The emitter is implemented as a fully connected layer that maps the pooled representation to the design space. The policy is trained using the Adam optimizer with an initial learning rate of 5 10-5, B = (0.8, 0.998), and an exponentially decaying learning rate, reduced by factor of 7 0.98 every 1000 epochs. In the Location Finding task, the model is trained for 10a epochs, and 10 contrastive samples are utilized in each training step for the estimation of the sPCE lower bound. Note that, for the CES task, we applied several adjustments, including normalizing the input, applying the Sigmoid and Softplus transformations to the output before mapping it the design space, increasing the depth of the network, and initializing weights using Xavier initialization [32]. However, DAD failed to converge during training in our experiments. Therefore, we report the results provided by Blau et al. [9].
vsOED [65] is an amortized BED method that employs an actor-critic reinforcement learning framework. It utilizes separate networks for the actor (policy). the critic (value function), and the variational posterior approximation. The reward signal for training the policy is the incremental improvement in the log- probabbibyyl of the ground-truth parameters under the variational posterior, which is estimated at each step of the experiment. Following the original implementation, distinct posterior network is trained for each design stage, while the actor and critic share a common backbone. For our implementation, the hidden layers of all networks are 3-layer MLPs with 256 units and ReLU activations. The posterior network outputs the parameters for an 8-component Gaussian Mixture Model (GMM). The input to the actor and critic networks is the zero-padded history of design-observation pairs, concatenated with one-hot encoding of the current time step. We train for 10t epochs with a batch size of 2 and a 108-sized replay buffer. The learning rate starts at 10-3 with 0.9999 exponential decay per epoch, and the discount factor is 0.9. To encourage exploration for the deterministic policy, Gaussian noise is added during training: the initial noise scale is 0.5
22

---

for the Location Finding task and 5.0 for the CES task, with decay rates of 0.9999 and 0.9998, respectively.
RI-BOEED [9] frames the design policy optimization as a Markov Decision Process (MDP) and employs reinforcement learning to learn the design policy. It utilizes stepwise reward function to estimate the marginal contribution of the t-th experiment to the sEIG.
The design network shares similar backbone architecture to that of DAD, with the exception that the deterministic output of the emitter is replaced by a Ganh-Gaussian distribution. The encoder comprises two fully connected layers with 128 units and ReLU activation, followed by output layer with 64 units and no activation. Training is conducted using Randomized Ensembled Double Q-learning (REDQ) [16], the full configurations are reported in Table A2.
Table A2: Additional hyperparameters used in RL-BOED [9].
Parameter
Location Finding
CES
Critics
2
Random subsets
2
2
Contrastive samples L
10
105
Training epochs
10
5 10
Discount factor y
0.9
0.9
Target update rate
10 10-3
5 10-3
Policy learning rate
10
3 10
Critic learning rate
3 10
3 1-10-
Buffer size
10
10c
C.2.3 Training and evaluation details
In the Location Finding task, the number of sequential design steps, T, to 30. For all evaluated methods, the sPCE lower bound is estimated using L = 10" contrastive samples. The ALINE is trained over 10 epochs with batch size of 200. The discount factor y for the policy gradient loss is set to 1. During the evaluation phase, the query set consists of 2000 points, which are drawn uniformly from the defined design space. For the CES task, each experimental run consists of T 10 design steps. The sPCE lower bound is estimated using L 10" contrastive samples. The ALINE is trained for 10. epochs with a batch size 200, and we use 2000 for the query set size.
C.3 Psychometric model
C.3.1 Model description
In this experiment, we use a four-parameter psychometric function with the following parameteriza-
tion:
. <(x) = 0y 04 + (1-81)FF(
02
where:
0, (threshold): The stimulus intensity at which the probability of a positive response reaches specific criterion. It represents the location of the psychometric curve. We use uniform prior U[-3,0 for 0,-
82 (slope): Describes the steepness of the psychometric function. Smaller values of 0g indicate a sharper transition, reflecting higher sensitivity around the threshold. We use uniform prior U [0.1, 2) for 02.
O3 (guess rate): The baseline response probability for stimuli far below the threshold, reflecting responses made by guessing. We use a uniform prior U(0,1, 0.9) for 03.
01 (lapse rate): The rate which the observer makes emoor independent of stimulus intensity, representing an upper asymptote on performance below We use uniform prior U00 0.5] for 04.
23

---

We employ Gumbel- type internal link function F = --- where 2 == =-61, Lastly, binary response y is simulated from the psychometric function (x) using Bernoulli distribution with probability of success p = (2).
C.3.2 Experimental details
We compare ALINE against two established Bayesian adaptive methods:
QUEST+ [70]: QUEST+ is an adaptive psychometric procedure that aims to find the stimulus that maximizes the expected information gain about the parameters of the psychometric function, equivalently, minimizes the expected entropy of the posterior distribution over the parameters. It typically operates on a discrete grid of possible parameter values and selects stimuli to reduce uncertainty over this entire joint parameter space. In our experiments, QUEST+ configured to infer all four parameters simultaneously.
Psi-marginal [58]: The Psi-marginal method is an extension of the psi method [43] that allows for efficient inference by marginalizing over nuisance parameters. When specific parameters are designated as targets of interest, Psi-marginal optimizes stimulus selection to maximize information gain specifically for these target parameters, effectively treating the others as nuisance variables. This makes it highly efficient when only a subset parameters is critical.
For each simulated experiment, true underlying parameters are sampled from their prior distributions. Stimulus values x are selected from a discrete set of size 200 drawn uniformly from the range [-5, 5].
D Additional experimental results
D.1 Active Exploration of High-Dimensiona. Hyperparameter Landscapes
To demonstrate ALINE'S utility on complex, high-dimensional tasks, we conduct a new set of experiments actively exploring hyperparameter performance landscapes. This experiment aims to efficiently characterize a machine learning model's 's overall behavior on a new task, allowing practitioners to quickly assess model family's viability understand its sensitivities. The task to actively query a small number of hyperparameter configurations to build a surrogate model that accurately predicts performance for a larger, held-out of target configurations. We use high-dimensional, real-world tasks from the HPO-B benchmark [2], evaluating on rpart (6D), svm (8D), ranger (9D), and xgboost (16D) datasets. ALINE is trained on their multiple pre-defined training sets. We then evaluate its performance, alongside non-amortized GP-based baselines and an amortized surrogate baseline (ACE-US), on the benchmark's held-out and entirely unseen test sets.
Table A3 shows the RMSE results after 30 steps, averaged across all test tasks for each dataset. First, both amortized methods, ALINE and ACE-US, significantly outperform all non-amortized GP-based baselines across all tasks. This highlights the power of meta-learning in this domain. GP-based methods must learn each new performance landscape from scratch, which is highly inefficient in high dimensions. In contrast, both ALINE and ACE-US are pre-trained on hundreds of related tasks, and their Transformer architectures meta-learn the structural patterns common these landscapes. This shared prior knowledge allows them to make far more accurate predictions from sparse data. Second, while ACE-US performs strongly due to its amortized nature, ALINE consistently achieves the best or joint-best performance. This demonstrates the additional, crucial benefit of our core contribution: the learned acquisition policy. ACE-US relies on standard heuristic, whereas ALINE'S policy trained end-to-end to learn how to optimally explore the landscape, leading to more informative queries and ultimately a more accurate final surrogate model.
D.2 Active learning for regression and hyperparameter inference
AL results on more benchmark functions. To further assess ALINE, we present performance evaluations on an additional set of active learning benchmark functions, see Figure A3. The results on Gramacy and Branin show that we are par with the GP baselines. For the Three Hump Camel, we see both AI.NEE and ACE-US showing reduced accuracy. This because the function's 's output value range extends beyond that of the GP functions used during pre- praitngin This highlights potential
24

---

Table A3: RMSE (4) on the HPO-B benchmark after 30 active queries. Results show the mean and
95% CI over all test tasks for each dataset.
GP-RS
GP-US
GP-VR
GP-EPIG
ACE-US
ALINE (ours)
1part (6D)
0.07 1 0.03
0.02
0.04 1.04100 0.02
0.05 0.02
0.01 1 0.00
0.0110.00
svm (8D)
0.22 0.11
0.11
0.12 0.07
0.15 0.08
0.04 0.04:0..00
0.03
ranger (9D)
0.10
0.07
0.08
0.08 0.02
0.02 0.01
0.02=0.01
xgboost (16D)
0.09 - 0.02
0.09 + 0.02
0.09 + 0.02
0.09 + 0.02
0.04 0.01
0.03=0.01
Gramacy 1D
Branin 2D
Three Hump Camel 2D
0.5
0.8
2.5
RMSB
0.3
0.6
2.0
1.5
0.4
0.2
1.0
0.2
0.5
0.0
0.0
0.0
20
25
30
10
20
30
50
10
20
30
40
50
GP-RS -
GP-US
GPVR
GP-EPIG
ACE-US
ALINE (ours)
Figure A3: Predictive performance in terms of RMSE three other active learning benchmark functioos. Results show the mean and 95% confidence interval (CI) across 100 runs. Notably, on the Three Hump Camel function, the performance of amortized methods like ALINE and ACE-US limited, as its output scale significantly differs from the pre-training distribution, highlighting a scenario of distribution shift.
area for future work, such as training ALINE on a broader prior distribution of functions, potentially leading to more universally capable models.
Acquisition visualization for AL. To qualitatively understand the behavior of our model, visualize the query strategy employed by ALINE for AL on a randomly sampled synthetic function Figure A4, This visualization illustrates how ALINE iteratively selects query points to reduce uncertainty and refine its predictive posterior.
Hyperparameter inference visualization. We now visualize the evolution of ALINE'S estimated posterior distributions for the underlying GP hyperparameters for a randomly drawn 2D synthetic GP function (see Figure A5). The posteriors shown after 1, 15, and 30 active data acquisition steps. As ALINE strategically queries more informative data points, its posterior beliefs about these generative parameters become increasingly concentrated and accurate.
assess computational efficiency ALINE, we report the inference times for the AL tasks in Table A4. The times represent the total duration to complete sequence of 30 steps for 1D functions and 50 steps for 2D functions, averaged over independent runs. As both ALINE and ACE-US perform inference via a single forward pass per step once trained, they are significantly faster compared to traditional GP-based methods.
Inference time.
Table A4: Comparison of inference times (seconds) for different AL methods on ID (30 steps) and 2D (50 steps) tasks. Values are averaged over runs (mean + standard deviation).
Inference time (s)
Methods
ID & 30 steps
2D & 50 steps
GP-US
0.62+10.09
1.72++0.23
GP-VR
1.41
4.03
GP-EPIG
1.34
3.224
ACE-US
0.08
0.19++0.02
ALINE
0.08 +0.00
0.19+
25

---

Step
Suep 2
Step
Slep 4
Step 5
Grrand Terght
Step
Step
Step
Step 9
Step 10
Step 11
Step 12
Stop 13
Stop 14
Stop 15
Step 16
Step
Slep
Step 19
Stop 20
Figure A4: Sequential query strategy of ALINE 1D synthetic GP funetion over 20 steps. As more points are queried, the model's prediction increasingly aligns with the ground truth, and the uncertainty is strategically reduced.
D.3 Benchmarking on Bayesian experimental design tasks
In this section, we provide additional qualitative results for ALINE'S performance on the BED benchmark tasks. Specifically, for the Location Finding task, we visualize the sequence of designs chosen by ALINE and the resulting posterior distribution over the hidden source's location (Figure A6). For the CES task, we present the estimated marginal posterior distributions for the model parameters, comparing them against their true underlying values (Figure A7). We see that ALINE offers accurate parameter inference.
D.4 Psychometric model
Demonstrations of flexibility. We conduct two ablations to explicitly validate ALINE'S flexible
targeting capabilities.
First, we test the ability to switch targets mid-rollout. We configure a single experiment where for the first 15 steps, the target is threshold & slope parameters, and at step 16, the target is switched to the guess rate lapse rate. As shown in Figure A8(a), ALINE'S acquisition strategy adapts immediately and correctly, shifting its queries from the decision threshold region to the extremes of the stimulus range gain maximal information about the new targets.
Second, we test generalization to novel target combinations. A single ALINE model is trained handle distinct targets separately: (1) threshold & slope and (2) guess & lapse rate. At deployment, we task this model with a novel, unseen combination: targeting all four parameters simultaneously. As shown in Figure A8(b), the resulting policy is sensible mixture of the two underlying strategies it has learned, strategically alternating queries between points near the decision threshold and points at the extremes. This confirms that ALINE can successfully compose its learned strategies to generalize to new inference goals at runtime.
Inference time. additionally assess the computational efficiency of each method proposing the next design point. The average per-step design proposal time, measured over the 30-step psychometric experiments across 20 runs, is 0.002 + 0.00s for ALINE, 0.07 + 0.00s for QUEST+, and 0.02 + 0.00s for Psi-marginal. Methods like QUEST+ and Psi-marginal, which often rely
26

---

t=1
t= 15
t= 30
Lengthscale Posterior
Lengthscale Posterior
Lengthscale Posterior
03
0.8
1.00
Valle
1.23
Density
Density
1.00
0.4
0.75
02
0.1
0.2
Lengthscale Value
Lengthscale Value
Tengthsrale Value
Lengthscale Posterior
Lengthscale 2 Posterior
Lengthscale 2 Posterior
0.6
1.2
1.0
0.3
Density
0.4
Density
0.8
0.3
04
0.4
Lengthscale 2 Value
Lengthscale 2 Value
Lengthscale Value
Scale Posterior
Scale Posterior
Scale Posterior
1.25
3.0
Density
0.75
Density
1.5
Densily
0.50
0.00
0.0
0.00
0.25
0.50
0.75
2.00
0.25
0.50
0.75
1.00
0.2
0.4
Scale Value
Scale Value
Scale Value
Figure A5: Estimated posteriors for the two lengthscales and the output scale obtained from ALINE after / 1, = 15, and = 30 active query steps. The posteriors progressively concentrate around the true parameter values as more data is acquired.
grid-based posterior estimation, face rapidly increasing computational costs the parameter space dimensionality or required grid resolution grows. ALINE, however, estimates the posterior via the transformer in single forward pass, making its inference time largely insensitive these factors. Thus, this computational efficiency gap is anticipated to become even more pronounced for more complex psychometric models.
E
Computational resources and software
All experiments presented in this work, encompassing model development, hyperparameter opti- mization, baseline evaluations, and preliminary analyses, are performed on GPU cluster equipped with AMD MI250X GPUs. The total computational resources consumed for this research, in- cluding all development stages and experimental runs, are estimated to be approximately 5000 GPU hours. For each experiment, it takes around 20 hours train an ALINE model for 10a epochs. The core code base is built using Pytorch (https:/ //pytorch org/ License: modi- fied BSD license). For the Gaussian Process (GP) based baselines, utilize Scikit-learn [56] ftttps:/sccikit
org/, License: modified BSD license). The DAD baseline is adapted from the original authors' publicly available code [23] (https: http//ttthuhbeeee/////// com ae foster / dad; MIT License). Our implementations of the RL-BOED and vsOED baselines are adapted from the of- ficial repositories provided by [6] (https://githut com//assirarrlas/L...... BOED; MIT License) and [65] /https://gitte thub. com/ugshen/vsOED MIT License), respectiveey We use questplus package (https: //githh wttm///occcebbbrrgrrruueuupppppp License: GPL-3.0) to implement QUEST+, and use Psi- sti-ssaarcass (https://gi cow//Michoo//Piiieeeeeeeoooae htt.://ssssstirrcsseeeee License: GPL-3.0) to implement the Psi-marginal method.
27

---

Location Finding
-300
23 Loje doy fotterso
0.6
900
0.2
1500
1800
Figure A6: Visualization of ALINE's design policy and resulting posterior for the Location Finding task. The contour plot shows the log posterior probability density of the source location (true location marked by blue star) after T = 30 steps. Queried locations, with color indicating the time step of acquisition, demonstrating concentration of queries around the true source.
40
2
p(0)
1
1.0
0.0
5.0
log(u)
Figure A7: ALINE'S estimated marginal posterior distributions for the parameters of the CES task after 7' 10 query steps. The dashed red lines indicate the true parameter values. The posteriors well-concentrated around the true values, demonstrating accurate parameter inference.
5
Stimuli Values
0
-5
10
20
30
10
20
30
Number of Steps f
Number of Steps r (b)
(a)
Figure A8: Demonstration of ALINE'S runtime flexibility on the psychometric task. (a) The acqui- sition strategy adapts after the inference target switched mid-rollout from (threshold & slope) to (guess & lapse rate). (b) When tasked with a novel combined target (all four parameters), the policy generalizes by mixing the two distinct strategies it learned during training.
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In this paper, propose ALINE, unified framework for amortized Bayesian inference and active data acquisition (Section 3), addressing the desiderata delineated in the abstract and introduction regarding simultaneous, flexible, and efficient data-driven learning. We explicitly state ALINE'S s key contributions, such as its targeted querying and capabilities, the end of the introduction. The experimental
joint pieerence-acquisiie results presented in Section 4 across various tasks (active learning, BED, psychometries) match and validate the claims made about ALINE'S performance and utility.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
It is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations of our work are discussed explicitly in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specificattin, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
29

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Complete proofs for the theoretical results are provided in Appendix A. Guidelines:
The answer NA means that the paper does include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proof' can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented
by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Details for all our experiments and the architecture used in our paper available in the Appendix B and C, as well as in the provided code.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should possible for other researchers have some path reproducing verifying the results.
30

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
We have included the code in the supplementary materials, and it will be publicly released.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nipe cc/
sei.iiiiisiioiiiooocoooci for more details. Bealliicguides/CceeSttiieee
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //iiss... wwwpcc/pppc/.aee.ec..eccce.cetle pppplc//guudde....p//ppppp /ttia//uuiides/oeuuubiiiiiiii /Ct////oddissiiiooooc ssiooooliiii/ for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
We describe the experimental settings for all experiments in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All the experiment results are accompanied by error bars. We also report the method for calculating the error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
31

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide an estimate of the resources needed for each experiment in Appendix E.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: h/terrip//////p...../ cc cctppblicc//ccccccccccccccctclct wawbllicEthicchuddelneeeee cstuidelineedde
Answer: [Yes]
Justification: Our work adheres the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The work is foundational research and not tied to particular applications. We believe there is no negative societal impact of the work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
32

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: For each asset (code, data) used in the paper provide citations, links and additional references as appropriate, in Appendix E.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
33

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. cow datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No specific additional assets introduced in the paper aside from research code, which will be released under a permissible MIT license.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
34

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard component.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy fttpp://neerips.t. htt.r/ppsccc/ooofeeenc///////230 for what should or should not be described.
35


Reviews:
[{"review_id": "EfBP1hQoap", "reviewer": "Reviewer_9JtZ", "summary": "This 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.", "questions": "1. **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.", "limitations": "Yes, the authors have discussed the limitations of their work.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Author's detailed rebuttal has addressed most of my questions, so I’m raising my score by 1.", "rebuttal": "Thank you for your positive assessment of our work and the points you raised. We address your points and questions below.\n\n> *1. The empirical study is confined to Gaussian-process benchmarks; broader experiments would strengthen the evidence.*\n\nWe would like to clarify that our evaluation already included diverse, non-GP settings, such as two classical BED benchmarks (Location Finding and CES) and a realistic psychometric modeling task with domain-specific baselines. We will mention more clearly in the paper.\n\nNevertheless, we fully agree that demonstrating scalability on even more complex problems is an excellent suggestion - a point also raised by other reviewers. To that end, we have conducted a new set of experiments on a practical high-dimensional task: **Active Exploration of Hyperparameter Performance Landscapes**. This experiment aims to efficiently characterize a machine learning model's overall behavior on a new task, allowing practitioners to quickly assess a model family's viability or understand its sensitivities. The task is to actively query a small number of hyperparameter configurations to build a surrogate model that accurately predicts performance for a larger, held-out set of target configurations.\n\nWe used the high-dimensional, real-world tasks from the **HPO-B meta benchmark** [1]. Specifically, we evaluated on the *rpart* (6D), *svm* (8D), *ranger* (9D), and *xgboost* (16D) datasets. ALINE was trained on their multiple pre-defined training sets. We then evaluated its performance, alongside non-amortized GP-based baselines and an amortized surrogate baseline (ACE-US), on the benchmark's held-out and entirely unseen test sets. This setup is explicitly designed to test the model's ability to generalize to new problems.\n\nThe table below shows the RMSE with 95% confidence interval after 30 steps, averaged across all test tasks for each dataset. First, **both amortized methods, ALINE and ACE-US, significantly outperform all non-amortized GP-based baselines** across all tasks. This highlights the power of meta-learning in this domain. GP-based methods must learn each new performance landscape from scratch, which is highly inefficient in high dimensions. In contrast, both ALINE and ACE-US are pre-trained on hundreds of related tasks and their Transformer architectures meta-learn the structural patterns common to these landscapes. This shared prior knowledge allows them to make far more accurate predictions from sparse data. Second, while ACE-US performs strongly due to its amortized nature, **ALINE consistently achieves the best or joint-best performance**. This demonstrates the additional, crucial benefit of our core contribution: the learned acquisition policy. ACE-US relies on a standard heuristic, whereas ALINE's policy is trained end-to-end to learn how to optimally explore the landscape, leading to more informative queries and ultimately a more accurate final surrogate model.\n\nWe believe this new experiment, alongside the existing BED and psychometric tasks, now provides a comprehensive and compelling demonstration of ALINE's practical utility. We will add these results to the revised manuscript.\n\n| RMSE ($\\\\downarrow$) | GP-RS | GP-US | GP-VR | GP-EPIG | ACE-US | ALINE |\n| :---- | :---- | :---- | :---- | :---- | :---- | :---- |\n| rpart (6D) | 0.07±0.03 | 0.04±0.02 | 0.04±0.02 | 0.05±0.02 | **0.01±0.00** | **0.01±0.00** |\n| svm (8D) | 0.22±0.11 | 0.11±0.05 | 0.12±0.07 | 0.15±0.08 | 0.04±0.01 | **0.03±0.01** |\n| ranger (9D) | 0.10±0.02 | 0.07±0.01 | 0.08±0.02 | 0.08±0.02 | **0.02±0.01** | **0.02±0.01** |\n| xgboost (16D) | 0.09±0.02 | 0.09±0.02 | 0.09±0.02 | 0.09±0.02 | 0.04±0.01 | **0.03±0.01** |\n\n\n> *2. 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\nWe appreciate the opportunity to elaborate on the concrete benefits of combining inference and active data acquisition. Our rationale is grounded in two core advantages:\n\n* **Significant Gains in Computational and Parameter Efficiency**: The most direct advantage of our unified framework is its efficiency. Both the inference task and the acquisition policy rely on the exact *same context history*. Our unified framework uses a *single, shared Transformer backbone* to process this context data. In contrast, a separate approach would lead to a near-doubling of model parameters, significantly increasing training costs and deployment latency, as two separate forward passes would be required for each step.  \n* **Improved Representation via Joint Learning**: Secondly, we believe our joint training framework could benefit from principles of multi-task learning (MTL) \\[2, 3\\]. The rationale is that the two tasks provide complementary learning signals that can mutually regularize the shared backbone. The inference objective pushes the model to learn underlying data patterns for accurate predictions, while the acquisition objective forces it to learn a robust representation of its own uncertainty to make informative queries. Forcing the model to learn a single representation that serves both objectives is a form of regularization that encourages a more generalizable understanding of the problem, compared to training on either task alone.\n\nWe will add the above discussion to both the method and the discussion section in the revised manuscript.\n\n\n> *3. Correct joint inference requires autoregressive prediction, yet this is deferred to future work. Why are the parameters and their inference not modeled autoregressively? An autoregressive treatment seems the principled approach.*\n\nGood point. Our decision to adopt a factorized approximation was a deliberate and critical **trade-off between modeling accuracy and computational tractability**. First, an autoregressive model would require sequential generation of each marginal dimension of the posterior, with each step conditioned on the previously generated ones. To compute the joint posterior's log-probability for $M$ targets, this would necessitate **$M$** **sequential forward passes** through the Transformer network. Since the reward for our policy agent is the improvement in this log-probability, calculating this reward at every single step of every training episode would become prohibitively expensive, as the cost would scale linearly with the number of targets. This would make training the acquisition policy **infeasibly slow**, especially for problems with many parameters or prediction targets. Therefore, ALINE adopts the mean-field approximation, which is a common and established practice in the Neural Process literature \\[4, 5, 6\\]. This allows for efficient, parallel computation of the marginals in a single forward pass. Also, our experiments demonstrate that even with this factorized assumption, ALINE achieves strong empirical performance, often outperforming competitive baselines. This suggests that for the tasks evaluated, the benefits gained from our joint amortization and flexible targeting strategy outweigh the limitations of the structural approximation of the posterior.\n\nWe agree that an autoregressive approach represents the most principled way to model the joint posterior distribution. We share your belief that developing more efficient methods for autoregressive inference in such models is a crucial and exciting direction for future research. As you rightly suggest, such advances are likely essential for scaling up to the next generation of foundation models in this field. For this work, however, we focused on establishing the core viability of the unified ALINE framework itself. We will add a discussion of this trade-off to the revised manuscript to make our design rationale clearer.\n\n> *4. 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\nThank you for raising this point. It is true that, in the general case, the posterior predictive outputs​ are not conditionally independent given the history, and thus the true objective should be based on the joint distribution. Our current Equation 6 is derived under the explicit assumption of a **factorized likelihood**, a choice we mention on line 144. This is a common and practical simplification in the Neural Process literature \\[4, 5, 6\\], adopted to maintain computational tractability. \n\nBut yes, the presentation could be clearer. We will revise to stating the objective first in a form that does not assume factorization, and then explicitly state that for practical implementation, we approximate this joint distribution with a factorized model.\n\n> *5. Consider relocating the Related-Work section to follow the Introduction to improve the paper’s narrative flow.*\n\nGood idea; we will update accordingly.\n\n**References**\n\n[1] Arango, et al., HPO-B: A Large-Scale Reproducible Benchmark for Black-Box HPO based on OpenML, NeurIPS Track on Datasets and Benchmarks (2021).\n\n[2] Liu, et al., Multi-task deep neural networks for natural language understanding, ACL (2019).\n\n[3] Yu, et al., Unleashing the Power of Multi-Task Learning: A Comprehensive Survey Spanning Traditional, Deep, and Pretrained Foundation Model Eras, arXiv preprint (2024).\n\n[4] Garnelo, et al., Conditional neural processes, ICML (2018).\n\n[5] Müller, et al., Transformers Can Do Bayesian Inference, ICLR (2022).\n\n[6] Chang, et al., Amortized Probabilistic Conditioning for Optimization, Simulation and Inference, AISTATS (2025)."}, {"review_id": "qCbsut5bT8", "reviewer": "Reviewer_gxPy", "summary": "This 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.", "questions": "Please see the comments in the Weaknesses part.", "limitations": "The authors discussed the limitations in the Section 6. They did not mention the Broader impacts but provided justification in the Checklist.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "__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.", "was_revised": true, "final_justification": "The authors responded to my concern by discussing the relevant related work and adding new comparison experiments. My main concern has been addressed, and I have accordingly increased my score.", "rebuttal": "Thank you for your constructive feedback. Below, we respond to the specific points raised.\n\n> *1. This paper has significant overlap with the already-published vsOED work.*\n\nThank you for pointing out this relevant vsOED [1] paper. There indeed are high-level similarities, particularly in using an RL framework with variational rewards for sequential design (similar to the RL-sCEE [2] we mentioned in our paper). \n\nHowever, we are confident that ALINE's core contributions are distinct and complementary to those of vsOED. To provide a direct and concrete comparison, **we have now added vsOED as a new baseline in our BED experiments.** Below, we outline the key distinctions and present the outcome of our new experiments.\n\n* The most critical distinction is ALINE's **dynamic targeting capability at runtime**. A single, pre-trained ALINE policy can instantly adapt its strategy to *any* user-specified target simply by changing an input to the model. vsOED, while flexible in its problem formulation (e.g., weighting between model and parameter EIG with its $\\\\alpha$ coefficients), would require training a new policy for each different design criterion. This on-the-fly adaptability is a central contribution of our work.  \n* ALINE introduces **a unified architecture** for both the inference and acquisition tasks. This design is not only computationally efficient by sharing a single Transformer backbone, but we also believe it benefits from jointly learning the two tasks (as elaborated in our response to Question #2 of Reviewer 9JtZ). The complementary objectives act as a form of mutual regularization, encouraging the model to learn a richer data representation. We believe this contributes to ALINE's strong performance, which is supported by our new empirical results showing that **ALINE outperforms vsOED on shared BED tasks**.  \n* While vsOED focuses on BED tasks, ALINE's unified framework naturally extends to active learning for predictive tasks. We introduce the **sequential Expected Predictive Information Gain (sEPIG)** and a variational lower bound for it. This is a specific methodological contribution aimed at the active learning domain.\n\n**New Empirical Comparison:** In our new experiments on the Location Finding and CES tasks, ALINE consistently achieves a higher PCE than vsOED (see the table below). For transparency, our re-implementation of vsOED achieves performance consistent with the results reported in their paper, which validates our comparison.\n\n| PCE ($\\\\uparrow$) | vsOED | ALINE |\n| :---- | :---- | :---- |\n| Location Finding | 7.30±0.06 | **8.91±0.04** |\n| CES | 12.12±0.18 | **14.37±0.08** |\n\nThank you again for prompting this important comparison. In our revised manuscript, we will add a summary of the above discussion of vsOED in the introduction and related work, and include it as a baseline in our results to make these distinctions clear.\n\n> *2. 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\nWe would like to respectfully clarify that this is a misunderstanding. ALINE **does not require an explicit or differentiable likelihood and naturally supports implicit likelihood models**, a key feature it shares with vsOED.\n\nOur reward signal, defined in **Eq. 10**, is derived from the log-probability of the ground truth parameters from the simulation, under our **variational posterior approximation** predicted by the neural network, i.e., $\\log q(\\theta|D)$. During training, we **never require evaluating the model's likelihood $p(y|\\theta,d)$**. The only interaction with the true generative process is the ability to **sample outcomes** ($y \\sim p(y|\\theta, x)$) to construct the experimental history. This is a fundamental property of methods based on a variational EIG lower bound.\n\nWe will clarify this in the paper by revising the methods section to explicitly state the likelihood-free nature of our approach, which allows ALINE to be applied to any simulation-based model [3].\n\n> *3. The paper lacks formal theoretical analysis regarding convergence properties, optimality guarantees, or theoretical relationships to information-theoretic objectives.*\n\nWe appreciate your emphasis on theoretical analysis. While our paper does include theoretical contributions, its focus is on the framework and its practical evaluation. In terms of the three points raised:\n\n* **Theoretical Relationships to Information-Theoretic Objectives:** Our framework is built directly upon and formally linked to these objectives. For parameter inference, our acquisition objective (Eq. 8\\) is a variational lower bound on the total Expected Information Gain (sEIG). This is a principled approach established in prior work [2, 4]. For predictive tasks, a key contribution of our work is the introduction of the sequential Expected Predictive Information Gain (sEPIG). In **Proposition 1 (with proof in Appendix A.2)**, we formally prove that our proposed objective (Eq. 9\\) is a valid variational lower bound on this new information-theoretic quantity.  \n* **Optimality Guarantees:** **Proposition 2** formally shows that the gap between our objective and the true information gain is precisely the KL divergence between our model's approximation and the true posterior. This means that in the limit of a perfect inference network (i.e., when the KL divergence is zero), our policy is optimized to maximize the **true sequential information gain** – which is the optimal criterion from an information-theoretic standpoint. While finding the globally optimal policy for this objective is generally intractable, our RL framework is a standard method for discovering high-performing policies that maximize this principled objective.  \n* **Convergence Properties:** The primary contribution of this paper is the unified and flexible framework for joint acquisition and inference, rather than developing a new reinforcement learning algorithm. Therefore, a formal convergence analysis of the specific RL algorithm is outside the scope of our contributions. We employed the standard REINFORCE algorithm for its simplicity and effectiveness. Its convergence properties can be referenced from existing works, such as [5].\n\nWe hope this clarifies the extent and focus of our theoretical contributions.\n\n> *4. (typo) Line 210: Should the \"locations\" here be deleted as not specifically discussing the Location Finding task here?*\n\nThank you for flagging this potential ambiguity. Our use of the term \"locations\" here was in the general sense of \"input points\" {$x^\\*\\_m$​}, a common term in the literature for points in a function's domain. However, we agree that this could be confusing given the specific \"Location Finding\" task discussed elsewhere. To prevent any misunderstanding, we will revise the manuscript and replace \"locations\" with the more general term **\"inputs\"** in this context.\n\n**References**\n\n[1] Shen et al., Variational sequential optimal experimental design using reinforcement learning, Computer Methods in Applied Mechanics and Engineering 444 (2025).\n\n[2] Blau et al., Statistically Efficient Bayesian Sequential Experiment Design via Reinforcement Learning with Cross-Entropy Estimators, arXiv preprint (2023).\n\n[3] Cranmer et al., The frontier of simulation-based inference, Proceedings of the National Academy of Sciences (2020).\n\n[4] Foster et al., Variational Bayesian Optimal Experimental Design, NeurIPS (2019).\n\n[5] Zhang et al., Sample efficient reinforcement learning with REINFORCE, AAAI (2021)."}, {"review_id": "ewEtImoIEZ", "reviewer": "Reviewer_xxdt", "summary": "Amortization 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.", "questions": "* 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?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\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).", "was_revised": true, "final_justification": "The reviewers answered a lot of the questions I had and ran several experiments on realistic and slightly higher dimensional tasks. The fundamental limitations in the current setup remain but the paper is stronger so I have increased my score.", "rebuttal": "Thank you for your positive feedback and questions. We address each of them in detail below.\n\n> *1. One of the main shortcomings of the paper is the limited scale of the empirical evaluations.*\n\nTo address the scalability concern, we have conducted a new set of experiments on a high-dimensional task: **Active Exploration of Hyperparameter Performance Landscapes.** The task is to actively query a small number of hyperparameter configurations to build a surrogate model that accurately predicts performance for a set of target configurations.\n\nWe use the real-world tasks from the HPO-B meta-benchmark [1]. Specifically, we evaluate on the *rpart* (6D), *svm* (8D), *ranger* (9D), and *xgboost* (16D) datasets. ALINE is trained on their multiple pre-defined training sets. We then evaluate it alongside non-amortized GP-based baselines and an amortized surrogate baseline (ACE-US), on the benchmark's entirely unseen test sets.\n\nThe table below shows the RMSE with 95% CI after 30 steps, averaged across all test sets. First, **both amortized methods significantly outperform all GP-based baselines** across all tasks. This highlights the power of meta-learning in this domain. GP-based methods must learn each new performance landscape from scratch, which is highly inefficient in high dimensions. In contrast, both ALINE and ACE-US are pre-trained on hundreds of related tasks and meta-learn the structural patterns common to these landscapes. This shared prior knowledge allows them to make more accurate predictions from sparse data. Second, while ACE-US performs strongly due to its amortized nature, **ALINE consistently achieves the best or joint-best performance.** This demonstrates the additional benefit of our core contribution: the learned acquisition policy. ALINE's policy is trained end-to-end to learn how to optimally explore the landscape, leading to more informative queries and ultimately a more accurate final surrogate model.\n\nWe believe this new experiment provides a compelling demonstration of ALINE's scalability. We will add a full description of these results to the revised manuscript.\n\n| RMSE ($\\\\downarrow$) | GP-RS | GP-US | GP-VR | GP-EPIG | ACE-US | ALINE |\n| :---- | :---- | :---- | :---- | :---- | :---- | :---- |\n| rpart (6D) | 0.07±0.03 | 0.04±0.02 | 0.04±0.02 | 0.05±0.02 | **0.01±0.00** | **0.01±0.00** |\n| svm (8D) | 0.22±0.11 | 0.11±0.05 | 0.12±0.07 | 0.15±0.08 | 0.04±0.01 | **0.03±0.01** |\n| ranger (9D) | 0.10±0.02 | 0.07±0.01 | 0.08±0.02 | 0.08±0.02 | **0.02±0.01** | **0.02±0.01** |\n| xgboost (16D) | 0.09±0.02 | 0.09±0.02 | 0.09±0.02 | 0.09±0.02 | 0.04±0.01 | **0.03±0.01** |\n\n> *2. Another weakness which is already mentioned in the paper is the choice of assuming factorized likelihood over the targets.*\n\nYou are correct, and we acknowledge this choice as a limitation in the paper. Our decision to adopt a factorized approximation was a deliberate and critical **trade-off between modeling accuracy and computational tractability**.\n\nAn autoregressive model would require sequential generation of each marginal dimension of the posterior, which is computationally prohibitive within our RL framework. Calculating the reward for $M$ targets would require $M$ **sequential forward passes** through the Transformer, making the training process **infeasibly slow**. Therefore, we adopted the factorized approximation - a common and efficient practice in the Neural Process literature [2, 3, 4] - which allows for computation in a single parallel forward pass. Our strong empirical results validate this trade-off between modeling accuracy and computational tractability. We will add a summary of this rationale to the revised manuscript.\n\n> *3. There is no discussion of the stability of the joint training procedure.*\n\nOur approach to ensuring stability is grounded in both the complementary nature of our learning tasks and a key practical implementation detail:\n\n* We can view the joint training as a multi-task learning problem where the two tasks are highly complementary, not conflicting. They mutually regularize the shared Transformer. A better inference model provides a more reliable reward signal for the policy, and a better policy gathers more informative data for the inference model. This synergistic relationship naturally promotes stable learning. We also find this observation in Bayesian optimization literature [5, 6].\n* From a practical standpoint, we found the model to be remarkably straightforward to train. A crucial trick for this is the **initial warm-up phase** where we first train only the inference network. By doing so, we ensure that the inference head provides a stable reward signal before we begin training the acquisition policy. This stable reward landscape is also why we were able to successfully use the basic REINFORCE algorithm, obviating the need for more complex RL methods.\n\nWe believe our strong empirical results provide substantial evidence for the method's viability. We agree that a deeper theoretical study would be a valuable direction for future work.\n\n> *4. 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\nALINE can be straightforwardly adapted to handle problems with **discrete inputs and outputs** with minor, standard modifications to the input embedding and output head layers, respectively. Recent Transformer-based neural processes, such as ACE [4], already demonstrate this capability.\n\n* **For discrete inputs:** Instead of using an MLP embedder, we could employ a learnable embedding layer. Each discrete value in the input space would be mapped to a unique, trainable embedding vector. These embedding vectors would then be processed by the Transformer backbone.\n* **For discrete outputs:** The modification would be in the inference head. Instead of parameterizing a GMM, the head would be adapted to output logits corresponding to each possible discrete value. A **Softmax function** would then be applied to these logits to produce a categorical probability distribution over the discrete outcomes.\n\n> *5. Is it possible to switch between predictive and parameter targets in the middle of a rollout?*\n\nYes, absolutely. This is a key feature of ALINE's flexible design. To switch targets, the user simply has to change this input specifier at the desired step.\n\nTo explicitly demonstrate this capability, we have run a new experiment in the psychometric task. We configured a rollout where, for the first 5 steps, the target was the **threshold & slope**, and at step 6, we switched the target to the **guess rate & lapse rate**. We present the sequence of queries in the table below as per the rebuttal format's limitations. The threshold value is 0.49 in this experiment. We observed that ALINE's acquisition strategy adapted immediately from querying the extremes to the values concentrated near the threshold. We will add this new experiment with full 30 steps and its plots to the revised manuscript.\n\n| Step | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 |\n| :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- | :---- |\n| Queried Stimulus Value | \\-3.93 | \\-4.92 | 4.70 | 3.04 | 3.06 | \\-0.31 | 0.86 | 1.60 | 1.03 | 0.3 |\n\n> *6. Could you elaborate on potential ways to make the method support inputs with arbitrary dimensionality?*\n\nThere is already some recent work that can support arbitrary dimensions. A standard approach would be to adopt a **bi-dimensional Transformer** [7, 8], the process would involve a **feature-wise Transformer** that operates across the different feature dimensions of each data point. Then we can use a **point-wise Transformer**, which operates across the sequence of data points as it does now.\n\nThis double-Transformer approach has been shown to be quite effective. Integrating this into ALINE is a very feasible direction for future research.\n\n> *7. What are the most important bottlenecks for scaling this to more complicated and higher dimensional problems?*\n\nWe see the most significant bottlenecks falling into three main categories:\n\n* **Architectural Flexibility and Scale:** Tackling more complex problems will necessitate significantly larger networks than the relatively small model used here.\n* **Quality of the Training \"Meta-Prior\":** The performance of ALINE is fundamentally dependent on the diversity and representativeness of the simulated problems used for training. A key bottleneck is the generation of a sufficiently rich training distribution that can bridge the \"sim2real\" gap and ensure robust generalization to real-world tasks that may differ from the training set.  \n* **Computational and Algorithmic Scalability:** Scaling up the model, data, and task complexity will demand substantial computational resources. Algorithmically, moving beyond tractable assumptions (like the factorized likelihood) to more principled but costly methods (like fully autoregressive models) will require significant advances in training efficiency to be feasible at a large scale.\n\nIn summary, we believe overcoming these challenges is a major focus for the entire field and a promising direction for future research.\n\n**References**\n\n[1] Arango, et al., HPO-B: A Large-Scale Reproducible Benchmark for Black-Box HPO based on OpenML, NeurIPS Track on Datasets and Benchmarks (2021).\n\n[2] Garnelo, et al., Conditional neural processes, ICML (2018).\n\n[3] Müller, et al., Transformers Can Do Bayesian Inference, ICLR (2022).\n\n[4] Chang, et al., Amortized Probabilistic Conditioning for Optimization, Simulation and Inference, AISTATS (2025).\n\n[5] Maraval, et al., End-to-end meta-bayesian optimisation with transformer neural processes, NeurIPS (2023).\n\n[6] Zhang, et al., PABBO: Preferential Amortized Black-Box Optimization, ICLR (2025).\n\n[7] Lee, et al., Dimension agnostic neural processes, ICLR (2025).\n\n[8] Qu, et al., Tabicl: A tabular foundation model for in-context learning on large data, ICML (2025)."}, {"review_id": "DgIpd1VIjq", "reviewer": "Reviewer_GfQV", "summary": "This 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.", "questions": "The 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?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "I think this is a solid paper and I vote for acceptance.", "rebuttal": "We sincerely thank you for your positive assessment of our paper. We appreciate your thoughtful suggestions and address your questions below.\n\n> *1. Whereas 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. ... 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\nExcellent point. We are indeed implicitly assuming that the model we are generating training data from is “well-specified”, which can be violated, thus requiring the sim2real gap to be bridged. We will state this assumption more prominently at the beginning of the method section.\n\n> *2. 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\nThank you for pointing out this reference.  The connection to our work is that both approaches leverage the power of Neural Processes to move beyond the limitations of traditional GPs in active learning. And we believe this interesting real-world application could provide a compelling potential use case for ALINE. We will gladly add this work to our related work section.\n\n> *3.  Figure 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\nWe fully agree. In fact, we did include such a baseline (ACE-US) in our active learning experiments. ACE (Amortized Conditioning Engine) [1] is a very recent, powerful **transformer-based neural process** that extends earlier work [2, 3]. We specifically chose it as the foundation for our baseline because, unlike many earlier NPs, it can perform amortized inference on both data (posterior predictive) and latent parameters (posterior), making it a fair and direct architectural counterpart to ALINE. We recognize and apologize that we did not make the role and nature of the ACE-US baseline clearer in our submission. We will revise the experimental section to explicitly mention ACE as an amortized neural process model and emphasize that this baseline serves this important ablative purpose.\n\n> *4. More generally, although I liked the experiments that were present in the paper, I thought the experiments were slightly on the light side.*\n\nThank you for the comment. To strengthen our evaluation, we have conducted a new set of experiments on a practical high-dimensional task: **Active Exploration of Hyperparameter Performance Landscapes**. The task is to actively query a small number of hyperparameter configurations to build a surrogate model that accurately predicts performance for a larger, held-out set of target configurations.\n\nWe used the high-dimensional, real-world tasks from the **HPO-B meta benchmark** [4]. Specifically, we evaluated on the rpart (6D), svm (8D), ranger (9D), and xgboost (16D) datasets. ALINE was trained on their multiple pre-defined training sets. We then evaluated its performance, alongside non-amortized GP-based baselines and an amortized surrogate baseline (ACE-US), on the benchmark's held-out and entirely unseen test sets. This setup is explicitly designed to test the model's ability to generalize to new problems.\n\nThe table below shows the RMSE with 95% confidence interval after 30 steps, averaged across all test tasks for each dataset. First, **both amortized methods, ALINE and ACE-US, significantly outperform all non-amortized GP-based baselines** across all tasks. This highlights the power of meta-learning in this domain. GP-based methods must learn each new performance landscape from scratch, which is highly inefficient in high dimensions. In contrast, both ALINE and ACE-US are pre-trained on hundreds of related tasks and their Transformer architectures meta-learn the structural patterns common to these landscapes. This shared prior knowledge allows them to make far more accurate predictions from sparse data. Second, while ACE-US performs strongly due to its amortized nature, **ALINE consistently achieves the best or joint-best performance.** This demonstrates the additional, crucial benefit of our core contribution: the learned acquisition policy. ACE-US relies on a standard heuristic, whereas ALINE's policy is trained end-to-end to learn how to optimally explore the landscape, leading to more informative queries and ultimately a more accurate final surrogate model.\n\nWe believe this new experiment, alongside the existing BED and psychometric tasks, now provides a comprehensive and compelling demonstration of ALINE's practical utility. We will add these results to the revised manuscript.\n\n| RMSE ($\\\\downarrow$) | GP-RS | GP-US | GP-VR | GP-EPIG | ACE-US | ALINE |\n| :---- | :---- | :---- | :---- | :---- | :---- | :---- |\n| rpart (6D) | 0.07±0.03 | 0.04±0.02 | 0.04±0.02 | 0.05±0.02 | **0.01±0.00** | **0.01±0.00** |\n| svm (8D) | 0.22±0.11 | 0.11±0.05 | 0.12±0.07 | 0.15±0.08 | 0.04±0.01 | **0.03±0.01** |\n| ranger (9D) | 0.10±0.02 | 0.07±0.01 | 0.08±0.02 | 0.08±0.02 | **0.02±0.01** | **0.02±0.01** |\n| xgboost (16D) | 0.09±0.02 | 0.09±0.02 | 0.09±0.02 | 0.09±0.02 | 0.04±0.01 | **0.03±0.01** |\n\n> *5. The 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\nWhen updating the acquisition policy, we **do stop the gradients** from the RL loss with respect to the parameters of the inference network. This is a deliberate implementation choice for the REINFORCE algorithm in this context, and it is crucial for stable training. The rationale is to ensure each component has a clear and distinct objective:\n\n* The **policy gradient** updates only the parameters of the **acquisition policy**, encouraging it to select actions that lead to higher rewards.  \n* The parameters of the **inference network**, in turn, are updated *only* by the **negative log-likelihood loss**, which has the sole objective of pushing to become a more accurate approximation of the target distribution.\n\nAllowing gradients to flow from the policy loss back into the inference network would create a problematic objective. The policy would learn not only to find informative points but also to directly manipulate the inference network to report higher rewards, potentially distorting the posterior estimates. We will add a sentence to the manuscript to explicitly clarify this important implementation detail.\n\n> *6. How well does the approach generalise to scenarios (e.g. combinations of targets) at deployment time?*\n\nTheoretically, ALINE is designed to generalize to novel target combinations at deployment time, even if those specific combinations were not seen during training. This is a direct benefit of our **query-target cross-attention mechanism**. It allows each candidate query point to compute its relevance and potential information gain with respect to each target embedding individually and in parallel. The final policy is then based on an aggregation of these individual relevance scores.\n\nTo empirically validate this, we ran a new experiment on the psychometric task.\n\n* **Training:** We trained a single ALINE model to handle two distinct targets separately: (1) the \"threshold & slope\" and (2) the \"guess & lapse rate\".  \n* **Testing:** At deployment, we tasked this model with a novel, unseen combination: to target **all four parameters simultaneously**.\n\nWe observed that the resulting acquisition policy was a sensible mixture of the two underlying strategies it had learned. It strategically alternated its queries between points near the decision threshold (optimal for the first target) and points at the extremes (optimal for the second). This behavior confirms that ALINE can successfully compose its learned strategies to generalize to new combinations of inference goals at runtime. We will include this experiment in the revised manuscript.\n\n**References**\n\n[1] Chang, et al., Amortized Probabilistic Conditioning for Optimization, Simulation and Inference, AISTATS (2025).\n\n[2] Müller, et al., Transformers Can Do Bayesian Inference, ICLR (2022).\n\n[3] Nguyen and Aditya, Transformer neural processes: Uncertainty-aware meta learning via sequence modeling, ICML (2022).\n\n[4] Arango, et al., HPO-B: A Large-Scale Reproducible Benchmark for Black-Box HPO based on OpenML, NeurIPS Track on Datasets and Benchmarks (2021)."}]

Author rebuttal:

Meta-review: The paper proposes ALINE, a unified framework for amortized Bayesian inference and active data acquisition built on a transformer backbone with joint inference–acquisition training. The method introduces a flexible target-specification mechanism and demonstrates strong empirical performance across regression-based active learning, Bayesian experimental design benchmarks, and psychometric tasks. Reviewers appreciated the clarity, the efficiency of the unified architecture, and the novelty of the flexible targeting strategy. The main concerns before the rebuttal phase centered on overlap with prior work (vsOED), reliance on a factorized likelihood approximation, and the sim2real training assumption. The authors addressed these with additional experiments on high-dimensional benchmarks, direct comparisons to vsOED, and clarifications of design trade-offs and limitations. 

Overall, I find the paper technically solid, with consensus among reviewers leaning toward acceptance. The authors are encouraged to strengthen the discussion of limitations and clarify distinctions from prior work in the final version.
Decision: Accept (spotlight)