Efficient Bayesian Experiment Design with Equivariant Networks
Authors: Conor Igoe
Tejus Gupta
Jeff Schneider
Venue: NeurIPS 2025 poster ()
Keywords: Bayesian Experiment Design
Bayesian Optimization
Active Search
Deep Learning
Deep Reinforcement Learning
Graph Neural Networks
Primary area: deep_learning

Abstract: Recent work in Bayesian Experiment Design (BED) has shown the value of using Deep Learning (DL) to obtain highly efficient adaptive experiment designs. In this paper, we argue that a central bottleneck of DL training for BED is belief explosion. Specifically, as an agent progresses deeper into an experiment, the effective number of realisable beliefs grows enormously, placing significant sampling burdens on offline training schemes in an effort to gather experience from all regions of belief space. We argue that choosing an appropriate inductive bias for actor/critic networks is a critical component in mitigating the effects of belief explosion and has so far been overlooked in the BED literature. We show how Graph Neural Networks are particularly well-suited for BED DL training due to their domain permutation equivariance properties, resulting in multiple orders of magnitude improvement to sample efficiency compared to naive parameterizations.

Principal among the motivations for the involvement of DL techniques is their potential to increase the scope of problems that admit practical BED solutions [IT] More specifically, practitioners typically trade off decision-theoretic performance with online computational costs when deploying BED policies real-world tasks. Deep Learning can amortize computationally costly but high- performing BED policies, allowing for improved decision-making, reduced compute burdens, or both. For example, recent work has shown that multistep lookahead- approaches for achieving non-myopic experiment designs involve considerable online computation in the face of nested inte- grals. Wu and Frazier [12] describe how these approaches can require up to an hour of compute for single decision in modest Bayesian Optimization (BO) problems. Although they propose novel method to reduce this compute overhead to at most several minutes, such online compute require- tasks, for example in remote Efficient Bayesian Experiment Design with Equivariant Networks
Conor Igoe*
Tejus Gupta' Robotics Institute Carnegie Mellon University te jujugecs cmu.
Jeff Schneider
Machine Learning Department Carnegie Mellon University ci cigoe@cs. cmu. edu
Robotics Institute
Cameegi Mellon University
schneide@cs. cmu edu
Abstract
Recent work in Bayesian Experiment Design (BED) has shown the value of us- ing Deep Learning (DL) obtain highly efficient adaptive experiment designs. In this paper, argue that central bottleneck of DL training for BED is be- lief explosion. Specifically, as an agent progresses deeper into an experiment, the effective number of realisable beliefs grows enormously, placing significant sam- pling burdens on offline training schemes in an effort to gather experience from all regions of belief space. We argue that choosing an appropriate inductive bias for actor/critic networks is a critical component in mitigating the effects of be- lief explosion and has far been overlooked in the BED literature. We show how Graph Neural Networks are particularly well-suited for BED DL training due their domain permutation equivariance properties, resulting in multiple orders magnitude improvement to sample efficiency compared to naive parameteriza- tions.
1 Introduction
Scientific inquiry depends on empirical observations to refine our understanding of the world. These observations generate data that be analyzed, interpreted, and used to test hypotheses and theo- ries. However, conducting experiments comes with costs, including time, financial resources, and logistical constraints, As result, determining which observations most valuable for given line of inquiry is a central concern in the Design of Experiments.
In particular, Bayesian Experiment Design, BED, has emerged
elegant formalism for un- derstanding the value of different experiment designs [ (2] Moreover, in recent years there has been a growing interest in adopting Deep Learning (DL) and Deep Reinforcement Learning (DRL)
techniques to obtain effective experiment designs for BED tasks B 4 5 8
ments are prohibitive in certain high- figgufreyeency
Equal contribution.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

9.
P4 C1
o(m))
(1
Figure 1: Illustration Domain Transformation Equivariance The figure on the left of the dotted line shows the posterior mean, μi, and marginal standard deviation, a. of a 2D Gaussian Process, along with the LogEl acquisition function used in Bayesian Optimization. This corresponds to the optimal 1-step lookahead landscape when cast the same posterior bclief 90%. Note that although these belief-action pairs appear different, they are related by the transformation & and represent equivalent decisions under an equivariant policy.
MDP, with the optimal action shown in red. The figure on the right shows and acquisition function under transformation that rotates the domain by
In contrast, amortized policies trained using DRL can achieve non-myopic designs while requiring several orders of magnitude less online computation 4 I0].
sensing [0] or edge computing [13]
Although the recent focus amortizing expert policies for BED has shown great promise, argument of this paper is that existing approaches have significantly high training compute burdens. For example, we find that these approaches take the order of days match the performance of simple experts on modest BED tasks.
We trace this issue to an explosion of possible posterior beliefs as inference progresses through an adaptive experiment, which we call "belief explosion". We observe that existing approaches struggle to learn in the presence of this large diversity of belief's at deeper timesteps. We argue that designing methods that can cope with this belief explosion efficiently should be a priority for the BED community, and our work proposes to leverage the equivariant structures of optimal BED policies to achieve this goal.
Specifically, we show how a large family of BED tasks admit optimal policies that are domain permutation equivariant, and that exploiting this structure leads to significantly more efficient learn- ing in the face of belief explosion. Moreover, we show in this work, these structural properties naturally suggest an inductive bias that is readily captured in Graph Neural Network (GNN) archi- tectures.
In particular, our contributions are follows:
we show that significant bottleneck when training policies for BED tasks is belief explo- sion, and show how standard architectures commonly found the literature are ill-suited to address this challenge;
we prove that a large family of BED tasks are domain permutation equivariant, which suggests natural inductive bias for policy and critic networks;
we demonstrate the utility of leveraging this domain cemmain-eermutationn equivariance on two prominent BED tasks, showing to two orders of magnitude improvement in sample effi- ciency compared to various architectures operating both belief-state and information-set representations;
we show that GNNs trained on small BED tasks can be used to generalise to significantly larger tasks, resulting in order of magnitude less online compute test time:
we show that these equivariances can be extended to information-set-bue based policies in tinuous BED tasks, allowing us to efficiently train performant policies for higher dimen- sional BO tasks.
2 Notation & Preliminaries
In order to provide a unified notation for our main theoretical and experimental results, we introduce the Bayesian Experiment Design and Markov Decision Process formalisms the following two subsections, followed by short overview of Graph Neural Networks.

---

2.1 Bayesian Experiment Design
Bayesian Experiment Design is general framework that helps agents make intelligent decisions to efficiently reduce their uncertainty over a set hypotheses. More specifically, BED models a sequential interaction process, or "experiment", between a decision-making agent and an environ- ment containing some latent variable of interest, denoted 0. BED task is defined by the tuple Popo, (0,X,) V. Po. p, r), where è defines space of possible values for the latent variable,
"hypothe- ses". X' defines a space of actions available to the agent, or "experiment designs". and y defines space of observations, or "experiment outcomes". The probabilistic models Po( and p(y|æ 0) define a prior over the possible hypotheses and the observation model, respectively.
At each time step t = 0, l..... the agent chooses the experiment design Ia based on the posterior belief state pt := p(0 {(#r, Vr (-)) The environment then yields the observation y/ ~ p(yl 0), which the agent adds to its information set Iz = {(Ir. }+)] The agent is interested in choosing informative experiment designs reduce the uncertainty about 0 as measured by some functional of the posterior r r(p(O[Z)))
The BED framework is flexible in its ability to model finite hypothesis classes well as contin- uous. For example, in this paper we focus on two tasks that feature prominently in the ML and robotics literature as case studies for experimental results: Bayesian Optimization (BO) and Active Search (AS). BO can be thought of as an instance of a BED task where 0 represents some unknown continuous black-box function that must be optimized. Similarly, AS can be thought of as a BED task where 0 is a binary-valued vector representing the location of objects of interest in a discredited real-world terrain. We return more detailed descriptions of the BO and AS tasks in their respective experimental sections.
2.2 Markoy Decision Processes
For ease of analysis, we now describe how to view a BED task as an instance of a Markov Decision Process (MDP).
We consider is continuous and the action space A is discrete, and the unknown state transition function P X x A + [0, oo) represents the probability density of the next state 8,41 E S given the current state s, E S and action a, E A. The environment generates reward ra E IR based on the current state S, and action a,. The environment generates trajectories of length T E N starting from an initial state drawn from initial state distribution p S - [0, oo). The standard MDP objective is to find a policy Te S + A by optimizing policy parameters to maximize the expected discounted return E
MDPs defined by the tuple (S. where the state space S
T- E Te /T]. Note that use & to denote the MDP timestep throughout this paper.
For the BED task, we define the MDP state as the posterior belief s; := Pt. and the action as the design a == The state definition satisfies Markovian dynamics assumptions by stochastically transitioning to state Pt+1 after incorporating (It, yt) using the inference equations, with stochastic- ity arising from the random variable y/- The reward is defined as a functional of pr, and can involve a variety of application-specifi quantities that incorporate the posterior belief state, as well as action costs and geometric quantities relating to the geometry of the underlying domain. In the experiment section, we describe some example reward functions that commonly found in the BED and related literature for our case study environments.
cntrmatiinnthhereee terms derived from
We note that is also possible to cast BED tasks in the MDP formalism by defining the state S; := I. The advantage of choosing information set state representations over belief-state representations that it allows for the training of end-to-end networks that amortize over the inference process mapping from I, to Pt, as well as the expected returns and optimal action. This may be attractive certain contexts, where inference itself is costly. Indeed, this the approach adopted in 3 g
Due to space constraints, in the main paper we present results that leverage discretized belief state representations following previous work 7 10 [6]. We include additional results on non-discretized information set methods in the Appendix E.2
14 15 g
3

---

2.3 Graph Neural Networks
Graph Neural Networks (GNNs) are a broad family of architectures designed to process data rep- resented graphs, where entities are modeled as nodes and their interactions as edges. Through iterative message passing, GNNs enable each node update its representation by aggregating infor- mation from its neighbors, thereby capturing relational structure and inductive biases that standard feedforward or convolutional architectures cannot. This makes them especially well-suited to do- mains where the underlying structure is non-Euclidean or relational.
Over the past several years, GNNs have been successfully applied across a wide range of machine learning and robotics subfields. In representation learning, GNNs underpin molecular property pre- diction 17. and structure-aware scene understanding [18] In reinforcement learning and control, GNNs have been used to model multi-agen interactions [19], physical dynamics [20], and robotic manipulation [2T]. leveraging their ability to encode symmetry and compositional structure. More broadly, GNNs have become a unifying framework for exploiting invariances and equivariances inherent in structured decision-making problems.
In our experiments, we build on these insights by leveraging GNNs to capture important equivari- ances in Bayesian Experiment Design (BED) tasks. As described in Section I our model architec- ture exploits the graph structure of the belief representation and design space to improve generaliza- tion across environments and hypothesis classes.
3 Belief Explosion as Policy Optimization Bottleneck in BED
In this section, we identify the belief explosion issue and identify how it affects policy optimization which the number of possible beliefs agent must consider becomes intractably large as the search progresses. This a well-known problem in the POMDP planning literature (22 23], but has received little attention in the contemporary BED and deep learning communities.
for BED tasks. More specifically, belief explosion refers to the phenomenon
Figure 2 visualizes the belief explosion for a ID BO task. The figure on the left shows four belief state trajectories, illustrating the increasing diversity of beliefs with time. The figure on the right illustrates this more precisely: As the agent progresses further into the trial, the average divergence between two realizable beliefs increases. This is shown by the blue curve, which measures the KL divergence between canonical representations of the beliefs the posterior mean vector and covariance matrix- in discretized Gaussian Process prior BO tasks. The orange curve shows that much of this growth is overstated: for each pair of realizable beliefs (s, 8), there exists a transformed belief in the set of all permuted beliefs B(š) which can have significantly lower divergence with compared to its original untransformed representation.
Stajocory
Thejetery
R.. [cokaas)
Increadig bettet state divexity
Figure 2: Belief Explosion
Next, we show that this rapid increase in S, with t renders current policy optimization methods for BED highly sample inefficient when using standard belief representations and network architecture choices. We consider a simple offline learning problem in a small 10-bin discretized ID BO task. The goal here is to imitate the I-step greedy expected information gain expert using behavior cloning with a fully connected network trained on concatenate(44 flatten (E))))n(E))) discretized belief-state representations. Figure shows cross entropy train-test loss curves with expert data set of 104 expert environment transition samples.
Mild overfitting is evident in the left plot of Figure 3 and is perhaps otherwise unremarkable. The fig- ure on the right in Figure shows the same test loss curve but decomposed over various timesteps. As example, the green curve shows the test loss when using test beliefs {pt} ) as input to ne.

