Technical Summary: Efficient Adaptive Data Acquisition via Pretrained Belief Representations (POLAR)

Problem Statement

Adaptive data acquisition encompasses sequential learning problems such as Bayesian Experimental Design (BED), Bayesian Optimisation (BO), and Active Learning (AL). In these settings, a learner must sequentially select queries (designs) to maximize a downstream utility (e.g., information gain, global optimum identification, or predictive accuracy) under a finite budget.

Existing approaches to learning acquisition policies generally fall into two problematic categories:

1. Posterior-based methods: These rely on fitting a probabilistic model to observed data to compute a posterior, which is then used to evaluate an acquisition function (e.g., Expected Improvement, Expected Information Gain). These methods are bottlenecked by the quality of the posterior approximation. Model misspecification or approximation errors can distort the acquisition rule, and these errors compound over sequential rounds.
2. Direct policy-learning methods: These map raw observation histories directly to decisions (e.g., via Reinforcement Learning). While they avoid explicit posterior inference, they must learn a complex mapping from raw data to decisions from scratch. This requires simultaneously learning how to represent the dataset and how to act on that representation, often leading to sample inefficiency and training difficulties.

Methodology: POLAR

The authors introduce POLAR (Policy Learning with Belief Representations), a framework that decouples representation learning from policy learning. The core insight is that the optimal sequential decision depends on the observation history only through a sufficient belief state. Rather than learning this state from scratch or relying on explicit posterior approximations, POLAR leverages pretrained predictive foundation models as belief-state encoders.

Architecture

1. Belief Encoder (Backbone): POLAR utilizes a Tabular Foundation Model (TFM), such as TabPFN or TabICL, pretrained on diverse synthetic tasks. Given a history of observations $D_{t-1}$ and a candidate design $x$, the TFM produces a candidate-conditioned hidden representation $r_\psi(x, D_{t-1})$.

   • Unlike the final predictive distribution output by the TFM, this hidden representation is treated as the belief representation. It encodes how predictions change with data (predictive uncertainty, local structure, context-dependent behavior) without committing to a specific posterior approximation.
   • The framework operates in a pool-based setting, where the policy selects from a finite set of candidates $C_t$. The backbone processes the history once to generate representations for all candidates in the pool.

2. Policy Head: A lightweight, learnable policy head $g_\eta$ maps the belief representation $r_\psi$ to a scalar logit for each candidate. A softmax over these logits defines the stochastic policy $\pi_\eta$. At deployment, the candidate with the highest score is selected.

3. Training Objective:

   • Policy Loss: The policy is trained via policy gradients on simulated rollouts. The objective maximizes a task-specific utility (e.g., Expected Information Gain for BED, terminal best value for BO). The authors use a local credit assignment strategy (one-step increments) to reduce variance, rather than a full long-horizon value function.
   • Backbone Adaptation (Optional): To align the pretrained backbone with the specific downstream task distribution, an auxiliary supervised prediction loss ($L_{pred}$) is applied to the backbone parameters.
   • Decoupling: A critical design choice is that policy gradients are not propagated into the backbone. The backbone is updated only via the supervised prediction loss, while the policy head is updated via the policy gradient. This prevents high-variance reinforcement learning signals from destabilizing the pretrained representations.

Theoretical Foundation

The paper provides a unified decision-theoretic view, showing that the Bayes-optimal policy $\pi^*$ is a function of the task-relevant posterior $p(w|D_{t-1})$. Since the posterior is a function of the history, any representation $e_t$ that allows for the recovery of the posterior (via a decoder) is a valid belief representation. POLAR posits that pretrained TFMs provide such representations that are rich enough to determine optimal actions but structured enough to make the mapping to the policy head easy to learn, avoiding the extremes of raw history (too complex) or explicit posterior (approximation error prone).

Key Contributions

1. Unified View: The authors formalize adaptive data acquisition as learning policies on top of belief representations, bridging the gap between raw history-based and posterior-based approaches.
2. Decoupled Architecture: They propose a simple architecture that uses pretrained foundation models as fixed (or finetuned) belief encoders and trains a separate policy head, separating representation learning from decision-making.
3. Sample Efficiency: Empirical results demonstrate that POLAR achieves state-of-the-art performance across diverse tasks (BED, BO, AL) while requiring significantly fewer training samples (up to $100\times$ fewer) compared to existing amortised methods.

Experimental Results

The authors evaluate POLAR on three main benchmarks:

1. Bayesian Experimental Design (BED):

   • Tasks: Location Finding (2D and 5D) and Constant Elasticity of Substitution (CES).
   • Results: POLAR significantly outperforms state-of-the-art amortised methods (DAD, RL-BOED, ALINE). In the 5D Location Finding task, POLAR surpasses the peak performance of RL-BOED and DAD using 10,000$\times$ fewer training samples and exceeds ALINE with 100$\times$ fewer samples.
   • Ablations:
     • Frozen vs. Finetuned Backbone: Finetuning the backbone with auxiliary prediction loss yields better performance than keeping it frozen, though even the frozen version is competitive.
     • Decoupled vs. Coupled Gradients: Allowing policy gradients to update the backbone (coupled) severely degrades performance, confirming the necessity of decoupling.
     • Pretrained vs. Scratch: Training the architecture from scratch converges much slower and performs worse, highlighting the value of the pretrained foundation model.

2. Bayesian Optimisation (Hyperparameter Optimisation):

   • Benchmark: HPO-B (Hyperparameter Optimisation Benchmark).
   • Results: POLAR achieves the lowest average regret and best average rank across six search spaces, outperforming classical GPs, Meta-GPs, and end-to-end amortised methods like NAP. It demonstrates that using a pretrained TFM backbone yields stronger performance than learning representations from scratch.

3. Molecular Optimisation:

   • Benchmark: DOCKSTRING (high-dimensional molecular docking).
   • Results: In a 512-dimensional discrete space, POLAR achieves the lowest regret compared to GPs with Tanimoto kernels, Random Forests, and a TabICL baseline with a hand-designed acquisition function. This confirms that learning the acquisition rule directly from the belief representation outperforms fixed acquisition strategies.

Significance and Claims

The paper claims that POLAR represents a significant step in the scalability and efficiency of amortised data acquisition. By shifting the paradigm from jointly learning representations and decisions to acting on pretrained belief representations, POLAR mitigates the sample inefficiency and training instability associated with direct policy learning and the approximation errors of posterior-based methods.

The authors emphasize that their approach is not about designing a new estimator or training algorithm, but rather reusing a pretrained predictive model as a belief encoder. This allows a small policy head to be trained with standard policy gradients, achieving superior performance with drastically reduced computational costs for training. The framework is presented as a unified solution applicable to BED, BO, and AL, differing only in the task-specific utility used for training.