1 Introduction↩︎

Many sequential learning problems can be posed as adaptive data acquisition: at each round, a learner chooses which query to make next to improve a downstream decision, based on the data observed so far, receives the resulting observation, and repeats until a budget is exhausted. This abstraction covers Bayesian experimental design (BED; ), Bayesian optimisation (BO; ), and active learning (AL; ). These settings differ in their downstream objectives, but share the same central challenge: mapping a growing observation history to acquisition decisions under a finite budget.

Learning effective acquisition policies in this setting is difficult. Existing approaches fall into two families: those that work through task-specific posterior, and those that map observation histories directly to decisions. Both are problematic in different ways. Posterior-based methods fit a probabilistic model from the observed data and use the resulting posterior quantities to evaluate an acquisition function. However, both classical [1]–[3] and amortised [4]–[6] methods are bottlenecked by the quality of the posterior. In practice, model misspecification or approximation errors can distort the acquisition rule—which typically requires well-calibrated uncertainties—and these errors can compound over the course of a sequential acquisition process [7], [8]. The second family maps observation histories directly to the next decision [9]–[12]. These methods avoid explicit posterior inference, but pay for this by having to learn a much more complex mapping from raw histories: they do not use architectures that exploit information in the model and must implicitly learn how the posterior changes with the data. In other words, they must simultaneously learn both how to represent the current dataset and how to act on that representation, which is often sample-inefficient and difficult to train.

The key insight that allows us to address these limitations is that, while the optimal sequential decisions depend only on the observation history through the corresponding posterior, a learned policy need not access this posterior explicitly: it may instead act on any representation that preserves the information needed to recover that task-relevant belief. This suggests a middle ground between raw observation histories and explicit posterior estimates. To exploit this insight, we therefore introduce an intermediate belief representation that captures underlying features in how the posterior varies with the data, without over-committing to a single posterior approximation.

Pretrained predictive foundation models provide a natural starting point for this belief representation. Namely, rather than learning a belief representation from observation histories from scratch, we can build on models whose pretraining has already shaped them to organise observed data into representations useful for prediction across a broad family of tasks. We instantiate this idea with Tabular Foundation Models (TFMs; [13]), such as TabPFN [14], [15] and TabICL [16], [17], which produce candidate-conditioned representations from a context set in a single forward pass.

Specifically, we introduce POlicy LeArning with Belief Representations (Polar), an amortised data acquisition framework that uses the hidden representations from a TFM as our belief representation and trains a policy head to select the next query from the encoding this provides. In doing so, Polar shifts amortised acquisition from jointly learning representations and decisions to learning how to act on pretrained belief representations. An overview of Polar is shown in 1.

Contributions. (1) We give a unified decision-theoretic view of adaptive data acquisition as learning acquisition policies on top of belief representations, rather than learning sequential decision-making directly from raw histories or acting through explicit posterior approximations. (2) We introduce a simple decoupled policy architecture that uses pretrained predictive foundation models as belief-state encoders and trains a policy head on top. (3) Across diverse adaptive data acquisition tasks, the learned policy substantially outperforms greedy acquisition on the same backbone and exceeds state-of-the-art amortised methods while requiring up to \(100\times\) fewer task-specific training samples.

2 Preliminaries↩︎

2.0.0.1 Adaptive data acquisition.

We consider a finite-horizon adaptive data acquisition problem in which a learner repeatedly selects designs based on the data observed so far. A latent world state \(z \sim p(z)\) determines observations through a model \(p(y\mid x, z)\), where \(x\in\mathcal{X}\) is a design and \(y\in\mathcal{Y}\) is the corresponding outcome. At round \(t\), the learner has observed \(\mathcal{D}_{t-1}=\{(x_i, y_i)\}_{i=1}^{t-1}\), chooses a feasible design \(x_t \in \mathcal{X}_t(\mathcal{D}_{t-1}) \subseteq \mathcal{X}\), observes \(y_t \sim p(\cdot \mid x_t, z)\), and appends \((x_t, y_t)\) to the history. An acquisition policy \(\pi=(\pi_1,\ldots,\pi_T)\) is a sequence of decision rules, \(\pi_t:\mathcal{H}_{t-1}\to \mathcal{X}\) satisfying \(\pi_t(\mathcal{D}_{t-1})\in \mathcal{X}_t(\mathcal{D}_{t-1})\). Adaptive data acquisition problems share this sequential structure, but differ in the downstream objective and which aspect of the world state matters for this. For example, in many BED problems, we seek information about unknown system parameters \(z\) [18]–[20]; in active learning, we want to learn to make effective future predictions [21]; and in BO, we seek evaluations that rapidly identify a global optimum [3].

Existing approaches often reduce adaptive acquisition to repeated posterior-based scoring. Given history \(\mathcal{D}_{t-1}\), they construct a posterior distribution, either by directly performing inference on the underlying model or using an amortised neural predictor [4], [6], [22], then select candidates with an acquisition function. Examples include Gaussian-process posteriors with Expected Improvement (EI) acquisition in BO [3] and parametric posteriors with Expected Information Gain (EIG) acquisition in BED [23], [24]. These methods rely on the quality of the posterior approximation exposed to the acquisition rule. Outside rare cases where exact inference is tractable, posterior approximation is hard [7], [8], and its errors can propagate into suboptimal acquisitions, especially across multiple rounds.

An alternative is to learn the acquisition policy. In this case, a parameterised policy is trained from simulated trajectories to maximise a task-specific sequential utility, without explicitly constructing a posterior or evaluating a hand-designed acquisition function. For example, amortised BED methods [9]–[11], [25]–[28], map observation histories directly to experimental designs, and related end-to-end policies have appeared for BO [12], [29] and AL [27], [30]. This amortises test-time decision-making, but requires the same model to learn both a representation of the current dataset and an acquisition strategy from sequential supervision alone.

2.0.0.2 Pretrained predictive models.

Recent work on Bayesian in-context learning has produced models that can approximate Bayesian prediction in a single forward pass, without iterative fitting or gradient updates. Transformer Neural Processes [31] and Prior-Fitted Networks [32] learn to map a context set of observed input-output pairs to predictive distributions for new inputs, after offline training on large families of synthetic tasks. More recent Tabular Foundation Models (TFMs), such as TabPFN [14], [15] and TabICL [16], [17], scale this idea to practical tabular prediction, achieving strong performance across a wide range of supervised learning problems. Given a dataset \(\mathcal{D}\) and a candidate input \(x\), these models produce a predictive distribution \(q_\psi(\cdot \mid x, \mathcal{D})\) for the corresponding outcome, together with internal representations \(r_\psi(x,\mathcal{D})\) that depend on both \(x\) and \(\mathcal{D}\). These representations are not designed for any particular sequential task, but can encode predictive uncertainty, local structure, and context-dependent function behaviour. Consequently, TFMs can be viewed as off-the-shelf amortised Bayesian predictors whose representations provide a natural substrate for acquisition policies.

3 A Unified Decision-Theoretic View of Adaptive Data Acquisition↩︎

To understand policy learning for adaptive data acquisition, we must first formalise what an ideal representation should capture. We do so by casting adaptive data acquisition under a unified Bayesian decision-theoretic framework.

3.0.0.1 Bayes-optimal policy.

Consider a finite-horizon sequential data acquisition problem. Once all data has been acquired, a downstream action, \(a \in \mathcal{A}\), is chosen, which incurs loss \(\ell(a, w)\), where \(w=\tau(z)\) represents the aspects of the overall world state that directly impact our loss. Depending on context, \(a\) might be an estimate, a predictive distribution, or a physical decision. Following Bayesian decision theory [24], [33]–[35] and defining \(p(w \mid \mathcal{D}_T)\) as the pushforward of our posterior onto \(w\), the Bayes-optimal acquisition policy is now: \[\pi^\star = \mathop{\mathrm{arg\,min}}_{\pi \in \Pi} \mathbb{E}_{p(\mathcal{D}_T;\pi)} \!\left[ \min_{a \in \mathcal{A}} \mathbb{E}_{p(w \mid \mathcal{D}_T)}[\ell(a,w)] \right]. \label{eq:general95design95obj}\tag{1}\] Here the expected loss of the Bayes action, \(\min_{a \in \mathcal{A}} \mathbb{E}_{p(w \mid \mathcal{D}_T)}[\ell(a,w)]\), is a function only of our posterior beliefs for a given loss function. We can thus rewrite 1 as an expected functional of the posterior: defining \(s(q,w) := \ell(a^\star(q), w)\) where \(a^\star(q):= \arg\min_{a \in \mathcal{A}} \mathbb{E}_{w \sim q}[\ell(a,w)]\), then \[\pi^\star =\mathop{\mathrm{arg\,min}}_{\pi \in \Pi} \mathbb{E}_{p(\mathcal{D}_T,w;\pi)} \!\left[ s\!\left(p_w(\cdot \mid \mathcal{D}_T), w\right) \right]. \label{eq:posterior95score95w}\tag{2}\] It turns out that this \(s\) is always a proper scoring rule for distributions on \(w\) [36], [37]. We can thus view 2 as an expected posterior uncertainty, \(\mathbb{E}_{p(\mathcal{D}_T;\pi)}[h_s[p_w(\cdot \mid \mathcal{D}_T)]]\), where our uncertainty measure is the generalised entropy \(h_s[q]=\mathbb{E}_{w\sim q}[s(q,w)]\) [38], [39].

