add references to BONSAI paper.

Summary: see title

Differential Revision: D92842590

Author: sdaulton

ax/api/client.py

@@ -244,8 +244,9 @@ def configure_generation_strategy(
                 initialization budget if more trials are needed.
             torch_device: The torch device to use for model fitting. If None, will
                 use the default device.
-            simplify_parameter_changes: Whether to simplify parameter changes in
-                arms generated via Bayesian Optimization by pruning irrelevant
+            simplify_parameter_changes: Whether to use BONSAI
+                [Daulton2026bonsai]_ to simplify parameter changes in arms
+                generated via Bayesian Optimization by pruning irrelevant
                 parameter changes.
         """
         generation_strategy = self._choose_generation_strategy(
@@ -1186,8 +1187,9 @@ def _choose_generation_strategy(
                 initialization budget if more trials are needed.
             torch_device: The torch device to use for model fitting. If None, will
                 use the default device.
-            simplify_parameter_changes: Whether to simplify parameter changes in
-                arms generated via Bayesian Optimization by pruning irrelevant
+            simplify_parameter_changes: Whether to use BONSAI
+                [Daulton2026bonsai]_ to simplify parameter changes in arms
+                generated via Bayesian Optimization by pruning irrelevant
                 parameter changes.
 
 

ax/api/utils/generation_strategy_dispatch.py

@@ -114,8 +114,8 @@ def _get_mbm_node(
             - "fast": Uses MBM defaults.
             - "custom": Uses the provided ``model_config``.
         torch_device: The torch device to use for the MBM node.
-        simplify_parameter_changes: Whether to simplify parameter changes in
-            the MBM node.
+        simplify_parameter_changes: Whether to use BONSAI [Daulton2026bonsai]_ to
+            simplify parameter changes in the MBM node.
         model_config: Optional model config to use for the MBM node.
             This is only supported when ``method`` is "custom".
         botorch_acqf_class: An optional BoTorch ``AcquisitionFunction`` class

ax/generation_strategy/dispatch_utils.py

@@ -399,9 +399,9 @@ def choose_generation_strategy_legacy(
             step and automatic selection will be skipped.
         use_input_warping: Whether to use input warping in the model. This is only
             supported in conjunction with use_saasbo=True.
-        simplify_parameter_changes: Whether to simplify parameter changes in
-            arms generated via Bayesian Optimization by pruning irrelevant
-            parameter changes.
+        simplify_parameter_changes: Whether to use BONSAI [Daulton2026bonsai]_ to
+            simplify parameter changes in arms generated via Bayesian Optimization
+            by pruning irrelevant parameter changes.
     """
     if experiment is not None and optimization_config is None:
         optimization_config = experiment.optimization_config

ax/generators/torch/botorch_modular/acquisition.py

@@ -6,6 +6,14 @@
 
 # pyre-strict
 
+"""
+References
+
+.. [Daulton2026bonsai]
+    S. Daulton, D. Eriksson, M. Balandat, and E. Bakshy. BONSAI: Bayesian
+    Optimization with Natural Simplicity and Interpretability. ArXiv, 2026.
+"""
+
 from __future__ import annotations
 
 import operator
@@ -782,7 +790,9 @@ def _prune_irrelevant_parameters(
         inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
         fixed_features: dict[int, float] | None = None,
     ) -> tuple[Tensor, Tensor]:
-        r"""Prune irrelevant parameters from the candidates.
+        r"""Prune irrelevant parameters from the candidates using BONSAI.
+
+        See [Daulton2026bonsai]_ for details.
 
         The method involves first optimizing the AF without any notion of irrelevance.
         Then, the irrelevant parameters are pruned via a sequential greedy algorithm