Title: Efficient Bayesian Experiment Design with Equivariant Networks
Venue: NeurIPS 2025
Accepted: yes
Keywords: Bayesian Experiment Design

Bayesian Optimization

Active Search

Deep Learning

Deep Reinforcement Learning

Graph Neural Networks
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.


Peer reviews:
comments: 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.

STRENGTHS AND WEAKNESSES:
## Strengths

- Novel, original approach towards better solutions on relevant problems
- Sound theoretical foundation
- Comprehensive quantitative evaluation
- Paper is clearly written, good presentation of results

## Weaknesses

- Rests on the assumption of well-specified models, which is commonly violated in practice
- Reliance on first- and second-order posterior moments; insufficient for handling hyperpriors (as discussed by the authors)

### Minor

- Readability could be improved by placing citations in brackets

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?
- For the continuous experiments, the GNN outputs weights to form a linear combination of points. Are there cases where this strategy falls short?
- 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
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-27

comments: 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.

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)

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
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-06-29

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:
+ 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).
+ 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).

Weaknesses:
+ 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.

QUESTIONS:
See "weaknesses"

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-04"}