Under this unified foundational framework of sequential Bayesian decision-making, we can thus equivalently think about defining adaptive data gathering problems through defining either downstream losses, scoring rules, or uncertainty measures, all of which also (implicitly) define the target variables of interest for our data gathering, \(w\). Most existing BED and active learning methods are explicitly based on defining uncertainty measures, which is typically taken to be entropy to yield the expected information gain [20], [23], though they can vary in what variables are targeted for learning, e.g., in AL, we can either take \(w\) to be model parameters [40] or downstream predictions [21]. In BO, it is more common to work directly with a downstream loss, with \(w\) taken to be the true function values at the queried inputs. For example, we can recover a strategy equivalent to targeting expected improvement by taking \(\ell_{\mathrm{EI}}(a,f)=-f(a)\) and then \(\mathcal{A}\) as the set of inputs we have queried (such that the minimisation over \(a\) trivially yields the point with the highest posterior mean).

3.0.0.2 From posterior beliefs to belief representations.

We can further rearrange 2 to make clear the dependency of the optimal policy on historical data as follows (where \(\mathcal{D}_{t:T} = \{(x_i,y_i)\}_{i=t}^{T}\)) \[\begin{align} \pi^\star =\mathop{\mathrm{arg\,min}}_{\pi \in \Pi} \mathbb{E}_{p(\mathcal{D}_{t-1} ; \pi)} \left[ \mathbb{E}_{p(w | \mathcal{D}_{t-1} ) p(\mathcal{D}_{t:T} | w ; \pi)} \!\left[ s\!\left(\frac{p_w(\cdot|\mathcal{D}_{t-1})p(\mathcal{D}_{t:T} | w ; \pi)}{\int p_w(w'|\mathcal{D}_{t-1})p(\mathcal{D}_{t:T} | w' ; \pi)dw'}, w\right) \right]\right]. \end{align}\] This yields a critical insight: The Bayes-optimal policy depends on the history only through the task-relevant posterior induced by that history (and more precisely the expected final posterior uncertainty that it induces). Equivalently, there exist maps \(f_{\mathcal{D}},\,f_{\mathrm{post}}\) such that \[\pi^\star(\cdot ; \mathcal{D}_{t-1}) \;=\; f_{\mathcal{D}}(\mathcal{D}_{t-1}) \;=\; f_{\mathrm{post}}\!\left(p_w(\cdot \mid\mathcal{D}_{t-1})\right). \label{eq:two95equivalent95forms}\tag{3}\] Thus, any two datasets that induce the same task-relevant posterior should lead to the same optimal next action. This invariance suggests that the right input to policy learning is not necessarily the raw dataset itself, but some representation that preserves the information in \(\mathcal{D}_{t-1}\) needed to recover the task-relevant belief. In particular, we want a representation that faithfully retains this information, while also providing the simplest possible mapping to the optimal policy decisions.

To this end, we break down \(f_{\mathcal{D}} = g \circ \mathrm{enc}\) into a belief encoder \(\mathrm{enc}\) and a policy head \(g\). We call \(e_t = \mathrm{enc}(\mathcal{D}_{t-1})\) a belief representation for \(w\) if there exists a decoder \(\mathrm{dec}: \mathcal{E}\to \mathcal{P}(\mathcal{W})\) such that \[p_w(\cdot\mid\mathcal{D}_{t-1}) = \mathrm{dec}(e_t). \label{eq:decoder}\tag{4}\] Whenever the decoding condition above holds, there also exists a map \(g\) for which \[\pi^\star(\cdot;\mathcal{D}_{t-1}) = g(e_t), \label{eq:h95t}\tag{5}\] because the task-relevant posterior is determined by \(e_t\). Two extreme choices of \(e_t\) recover the two dominant approaches in the literature:

• Raw history: \(e_t = \mathcal{D}_{t-1}\). The raw history is trivially a belief representation, since the posterior can be recovered from it. This is the input used by direct policy-learning methods, but it leaves the policy head \(g\) with the entire job of capturing the dependency of the posterior on the data, significantly complicating practical training.

• Posterior: \(e_t = p_w(\cdot \mid\mathcal{D}_{t-1})\). This is the belief state which yields a trivial decoder (the identity function) and thus incorporates all desired posterior invariances (though there may be further invariances in the mapping from intermediary posteriors to expected future uncertainties), thus making the training of \(g\) much easier. In practice, however, this belief state cannot usually be calculated exactly and is replaced by an approximation \(\hat{p}_w(\cdot\mid\mathcal{D}_{t-1})\), breaking 4 and exposing \(g\) to the approximation error of the inference step.

Rather than going with either of these extremes, our key idea is to use an encoder that avoids needing an overly complex \(g\) while still preserving the required information in \(\mathcal{D}_{t-1}\) that determines optimal actions (and in particular avoiding the errors that build up from direct posterior approximation). We achieve this by using encodings that provide effective representations of the posterior, without over-committing to using precise posterior approximation. 4 shows how predictive foundation models provide an appropriate and practical encoder as a by-product of their training.

4 Policy Learning with Belief Representations↩︎

Section 3 argued that acquisition policies should act on representations of the task-relevant posterior: representations that retain the information in observation histories needed for optimal decisions, while structured enough to make the map to acquisitions easy to learn. Pretrained predictive foundation models are natural candidates for this role: to predict well across diverse tasks from in-context observations alone, such models must internally summarise the observed data into representations that track how beliefs change as data accumulates, precisely the information an acquisition policy needs. We now describe how Polar instantiates this principle using Tabular Foundation Models (TFMs) as practical belief encoders.

4.0.0.1 Foundation models as belief encoders.

A pretrained TFM \(F_\psi\) is trained to map context sets to predictive distributions across a broad family of synthetic tasks [14]–[16]. Given a dataset \(\mathcal{D}_{t-1}\) and a candidate \(x\in\mathcal{X}\), a forward pass returns (i) a predictive distribution \(q_\psi(\cdot\mid x,\mathcal{D}_{t-1})\) over the outcome \(y\), and (ii) a candidate-conditioned hidden representation \(r_\psi(x,\mathcal{D}_{t-1})\in\mathbb{R}^{d_r}\) from which this predictive distribution is decoded. The predictive distribution \(q_\psi\) is one particular output decoded from \(r_\psi\) through a fixed prediction head. It lives in outcome space, approximating a posterior predictive over \(y\) at a candidate \(x\), whereas the decision-theoretic object in 3 is the posterior over the task-relevant quantity \(w\). Thus, although \(q_\psi\) is useful for prediction and for some task-specific surrogates, it is not in general itself a belief representation for \(w\).

We therefore use the hidden representation \(r_\psi\), rather than \(q_\psi\) alone, as the interface for policy learning, so that the acquisition policy can act on the candidate-conditioned state before it is compressed into the final predictive output. Although predictive pretraining optimises the decoded distribution \(q_\psi\), it does so through hidden representations \(r_\psi\) that must support prediction across a broad family of synthetic priors and variable-size target sets. This training pressure encourages \(r_\psi\) to organise \(\mathcal{D}_{t-1}\) along directions that determine how predictions change with the data—predictive uncertainty, local structure, and context-dependent function behaviour—which are precisely the kinds of information an acquisition policy needs in order to decide where to query next. We therefore treat \(r_\psi\) as a belief representation in the sense of 3. As with any neural encoder, \(r_\psi\) is not an exact sufficient statistic for \(w\), but predictive pretraining shapes it to retain the information that determines optimal acquisitions, and we validate its adequacy empirically in 6.

While the framework of 3 applies to general design spaces, TFMs are far more efficient when the context \(\mathcal{D}_{t-1}\) is processed once and reused across a finite candidate set. We therefore work in a pool-based setting: at each round, the policy chooses from a feasible pool \(\mathcal{C}_t \subseteq \mathcal{C}\). A single batched pass over \(\mathcal{C}_t\) returns the candidate-conditioned representations \(\{r_\psi(x,\mathcal{D}_{t-1}) : x\in\mathcal{C}_t\}\). This setting matches many practical problems: hyperparameter optimisation over discrete configuration grids, molecular optimisation over finite molecule libraries, and active learning over finite unlabelled pools. For continuous design spaces, the same architecture can be combined with a candidate generator or evaluated over sampled candidate sets.

4.0.0.2 Policy head.

We instantiate the policy head \(g\) as a function \(g_\eta:\mathbb{R}^{d_r}\to\mathbb{R}\) with learnable parameters \(\eta\). For each candidate \(x\in\mathcal{C}_t\), \(g_\eta\) maps \(r_\psi(x,\mathcal{D}_{t-1})\) to a scalar logit, and normalising these logits across the candidate pool with a softmax gives the stochastic policy \(\pi_\eta(\cdot\mid\mathcal{D}_{t-1},\mathcal{C}_t)\). We parameterise \(\pi_\eta\) as a distribution over candidates during training, which provides exploration and admits standard stochastic-gradient policy training. The Bayes-optimal policy is deterministic at convergence [24], so at deployment, we select the candidate with the largest learned score.

4.0.0.3 Task scoring rules.

The architecture is shared across all tasks; only the task-specific scoring rule changes. For each task, we define \(\widehat{\Phi}(\mathcal{D}_T,w)\) as a tractable surrogate of \(-s(p_w(\cdot | \mathcal{D}_T),w)\). In some tasks, this score is directly observable, while in others a surrogate estimator is needed. In BED, the natural scoring rule is the log loss, whose expected score is the EIG. Since the EIG is intractable, we approximate it with the sPCE lower bound [9]. Other EIG estimators could be substituted without further change. In BO, the queried function values \(y_t\) are directly observed, so the score reduces to a tractable function of the trajectory. We use the terminal best observed value, \(\widehat{\Phi}_{\mathrm{BO}}(\mathcal{D}_T,w)=\max_{t\leq T} y_t\), for which no posterior approximation is required. In AL, when the goal is to make future predictions [21], [35], the terminal score is the predictive log-density on a target set, which is also intractable under the true posterior, so we use a variational surrogate built on \(q_\psi\).

4.0.0.4 Policy training.

We train the policy by rolling out \(\pi_\eta\) on simulated tasks and applying a score-function policy-gradient objective. Given a task-specific score \(\widehat\Phi(\mathcal{D}_T,w)\), to provide dense training signals, we define local utility increments \(u_t=\widehat\Phi(\mathcal{D}_t,w)-\widehat\Phi(\mathcal{D}_{t-1},w)\), so that \(\sum_{t=1}^T u_t=\widehat\Phi(\mathcal{D}_T,w)-\widehat\Phi(\mathcal{D}_0,w)\). With \(\tau=(x_{1:T},y_{1:T})\) denoting a sampled trajectory, the policy loss is \[\mathcal{L}_{\mathrm{pol}}(\eta)=-\mathbb{E}_{\tau\sim \pi_\eta}\left[\sum_{t=1}^T u_t \log \pi_\eta(x_t;\mathcal{D}_{t-1})\right]. \label{eq:policy95loss}\tag{6}\] Here, following [12], [27], we use the one-step increment rather than the full reward-to-go. This objective should be understood as a low-variance local-credit surrogate for the terminal-score objective, rather than an unbiased policy-gradient estimator of the full non-myopic objective.

4.0.0.5 Backbone adaptation.

The pretrained TFM backbone \(F_\psi\) is trained on synthetic priors that may not perfectly match the downstream task distribution. We therefore optionally adapt \(\psi\) using an auxiliary supervised prediction loss, \(\mathcal{L}_{\mathrm{pred}}(\psi)=\sum_{t=1}^T\sum_{m=1}^M\ell_{\mathrm{pred}}\!\big(q_\psi(\cdot\mid x_m^\star,\mathcal{D}_{t-1}),y_m^\star\big)\), where \(\{(x_m^\star,y_m^\star)\}_{m=1}^M\) are target points sampled from the same simulated task during training. The form of \(\mathcal{L}_{\mathrm{pred}}\) depends on the prediction head: for density-based heads [41], it is negative log-likelihood, while for quantile-based heads [17], it is quantile regression loss. A key design choice is that policy gradients are not propagated into the pretrained backbone: \(\mathcal{L}_{\mathrm{pol}}\) updates only the policy head, while \(\mathcal{L}_{\mathrm{pred}}\) is the only loss that may update \(\psi\). This lets the backbone adapt through supervised prediction signals while preventing high-variance policy-gradient updates from destabilising the pretrained model. Thus, representation refinement and decision learning remain decoupled.

4.0.0.6 Limitations and possible extensions.

We close this section by noting the scope of the particular realisation of Polar described above, and the natural extensions it suggests. First, our policy acts over finite candidate pools. This covers the pool-based settings considered in this work, but problems requiring highly precise continuous designs may benefit from continuous-action variants, for example, by optimising the learned score over generated candidates or by introducing a global summary token that feeds a continuous policy head. Second, our policy-gradient objective uses local credit assignment rather than an explicit long-horizon value function. This keeps training simple and low variance, but tasks with strongly delayed information gains may benefit from value-based extensions or reward-to-go estimators. Third, when the downstream task distribution differs substantially from the synthetic priors used to pretrain the backbone, auxiliary prediction losses may be needed to realign \(r_\psi\), adding training cost. Parameter-efficient adaptation strategies, such as LoRA-style adapters [42] or partial-layer finetuning, could reduce this overhead.

5 Related Work↩︎

Amortised policy-based BED was first proposed by DAD [9], which learns a design policy offline and deploys it in a single forward pass, with iDAD [10] extending the framework to implicit likelihood models. A parallel line recasts the problem as reinforcement learning [25], [43], and [11] amortise designs against an explicit downstream decision utility. [44] introduce a semi-amortised variant that combines an offline policy with online refinement, while [45] consider constraint-aware BED via online planning, and a recent thread improves training signals through alternative EIG estimators [26], [28], [46], [47]. [27] show that prediction-driven supervision can support policy learning, but its shared backbone is trained from scratch under coupled prediction and policy objectives, making it sample-inefficient. Our work is complementary to all of these, rather than designing a new estimator or training algorithm, we reuse a pretrained predictive model as the belief encoder so that a small policy head can be trained with standard policy gradient.

In BO, one line of work retains the classical surrogate-plus-acquisition template while replacing the surrogate with increasingly powerful neural predictors, including PFN- and transformer-based models [4], [6], [22], [48], [49]. A second line instead learns the acquisition policy end-to-end, mapping observed trajectories directly to the next query [12], [29], [50]–[53]. Polar sits between these threads: it inherits the strong representations of a pretrained surrogate, but learns a policy on top of those representations rather than committing to a fixed acquisition rule.

Our work is also related to the broader family of meta-learned predictors that map context sets to predictive distributions. Originating from the Conditional Neural Processes [54], these architectures evolved into Transformer Neural Processes [31] and early PFNs [32]. Recent tabular foundation models such as TabPFN [14], [15], [41] and TabICL [16], [17] scale this paradigm to broad supervised-learning regimes and have stimulated a growing literature [13], [55]–[57]. Finally, since our backbone is adapted with an auxiliary prediction loss, our work is also related to recent studies on fine-tuning TFMs to better match downstream data distributions [58]–[60].

6 Experiments↩︎

We now empirically evaluate Polar across a range of tasks. 6.1 benchmarks Polar on two standard BED tasks, location finding and constant elasticity of substitution, where we also report a series of ablations isolating the contribution of each component of our method. 6.2 evaluates Polar for BO on a hyperparameter optimisation benchmark, and 6.3 turns to a real-world high-dimensional molecular optimisation task. We additionally study loss-driven active learning in Appendix 10, which demonstrates that Polar extends naturally to settings defined by user-specified predictive losses. Unless otherwise stated, all experiments use TabICL v2 [17] as the backbone; additional results with TabPFN v2.5 [15] are reported in Appendix 11.

6.1 Benchmarking on Bayesian experimental design tasks↩︎

We begin with two standard BED benchmarks, location finding and constant elasticity of substitution (CES), which have been used extensively in prior work [25], [27], [61], [62]. The goal of location finding [63] is to infer the locations of \(K\) hidden sources in \(d_x\)-dimensional space from noisy distance measurements. We consider two variants to test scalability: a standard 2D setting (\(d_x=2, K=2, T=30\)) and a higher-dimensional 5D setting (\(d_x=5, K=2, T=50\)). CES [64] involves eliciting the parameters of an economic utility function from a series of pairwise preference queries. We compare against Random, DAD [9], RL-BOED [25], and ALINE [27]. For both tasks, we report the EIG lower bound estimated by sPCE, and the upper bound estimated by sNMC [9]. Full details on the task and training setups are provided in Appendix 9.1.

6.1.0.1 Main results.

To explicitly assess the training efficiency, we report performance as a function of the total number of training samples (see Appendix 9.1 for details on how we calculate this metric for a fair comparison). For location finding, Figures 2a and 2b show that Polar is dramatically more sample-efficient than prior amortised BED policies. In the standard 2D task, our method slightly outperforms ALINE, the current state-of-the-art, using approximately 100\(\times\) fewer training samples, and surpasses the best reported results of both RL-BOED and DAD using 1,000\(\times\) fewer samples. This dramatic efficiency gap further widens in the more complex 5D environment, where our method eclipses the peak performance of RL-BOED and DAD with 10,000\(\times\) fewer samples, and exceeds ALINE with 100\(\times\) fewer samples. This sample efficiency is particularly valuable in settings where simulation or real data collection is expensive, as is typical in scientific applications of experimental design. 2c provides visualisations of design trajectories. The sNMC upper bound results are provided in Appendix 11.1. The conclusion remains consistent for the CES task. ¿tbl:tab:ces? shows that Polar outperforms all amortised baselines at matched training-sample budgets. Results at full convergence for each method are reported in Appendix 11.2, where Polar also retains its lead. Owing to the cost of evaluating a foundation-model backbone, our method is slightly slower to deploy than the other policies, but the absolute overhead remains negligible in practice: at roughly 0.03 seconds per decision, the latency is still effectively imperceptible at the human timescale.

To better understand where these gains come from, we conduct three ablations on the 2D location finding task, all reported in 3. Additional ablations on the choice of pretrained backbone and the representation layer used as the belief state are reported in 8, which shows Polar is not tied to a single backbone, and later backbone layers provide more useful belief states for acquisition.

6.1.0.2 Frozen vs. finetuned backbone.

We first compare our default setting, in which the backbone is adapted by the supervised prediction loss, against a frozen-backbone variant. Even without any weight updates, the frozen variant remains competitive with the peak performance of RL-BOED. However, it clearly falls short of our finetuned variant. Because the backbone is a general-purpose foundation model pretrained on a vast array of synthetic distributions, its representations are not tailored to a specific task. The auxiliary prediction loss acts as an alignment signal, finetuning the backbone to adapt its predictive representations to the target domain, which in turn provides a better belief state for the policy head to act upon.

6.1.0.3 Decoupled vs. coupled policy gradients.

Our default method updates the backbone only through the prediction loss. We compare this against a coupled variant in which policy gradients are also allowed to update the backbone. The coupled variant severely degrades performance. Once policy gradients are allowed to modify the backbone, the same parameters must simultaneously satisfy two rather different objectives: a dense, relatively stable prediction objective and a sparse, high-variance reinforcement-learning objective. In practice, this appears to destabilise the learning and weaken the benefits of pretraining. By decoupling the gradients, we effectively assign the optimal learning objective to each component: low-variance supervised learning for belief state estimation, and policy gradients for decision-making on a stable representation manifold.

6.1.0.4 Pretrained backbone vs. training from scratch.

We next compare our method against the same architecture trained from scratch, without loading the pretrained checkpoint. The difference is substantial: training from scratch converges much more slowly and remains clearly worse within the same sample budget. This supports a central message of the paper: learning underlying belief representations is a critical part of the overall policy training for amortized adaptive design methods. A pretrained foundation model has already significantly alleviated this problem, and starting from it confers the bulk of our sample efficiency.

6.2 Hyperparameter optimisation benchmarks↩︎

We next evaluate our method on BO using HPO-B [65], a large-scale hyperparameter optimisation benchmark containing the evaluations of a wide range of machine learning models across thousands of hyperparameter configurations and datasets. HPO-B is pre-partitioned into multiple search spaces, each corresponding to a model family with its own hyperparameter parameterisation. Following prior work [6], [12], we report results on six search spaces, spanning input dimensionalities from 2 to 16 and a representative range of model families. For each search space, we train our policy on the provided meta-training tasks and evaluate on the fixed test tasks and seeds defined by the benchmark.

We compare against baselines spanning both classical and amortised BO. Random samples configurations uniformly from the candidate pool. GP fits a Gaussian process [66] directly on the test dataset. Meta-GP, following [12], pretrains the kernel hyperparameters on the meta-training datasets and initialises the GP with these parameters at test time. PFNs4BO [22] uses a prior-fitted network [32] as a surrogate. NAP [12] is an end-to-end fully amortised method based on transformer neural processes [31] that maps observation histories directly to next queries. Finally, TabICL uses the same finetuned TabICL backbone as Polar, but pairs it with a hand-designed acquisition function rather than a learned policy head. All surrogate-based baselines use EI as the acquisition strategy. Additional details on baseline configurations and the dataset are provided in Appendix 9.2.

6.2.0.1 Results.

4 shows the average normalised regret and average rank aggregated across all search spaces. Our method achieves the lowest regret and the best average rank throughout the acquisition trajectory, with the gap to the strongest baselines opening within the first ten queries and persisting through the full budget. Notably, compared with NAP, our method shows that using pretrained TFMs can yield stronger performance without requiring the policy to learn its own belief representation from scratch. Additional per-search-space regret and rank curves are reported in Appendix 11.3.

6.3 Molecular docking optimisation↩︎

Finally, we evaluate our method on a real-world high-dimensional molecular optimisation benchmark, Dockstring [67], which provides docking scores for over 260,000 molecules against a panel of protein targets. Since the physicochemical interactions that govern binding affinity are inherently correlated across these structurally related targets, this dataset serves as a natural testbed for meta BO.

We formulate the task as identifying the molecule that minimises the docking score for a given protein target. We train our policy on a meta-training set of 8 protein targets and evaluate its zero-shot transfer performance on a held-out test set of 6 distinct targets; full target lists are provided in Appendix 9.3. Each molecule is represented by a 512-bit Morgan fingerprint computed from its SMILES string.

Operating in a 512-dimensional space presents significant challenges for traditional BO. Standard GPs equipped with continuous kernels (e.g., RBF) notoriously struggle in high-dimensional, discrete spaces. Instead, we equip the GP with a Tanimoto kernel, which measures set overlap between binary fingerprints and is the standard choice for fingerprint-based molecular BO [68]. We additionally compare against a Random Forest surrogate, a standard baseline in cheminformatics that is commonly paired with molecular fingerprints [69], [70]. For the amortised surrogate baseline, we use exactly the same finetuned TabICL backbone as in our method.

6.3.0.1 Results.

5 reports average regret across the six test targets. Polar achieves the lowest regret throughout the trajectory. Both Polar and TabICL substantially outperform the non-amortised GP and Random Forest baselines, indicating that meta-training across targets allows the backbone to capture shared structure. The further gap between Polar and TabICL confirms that learning the acquisition rule directly from the belief representation outperforms a hand-designed acquisition.

7 Conclusions↩︎

We introduced Polar, an amortised framework for adaptive data acquisition based on a decision-theoretic view: policies need not act directly on raw histories or explicit posterior approximations, but can instead act on task-relevant belief representations. By instantiating these representations with pretrained predictive foundation models, Polar decouples belief formation from decision-making and trains an acquisition policy on them. Across tasks, Polar surpasses classical non-amortised methods and state-of-the-art amortised baselines while requiring up to \(100\times\) fewer training samples.

The appendix is organized as follows:

• In Appendix 8, we provide additional details of Polar, including the backbone architectures, the policy head, and the training details.

• In Appendix 9, we describe the experimental setups in detail, covering task specifications, baseline implementations, and evaluation protocols for the Bayesian experimental design, hyperparameter optimisation, and molecular docking benchmarks.

• In Appendix 10, we extend Polar to a loss-driven active learning task.

• In Appendix 11, we report additional experimental results, including sNMC upper bounds for location finding, a backbone ablation with TabPFNv2.5, best-performance comparisons on CES, and per-search-space breakdowns on HPO-B.

• In Appendix 12, we provide an overview of the computational resources and software dependencies used in this work.

8 Additional details about Polar↩︎

8.1 Backbone↩︎

The default backbone in our experiments is TabICLv2 [17], the regression variant of the TabICL family. TabICLv2 is a transformer-based in-context learner pretrained entirely on synthetic tabular data. Architecturally, TabICLv2 inherits the three-stage compress-then-ICL design of TabICLv1 [16]: a column-wise transformer that embeds each feature in isolation, a row-wise transformer that aggregates feature embeddings into a single fixed-dimensional representation per row, and a dataset-wise transformer that performs in-context learning across the entire context-query set. TabICLv2 extends this design with target-aware row embeddings, repeated feature grouping, and a query-aware scalable softmax that improves generalisation to longer contexts than those seen during pretraining; the model is pretrained with the Muon optimiser on a large corpus of synthetic datasets generated to maximise prior diversity. We use the official regression checkpoint released by the authors, which is trained to predict 999 conditional quantiles of \(y^\star\) via the aggregated pinball loss. As a robustness check, we additionally evaluate Polar with the TabPFNv2.5 backbone [15], using the public tabpfn-v2.5-regressor-v2.5_default.ckpt checkpoint; results are reported in Appendix 11.1.

The auxiliary prediction loss aligns the backbone with the current task distribution by supervising its predictive head on target points sampled within each trajectory. The exact form depends on which quantity the backbone is trained to predict.

TabPFNv2.5 outputs a categorical distribution over discretised output bins, which we treat as an approximate predictive density. The auxiliary loss is the standard negative log-likelihood \[\ell_{\mathrm{pred}}^{\mathrm{NLL}} = - \log q_\psi(y_m \mid x_m, \mathcal{D}_{t-1}).\]

TabICLv2 outputs \(K = 999\) conditional quantile estimates \(\hat{q}_\psi^{(\tau_k)}(x_m, \mathcal{D}_{t-1})\) for evenly spaced quantile levels \(\tau_k = k / (K+1)\), \(k=1, \ldots, K\). We follow the original training objective and adopt the aggregated pinball loss \[\ell_{\mathrm{pred}}^{\mathrm{pinball}}\bigl(q_\psi(\cdot \mid x_m, \mathcal{D}_{t-1}), \, y_m\bigr) = \frac{1}{K} \sum_{k=1}^{K} \rho_{\tau_k}\!\bigl(y_m - \hat{q}_\psi^{(\tau_k)}(x_m, \mathcal{D}_{t-1})\bigr),\] where \(\rho_\tau(u) = \max(\tau u, (\tau - 1) u)\) is the standard pinball (tilted absolute) loss for quantile level \(\tau\). In all experiments, the policy loss in 6 and the auxiliary prediction loss are weighted equally, \(\lambda_{\mathrm{pol}} = \lambda_{\mathrm{pred}} = 1\). For density-based rewards and NLPD evaluation in Appendix 10, following [17], we convert the predicted quantiles into a continuous distribution by enforcing monotonicity, linearly interpolating the quantile function between adjacent levels, extrapolating exponential tails beyond the extreme predicted quantiles, and then evaluating \(\log q_\psi(y\mid x, \mathcal{D}) = -\log \partial_\tau Q_\psi(\tau\mid x, \mathcal{D})\vert_{\tau = F_\psi(y\mid x, \mathcal{D})}\); the predictive mean is computed as \(\mu_\psi(x, \mathcal{D})=\int_0^1 Q_\psi(\tau\mid x, \mathcal{D})\,d\tau\).

8.2 Policy head↩︎

The default policy head used throughout the main experiments is a lightweight MLP with two hidden layers, hidden width 128, and GELU activations. It maps each candidate-conditioned representation to a scalar logit, and the policy distribution over candidates follows the softmax.

For the loss-driven active learning experiments in Appendix 10, where the goal is to optimise predictive performance over a user-specified target set, we use a target-aware variant of the policy head. This variant uses a single transformer encoder layer (4 attention heads, hidden width 1024, GELU activation), allowing each candidate token to attend to the target-point representations produced by the backbone.

During training, we optimise the policy head and the backbone using AdamW. Gradients are clipped to unit global norm. We sample candidates stochastically from the policy distribution to ensure sufficient exploration during training. At test time, we instead select the argmax candidate.

9 Experimental details↩︎

9.1 Benchmarking on Bayesian experimental design tasks↩︎

Location finding [63] is a widely adopted benchmark for sequential BED [9], [10], [27], [61]. The objective is to recover the positions of \(K\) unknown sources in \(\mathbb{R}^d\), denoted \(\theta = \{\theta_k \in \mathbb{R}^d\}_{k=1}^K\), by adaptively choosing measurement locations \(x \in \mathbb{R}^d\) at which noisy signal intensities are recorded. Each source radiates a signal whose magnitude attenuates with distance under an inverse-square law, so that a measurement taken at location \(x\) aggregates contributions from every source: \[\mu(\theta, x) \;=\; b \;+\; \sum_{k=1}^K \frac{\alpha_k}{m + \lVert \theta_k - x \rVert^2},\] where the constants \(\alpha_k\) specify the strength of each source, while \(b > 0\) and \(m > 0\) govern the background offset and the saturation behaviour of the signal at short range, respectively. The experimenter does not observe \(\mu\) directly; instead, the log-intensity is observed under additive Gaussian noise: \[\log y \mid \theta, x \;\sim\; \mathcal{N}\!\big(\log \mu(\theta, x), \sigma^2\big).\] Following standard practice, we instantiate the task with \(K=2\), \(\alpha_k=1\), \(b=0.1\), \(m=10^{-4}\), and \(\sigma=0.5\). Designs are constrained to the unit hypercube \(x \in [0,1]^d\), and the prior over each source places independent uniform distributions on every coordinate, \(\theta_{k,j} \sim \mathrm{Unif}[0,1]\) for \(j=1,\dots,d\).

Constant elasticity of substitution (CES; [64]) originates in economic theory and models how a consumer values bundles of goods. In the experimental design formulation, the experimenter aims to identify a participant’s latent preference structure by repeatedly presenting them with pairs of baskets and recording a (potentially noisy) judgement about their relative desirability. A basket \(z \in [0,100]^K\) encodes the quantities held of each of \(K\) goods. The participant’s preferences are captured by latent parameters \(\theta=(\rho,\boldsymbol{\alpha},u)\), where \(\rho \in (0,1)\) determines the degree of substitutability between goods, \(\boldsymbol{\alpha} \in \Delta^{K-1}\) assigns simplex-valued weights to the goods, and \(u>0\) modulates the overall responsiveness of the participant’s reports. A single experimental design takes the form of a basket pair \(x=(z,z') \in [0,100]^{2K}\), and elicits a bounded preference rating \(y \in [0,1]\).

The utility assigned to a basket follows the canonical CES form: \(U(z) \;=\; \left(\sum_{i=1}^{K} z_i^{\rho}\,\alpha_i\right)^{\frac{1}{\rho}}\). We adopt the following prior specification over the latent parameters: \[\rho \sim \mathrm{Beta}(1,1),\qquad \boldsymbol{\alpha} \sim \mathrm{Dirichlet}(\mathbf{1}_K),\qquad \log u \sim \mathcal{N}(1,3^2).\] For a given query \(x=(z,z')\) and parameter setting \(\theta\), the underlying (unobserved) utility differential is generated from \[\eta \sim \mathcal{N}\!\Big(u\,(U(z)-U(z')),\; u^2\,\tau^2\,(1+\lVert z-z'\rVert)^2\Big),\] and then transformed through a sigmoid link and clipped to a bounded interval to yield the reported outcome: \(y \;=\; \mathrm{clip}\big(\sigma(\eta),\;\epsilon,\;1-\epsilon\big)\), with \(\sigma(\cdot)\) the standard sigmoid. We use \(K=3\), \(\tau=0.005\), \(\epsilon=2^{-22}\), and query budget \(T=10\) in all experiments. Since TFMs are pretrained on standardised inputs and outputs, we standardise the features before feeding them to the backbone.

9.1.1 Baselines↩︎

Deep Adaptive Design [9] learns an amortised design policy by directly maximising the sPCE lower bound on the total EIG over \(T\)-step trajectories. The policy network consists of a two-layer MLP encoder with hidden width that maps each design-observation pair \((x_i, y_i)\) to a 16-dimensional embedding, a permutation-invariant pooling operation that aggregates these embeddings into a history representation, and a one-layer MLP emitter that maps the pooled representation to the next design. We train with Adam (learning rate \(5\times 10^{-5}\), \(\beta=(0.8, 0.998)\)), gradient clipping at 1.0, and an exponential learning-rate decay with factor 0.98 every 1000 epochs. For location finding, we use \(L=2 \times 10^4\) contrastive samples for sPCE estimation during training. For CES, we use \(L=10^5\).

RL-BOED [25] formulates sequential BED as a hidden-parameter Markov decision process and learns an amortised design policy using reinforcement learning. The reward is defined stepwise as the marginal contribution of the newly acquired design-observation pair to the sPCE. We use Randomized Ensembled Double Q-learning (REDQ) to train the continuous-design policy. The history encoder follows a DAD-style permutation-invariant architecture: each concatenated design-observation pair \((x_i,y_i)\) is passed through a two-layer MLP with 128 hidden units and ReLU activations, followed by a 64-dimensional linear output layer; the resulting embeddings are summed to form the history representation. The policy emitter outputs the mean and log-variance of independent Tanh-Gaussian distributions over the design dimensions. The critic networks use the same history encoder and concatenate the encoded history with the candidate design before passing it through a two-layer 128-unit ReLU MLP. We train with \(L=10^5\) contrastive samples for the sPCE reward. The remaining RL hyperparameters are listed in 1.

ALINE [27] jointly amortises sequential design and posterior inference with a shared masked transformer architecture. Each design \(x\) and observation \(y\) is embedded by separate two-layer MLPs with hidden width 128 into 32-dimensional representations; context tokens are formed by summing the design and observation embeddings, candidate-query designs are represented by design embeddings alone, and one learnable target token is introduced for each latent parameter dimension. These tokens are processed by a 3-layer transformer encoder with model width 32, feedforward width 128, and 4 attention heads. A one-hidden-layer MLP acquisition head maps the query representations to a categorical distribution over the candidate design pool, while a separate one-hidden-layer MLP posterior head predicts a 10-component Gaussian mixture for each target parameter token. We train ALINE with a joint objective consisting of a negative log-likelihood term for posterior prediction and a REINFORCE-style design loss based on one-step improvements in target log-likelihood, weighted equally with \(\alpha=1\) and discount factor \(\gamma=1\). For the sample-efficiency experiments, we use AdamW with cosine annealing and layer-wise learning rates of \(10^{-3}\) for the acquisition and posterior heads and \(2\times 10^{-4}\) for the shared embedder and transformer.

9.1.2 Evaluation details↩︎

Following [9], we evaluate all amortised policies using the sequential Prior Contrastive Estimation (sPCE) lower bound and the sequential Nested Monte Carlo (sNMC) upper bound on the total expected information gain. Given a sampled trajectory \((\theta_0, h_T) \sim p(\theta, h_T \mid \pi)\) together with \(L\) contrastive samples \(\theta_{1:L} \sim p(\theta)\) drawn independently from the prior, the sPCE and sNMC estimators are defined as \[\mathcal{L}_T(\pi, L) = \mathbb{E}\!\left[\log \frac{p(h_T \mid \theta_0, \pi)}{\frac{1}{L+1}\sum_{\ell=0}^{L} p(h_T \mid \theta_\ell, \pi)}\right], \quad \mathcal{U}_T(\pi, L) = \mathbb{E}\!\left[\log \frac{p(h_T \mid \theta_0, \pi)}{\frac{1}{L}\sum_{\ell=1}^{L} p(h_T \mid \theta_\ell, \pi)}\right],\] where the outer expectation is taken over \(\theta_0, h_T \sim p(\theta, h_T \mid \pi)\) and \(\theta_{1:L} \sim p(\theta)\). Both bounds become tight as \(L \to \infty\) at a rate \(O(L^{-1})\) [9]. For training, Polar uses \(L=10^5\) contrastive samples for both tasks. For the evaluations, we use \(L=10^6\) contrastive samples for location finding and \(L = 10^7\) for CES.

We report performance against the total number of training samples consumed during policy learning, where a "sample" denotes a single design-observation pair \((x, y)\) produced by the simulator. This convention provides a uniform unit of comparison that is independent of architectural choices such as batch size, replay-buffer capacity, or auxiliary loss terms.

For DAD and ALINE, every gradient step consumes one fresh batch of trajectories of length \(T\), so the cumulative number of training samples after \(E\) epochs is \(E \cdot B \cdot T\), where \(B\) is the batch size and \(T\) is the experiment horizon. For RL-BOED, which maintains a replay buffer, only the trajectories newly synthesised at each epoch contribute to the sample count; the batch size of the RL update governs how many transitions are drawn from the buffer per gradient step, but does not generate new simulator data, and is therefore excluded from the accounting. Concretely, if RL-BOED rolls out \(B_{\mathrm{env}}\) new trajectories per epoch, the cumulative count after \(E\) epochs is \(E \cdot B_{\mathrm{env}} \cdot T\).

For Polar, the auxiliary prediction loss requires evaluating the backbone on \(M\) target points per trajectory in addition to the \(T\) design-observation pairs visited by the policy. Each target point corresponds to a fresh sample drawn from the simulator and therefore contributes to the total sample budget. We thus account for both sources, giving a per-epoch consumption of \(B \cdot (T + M)\) and a cumulative count of \(E \cdot B \cdot (T + M)\) after \(E\) epochs. Our sample-efficiency comparison counts only the simulator samples consumed during policy learning for the target task; the cost of pretraining the TFM is not included. We consider this accounting appropriate because the backbones are pretrained entirely on synthetic tabular data that is task-agnostic, cheap to generate, and shared across all downstream applications. Once pretrained, the same checkpoint is reused without modification across BED, BO, and AL tasks, so the pretraining cost is amortised across all downstream uses rather than attributable to any single task, analogous to how downstream evaluations of ImageNet-pretrained vision models or pretrained language models typically report finetuning cost rather than pretraining cost. Polar inherits this off-the-shelf reusability from the broader tabular foundation model paradigm.

Both Polar and ALINE operate over a finite candidate pool at each acquisition step, whereas DAD and RL-BOED produce continuous designs directly. To match the original formulations of all baselines while keeping the pool-based methods well-resourced, we use a candidate pool of \(|\mathcal{C}| = 2000\) designs sampled uniformly from the design space at each step for location finding, and \(|\mathcal{C}| = 20000\) for CES.

9.2 Hyperparameter optimisation↩︎

9.2.1 Task description↩︎

HPO-B [65] is a large-scale meta-dataset for black-box hyperparameter optimisation (HPO) assembled from the OpenML repository. It contains 6.4 million hyperparameter evaluations. We use the curated HPO-B-v3 split, which retains the 16 search spaces and provides predefined meta-train, meta-validation, and meta-test partitions, together with five fixed initialisation seeds per test task to support reproducible comparisons. Each search space corresponds to the hyperparameter space of a particular machine-learning algorithm (e.g., SVM, XGBoost, glmnet), and each meta-dataset within a search space corresponds to evaluations on a specific tabular dataset. The hyperparameter ranges are normalised to \([0,1]^d\) and categorical hyperparameters are one-hot encoded. Following [12], we evaluate on six search spaces that together cover a representative range of underlying model families and input dimensionalities: glmnet (search space 5860, \(d=2\)), rpart.preproc (search space 4796, \(d=3\)), rpart (search space 5859, \(d=6\)), ranger (search space 5889, \(d=6\)), svm (search space 5527, \(d=8\)), and xgboost (search space 5906, \(d=16\)). We train on the provided meta-training datasets and evaluate on the held-out meta-test tasks. We adopt the recommended evaluation protocol of [65]: each test task is run with the five fixed initialisation seeds for \(T=50\) acquisition steps, and we report the average normalised regret and average rank as defined in [65].

9.2.2 Baselines↩︎

9.2.2.1 GP.

A standard Gaussian process surrogate fitted directly on the test task, with no use of meta-training data. We use an exact GP with a constant mean, a scaled RBF kernel, and a Gaussian likelihood. The kernel and noise hyperparameters are optimised by maximizing the exact marginal likelihood on the observed configurations after each BO round.

9.2.2.2 Meta-GP.

Following [12], a meta-trained Gaussian process pretrains the RBF kernel hyperparameters on the meta-training datasets of the corresponding search space and uses these pretrained values to initialise the GP at test time. Subsequent acquisitions then refit the kernel hyperparameters on the growing test-time observation set, starting from this meta-learned initialisation rather than from a generic prior.

9.2.2.3 PFNs4BO [22].

PFNs4BO uses Prior-data Fitted Networks [32] as transformer-based surrogates for BO: a single transformer is meta-trained offline to approximate the posterior predictive distribution of a chosen prior over functions, after which inference for any new context set is performed in a single forward pass without further fitting. The strongest variant of the model is trained on the so-called HEBO+ prior, which augments the HEBO Gaussian process prior [71] with input and output warping and additional hyperparameter randomisation, yielding a surrogate that has been shown to outperform standard GP surrogates on the HPO-B benchmark. We use the publicly released hebo-plus-model checkpoint from the official PFNs4BO repository without any further finetuning on the meta-training tasks.

9.2.2.4 NAP [12].

NAP is an end-to-end differentiable meta-BO framework that jointly meta-learns a transformer-based neural process [31] surrogate and an acquisition function via reinforcement learning, with an auxiliary supervised likelihood loss that stabilises training and provides an inductive bias toward valid probabilistic predictions. NAP is the closest amortised baseline to Polar, in that both methods learn the acquisition rule rather than relying on a hand-designed one. Since the HPO-B search spaces, test tasks, and initialisation seeds used by [12] are identical to ours, we directly use the per-step regret values reported in their official release for fair comparison.

9.2.2.5 TabICL [17].

A surrogate-based ablation that uses the same finetuned TabICL backbone as Polar, but discards the learned policy head and instead pairs the backbone’s predictive distribution with the EI acquisition function. This baseline isolates the contribution of learning the acquisition rule from the contribution of the meta-learned belief representation: the difference between TabICL and Polar can be attributed entirely to replacing a hand-designed acquisition with one trained on top of the same belief representation.

9.2.3 Evaluation details↩︎

For each HPO-B test task, all methods start from the same official initialisation and are evaluated on the same candidate pool. We report the normalised regret \[r_t = \frac{f^\star - f_t^{\mathrm{best}}}{f^\star - f_{\min}},\] where \(f_t^{\mathrm{best}}\) is the best objective value observed by the method up to step \(t\), \(f^\star\) is the best value in the full task pool, and \(f_{\min}\) is the worst value in the same pool. We average normalised regret over test datasets and initialisation seeds.

We also report the normalised rank to compare methods across tasks with different regret scales. For each task, seed, and optimisation step, methods are ranked by their normalised regret, with lower regret receiving a better rank and ties assigned the average rank. These ranks are then averaged over datasets, seeds, and search spaces.

9.3 Molecular docking optimisation↩︎

9.3.1 Task description↩︎

Dockstring [67] is a benchmark for molecular optimisation built around AutoDock Vina [72] docking simulations against a curated panel of medically relevant protein targets. The full dataset comprises docking scores and binding poses for over 260000 drug-like molecules evaluated against 58 targets, where each docking score quantifies the predicted binding affinity between a ligand and a target. Lower scores indicate stronger predicted binding, so the optimisation objective is to identify molecules that minimise the docking score for a given target. We frame this as a pool-based BO problem. Each molecule is represented by its SMILES string and converted to a 512-bit Morgan fingerprint with radius 2, matching the featurisation used by [67] for their property-prediction baselines. We randomly subsample 10000 molecules from the full dataset and use this fixed subset throughout all experiments. We split the 14 protein targets used in our experiments into a meta-training set of 8 targets, JAK2, KIT, LCK, SRC, IGF1R, ABL1, MET, EGFR, and a held-out test set of 6 targets, PTK2, FGFR1, CSF1R, CDK2, MAPK14, KDR. All 14 are protein kinases, which share a conserved ATP-binding pocket and substantial structural homology across the family. This shared geometry means that the physicochemical features predictive of strong binding are correlated across targets, making the kinase family a natural setting for meta-learning.

9.3.2 Baselines↩︎

9.3.2.1 GP with Tanimoto kernel.

Standard continuous kernels such as the RBF or Matérn kernels are poorly suited to high-dimensional binary fingerprint inputs, where Euclidean distance is a weak measure of chemical similarity. We instead equip the GP with the Tanimoto kernel, a standard choice for fingerprint-based molecular BO [68]. For two binary fingerprints \(x, x' \in \{0,1\}^d\), the Tanimoto kernel is defined as \[k_{\mathrm{Tanimoto}}(x, x') = \sigma^2 \cdot \frac{\langle x, x' \rangle}{\|x\|^2 + \|x'\|^2 - \langle x, x' \rangle},\] where \(\sigma^2\) is a learnable output scale and \(\langle \cdot, \cdot \rangle\) denotes the inner product. The kernel returns 1 for identical fingerprints and 0 for fingerprints with no shared bits, providing a meaningful similarity signal for sparse binary representations. We refit the kernel hyperparameters by maximum marginal likelihood after each acquisition step.

9.3.2.2 Random Forest.

Random forests are a standard surrogate in cheminformatics and are commonly paired with molecular fingerprints [69], [70]. We use an ensemble of 256 regression trees with bootstrapped training sets and the standard \(\sqrt{d}\) feature-subsampling rule at each split. The surrogate is refit on the full set of observed molecules at each acquisition step.

9.3.2.3 TabICL.

We use the same setup as we used in our HPO-B experiments, where the same finetuned TabICL backbone as Polar is used but replaces the learned policy head with a hand-designed acquisition function.

9.3.3 Evaluation protocol↩︎

For each of the 6 held-out test targets, we evaluate all methods under the same protocol: 10 randomly sampled molecules form the initial context, after which the policy selects \(T = 50\) acquisitions, with a candidate pool of size \(|\mathcal{C}| = 2000\). We report the average regret and aggregate by averaging across the 6 test targets and 5 random seeds for each test target.

10 Loss-driven active learning↩︎

In this section, we evaluate Polar on a loss-driven active learning task, where the goal is to select query points that improve predictive performance on a downstream task distribution \(p^\star\) under a user-specified loss.

Active learning seeks to acquire labelled data that maximises the predictive performance of a downstream model. In practice, a clinician may care primarily about predicting high-risk patients accurately, a structural engineer about responses near a failure threshold, and a drug discovery pipeline about molecules with high binding affinity. Standard information-theoretic acquisition objectives such as BALD [40] and EPIG [21] treat all outcomes symmetrically and therefore cannot directly target such loss-weighted criteria. [35] formalise this gap and derive a principled extension via weighted Bregman divergences, yielding myopic acquisition rules such as weighted variance reduction (GP-VR\(_w\)) that explicitly target a user-specified weighted loss. We adopt the same loss-driven framing but instantiate it in our amortised setting.

Concretely, we consider pool-based active learning under a weighted loss \(\ell_\omega(y, q) =\omega(y) \ell(y,q)\), where \(\omega: \mathcal{Y} \to \mathbb{R}_+\) encodes which outcomes matter most for the downstream task. At round \(t\), the learner observes a context \(\mathcal{D}_{t-1} = \{(x_i, y_i)\}_{i=1}^{t-1}\), selects a query \(x_t\) from a finite pool \(\mathcal{C}_t\), and observes \(y_t \sim p(y \mid x_t, z)\) where \(z\) is the latent function. After \(T\) rounds, performance is evaluated on a held-out target set drawn from the target input distribution \(p_\star(x^\star)\).

Polar instantiates this setting using the same two-component architecture as in our other experiments. TFM plays the role of an amortised surrogate model: given the current context \(\mathcal{D}_{t-1}\) and a query input \(x^\star\), a single forward pass yields the approximate posterior predictive \(q_\psi(y^\star \mid x^\star, \mathcal{D}_{t-1})\), replacing the GP surrogate that classical AL methods refit at every round. The policy head then scores each candidate in \(\mathcal{C}_t\) from the same forward pass and samples the next query \(x_t\). For the utility, we use the weighted one-step improvement in predictive log-density on the target set, \[\label{eq:al-reward} u_t = \mathbb{E}_{x^\star \sim p^\star}\big[\omega(y^\star)\big(\log q_\psi(y^\star \mid x^\star, \mathcal{D}_t) - \log q_\psi(y^\star \mid x^\star, \mathcal{D}_{t-1})\big)\big],\tag{7}\] which measures the weighted improvement in the backbone’s own predictive log-density on the target set after observing the new query, and which generalises the dense sEPIG reward of [27] from the unweighted to the weighted setting. In our experiments, we adopt the exponential weighting \(\omega(y) = \exp(\beta y)\) with \(\beta = 10\), prioritising regions of high outcome value. At deployment, no GP fitting or acquisition optimisation is required: the next query is produced in a single forward pass per round.

10.1 Training distribution↩︎

We train Polar on synthetic functions drawn from a distribution over GPs with randomised kernels and hyperparameters, following a procedure similar to ALINE [27]. Each task is sampled by:

1. Drawing a kernel uniformly from {RBF, Matérn-3/2, Matérn-5/2};

2. Drawing a length-scale \(\ell \sim \text{LogUniform}(0.1, 2.0)\) and an output scale \(\sigma_f \sim \text{Uniform}(0.1, 1.0)\);

3. Sampling a function \(f \sim \mathcal{GP}(0, k_{\ell, \sigma_f})\).

Before each forward pass, we normalise each function using task-level statistics computed independently of the observed context: for inputs, we use the known domain bounds, and for outputs, we estimate the normalisation mean and variance from a large set of reference points sampled from that function. The horizon is \(T = 20\), with an initial context of 2 points and a candidate pool of size \(K = 200\).

10.2 Baselines↩︎

We compare against four non-amortised GP baselines.

GP-RS (random sampling) selects \(x_t\) uniformly at random from the pool.

GP-US (uncertainty sampling) selects the candidate with the highest posterior predictive variance, \[\text{US}(x) = \mathbb{V}[y \mid x, \mathcal{D}_{t-1}].\]

GP-VR (variance reduction) selects the candidate that maximally reduces total predictive variance over the target set \(\{x^\star_m\}_{m=1}^{M}\): \[\text{VR}(x) = \sum_{m=1}^{M} \frac{\text{Cov}_{t-1}(x^\star_m, x)^2}{\mathbb{V}[y \mid x, \mathcal{D}_{t-1}]},\] where \(\text{Cov}_{t-1}(x^\star, x)\) is the GP posterior covariance between the latent function values at \(x^\star\) and \(x\) given \(\mathcal{D}_{t-1}\).

GP-VR\(_\omega\) (weighted variance reduction; [35]) is the loss-matched counterpart of GP-VR, derived from the same weighted-Bregman framework as our reward. It targets the weighted predictive variance under the reweighted predictive \(p_\omega(y \mid x^\star) \propto \omega(y) p(y \mid x^\star)\) rather than the unweighted variance.

10.3 Evaluations↩︎

We evaluate on two distributions of test functions.

GP synthetic. We sample 100 functions from the same generative procedure used during training. This measures in-distribution performance.

Benchmark functions. We additionally evaluate on four standard regression benchmarks to measure out-of-distribution generalisation: Forrester, Gramacy-Lee, Higdon, and Sine-Gaussian Bumps. The functional forms are: \[\begin{align} f_{\text{Forr}}(x) &= (6x - 2)^2 \sin(12x - 4), \\ f_{\text{GL}}(x) &= \frac{\sin(10\pi x)}{2x} + (x - 1)^4, \\ f_{\text{Higdon}}(x) &= \sin(2\pi x / 10) + 0.2 \sin(2\pi x / 2.5), \\ f_{\text{SGB}}(x) &= 2\sin(2x) + 8 \phi_{2.5, 0.5}(x) + 10 \phi_{7.5, 0.25}(x) - 6 \phi_{-4.5, 0.5}(x). \end{align}\] where \(f_0(x) = \sin(2x)\) and \(\phi_{\mu, \sigma}\) is a Gaussian density with mean \(\mu\) and standard deviation \(\sigma\). Before evaluation, we apply a fixed per-function normalisation: inputs are centred and scaled using the benchmark domain, and outputs are standardised using statistics estimated from Sobol samples over the same domain. All methods, including GP and Polar, are evaluated on these same normalized benchmark versions.

For each test function (whether GP-sampled or benchmark), we run 100 independent seeds. Each seed draws an initial context of 2 points, a candidate pool of 200 points, and a target set of 100 points; all three sets are sampled uniformly. Performance is reported as weighted RMSE and weighted NLPD on the target set: \[\text{RMSE}_\omega = \sqrt{\frac{\sum_{m=1}^{M} \omega(y^\star_m)(y^\star_m - \mu(x^\star_m))^2}{\sum_{m=1}^{M} \omega(y^\star_m)}}, \qquad \text{NLPD}_\omega = -\frac{\sum_{m=1}^{M} \omega(y^\star_m)\log q(y^\star_m \mid x^\star_m)}{\sum_{m=1}^{M} \omega(y^\star_m)},\] where \(\mu(x^\star)\) is the predictive mean under the surrogate at the end of the acquisition trajectory.

10.4 Results↩︎

2 reports final-step weighted RMSE and weighted NLPD on both evaluation distributions. First, the two loss-aware methods GP-VR\(_\omega\) and Polar substantially outperform the loss-agnostic baselines (GP-RS, GP-US, GP-VR) on both metrics and both evaluation distributions. This confirms the central message of [35]: when the downstream loss assigns non-uniform weight to outcomes, acquisition rules that ignore the weighting are systematically suboptimal. Second, Polar consistently outperforms GP-VR\(_\omega\) across both evaluation distributions, suggesting that learning the acquisition rule directly from data outperforms a hand-designed one. 6 visualises the acquisition behaviour of GP-VR, GP-VR\(_\omega\), and Polar on the Sine-Gaussian Bumps function from [35].

11 Additional experimental results↩︎

11.1 Location finding↩︎

7 and 8 report complementary results on the location finding task. 7 (a) and 7 (b) show the sNMC upper bound on EIG against total training samples in the 2D and 5D settings, respectively. The relative ordering of methods mirrors the sPCE comparison in 2.

8 (a) examines whether Polar’s gains depend critically on the choice of pretrained backbone by replacing TabICLv2 with TabPFNv2.5 on the 2D task. The TabPFN-backed policy converges more slowly during the early stages of training, but the two variants reach essentially the same EIG at convergence. This suggests that Polar is not tied to a particular backbone: as long as the pretrained model produces representations from which the task-relevant belief can be approximately decoded, a lightweight policy head can turn them into an effective acquisition policy.

Another natural question is which layer of the backbone supplies the most informative belief representation for policy learning. To test this, we re-train Polar on the 2D task using representations taken from the last, second-to-last, fourth-to-last, and eighth-to-last transformer layers of the TabICLv2 backbone, keeping all other components fixed. 8 (b) shows the last and second-to-last layers perform comparably throughout training, with the second-to-last layer slightly ahead in the low-sample regime and the last layer marginally ahead at convergence. Earlier layers degrade performance monotonically. A natural concern is that the layer ordering in 8 (b) reflects the effect of finetuning rather than the structure of the pretrained backbone. 8 (c) repeats the comparison with a frozen backbone and shows the same ordering, confirming that later layers provide more informative belief representations regardless of whether the backbone is adapted. Therefore, we adopt the last layer as the default in all other experiments.

11.2 Constant elasticity of substitution↩︎

3 reports each method’s best-attained performance. Polar achieves the strongest sPCE among all methods, retaining its lead from the matched-budget comparison in ¿tbl:tab:ces?. During training, we observed that both DAD and RL-BOED were prone to instability: their training losses diverged in the late stages, with sPCE peaking and then degrading. We therefore applied early stopping based on validation sPCE for both baselines and report the checkpoint at which validation sPCE was highest. A separate observation concerns the sPCE–sNMC gap. For RL-BOED, this gap is markedly larger than for the other methods. A wide gap at finite contrastive samples indicates that the EIG bounds are loose for this particular policy, rather than that the true EIG is correspondingly large.

11.3 Hyperparameter optimisation↩︎

9 and 10 break down the aggregate results of 6.2 by search space, reporting the normalised regret and the average rank, respectively, for each of the six HPO-B search spaces. Polar is consistently among the top performers across search spaces, achieving the lowest regret on most of them and remaining competitive on the others. The advantage over the strongest baselines is most pronounced on the higher-dimensional search spaces (5527 and 5906), where the meta-learned belief representation appears to confer the largest benefit. On the lower-dimensional search spaces the comparison is closer, but Polar retains either the best or second-best regret throughout the acquisition trajectory, and the average-rank curves in 10 confirm that this consistency holds across tasks and seeds rather than being driven by a small number of favourable cases.

12 Computational resources and software↩︎

All experiments presented in this work, including model development, hyperparameter tuning, baseline evaluations, and preliminary analyses, were performed on a GPU cluster equipped with NVIDIA H200 GPUs. The total computational resources consumed for this research, across all development stages and experimental runs, are estimated to be approximately 2000 GPU hours. The wall-clock cost of training Polar varies with the task: on the HPO-B benchmark, a strong policy is typically obtained within 30 minutes per search space; on the location finding, policy training takes around 5 hours; on the CES, it takes about 16 hours to get the final best performance; on the Dockstring benchmark, the high input dimensionality increases the per-step cost of the backbone forward pass, and a full training run takes approximately one day. The core code base is implemented in PyTorch (https://pytorch.org, License: modified BSD license), which also underlies our policy head as well as the TabPFN and TabICL backbones. We use the publicly released TabPFNv2.5 backbone from the PriorLabs repository (https://github.com/PriorLabs/TabPFN, License: Apache 2.0), with the corresponding checkpoint hosted on Hugging Face (https://huggingface.co/Prior-Labs/tabpfn_2_5). For the TabICLv2 backbone, we use the official implementation released by the authors (https://github.com/soda-inria/tabicl, License: BSD 3-Clause), with the regression checkpoint hosted on Hugging Face (https://huggingface.co/jingang/TabICL-clf). BED baselines are adapted from the original authors’ publicly available code. For the HPO-B baselines, we use the publicly released hebo-plus-model checkpoint from the official PFNs4BO repository (https://github.com/automl/PFNs4BO, License: Apache 2.0), and we use the per-step regret values released as part of the official NAP codebase (https://github.com/huawei-noah/HEBO/tree/master/NAP) for the NAP comparison. The Gaussian process baselines are implemented using GPyTorch (https://github.com/cornellius-gp/gpytorch, License: MIT). The HPO-B benchmark is obtained from its official repository (https://github.com/releaunifreiburg/HPO-B, License: MIT), and the Dockstring benchmark is obtained from its official repository (https://github.com/dockstring/dockstring, License: MIT).

References↩︎