Sample-Then-Optimize Batch Neural Thompson Sampling
Zhongxiang Dai, Yao Shu, Bryan Kian Hsiang Low, Patrick Jaillet
Abstract
Bayesian optimization (BO), which uses a Gaussian process (GP) as a surrogate to model its objective function, is popular for black-box optimization. However, due to the limitations of GPs, BO underperforms in some problems such as those with categorical, high-dimensional or image inputs. To this end, recent works have used the highly expressive neural networks (NNs) as the surrogate model and derived theoretical guarantees using the theory of neural tangent kernel (NTK). However, these works suffer from the limitations of the requirement to invert an extremely large parameter matrix and the restriction to the sequential (rather than batch) setting. To overcome these limitations, we introduce two algorithms based on the Thompson sampling (TS) policy named Sample-Then-Optimize Batch Neural TS (STO-BNTS) and STO-BNTS-Linear. To choose an input query, we only need to train an NN (resp. a linear model) and then choose the query by maximizing the trained NN (resp. linear model), which is equivalently sampled from the GP posterior with the NTK as the kernel function. As a result, our algorithms sidestep the need to invert the large parameter matrix yet still preserve the validity of the TS policy. Next, we derive regret upper bounds for our algorithms with batch evaluations, and use insights from batch BO and NTK to show that they are asymptotically no-regret under certain conditions. Finally, we verify their empirical effectiveness using practical AutoML and reinforcement learning experiments.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers17
- Large Language Models to Enhance Bayesian OptimizationTennison Liu, Nicolás Astorga, Nabeel Seedat, Mihaela van der SchaarICLR 2024 · 143 citations
- A Study of Bayesian Neural Network Surrogates for Bayesian OptimizationYucen Lily Li, Tim G. J. Rudner, Andrew Gordon WilsonICLR 2024 · 59 citations
- PINNACLE: PINN Adaptive ColLocation and Experimental points selectionGregory Kang Ruey Lau, Apivich Hemachandra, See-Kiong Ng, Bryan Kian Hsiang LowICLR 2024 · 43 citations
- Unifying and Boosting Gradient-Based Training-Free Neural Architecture SearchYao Shu, Zhongxiang Dai, Zhaoxuan Wu, Bryan Kian Hsiang LowNeurIPS 2022 · 41 citations
- Use Your INSTINCT: INSTruction optimization for LLMs usIng Neural bandits Coupled with TransformersXiaoqiang Lin, Zhaoxuan Wu, Zhongxiang Dai, Wenyang Hu et al.ICML 2024 · 26 citations
Builds on24
- Neural Contextual Bandits with UCB-based ExplorationDongruo Zhou, Lihong Li, Quanquan GuICML 2020 · 329 citations
- Neural Thompson SamplingWeitong Zhang, Dongruo Zhou, Lihong Li, Quanquan GuICLR 2021 · 152 citations
- Federated Bayesian Optimization via Thompson SamplingZhongxiang Dai, Bryan Kian Hsiang Low, Patrick JailletNeurIPS 2020 · 144 citations
- Bayesian Deep Ensembles via the Neural Tangent KernelBobby He, Balaji Lakshminarayanan, Yee Whye TehNeurIPS 2020 · 136 citations
- Fault-Tolerant Federated Reinforcement Learning with Theoretical GuaranteeFlint Xiaofeng Fan, Yining Ma, Zhongxiang Dai, Wei Jing et al.NeurIPS 2021 · 102 citations
Related papers
- Thompson Sampling in Function Spaces via Neural OperatorsRafael Oliveira, Xuesong Wang, Kian Ming A. Chai, Edwin V. BonillaNeurIPS 2025 · 2 citations
- Bayesian Optimization for Categorical and Category-Specific Continuous InputsDang Nguyen, Sunil Gupta, Santu Rana, Alistair Shilton et al.AAAI 2020 · 59 citations
- On Regret Bounds of Thompson Sampling for Bayesian OptimizationShion Takeno, Shogo IwazakiICML 2026 · 3 citations
- Bayesian Optimization over Permutation SpacesAryan Deshwal, Syrine Belakaria, Janardhan Rao Doppa, Dae Hyun KimAAAI 2022 · 27 citations
- BayeSQP: Bayesian Optimization through Sequential Quadratic ProgrammingPaul Brunzema, Sebastian TrimpeNeurIPS 2025 · 7 citations
