Skip to main content

Epistemic Nearest Neighbors

Project description

Epistemic Nearest Neighbors

A fast, alternative surrogate for Bayesian optimization

ENN estimates a function's value and associated epistemic uncertainty using a K-Nearest Neighbors model. Queries take $O(N lnK)$ time, where $N$ is the number of observations available for KNN lookups. Compare to an exact GP, which takes $O(N^2)$ time. Additionally, measured running times are very small compared to GPs and other alternative surrogates. [1]

Contents

  • ENN surrogate, EpistemicNearestNeighbors [1]

  • TuRBO-ENN optimizer via create_optimizer with config factories

    • turbo_one_config() - TuRBO [2], matching the reference implementation.
    • turbo_enn_config() - Uses ENN instead of GP.
    • turbo_zero_config() - No surrogate
    • lhd_only_config() - LHD design on every ask(). Good for a baseline and for testing. The optimizer has an ask()/tell() interface. All turbo_*() methods follow TuRBO:
    • Generate candidates with RAASP [3] sampling.
    • Select a candidate with Thompson sampling (TuRBO-one), UCB (TuRBO-ENN), or randomly (TURBO-zero).
  • Overview of algorithms: algos.pdf

[1] M. Bafna, Jadhav, S. a., & Sweet, D., (2025). Taking the GP Out of the Loop. arXiv preprint arXiv:2506.12818. https://arxiv.org/abs/2506.12818 [2] Eriksson, D., Pearce, M., Gardner, J. R., Turner, R., & Poloczek, M. (2020). Scalable Global Optimization via Local Bayesian Optimization. Advances in Neural Information Processing Systems, 32. https://arxiv.org/abs/1910.01739 [3] Rashidi, B., Johnstonbaugh, K., & Gao, C. (2024). Cylindrical Thompson Sampling for High-Dimensional Bayesian Optimization. Proceedings of The 27th International Conference on Artificial Intelligence and Statistics (pp. 3502–3510). PMLR. https://proceedings.mlr.press/v238/rashidi24a.html

Installation

pip install ennbo[with-deps] or cargo add ennbo

Demonstration

demo_enn.ipynb - Shows how to use EpistemicNearestNeighbors to build and query an ENN model. demo_turbo_enn.ipynb - Shows how to use TurboOptimizer to optimize the Ackley function.

Installation, MacOS

On my MacBook I can run into problems with dependencies and compatibilities.

On MacOS try:

micromamba env create -n ennbo -f admin/conda-macos.yml
micromamba activate ennbo
pip install --no-deps ennbo
pytest -sv tests

You may replace micromamba with conda and this will probably still work.

The commands above make sure

  • You use the MacOS-specific PyTorch (with mps).
  • You avoid having multiple, competing OpenMPs installed PyTorch issue faiss issue.
  • You use old enough versions of NumPy and PyTorch to be compatible with faiss faiss issue.
  • Prevent matplotlib's installation from upgrading your NumPy to an incompatible version.
  • ennbo's listed dependencies do not undo any of the above (which is fine b/c the above commands set the up correctly).

Run tests with

pytest -x -sv tests

and they should all pass fairly quickly (~10s-30s).

If your code still crashes or hangs, try this hack:

export KMP_DUPLICATE_LIB_OK=TRUE
export OMP_NUM_THREADS=1

I don't recommend this, however, as it will slow things down.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

ennbo-0.3.2.tar.gz (103.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

ennbo-0.3.2-cp311-cp311-macosx_11_0_arm64.whl (1.3 MB view details)

Uploaded CPython 3.11macOS 11.0+ ARM64

File details

Details for the file ennbo-0.3.2.tar.gz.

File metadata

  • Download URL: ennbo-0.3.2.tar.gz
  • Upload date:
  • Size: 103.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: maturin/1.12.5

File hashes

Hashes for ennbo-0.3.2.tar.gz
Algorithm Hash digest
SHA256 0b66fae648f6b1c7cade612d64939d75839489dee4a72ea25f6a2447cdbe9ef9
MD5 4b9f9e2c4fe771d3d31a0b98cdf3a59a
BLAKE2b-256 45785c8a69563b490020caf76b736afdbd05967a378dbd9b3bf326a7397c2647

See more details on using hashes here.

File details

Details for the file ennbo-0.3.2-cp311-cp311-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for ennbo-0.3.2-cp311-cp311-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 073b36e9b1ad17ac3f2e76d521e95b4d3842c38c2206b942c52de3431594c9c5
MD5 11b94926324ab59cf137a79c40755d61
BLAKE2b-256 443feb32cdec3aaa3a4993b6f72d5e1fefe4a864127eeeecdf28951ccc8435a2

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page