---

0.06
0.06
Twai
Text
0.04
Tess Loss
0.04
0.02
0.02
0.00
0.00
5000
10000
10000
D
5000
10000
15000
Gradient Steps
Figure 3: Test Loss Temporal Decomposition Note that each curve is averaged over seeds. The shaded
region shows =2 standard error..
make the following observations about Figure
The best test loss achieved by Ti increases monotonically as a function until timestep ( = 6, after which the best test loss stagnates. That is, for "shallow" beliefs, Tip general- izes well, while for "deep" beliefs, generalization performance plummets. Given that each timestep is associated with the same number of training samples, is common in stan- dard rollout-based DRL data collection strategies, Figure 3 is a clear illustration of belief explosion in full effeect;
2.
While the aggregate test curve in left plot of Figure 3 shows mild overfitting, decomposing the test loss over the timesteps tells a very different story: namely, that for shallow beliefs, our network hasn't 't yet converged to the best test loss, while for deep beliefs, we experience drastic overfitting.
The key issue that a fully connected network fails to incorporate the structural properties of the BED task, and as result, it struggles to generalize to unseen belief states at test time. This limitation poses challenges even in very small BO instances, such as this task with only 10 discrete bins.
We argue that embedding BED-specific inductive biases into the model architecture is essential for efficient learning. Our key insight is that naive network parameterizations experience belief explosion along the blue curve of Figure 2 while networks that leverage property we call domain permutation equivariance (introduced in the next section) experience significantly tamer growth of belief diversity according to the orange curve.
4 Domainn-ermmuation equivariance in BED and GNNs
In this section, we present some intuition about the structural properties of BED tasks that can be exploited for efficient learning. Figure iillustrates how we might intuitively expect the acquisition values of belief Pt to be equivariant with respect to a simple transformation formalize this intuition in our main theoretical result below about the Nomain-permutation equiv- ariance of optimal policies in BED tasks and use it to motivate the use of GNNs as an appropriate inductive bias.
the domain X.
Specifically, for case of exposition, we focus our attention on BED tasks in which the posterior belief over the unknown function is fully characterized by the first two moments, as is the case in Bayesian Optimization and Active Search.
Let f Rd + be the unknown function which the active sensing agent is querying. We consider BED tasks over a discretised space Xaisc {x 22 ,.,", The posterior distribution over (belief-statee) given the information set I, is therefore a tuple of finite dimensional moments s, (n,, Σx) with:
2...,,,...
#4 (Eif(>')
RM and E, (E((f(=*) 14) - -()
M
SM
Consider set of permutations 4, where each permutation e E + permutes the belief state s,, More specifically, each permutation is a bijection Φb 2 M) - {1,2...., M) that reorders the indices of the M input points in Xise. Thus, (s(:) (Ju,, Ex,e. where:
(z)
d(1)
Ps.1
i=1 1=1
5

---

and
- (B((f(=0()))
as x (f(x - -u ( [zol
34 i,j=1 =(2 (211)))000)
E5 =
fi=1"
We overload the notation @ to also map actions to permuted actions. In our active sensing setup, each action a senses some subset of & i.e., N' S Kdise= and we define Φ(a) {2((0) x' E Xa)-
Theorem 1 (Domain Permutation Equivariance). Let D be set of permutations. Suppose a Bayesian Experiment Design (BED) task satisfies, for all Φ &b and all state-action-eex--ttai transitions 8,8' E, E
(i) Reward invariance:
r(s,a) ((o(s),
(i) Transition invariance:
p(s' s, s, p(o(s') @(s), 6(a))
Then, for all & E Φ all a S, a, A, and 7 E [0, 1]:
(o)( (s))
(optimal policy equivariance)
(((0(s))(aa) Q (s. a) (optimal Q-value invariance).
This reward and transition equivariance property holds for wide range of BED tasks. For example, with the BO task, a practitioner might choose the reward as the negative posterior argmax entropy. Since the transition is the GP inference, both reward and transition invariance hold, since neither depend explicitly on the ordering of states in the domain. These conditions also hold for the AS task with more constrained set of permutations, which is described in detail in the appendix.
One way to interpret this result is that using mean vectors and covariance matrices as belief repre- sentations during learning implicitly imposes fixed ordering on the underlying random variables. This is problematic, as the BED task ultimately requires reasoning about sets of jointly distributed random variables, which are inherently unordered. By designing policies that are equivariant to per- mutations of these variables, we enable agents to generalize optimal behaviors across a broader class of decision-theone theoretically equivalent belief states, thereby substantially improving sample efficiency.
We emphasize that this theorem complementary to Theorem by Foster et al. in B,, which concerns observation-sel permutations. Whereas observation-set permutation invariance focuses on temporal structure, domaiin-permutation equivariance is concerned with the geometric structure of the belief state itself.
The significance of this result is that we can easily generalize between permuted belief states during policy optimization. We use a straightforward way of doing so: using graph representations of the posterior belief as our state, with Graph Neural Networks as our policy class, which are permutation equivariant by design.
Input to MLP
Imput to n
GNN
Undenil
oflhy Ro)
<260
oft
<20r.0
of
Figure 4: Dense Network & GNN Comparison: Input Representations Note that although the graph input ONN; is complete graph, we omit all edges from the figure for ease of illustration.
to ni
We use complete undirected graph to represent our belief, with M nodes representing the un- known function value at points in the discretized domain Maice and M(M - 1)/2 edges rep- resenting our joint uncertainty between points in Xdise- For example, in the Bayesian Optimi- sation setting, each node in the graph is labeled with a 2d vector containing posterior statistics (f(())I,,) E IR? Each edge of the graph is labeled with the posterior covariance Cov(J(x). [(r') E R+. We then use a TransformerConv architecture 24 for the GNN which transforms labels on the nodes by combining information from the labels neighboring nodes and edges using learned soft attention. Figure illustrates the input represen-
(pe(z) [Z)()) o2(x)
tation to the GNN and the MLP.

---

Experiments
Our experiments aim to answer the following questions:
1. Do GNNs offer an effective inductive bias for efficiently learning BED policies using data from an expensive oracle?
2. Can GNNs be used to learn non-myopic BED policies using reinforcement learning?
3. How well do GNN-based policies transfer across BED problems, and do they scale to larger search problems?
Note that due to space constraints, we leave our reinforcement learning, generalization, and contin- uous BO experiments Appendix E.2
5.1
Setup
Tasks: We use two BED tasks: Bayesian Optimization (BO) and Active Search (AS).
Bayesian Optimization models the problem of efficiently locating the argmax of an unknown func- tion using black-box queries. Our latent variable of interest 0 is an unknown function f X - IR defined the domain X' C Rd. The action space A consists of point locations in the domain X, and the observation model p(y)a jr, 9) is black-box evaluation at the detection location with Gaussian noise Paoise (0, of n. The prior p(0) defined using a zero-mean Gaussian prior with a chosen kernel.
For our experiments, we use a
kernel k(x,, x') = exp
cxcxp (- I0- 22
We
also discretize the domain X' into finite set of points {x Xaisc C X' allowing for exact closed-form inference over the set of random variables {f(**) using standard Gaussian infer- equations [25] We also restrict the agent's actions to Xasse-
There are several candidate choices for the MDP reward function, but a natural choice 26 27 for BO tasks is to define r; ). Here we define P; as the posterior over the arg max of and we use H. to denote the Shannon entropy. With this reward definition, we incentivize designs that improve the agent's understanding of the location of the maximum of f. equivalently, designs that reduce the agent's 's posterior uncertainty about the location of an optimal function value. We note that, unlike posterior inference, there therefore use Monte Carlo sampling to approximate (x(")) with the following estimator:
closed-form expressions available for H(p)). We
1 N
pc(x*)
II arg max f((x) = =x"]
where Si ~ pe(f) and then use r. = =-((((<=")).
We use ID domain with M 32 grid points in region length 8 with é = 1, ouais 0.1 and o² =1 1. These parameters were chosen such that draws ~ Pu typically have 3-4 local optima, and the smoothness is qualitatively discernible. We set the episode length to T = 8. As we shall see, although this toy problem instance may seem trivial, it poses significant challenges when training with standard dense neural networks. In the appendix we explore fully continuous non-discretized BO tasks.
Bayesian Optimisation
Active Search
M(N((r") =
Jujas call
after robeeretiies
Postorke datchuation offer
X-
P.F.A. X
Figure 5: Overview of the Bayesian Optimization and Active Search Tasks
7

---

Active Search models the problem of efficiently locating sparse objects in an unknown environment by actively taking sensing actions given all observations thus far 29 illustrates an Active Search problem using aerial remote sensing agent. We adopt the environment model from 10]
30
31
32
Figure
The Active Search problem can be thought of as an instance of the standard Bayesian Optimization problem with some modifications, the key differences being:
1. BO: f smooth; AS: f sparse
2. BO: pointwise sensing of f: AS: contiguous interval sensing of
3. BO: homoskedastic noise; AS: heteroskedastic noise based on sensed interval Kz
Here, f: X {0,1] 1] denotes object presence, with X' representing a ID, 2D, 3D search space. The prior p(0) assumes sparsity. Instead of selecting single points, the agent senses contiguous intervals Ng S X, receiving noisy observation f(G).. Wider intervals yield more noise, reflecting the trade-off between resolution and coverage (e.g., higher altitudes capture more area but with less precision).
We adopt the recovery reward, which selects actions to maximize expected detections under the posterior 33 10 34).
Experiments use a ID grid with M 32 and episode length T 32. Actions span up to I 4 points, with noise increasing with interval size-inducing a trade-off between exploration and precision.
Learning-based baselines: We compare GNNs against three standard parameterizations the pol- icy and critic networks used in previous work: FCNs, CNNs, and Transformers on the information set.
FCN or Dense networks are a naive, but common choice in the DRL. literature 35 36 Multiple BED DRL papers have used similar architectures while focusing primarily on other dimensions of the BED problem 7 [0]
CNNs, on the other hand, provide stronger inductive bias for BED tasks They are set up to predict the Q value or the action probability for each point in the domain using marginal statistics that point. Although this architecture can be efficient learning certain functions, it completely ignores the correlations between different points in the domain and, therefore, is insufficient for representing the multistep value function or optimal policy for BED tasks.
Transformers are another promising alternative 37., capable of processing entire information sets directly. Though there is similarity in transformers and GNNs, they operate different state representation and offer distinct invariances: transformers directly operate on the information set and provide information permutation invariance, while the GNNs operate on the belief state and additionally provide domain permutation-set equivariance. Note that in the appendix we include more experiments comparing Transformers and information-set-bue GNNs without requiring any discretization strategies.
Non-amortized baselines: For Bayesian Optimization, compare our learned policies with the greedy policies w.r.t. the Expected Improvement (38) and UCB (39] acquisition functions, as well Thompson Sampling [4). For Active Search, compare with random waypoint selection and a linear sweeping of the search domain 288
5.2 Behavior Cloning Experiments
We begin our experiments by comparing different policy architectures for behavior cloning an ex- pensive oracle and validating that GNNs offer a helpful inductive bias.
Our goal to train policies clone the 1-step greedy expert т Ti-step = arg max, E(r++1 Pe, a = a). with negative posterior argmax entropy reward function. We use Monte Carlo rollouts to generate an offline dataset of beliefs encountered by the 1-step greedy expert, along with the expert actions: D = {(rt (t)) where p denotes the posterior belief state at timestep / in trial i when following policy Ti.step- other words, is dataset of (state, - l-step-expert) collected from the 1-step lookahead policy. The Active Search task uses 32 timesteps per rollout, and the Bayesian

---

Active Search
Bayolia Optiniiationn
Duttr (bat uacaC
-2:00
2/0
Epods
Kpecks
Lipoda
1godta
CSS: Indn
CNC Imis CNC ml
Cnunfattes: Iniui
PUN vol
GAN: vol
Figure 6: Comparison of Inductive Biases for Behavior Cloning 5,000 samples: Training (solid) and test (dotted) bchavior-cloning loss, and expected return for various networks on the Active Search (left) and Bayesian Optimization (right) tasks. The dotted black line denotes the oracle's expected return. During training, we keep track of the best-performing policy the validation set using the behavior cloning loss (i.e., the rolling best policy). The expected return is then computed evaluating this best-val policy. Note that each averaged over 5 seeds. The shaded region shows #2 standard errors.
Optimisation task higher dimensional continuous BO tasks using 32 timesteps in the appendix.
We note that include behavior cloning results on BO instances for
Figure 6 shows the standard train-test loss curves and online performance for different pol- icy networks trained using behavior cloning on AS and BO respectively. We train these net- works 80:20 train:test split of data from D using Adam, with varying dataset size |D] (50, 500, 5000, 50000} (full results presented in the appendix). The dotted black lines on the right subfigures show the online performance of Tsssee
Across both tasks and all dataset sizes, we observe that the GNN shows excellent generalization on the test set, while the other networks overfit and perform poorly on the test set. Moreover, the GNN able to match the expert's rollout performance using 5,000 samples on both tasks.
Looking at the baselines, we see that the FCN and Transformer are very data-hungry and require prohibitive number of samples even for small tasks. CNNs improve upon FCNs and Transformers and can learn much more efficiently. However, they are limited by their inability to use cross- correlation terms the belief state, which essential for doing well on many BED tasks.
Terestal Satrgg
17.5
12.5-
7.5-
Bonku
CS0 tueado
Figure 7: Comparison with Non-Amortized baselines BO
Figure compares the trained GNN with non-learning based baselines: Upper Confidence Bound (UCB), Thompsons Sampling (TS), Expected Improvement (EI) and Expected Information Gain (EIG). We report inference time and performance, measured by the terminal posterior argmax entropy, that is, the entropy of the posterior argmax implied by the belief state at the final timestep of the rollout = T' = 8). We that the GNN can match the oracle's performance while being significantly faster. Acquisition functions such as and UCB are faster to compute, but don t perform well.
Our experiments in this section validate our claim that the choice of inductive bias a bottleneck for efficiently learning BED policies. We show that GNNs are more sample efficient than other networks by at least an order of magnitude on all tasks, confirming the benefit of exploiting the Nomain--ermutain equivariance structure for BED tasks.

---

5.3 Reinforcement Learning Experiments
In this section we provide results on learning non-myopic policies using reinforcement learning. We use DDQN with nonstationary Q-values. integrated Bellman targets and various critic architectures learn policies described in Section As with the behavior cloning experiments, we compare the performance of the Graph Neural Networks with the Dense Networks, CNNs, and Transformers.
For the BO enviornment, the AS environment, we use the full recovery rate reward [10].
the negative posterior argmax entropy the reward function. For
Figure g shows the training curves for our reinforcement learning experiments. The black dotted line shows the performance of the greedy policy, As in the offline training setting, we see significant improvements in training efficiency when using GNNs. In fact, the GNNs are the only policy that can learn non-myopic policies in either task. All other networks train very slowly and we expect them to require at least order of magnitude more training steps even reach the myopic policy's performance.
Active Search
Bayesian Optimization
-14
Mean Return
2.5
Moan Return
-16
-18
Greedy Policy
1.5
FCN
GNN
Transformer
22
0K
150K
300K
450K
600K
750K
Env Steps
Env Steps
Figure 8: Comparison of Inductive Biases for Reinforcement Learning: RL training curves for various architecture the AS (left) and BO (right) tasks. The dotted black line denotes the oracle's expected return. Note that each curve is averaged over seeds. The shaded region shows standard errors.
Conclusion
In this work, we identified belief explosion as bottleneck policy optimization for BED tasks, and demonstrated that designing networks with the right equivariance structure is an effective strategy taming this issue. We empirically show that our method can amortize expensive oracles, learn non-myopic policies, and generalize to larger scale BED tasks at test time. We also showed how GNNs can be leveraged both discrete and continuous contexts, with significantly increased sample efficiency gains in higher-dimensionas BO tasks. An exciting future direction involves discovering for a broader range of BED tasks such as SIR parameter identification, as well as expanding our approach to misspecified BED tasks.
Acknowledgments and Disclosure of Funding
This work was supported by the U.S. Army Research Office and the U.S. Army Futures Command under Contract No. W911NF-20-D-00
We would also like to thank Elissa Wu for her continuous support and encouragement during this project.
10

---

References
[1] Dennis V Lindley. On a measure of the information provided by an experiment. The Annals of Mathematical Statistics, 27(4))86--105, 1956.
[2] Kathryn Chaloner and Isabella Verdinelli. Bayesian experimental design: A review. Statistical Science, pages 273-304, 1995.
[3] Adam Foster, Desi R Ivanova, Ilyas Malik, and Tom Rainforth. Deep adaptive design: Amor- tizing sequential bayesian experimental design. In International Conference on Machine Learning, pages 3384-3395, PMLR, 2021.
[4] Adam Foster, Martin Jankowiak, Elias Bingham, Paul Horsfall, Yee Whye Teh, Thomas Rain- forth, and Noah Goodman. Variational bayesian optimal experimental design. Advances in Neural Information Processing Systems, 32, 2019.
[5] Adam Foster, Martin Jankowiak, Matthew O'Meara, Yee Whye Teh, and Tom Rainforth. A unified stochastic gradient approach to designing bayesian-optimal experiments. In Interna- lional Conference on Artificial Intelligence and Statistics, pages 2959-2969. PMLR, 2020.
[6] Takashi Goda, Tomohiko Hironaka, Wataru Kitade, and Adam Foster. Unbiased mime stochas- tic gradient-basee optimization of bayesian experimental designs. SIAM Journal on Scientific Computing, 14(1):A286-A33) 2022.
[7] Michael Volpp, Lukas P Fröhlich, Kirsten Fischer, Andreas Doerr, Stefan Falkner, Frank Hut- ter, and Christian Daniel. Meta-learning acquisition functions for transfer learning in bayesian optimization. arXiv preprint
2019.
[8] Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian- inference by meta-learning on prior-data. Fifth Workshop on Meta-Learning at the Conference on Neural Information Processing Systems, 2021.
[9] Vincent Lim, Ellen Novoseller, Jeffrey Ichnowski, Huang Huang, and Ken Goldberg. Policy- based bayesian experimental design for non-differentiable implicit models. arXiv preprint arXiv:2203.04272. 2022.
[10] Conor Igoe, Ramina Ghods, and Jeff Schneider. Multi-agent active search: A reinforcement learning approach. IEEE Robotics and Automation Letters, 7(2):754-761, 2022.
[11] Tom Rainforth, Adam Foster, Desi R Ivanova, and Freddie Bickford Smith. Modern bayesian experimental design. Statistical Science, 39(1):100-114, 2024.
[12] Jian Wu and Peter Frazier. Practical two-step lookahead bayesian optimization. Advances in neural information processing systems, 32, 2019.
[13] Mahadev Satyanarayanan. The emergence of edge computing. Computer, 50(1):30-39, 2017.
[14] Tom Blau, Edwin V Bonilla, Iadine Chades, and Amir Dezfouli. Optimizing sequential ex- perimental design with deep reinforcement learning. In International conference on machine learning, pages 2107-2128 PMLR, 2022.
[15] Samuel Müller, Matthias Feurer, Noah Hollmann, and Frank Hutter. Pfns4bo: In-context learning for bayesian optimization. In International Conference on Machine Learning, pages 25444-25470. PMLR, 2023.
[16] Anindya Sarkar, Nathan Jacobs, and Yevgeniy Vorobeychik. partially-supervssed reinforce- ment learning framework for visual active search. Advances in Neural Information Processing Systems, 36:
2023.
[17] Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neu- ral message passing for quantum chemistry. In International conference on machine learning. pages 1263-1272. Pmlr, 2017.

---

[18] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalee Vinicius Zam- baldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprini arXiv: 18.6.180.6.666 2018.
[19] Philip Kwaku Adjei, Qin Zhiguang, Isaac Amankona Obiri, Ansu Badjie, Christian Nii Aflah Cobblah, Ali Alqahtani, Yeong Hyeon Gu, and Mugahed A Al-Antari. A graph attention network-based multi-agent reinforcement learning framework for robust detection of smart contract vulnerabilities. Scientific Reports, 15(1):29810, 2025.
[20] Alvaro
Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Ried- miller, Raia Hadsell, and Peter Battaglia. Graph networks learnable physies engines for inference and control. In International conference on machine learning, pages 4470-4479.
PMLR, 2018.
[21] Yixin Lin, Austin S Wang, Eric Undersander, and Akshara Rai. Efficient and interpretable robot manipulation with graph neural networks. IEEE Robotics and Automation Letters, 7(2):2740-2744) 2022.
[22] Adhiraj Somani, Nan Ye, David Hsu, and Wee Sun Lee. Despot: regularization. Advances neural information processing systems, 26, 2013.
Online pomdp planning with
[23] Yuanfu Luo, Haoyu Bai, David Hsu, and Wee Sun Lee. Importance sampling for online plan- ning under uncertainty. The International Journal of Robotics Research, 38(2-3):162-181) 2019.
[24] Yunsheng Shi, Zhengjie Huang, Shikun Feng. Hui Zhong, Wenjin Wang, and Yu Sun. Masked label prediction: Unified message passing model for semi-supervised classification. arXiv preprint arXiv:2009.00 03509, 2020.
[25] Carl Edward Rasmussen. Gaussian processes in machine learning. Springer. 2004.
[26] Daniel Hermández-Lobato, Jose Hernandez-Lobato, Amar Shah, and Ryan Adams. Predic- tive entropy search for multi-objective bayesian optimization. In International conference on machine learning, pages 1492-1501. PMLR, 2016.
[27] Zi Wang and Stefanic Jegelka. Max-value entropy search for efficient bayesian optimization. In International Conference on Machine Learning, pages 3627-3635. PMLR, 2017.
[28] Nikhil Angad Bakshi, Tejus Gupta, Ramina Ghods, and Jeff Schneider. Guts: Generalized uncertainty-aware thompson sampling for multi-agent active search. In 2023 IEEE Interna- lional Conference on Robotics and Automation (ICRA), pages 7735-7741, IEEE, 2023.
[29] Arundhati Banerjee, Ramina Ghods, and Jeff Schneider. Multi-agent active search using de- tection and location uncertainty. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 7720-7727. IEEE, 2023.
[30] Roman Garnett, Yamuna Krishnamurthy, Donghan Wang, Jeff Schneider, and Richard Mann. Bayesian optimal active search on graphs. Ninth Workshop on Mining and Learning with Graphs. Citeseer, 2011.
[31] Roman Garnett, Yamuna Krishnamurthy, Xuchan Xiong, Jeff Schneider, and Richard Mann. Bayesian optimal active search and surveying. arXiv preprint arXiv:1206.64400 2012.
[32] Robin R Murphy. Human-robot interaction in reseue robotics. IEEE Transactions on Systems, Man, and Cybernetics, Part C (Applications and Reviews), 34(2):138-153 2004.
[33] Ramina Ghods, Arundhati Banerjee, and Jeff Schneider. Decentralized multi-agent active search for sparse signals. In Cassio de Campos and Marloes H. Maathuis, editors, Proceed- ings of the -eirty-Seventh Conference on Uncertainty Artificial Intelligence, volume 161 of Proceedings of Machine Learning Research, pages 696-706. PMLR, 27-30 Jul 2021.
[34] Nikhil Angad Bakshi, Tejus Gupta, Ramina Ghods, and Jeff Schneider. Guts: Generalized uncertainty- thompson sampling for multi-agent active search. In 2023 IEEE Interna- tional Conference on Robotics and Automation (ICRA), pages 7735-7741, 2023.
12

---

[35] Marcin Andrychowicz, Anton Raichuk, Piotr Starczyyk, Manu Orsini, Sertan Girgin, Raphael Marinier, Léonard Hussenot, Matthieu Geist, Olivier Pietquin, Marein Michalski, et al. What matters in on-policy reinforcement learning? large-scale empirical study. arXiv preprint wwXiv:2006.059900 2020.
[36] Conor Igoe, Swapnil Pande, Siddarth Venkatraman, and Jeff Schneider, Multi- soft actor- critic: Overcoming stochastic biases in maximum entropy reinforcement learning. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 7162-7168. IEEE, 2023.
[37] Lei Song, Chenxiao Gao, Ke Xue, Chenyang Wu, Dong Li, Jianye Hao, Zongzhang Zhang, and Chao Qian. Reinforced in-context black-box optimization. arXiv preprini arXiv:2402. 2024.
[38] Donald R Jones, Matthias Schonlau, and William Welch. Efficient global optimization of expensive black-box functions. Journal of Global optimization, 13:455 492,, 1998.
[39] Niranjan Srinivas, Andreas Krause, Sham M Kakade, and Matthias Seeger. Gaussian pro- cess optimization the bandit setting: No regret and experimental design. arXiv preprint arXiv:0912.3995, 2009.
[40] Daniel Russo and Benjamin Van Roy. Learning optimize via posterior sampling. Mathe- matics of Operations Research, 39(4): 2014.
[41] Maximilian Balandat, Brian Karrer, Daniel R. Jiang, Samuel Daulton, Benjamin Letham, An- drew Gordon Wilson, and Eytan Bakshy. BoToreh: A Framework for Efficient Monte-Carlo Bayesian Optimization. In Advances in Neural Information Processing Systems 33, 2020.
[42] Sebastian Ament, Samuel Daulton, David Eriksson, Maximilian Balandat, and Eytan Bakshy. Unexpected improvements to expected improvement for bayesian optimization. Advances in Neural Information Processing Systems, 36:20577-20612, 2023.
[43] José Miguel Hernández-Lobite Matthew W Hoffman, and Zoubin Ghahramani. Predictive entropy search for efficient global optimization of black-box functions. Advances in neural information processing systems, 27, 2014.
13

---

A Limitations
This work has several limitations. First, all experiments are conducted on synthetic Bayesian ex- periment design task instances. While these are widely used benchmarks, they do not capture the full complexity of real-world applications. Second, we assume well-specified models throughout, which may not hold in practice. Finally, our framework presumes access to exact posterior moments (e.g., means and covariances), which appropriate under fixed priors. However, many practical BED settings involve hyperpriors (such as distributions over kernel lengthscales) that induce pos- terior distributions over parameters for which first- and second-order moments may no longer be sufficient. This complicates the construction of GNN input representations and poses challenge for generalizing our approach.
B Compute Resources
All experiments in this paper were conducted on cluster of 8 NVIDIA 2080 Ti GPUs. The longest- running experiments involved training the 10 Transformer model seeds on the 8D Bayesian Opti- mization continuous behavior cloning task, which took approximately one week.
C Proof of Theorem 1
Theorem (Domain Permutation Equivariance). If our BED task has the following two proper- fies for all permutations Φ E 4:
1. r(s, r((()),(
(reward invariance)
2. a=) o(a))
(transition invariance)
then, for all permutations & E 1 for all beliefs s for all discounts 7,, =;(ő(s)) = ()(T (optimal policy equivariance)
Proof: We'd like to prove that (s) = o ( Vs. This equivalent showing that
arg max q' (8, -
(arg max (o(+), Vs
We will also show that
9- (8,a)= (d)()) ((a))) Vs, Va
(2)
We will prove both Eq and Eq using induction.
Let's start with the last timestep T. Eq Eqg directly holds from reward equivariance.
Proof for Eq at the last timestep:
o- (arg max, (o(s), a)) = o- o max, r(o(s), a)) (the action-value function timestep T is equal to the reward)
l (arg max, 47 (p(s), arg max, r (()(s), (ba)) (using arg max, f(o(z)) = (arg max, f(x) o)
l (arg max. 4 (0(), a)) = arg max, r(s. a)
(using reward invariance)
o (arg max, (p(s), a)) arg max, of (s, a).
Now we'd like to show that Bq and Eq holds for timestep /. 1, also true for timestep
14

---

の (arg max (d(s), a)) = x- " max r(o(s), max)m(x)))) a) + tB a(s),a [max l.41 (5', a')l)
(Bellman equation)
- max haan:n/eeeeooooooooo//o//////--- (ma) mrgmax(()(s) (o()) |max g:t1 (s', a')]
(using arg max [(((x)) o "arr max (x)))
(arg max (ox)) arg maxr(s.) , (((o)) [max z (s' a')l
(reward invariance)
x (arg max (d(s), a)) = arg max r(),a) + Bl [max ze (o(s')
of
(transition invariance)
o (arg max ( arg magrax((s. + o mmax (s))
(change of variables: -
の- (arg max (d(s), a)) arg max r(s, a) + > l [ (4 (d))
(transition invariance)
o (arg max (o((s)), arg maxr a) + I a.a [max
a' (0 (
(invariance of max under Φ)
o (arg max m (o(s), arg maggaasssso a) + / [max 1 as')
(inductive assumption: Eq. 2 holds for + 1)
ক (arg max of (o(+), a)) arg argaaggann max q" (=(saa)
This completes our proof by induction.
15

---

D Discrete Belief Representations and Policy Inputs
E Further Experiments
E.1 Full Behavior Cloning Experements
In this section, we present the full results from our behavior cloning experiments described in section
5.2
semples
samples
11.00
Stme Seire 1am (awed))
5
2/6
1.75
+0008
100
5000 samples
500 complex
4200
MoetBree (hat medeel
+20000
925
2.00
0.55
2.30
Epech
1p-db
1pm-30
00000 auupher
0000 sumple
2018
110-
11:03
2.00
2.25
Fpedk
Tpodde
Rpedk
(a) Active Search
(b) Bayesian Optimization
BON Inde
DNX Bids
ONN Cods
Tixe Shunc bonk
Figure 9: Comparison of Inductive Biases for Behavior Cloning: We plot the (i) train (solid lines) and test (dotted lines) behavior-cloning loss and (ii) expected return for various networks the Active Search (left) Bayesian Optimization (right) tasks using increasing dataset size (top bottom). dotted black line denotes the oracle's s expected return. Note that each curve is averaged over seeds. The shaded region shows the 95% confidence interval for the mean.
16

---

E.2 Generalization and Scaling Experiments
In this section, we evaluate whether GNN- policies can scale to large search domains where established acquisition functions are prohibitively expensive to compute. We show results on three transfer setups in Table 2-4, training the GNN policy on a smaller problem and evaluating transfer performance on larger versions of the problem in each setup.
show the performance of GNN trained using hehavior cloning a greedy expert on a 8 x 8 active search environment, and evaluated on a 32 x 32 environment. The GNN almost matches the performance of the expert on the 8 grid, and allows us to scale 32 x 32 grid. In contrast, it's computationally intractable to run the greedy search on the larger grid because requires us to estimate the expected information gain for each candidate action, which in turn involves nested integrals. Table 4 and Similalryy similarly show that GNN- based policies can scale to higher- dimensional tasks and larger grids, and match or outperform expert performance.
For example, Table
Table 1: Active Search Scaling
Table 2: Active Search Scaling
(Train: 8 x 8 grid)
(Test: 32 x 32 grid)
Method
Performance
Time (ms)
Method
Performance
Time (ms)
Random
15.28 + 0.32
< 0.01
Random
3.72 + 0.08 3.16 + 0.07 N.A. 7.15 + 0.12
< 0.01
Simple Search
< 0.01
Simple Search
< 0.01
Greedy Search
35.72 + 0.40
265.4
Greedy Search
> le4
GNN
34.13 + 0.49
3.34
GNN
67.50
F Continuous BO Experiments
Here we present results of experiments comparing the simple regret performance of Transformer and GNN models trained to clone the behavior of the LogEl policy acting on various continuous BO tasks (2D, 4D, 8D), each rollout using T 32 timesteps. In these experiments, no discretization is used: both the Transformer and GNN take the information set I, as input and output a point directly in continuous design space, a, E X' S Rd. This is desirable for contexts where discretization is inappropriate for the level of precision demanded by the application or where it is necessary to additionally amortize inference as well acquisition estimation and optimization.
We use the BoTorrh [1 implementation of LogEl as our expert oracle with default hyperparameters recommended in their documentation + This oracle is numerically robust implementation of the canonical Expected Improvement acquisition function, as described in We use a zero mean Gaussian noise assumption with noise variance of 10-6
We train both the Transformer and the GNN models on expert trajectories using standard behaviour cloning losses. Table shows the key architectural details of both models.
Table 3: Transformer and GNN hyperparameters for continuous BO experiments
Model Trainable Params
Hidden Dim
# Heads
# Layers
Transformer GNN
10?
40
8 10
1 x 10?
64
8
10
As in the discrete case, the GNN input a graph constructed from the information set Z,. Specifi- cally, we use the same basic architecture as in the discrete experiments, but here our graph contains I nodes at timestep instead of fixed number of nodes for all timesteps to match a desired level of discretization. The nodal features are the observed values {4r )-=1 and the edge feature for edge i is the kernel evaluation k(z,, I.).. In the discrete experiments, the GNN output logits for the discrete set of actions. This allowed the network to naturally take advantage of domain transfor- mation equivariance. In these continuous experiments, we transform the nodal features to a final set
hatps://botorch.org/
17

---

of linear weights {w-) which are used to combine the points in the information set to define an action distribution for the next timestep:
V E= W+IT: o²) where {w+) = 16
PUNNN (Ix).
This parameterization is equivariant to the same of transformations & that the kernel A: invariant to. Note that in the discrete context, Φ 3 Φ [N] + [N] is a permutation of the discretized domain, while in this continuous setting we abuse notation and use Φ D Φ Rd + Rd to denote pointwise invertible transformations of the domain. In our experimental results below, again kernel with lengthscale of 0.2, 0.3, 0.4, for d = 2, 4, 8, over an optimization domain refined to the disk X {ar E Rd ||>||| - 0.5}. This results in 4 O(d), where O(d) denotes the orthogonal group, the group of distance-preserving transformations of d-dimensional Euclidean space. We use fixed variance actor while training both the GNN and Transformer, and during simple-regret rollouts we act deterministically by following the mean policy a, Σi= WyFC7. Finally, note that we warm-start the information sets in every rollout with d + 1 uniformly at random chosen points in X' for all methods. Timestep 0 in the figure below refers to the first meaningful decision after the warm-start points have already been collected.
RBF
The figures below show that GNNs significantly outperform Transformers when trained with the same dataset, with increased delta as BO task dimension increases. The GNN can come close to the simple regret performance of the expert in each task with at least an order of magnitude less data than the Transformer. Moreover, both GNN and Transformer save roughly 1 order of magnitude test-time compute by amortizing the inference, acquisition evaluation and optimization necessary for the BoTorch LogEl expert.
These results demonstrate that primary equivariance insight is not just restricted to discrete domains. This is important, since real-world BED tasks often involve continuous modeling choices enable precise experiment designs.
18

---

Medban Menital tae Obverval Ware with 100000
2.00
1.50-
1.25
1.00-
877-
0.0-
0.2-
ONN Halp
Tram Eiep
GNn Asoa
Thme Estaus
OWN Helbon
Trans
Terminal Best Observed Value 4D
Stettnn sertited Seet Obverved Waan with 000
25
1.5
Thost 10op
OnN 1509
GND 1i000
Trmme Hopo00
ONN Isondt
Omede
Tram
Best
Mettan temmuaa teet Cheervve Satue vod 0075
3.0-
Treas Usep
oNN 1860
Triss Loeno
ONW Issl
Hoose
GyN toonn
Oreale
Figure 10: Continuous BO experimental results. We 1,000 rollouts to evaluate the terminal best observed value. Note that optimizing this quantity is equivalent to optimizing simple regret. Here median and 1,000-seed bootstrap confidence intervals is standard in continuous BO literature due to the heavy-tailed nature of simple regret curves in continuous BO, such as in [43. Both the GNN and Transformer models were trained until overfitting the indicated number of expert rollout trajectories, with 200 trajectories held out for validation. The simple performance is then computed using the models that achieved the lowest behavior cloning validation loss.
G Additional Generalisation Results
In this section, we evaluate the generalization capabilities of our amortized BO policies across search domain dimensions and sizes. Note that these results use discrete belief representations in the main paper.
Table reports results for policy trained on a 2D BO task and then evaluated on set of higher- dimensional tasks. We use an adaptively discretized domain inspired by Volpp al. 7 (ex- plained below) to keep the size of the belief state tractable. The GNN was trained by behavior cloning the Expected Improvement (EI) policy using 5,000 samples, and was evaluated zero-shot on the 3D and 5D tasks. We compare its performance with random policy, Thompson Sampling
19

---

1.0
0.8
Random
Trans 1500
GNN 1500
Time (s1
15000
GNN 15000
Trans 150000
0.1
GNN 150000
MetaBO 6k
0.2
0.0
30
80 Step
Figure 11: Continuous BO timing results. We additionally include MetaBO variant of the LogEI expert that restricts the acqusition function optimisation to the adaptive discretisatino grid proposed in MetaBO paper 77.. Note that amortized methods significantly outperform non-amortized methods on wall clock time while maintaining comparable performance in terms of simple regret.
(TS), and the El policy restricted to the discrete domain points. We note that the GNN matches the El policy's performance on the 2D training task and generalizes well to the 3D and 5D domains.
Table presents another generalization experiment: we train the GNN by behavior cloning the Expected Information Gain (EIG) policy on a domain [0, 4] discretized into bins using 5,000 samples, and tested on larger domain of [0. 128] discretized into 1,024 bins. Once again, the GNN shows strong imitation and zero-shot generalization. Note that, due to the computational cost of the EIG policy (exceeding 10 seconds per decision). direct evaluation the large domain is infeasible. This underscores the value of generalizable amortized policies.
Table 4: Scaling Amortized BO Policies with Search Dimension
Train: 2D
Test: 3D
Test: 5D
Method
Performance
Time (ms)
Performance
Time (ms)
Performance
Time (ms)
Random
1.07 = 0.05
0.02
1.13 0.05
0.05
0.02
1.13 +: 0.05
0.02
TS
1.25 = 0.05
1.13
1.29 + 0.05
1.86
1.23 0.05
1.04
EI
1.31 + 0.05
0.29
1.48 + 0.06
0.30
1.48 + 0.05
1.94
GNN
1.28 + 0.05
67.01
1.49 + 0.06
56.32
1.47 + 0.05
57.46
Table 5: Scaling Amortized BO Policies with Grid Size
Train: N
32
Test: N 1024
Method
Performance
Time (ms)
Performance
Time (ms)
Random
-15.43 + 0.17
0.05
-33.79 + 0.27
0.11
TS
11.44 0.20
0.21
-29.21 + 0.38
29.11
UCB
-13.74 + 0.17
0.15
-99.10 + 0.33
0.24
EIG
-11.95 + 0.20
1067.4
N.A.
> 1e4
GNN
-11.60 - 0.20
9.83
277.94 + 0.36
70.63
Sample n points uniformly at random the domain, choose the top & using the El acquisition function, and then sample points each of these k best points. The candidate set consists of n + km points. We use n 100, A: = 5, m 20.
Adaptive discretization method:
Evaluation procedure: For all the experiments in this section, we evaluate the mean return of the policy using a fixed episode length of 8, averaged over 200 trials, using the reward for Table and EIG reward for Table Since the El reward requires previous max-value, instead of assuming
20

---

fixed initial max value, we choose to initialize each trial with single point chosen uniformly at random the domain. This point does not contribute the timestep count in our experiments.
H Generalization through data augmentation
In this experiment, we compare two different strategies for leveraging the domain permutation equiv- ariant structure in Bayesian Optimization policies for efficient Sehavior-cloning (i) using network with suitable inductive bias, i.e., the Graph Neural Network described Section 4, and (ii) apply- ing data augmentation during training. This data augmentation corresponds to the set of permuta- tions the policy is equivariant to.
Figure 12 shows the behavior-cloning training and validation loss alongside the resulting policy performance across several network architectures. Additionally, it shows outcomes from training fully connected network (FCN) using data augmentation. The results indicate that data augmentation significantly improves the generalization capability of the FCN, as evidenced by the near- pear-ctrfect alignment of training and validation loss curves (depicted in purple). However, this improvement comes at a substantial computational cost- trainin times are considerably prolonged. Notably, the GNN achieves expert-level performance within approximately 30 epochs, whereas the FCN augmented with permutation-bant based data augmentation fails to match this performance even after 7, 7,000 epochs. This disparity is expected since the equivariant permutation set is very large, equal to n! for re-bin discrete BO task. This the number of permutation the GNN is naturally equivariant to, but the FCN needs explicit data augmentation for during training. For this particular experiment, number of permutations is n! = 32!
1035.
Behavior Cloning for Bayesian Optimization with 5,000 samples
-12.00
Mean Return (best modal)
-12.50
-13.00
1,0s
-13.50
-14.00
10*
-14.30
-15.00
+15.50
1.0
10t
102
10a
1000
2000
3000
4000
5000
6000
7000
Epochs
Epochs
FCN: train
GNN: train
CNN: train
FCN with data augmentation (permute): train FCN with dala sugmentation (permule]:
FCN:
GNN: val
CNN:
Figure 12: Comparison of Inductive Biases and Data Augmentation for Behavior Cloning: We plot the (i) train (solid lines) and test (dotted lines) hehavior-cloning loss and (ii) expected return for various networks on the Bayesian Optimization using dataset size of 5,000 samples. The oracle's expected return is indicated by the dotted black line. Regular training was terminated at 1,000 epochs after performance had stagnated. Note that each curve averaged over 5 seeds. The shaded region shows the 95% confidence interval for the mean.
I
Generalization and Scalability for Bayesian Optimization
In this section, we discuss the generalizability and scalability of our amortized approach Bayesian Optimization (BO).
I.1 Generalizability
In all but the scaling experiments section our experimental setup assumes well-specified prior: both training and test functions are sampled from a known Gaussian Process (GP) prior.
21

---

While this assumption facilitates controlled evaluation, it does not always hold in practice. Real- world applications of BO often involve misspecified priors, where the true function does not match the assumed GP prior. In these cases, surrogate models are fit by estimating GP hyperparameters from data.
To generalize our amortization approach beyond the well-specified regime, the learned policy should (i train distribution BED tasks induced by our prior distribution over hyperparameters, and (ii) condition on the posterior distribution over the full set of random variables in the BED task, which now includes the hyperparameters.
1.2 Scalability
Our approach can help with scalability by amortizing the computationally expensive BO methods via neural policy that maps belief states actions, thereby reducing online computational burden. However, special care is required in representing the belief state. Our method relies on an explicit representation of the posterior over the search space, which can become prohibitively expensive in high-dimensional settings. In particular, a naive uniform discretization leads to a grid whose size grows exponentially with the number of input dimensions.
To mitigate this, we adopted simple adaptive discretization strategy inspired by MetaBO 7 for our higher dimensional BO experiments in section G By concentrating representational capacity in high-mean or high- high-uncertainty regions, we retain a tractable belief representation. Nonetheless, the choice of discretization remains a critical factor, and more principled exploration of this design choice is left to future work.
Network Hyperparamters for Discrete Experiments
Table 6: Hyperparameters for FCN.
Hidden Layers
[512, 512, 512]
Table 7: Hyperparameters for CNN.
Number of Layers 6
Kernel size 9
Table 8: Hyperparameters GNN.
Hidden Dim
64
Number of Heads
8
Number of Layers 2
Table 9: Hyperparameters for Transformer.
Number of Layers 5
Number of Heads
16
Embedding Size
64
22

---

K
Training Hyperparamters
Table 10: Behavior Cloning Hyperparameters
Learning Rate
3e-4
Batch Size
32
Optimizer
Adam
Table 11: Reinforcement Learning Hyperparameters
Algorithm
DDQN
Exploration
e-greedy
Discount factor
0.95
Learning Rate
3e-4
Batch Size
32
Optimizer
Adam
We also non-stationary Q-functions with DDQN since we model BED tasks finite-horizon MDPs. In addition, we use integrated Bellman targets to reduce noise in training.
L
Visualisation of Evolution of GNN & FCN Critic
--)-0
[9]
-
[98
[2t]
[S]]
|
1=1
1-4
1=5
1=7
1-8
Figure 13: Evolution of BO GNN policy through RL training. Each column illustrates a rollout after given number of collected MDP transitions (indicated in the column title). The left figure in each column depicts the belief state Pt and ground-truth hidden function, and the right figure illustrates the estimated g-values. Note that the GNN quickly learns myopic search behavior (first check both corners of the domain). After gathering more experience, instead learns to first query the internal point the domain to act less greedy while maximizing return. Note also that the GNN from initialisation captures symmetry and smoothness priors, and that after 8, 8,192 transitions has learned to represent both sharp discontinuities and smoothness, which is essential when modeling the drop in value for re-sensing previously queried point, but sharp increase in value for querying neighboring points the observed point is high.
23

---

|-1=0
2
1024
2 2048
3072
[
4096
1=0
1=1
1=2
NWW
1=3
NWW
www
1=4
www
1=5
1=6
1=7
1=8
Figure 14: Evolution of BO FCN policy through RL training. Each column illustrates a rollout after a given number of collected MDP transitions. The left figure in each column depicts the belief state Pt, and the right figure illustrates the estimated vavauess. Note that no qualitative improvement the q- valuue visible after 4,096 MDP transitions, contrast to the GNN which achieves high rewards and clear structure in the g-estimates.
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: The main claims in the abstract and introduction concern the belief explosion problem as a bottleneck for sample efficient BED DL training, the role of domain permu- tation equivariance as an exploitable property of BED tasks, and the role of GNNs as an inductive bias to leverage this.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it 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: [es]]

---

Justification:
We have included a limitations section in the appendix of the paper.
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 encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors 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 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 ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might be used reliably provide closed captions for online lectures because fails to 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 ad- dress 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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
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: We include a full proof in the appendix for our theorem concerning optimal policy equivariance.
Guidelines:
The answer NA means that the paper does not 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 in the statement of any theo- rems.
The proofs can either appear in 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 comple- mented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: algorithmic and BED environment hyperparameter details, as well as model references are provided in the paper and appendix.
25

---

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 a 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 to 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 to replicate the results, access to a hosted model (e.g., in the case of a 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 sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If 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 way produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recogmize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibiliiy. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: 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?
Answer: [No]
Justification: We have not provided full codebase for this paper. However, our contribu- tion is in articulating specific problem independent of specific code implementation, and our results rely on open source widely used GNN, Transformer, CNN and FCN models.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
Polooooooolyoooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not 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:
www.Summiiiiiooiiiiiiiiii onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to 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.
26

---

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 is permitted.
6. Experimental setting/details
Question: 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?
Answer: [Yes]
Justification:
We have provided these details in the appendix.
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, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Ys]]
Justification:
We report standard errors in all figures and describe our approach in detail.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
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 (c.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably 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 of range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the 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 have described the compute resources used in all of our experiments in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
27

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount 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 t make it 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 fttngh/trrrpe./errrro.../opo/ooa
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethics and believe that we conform in every respect.
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.. if there is a special consid- eration due to laws or regulations 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: Our paper studies specific bottleneck in a purely synthetic context that is of broad interest to the BED community but which we believe directly related any clear broader social impact.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML..
11. Safeguards
Question: 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)?
28

---

The answer NA means that the paper poses such risks. Answer: [NA]
Justification: Our paper studies a specific bottleneck in a purely synthetic context that is of broad interest to the BED community but which we believe does not raise any significant safety issues.
Guidelines:
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 re- quiring that users adhere to usage guidelines or restrictions 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 require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses existing assets
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?
Question:
Answer: [Yes]
Justification:
We have credited all models, enviormments and libraries that we used in our experiments.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URI
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.
assets are released, the license, copyright information, and terms of use pack- age should be provided. For popular datasets, ppapesswthcooe com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [NA]
Justification:
We have not provided any new assets in this paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details 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 used.
At submission time, remember to anonymize your assets 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 pa- per include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We have not used any crowdsurcing or human subjects in this research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should paid at least the minimum wage 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 equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We have not used any crowdsurcing or human subjects in this research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to 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 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 methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [No]
Justification:
This research does not concern LLMs.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https://neurips.com hett:p//meerrrpssccc///oeeeerree httprrrennnne//22/2L11/ for what should should not be described.
30


Reviews:
[{"review_id": "mfTNDW3R42", "reviewer": "Reviewer_sQcj", "summary": "This paper studies Bayesian Experiment Design (BED) and proposes a way to reduce the sample requirements when using deep learning for BED. Specifically, it proposes to leverage graph neural networks to exploit problems that exhibit permutation equivariance in the reward and transition. By taking advantage of the equivariance in the belief state, GNNs generalize better over permuted belief states. \n\nThe experiments study the Bayesian Optimization and Active Search tasks, and compare the efficacy of the GNN architecture compared to fully connected networks, convolutional neural networks, and transformers. They also compare to more traditional baselines for BO and AS, including UCB and Thompson Sampling. The results show that the GNNs generalize better than other deep learning architectures, and match the oracle better than the non-learning baselines.", "questions": "* How does this differ from prior works that use GNNs to develop deep reinforcement learning algorithms? Does your implementation make specific assumptions for the BED case? For example:\n    * Almasen et al. Deep Reinforcement Learning meets Graph Neural Networks: exploring a routing optimization use case.\n    * Gammelli et al. Graph Neural Network Reinforcement Learning for Autonomous Mobility-on-Demand Systems.\n* To address the issue described in Lines 154-155, i.e., that the network underfits for shallow beliefs but overfits for deep beliefs, would it make sense to have a different network/policy for each time-step?\n* Could we see what Fig. 3 (right) would look like when using GNNs?\n* In Fig. 6, are the “Mean Return (best val model)” plots showing the best validation return seen so far? How is the validation set defined? I’m also curious about the return achieved by the model at that point in training. Could we see that plot? That would show how much fluctuation in training there is.\n* In Fig. 6, do the vastly different losses just indicate that training parameters need to be tuned and that some basic regularization needs to be applied?\n* In Fig. 7, what is terminal posterior argmax entropy? Is that the same as the one-step reward or is it the reward at the final time-step? What would the return look like?\n* What are the error bars in the figures computed over?", "limitations": "Yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n* The paper focuses on the problem of Bayesian optimal experimental design, and investigates the use of graph neural networks to reduce the sample requirements that traditional methods require.\n* Although the experiments only study simple toy tasks, they seem thorough and well-designed, and do the intended job of showing that the use of GNNs improves the sample efficiency over baselines for different BED tasks.\n\nWeaknesses:\n* The introduction helps provide context into BED and the prior approaches to this problem, covering both deep learning and non-deep learning techniques. However, it doesn’t cover GNNs or their prior applications into RL/other sequential decision-making problems. How do those prior works and the conclusions from them relate to this paper?\n* Nit: I think more effort into putting the major experiments in the main paper would improve the readability of the paper. Currently, only one preliminary experiment (behavior cloning) is included in the main paper while the more central experiments are in the appendices.", "was_revised": true, "final_justification": "I've read the authors' responses and maintain my recommendation to accept.", "rebuttal": "We would like to thank the reviewer for their time and valuable feedback. We agree with the reviewer on both of their main points of feedback:\n\nI. As the reviewer points out, GNNs have been commonly used as policy representation to naturally capture the relational structure among components in a system, such as joints in a robot [1] or agents in a multi-agent environment [2], and the paper would benefit from describing this prior use of GNNs and comparing it to our work. The main difference is that prior work on GNN-based policies are motivated by their adaptivity to the size of the input graph and structure (corresponding to varying robot topology and number of agents in these examples). In contrast, our work uses GNNs to leverage policy equivariance that we show exist in (any fixed-sized) BED problems specifically.\n\n1. Almasan et al., \"Deep reinforcement learning meets graph neural networks: Exploring a routing optimization use case.”\n2. Nurbek et al., \"Exploring Graph Neural Networks in Reinforcement Learning: A Comparative Study on Architectures for Locomotion Tasks.”\n\nII. We also agree that readability would improve by moving some continuous and high-dimensional BO experimental results from the appendix to the main body of the paper. To accommodate this, we will condense the behavior cloning experiments by presenting plots for only one representative dataset size per problem type in the main text, shifting additional dataset size results to the appendix. This will allow the reader a greater opportunity to appreciate the range of experiments included in this work.\n\n---\n\nAnswers to reviewer’s questions:\n\n**Q. How does this differ from prior works that use GNNs to develop deep reinforcement learning algorithms? Does your implementation make specific assumptions for the BED case?**\n\nWe have provided an answer to this in our response above, but to summarize concisely: while prior works often use GNNs for problems with variable graph structures, our key insight is using a GNN to exploit the domain-permutation equivariance inherent in BED problems. This allows our policy to generalize efficiently across a vast space of possible beliefs, which is an assumption specific to the BED case that we leverage.\n\n**Q. To address the issue described in Lines 154-155, i.e., that the network underfits for shallow beliefs but overfits for deep beliefs, would it make sense to have a different network/policy for each time-step?**\n\nThat is an interesting suggestion. However, using a different policy for each timestep would not address the root of the problem. The main bottleneck is that the learning problem becomes fundamentally harder at later timesteps due to the \"belief explosion\"—the much wider diversity of belief-states the policy encounters.\n\nWhile it is not the focus of our work, we also expect a single policy network to have some inter-timestep generalization which we would lose by having a different policy network for each timestep.\n\nA related brute-force approach might be to collect significantly more training data for these later timesteps. In contrast, our paper proposes a more sample-efficient solution.\n\n**Q. Could we see what Fig. 3 (right) would look like when using GNNs?**\n\nWe observed that the GNN clearly generalizes well across all the nine timesteps in that example. We didn’t include in our original submission due to space constraints and are not able to share it in our response due the NeurIPS rebuttal policy, but we will include it in our final manuscript in the appendix.\n\n**Q. In Fig. 6, are the “Mean Return (best val model)” plots showing the best validation return seen so far? How is the validation set defined? I’m also curious about the return achieved by the model at that point in training. Could we see that plot? That would show how much fluctuation in training there is.**\n\nThe validation set consists of a 20% held-out portion of the offline dataset of expert trajectories.\n\nThe plot shows the mean return of the policy checkpoint that has achieved the lowest behavior-cloning validation loss up to that epoch.\n\nRationale: The return of the latest policy has a U-shape, improving initially and then degrading as the policy overfits on the behavior cloning data. We chose this presentation to more clearly illustrate the best achievable online performance for each architecture as a function of training data, factoring out the noise from overfitting. We can add the plots of the return achieved by the model at that point in training to the appendix for completeness.\n\n**Q. In Fig. 6, do the vastly different losses just indicate that training parameters need to be tuned and that some basic regularization needs to be applied?**\n\nOur goal in this experiment was to isolate and compare the effectiveness of the inherent inductive bias of each architecture. Therefore, we intentionally evaluated the models on an equal footing, using a similar number of parameters and without architecture-specific regularization. We chose the simplest parameterization of each network type, choosing the default GNN implementation in the pytorch-geometric library for example, and did not perform any hyperparameter tuning, apart from making sure that all the networks had a similar number of parameters. We will include our motivation for the policy network parameterizations in the appendix.\nFinally, we note that all the policy networks we used were capable of overfitting in these experiments and achieve similar training loss; their validation loss and return vary vastly primarily because of their different generalization ability.\n\n**Q. In Fig. 7, what is terminal posterior argmax entropy? Is that the same as the one-step reward or is it the reward at the final time-step? What would the return look like?**\n\nThe \"terminal posterior argmax entropy\" is the entropy of the belief over the function's maximum location at the final timestep of the episode. In our BO experiments, the per-step reward is the expected reduction in this entropy, also known as Expected Information Gain. This is the same objective that popular entropy-search methods such as PES [4] optimizes.\n\nMaximizing the cumulative return is equivalent to minimizing this terminal entropy, i.e., over the entire trial, not just the next timestep. The bar plot in Figure 7 shows this terminal entropy.\n\n[3] José Miguel Hernández-Lobato et. al. Predictive Entropy Search for Efficient Global Optimization of Black-box Functions\n\n**Q. What are the error bars in the figures computed over?**\n\nThe error bars in all training curve plots represent ±2 standard errors of the mean, computed over 5-10 independent runs with different random seeds. We used 5 seeds for the discrete BO and AS experiments, and used 10 seeds for the continuous BO experiments. We can include links to the complete comet training logs for all our experiments with the final manuscript.\n\nFor policy evaluation in the continuous BO task, they represent 95% confidence intervals estimated with a hierarchical bootstrap, a standard practice for that domain. Specifically, there are two sources of randomness for the continuous BO Behavior Cloning task: the 10 seeds of model training, and the 1,000 seeds of model evaluation. We use a hierarchical bootstrap estimate of the 95% CI for the statistics we report in the figures in the appendix."}, {"review_id": "e4sG4wXHc9", "reviewer": "Reviewer_Dbo7", "summary": "This paper tackles Bayesian Experiment Design (BED) with deep learning by mitigating the belief explosion issue (similar to low sample efficiency issue in reinforcement learning). The authors propose to leverage an appropriate inductive bias for the deep neural networks as a solution. Specifically, they adopt graph neural network due to their domain permutation equivariance properties. Their experiments show that both this design outperform other alternatives (CNN, Transformer, etc.) on two instances of BED tasks, namely Bayesian Optimization and Active Search.", "questions": "See \"weaknesses\"", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\n+ The paper is well-motivated and explain the belief explosion issue relatively well, which seems to be overlooked in the literature (as what the authors suggest).\n+ The experiments show adequate evidence as it draws comparison with several learning-based baselines as well as non-learning-based ones (e.g., Thompson Sampling).\n\nWeaknesses:\n+ While GNN is reasonable here, there might be finer-grained network design choices to be studied. I wonder within a generalized class of GNN (or say permutation-equivariant functions), is there specific designs that are better than others. For instance, network structures in other relevant fields has been studied before such as PointNet, Deep Sets, and Set Transformer.", "was_revised": true, "final_justification": "I remain my original rating, i.e, learning towards acceptance.", "rebuttal": "We would like to thank the reviewer for their time and for raising this point about our specific choice of GNNs.\n\nWe agree with the reviewer that that a study of finer-grained network design choices would be valuable, and that our work suggests GNNs as just one way to learn domain permutation-equivariant functions. However, we would like to point out that while other existing networks like PointNet, Deep Sets and Set Transformers do handle permutation symmetries, they are not suitable to leverage the restricted class of domain permutation-equivariances in BED problems. This distinction arises from the **type of input representation** and the **corresponding type of symmetry** we can leverage. There are two primary ways to formulate a BED policy:\n\n**Option A: Map directly from the** **information set to the action distribution $a_t \\sim \\pi(\\{(x_{\\tau}, y_{\\tau})\\}_{\\tau=0}^{t-1})$.** In this setup, the policy is invariant to all permutations of the information set (i.e. it does not matter whether we observe the sequence $(x_a,y_a), (x_b,y_b), (x_c,y_c)$ or $(x_b,y_b), (x_c,y_c), (x_a,y_a)$ , the next optimal action is the same). Any network designed to learn functions over such sets is appropriate, including PointNet, Deep Sets and Set Transformers. Our Transformer baseline (with no positional encoding mechanism) is an example of such a model, and as our main results show, only leveraging information set permutation invariance is insufficient to combat the belief explosion issues we observe.\n\n**Option B: Perform Bayesian inference to compute the belief-state and then** **map from the belief state to the action distribution $a_t \\sim \\pi(p(\\theta|\\mathcal{I}_t))$.** This two-step process automatically handles the information-set permutation invariance from the first setup (since inference results in the same posterior regardless of data order). However, it also introduces a new, more structured symmetry: **domain permutation equivariance**. Figure 1 in the paper illustrates this concretely: if we rotate the belief-state, the optimal action distribution should also rotate. GNNs are especially suitable for this task because the belief state (i.e., the posterior mean and covariance) has a natural graph structure as it basically models a set of random variable with marginal and joint statistics. PointNet or Deep Sets would treat the the marginal mean and variances as an unordered set of values, completely ignoring the rich relational information encoded in the **joint** statistics such as the covariance matrix. GNNs, on the other hand, can preserve the relationship between points in the domain by encoding the joint statistics (covariance information) in the graph edges, while being equivariant to domain permutations such as rotations.\n\nWe thank the reviewer again for prompting this important discussion. We will add a paragraph to Section 4 to explicitly clarify this distinction and better contextualize our choice of GNNs with respect to these other architectures."}, {"review_id": "tKBLd3PA5f", "reviewer": "Reviewer_Hk2o", "summary": "This paper discusses the (modelling) choice of a (deep) policy network for Bayesian Experiment Design (BED) tasks (i.e., a generalization of Bayesian Optimization and/or Active Search). It argues that the belief state undergoes an exponential explosion with as the time horizon increases (a well known result since the 1960's which led to the term \"curse of dimensionality\" by Bellmann). The paper argues that in BED tasks which have a state-permutation invariance in both the MDP model and reward function, a Graph Neural Network architecture provides the right inductive bias to reduce the number of necessary samples as well as avoids overfitting. The claim is empirically validated on a toy problem.", "questions": "How exactly is the state space mapped to the nodes of a GNN and CNN for both the Bayesian Optimization and Active Search example in Section 5?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "A strength of the paper is the empirical evaluation on a rather small, discrete state-space example which already highlights the substantial differences in generalization performance - also compared to non DNN approaches such as UCB or Thompson sampling. For me, a major weakness is the vague description how the neural network architectures are mapped to the policy network; I still find it hard to understand how exactly the GNN is mapped to this problem to achieve the improved generalization. As this is the main contribution of the paper, I would have preferred more clarity there (Figure 4 is too vague for my taste)", "was_revised": true, "final_justification": "Thanks for the clarification by the authors in the rebuttal - I have adjusted my scores accordingly", "rebuttal": "We would like to thank the reviewer for their time and their valuable feedback. We appreciate the positive comments on our empirical evaluation. We would like to mention that we organized the paper to contain the most pedagogically understandable examples in the main paper but include significantly more experiments in the appendix, including fully continuous higher-dimensional (4D and 8D) Bayesian Optimization tasks.\n\nThe reviewer is right to point out that the description of how the GNN is mapped to the policy network could be clearer, and we are happy to provide a more detailed description below.\n\nAs a reminder of our setup, our policy network maps from the belief-state to a distribution over actions, and our goal is design a domain-permutation equivariant policy. For example, in Bayesian Optimization, this belief-state corresponds to the mean vector $\\mu$ and covariance matrix $\\Sigma$, and the simplest example of a domain permutation we want the policy to be equivariant to is flips for 1D problems (generalized to rotations for higher-dimensional BO problems).\n\nOur parameterization of the input graph to the GNN is directly motivated by these domain-permutation equivariances. Each node of the graph corresponds to a point in the domain; the node features are the marginal mean and variance, and the edge features are the covariance terms. The GNN processes this input graph and the final layer outputs a single logit for each node in the graph. Since each node corresponds to a point in the domain, this gives us a logit for each possible discrete action. We then apply a softmax function over these logits to produce the final probability distribution for our policy.\n\nThis specific parameterization ensures that the desired equivariance is built into the architecture. Any permutation of the graph's nodes (e.g., a rotation or flip of the domain) will result in an identical permutation of the output logits, thereby making the policy exactly equivariant. This structure allows the network to generalize far more efficiently, as it learns a single underlying mapping to action distributions that applies across all symmetric belief states.\n\nWe describe the design for a similarly motivated input graph for the GNN for fully-continuous BO tasks in Section E of the appendix.\n\nWe thank the reviewer again for this constructive feedback, and we will update our manuscript to include a detailed step-by-step description of the GNN.\n\nWith regards to the Bellman curse of dimensionality observation, we agree with the reviewer that our observation about belief explosion is strongly related to Bellman’s original observation. We acknowledge this when introducing the concept of belief explosion in our paper on line 125 by connecting it to the literature on POMDP planning. However, while Bellman’s original observation concerned a broad class of decision-making under uncertainty problems, part of our contribution is to highlight belief explosion for BED problems specifically and how existing contemporary approaches fail to recognise this as a training bottleneck. Moreover, our central contribution is to **propose a specific solution for certain BED problems** which tames this problem significantly. \n\nWe thank the reviewer for the formatting and notation issues, we will address all of these issues in the final manuscript."}, {"review_id": "3GargxxBrv", "reviewer": "Reviewer_mCGY", "summary": "The authors prove that many Bayesian Experimental Design (BED) problems are domain permutation invariant. To exploit this property, they propose to use Graph Neural Networks (GNNs), which are inherently permutation equivariant. Through extensive experiments, they show that leveraging this architectural bias can dramatically improve sampling efficiency, unlocking significantly larger tasks, making progress towards addressing the \"belief explosion\" issue.", "questions": "- The limitations section points out that handling hyperpriors is an open challenge. Can you elaborate on how the framework might be extended to handle this?\n- For the continuous experiments, the GNN outputs weights to form a linear combination of points. Are there cases where this strategy falls short?\n- As acknowledged in the limitations section, in real-world applications, model misspecification will be common. How robust is the method to model misspecification? What might be ways to improve robustness?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n- Novel, original approach towards better solutions on relevant problems\n- Sound theoretical foundation\n- Comprehensive quantitative evaluation\n- Paper is clearly written, good presentation of results\n\n## Weaknesses\n\n- Rests on the assumption of well-specified models, which is commonly violated in practice\n- Reliance on first- and second-order posterior moments; insufficient for handling hyperpriors (as discussed by the authors)\n\n### Minor\n\n- Readability could be improved by placing citations in brackets", "was_revised": true, "final_justification": "I appreciated the thorough response of the authors, which resolved my questions, and maintain my recommendation to accept the paper.", "rebuttal": "We would like to thank the reviewer for their time and insightful questions. We appreciate the positive comments, and we respond to the individual questions below.\n\n**For the continuous experiments, the GNN outputs weights to form a linear combination of points. Are there cases where this strategy falls short?**\n\nYes, this strategy can fall short if the information set does not sufficiently span the design space. Specifically, if the information set contains points that lie within a low-dimensional subspace of $\\mathcal{X}$, then any linear combination of those points will remain within that subspace, limiting the exploration capabilities of the policy.\n\nTo address this, we initialize each trial with $d+1$ points sampled uniformly at random from $\\mathcal{X}$, where $d$ is the dimensionality of the design space. This ensures (almost surely) that the information set spans $\\mathcal{X}$. This warm-start practice is common in continuous Bayesian optimization and aligns with protocols in the literature. For example, Astudillo & Frazier (2019) write: “For all problems and methods, an initial stage of evaluations is performed using $2(d + 1)$ points chosen uniformly at random over $\\mathcal{X}$.” [1]\n\nThus, while the linear combination strategy has inherent limitations, we mitigate them through principled initialization that ensures sufficient coverage of the design space.\n\n[1] Astudillo, Raul, and Peter Frazier. \"Bayesian optimization of composite functions.\" International Conference on Machine Learning, 2019.\n\n**The limitations section points out that handling hyperpriors is an open challenge. Can you elaborate on how the framework might be extended to handle this?**\n\nIn the continuous setting, one natural extension is to train the GNN on information sets generated using a distribution over hyperparameters e.g., sampling kernels or other GP hyperparameters at training time. This effectively averages over a richer space of data-generating processes, helping the model learn representations that are robust to different prior assumptions.\nFor the discrete setting, a viable alternative is to use MAP estimates of kernel hyperparameters (obtained via standard hyperprior inference) to compute GP posterior moments. These moments can then be encoded as input graphs, as done in our current framework. While this collapses the full Bayesian belief state into a second-order approximation, we hypothesize that it suffices for many realistic forms of misspecification.\n\nNonetheless, we agree this is an important and nuanced direction, and a dedicated study of hyperprior graph representation is warranted.\n\n**As acknowledged in the limitations section, in real-world applications, model misspecification will be common. How robust is the method to model misspecification? What might be ways to improve robustness?**\n\nWe appreciate the reviewer’s emphasis on this important issue. In general, amortized policies (such as those learned by our GNN) are inherently sensitive to discrepancies between training-time and test-time data-generating processes. That said, robustness can potentially be improved by diversifying the training data.\n\nFor instance, if even a few representative real-world functions are available, one could augment the training set with these alongside synthetic GP samples. This could enable the policy to implicitly learn priors over a broader and potentially non-Gaussian class of functions. Such a strategy would allow the policy to adapt beyond the GP framework without requiring explicit changes to the architecture or the inference mechanism.\n\nFinally, we believe that there are many competing reasonable notions of misspecification and robustness that warrant careful study. We agree that a dedicated study of robustness to model mispecification is an important future direction of this work."}]

Author rebuttal:

Meta-review: This paper studies Bayesian Experiment Design and addresses the “belief explosion” problem, where the number of realizable beliefs grows exponentially with experiment length, creating severe sample inefficiency. The authors propose using Graph Neural Networks to exploit the domain permutation equivariance in BED problems, enabling policies that generalize efficiently across symmetric belief states. Experiments on Bayesian Optimization and Active Search show that GNN-based policies outperform alternative architectures (MLPs, CNNs, Transformers) and match or exceed traditional baselines like UCB and Thompson Sampling.

The paper is well-motivated, technically sound, and clearly written, with a novel use of GNNs that is distinct from prior RL applications. Reviewers found the empirical evaluation convincing, and the authors’ rebuttal addressed concerns about architectural design, validation procedures, error bars, and the mapping of belief states to GNN nodes. Limitations include reliance on model assumptions, sensitivity to hyperpriors, and evaluation mostly on toy tasks, but these are acknowledged and discussed.

Overall, the paper makes a solid contribution by combining a principled architectural choice with empirical evidence of improved sample efficiency. Given the technical quality, originality, and thorough responses during the rebuttal, the recommendation is to accept.
Decision: Accept (poster)