Neural Operator Thompson Sampling (NOTS)

Code repository for research paper "Thompson Sampling in Function Spaces via Neural Operators" presented at NeurIPS 2025 in San Diego, CA. A preprint is available on arXiv at: https://arxiv.org/abs/2506.21894

Dependencies

To run the code, the dependencies should be first installed. Preferrably a python virtual environment should be created and activated before running the experiments. On a Unix environment, one can do so by running:

$ python -m venv .venv
$ source .venv/bin/activate
$ pip install -r requirements.txt

Running experiments

The main experiment scripts can be found on the root folder. Instructions on how to run them can be found by running them with the --help command line option, e.g.:

python pool_nots.py --help

To launch experiments with multiple algorithms, please use run_trials.py, which allows for specifying multiple algorithms to run with the same problem settings and launching them as separate processes. For example, the following command will run NOTS and the baselines on the shallow water problem in the paper:

python run_trials.py --dataset=spherical_swe_long --functional=inverse_problem --n-iterations=300 --n-trials=10 --n-init=10 --noise=0.1 --seed=42

If no option is passed with --algorithms, all the methods are run by default. To specify only particular methods to run, one can append to the command above, for example, --algorithms nots bo gpts bfo to run only NOTS, BO (with expected improvement algorithm), GP-TS and Bayesian functional optimization.

To run on GPU or other PyTorch-compatible devices, one should use the --device option on any of the experiment scripts, e.g., --device=cuda, selects NVIDIA GPUs, and --device=mps, selects Apple Metal GPUs as device to run the algorithms on. By default, the CPU device is used.

Darcy flow data

The Darcy flow PDE dataset is provided under the data/ directory as a compressed archive darcy_16.tgz (16-by-16 grid resolution used in our experiments) and should be extracted before attempting to run experiments with that dataset, specified via the command-line option --dataset=darcy_flow_small on the experiment scripts. This data file originates from the Neural Operator project (https://github.com/neuraloperator/neuraloperator) and is redistributed in this repository under the terms of its original MIT license, which is included alongside the data. The upstream project is also used as a software dependency.

References

Please cite the paper in your publications as:

@inproceedings{
oliveira2025thompson,
title={Thompson Sampling in Function Spaces via Neural Operators},
author={Rafael Oliveira and Xuesong Wang and Kian Ming A. Chai and Edwin V. Bonilla},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
address={San Diego, CA, USA},
url={https://openreview.net/forum?id=F54u4NkFvS}
}