Update the Eagle Strategy to support multiplicative continuous feature perturbation.

PiperOrigin-RevId: 723098266

Author: vizier-team

vizier/_src/algorithms/optimizers/eagle_strategy.py

@@ -34,7 +34,7 @@
 in addition to the "attraction" forces, meaning fireflies move towards the
 bright spots as well as away from the dark spots. We also support non-decimal
 parameter types (categorical, discrete, integer), and treat them uniquely when
-computing distance, adding pertrubation, and mutating fireflies.
+computing distance, applying pertrubation, and mutating fireflies.
 
 For more details, see the linked paper.
 
@@ -68,6 +68,7 @@
 # Run the optimization.
 trials = optimizer.optimize(problem_statement, objective_function)
 """
+
 import dataclasses
 # pylint: disable=g-long-lambda
 
@@ -97,6 +98,16 @@ class MutateNormalizationType(enum.IntEnum):
   UNNORMALIZED = 2
 
 
+@enum.unique
+class ContinuousFeaturePerturbationType(enum.IntEnum):
+  """The type of perturbation to apply to continuous features."""
+
+  # Add the perturbation to the feature value.
+  ADDITIVE = 0
+  # Multiply the feature value by exp(perturbation).
+  MULTIPLICATIVE = 1
+
+
 @struct.dataclass
 class EagleStrategyConfig:
   """Eagle Strategy optimizer config.
@@ -106,6 +117,8 @@ class EagleStrategyConfig:
     gravity: The maximum amount of attraction pull.
     negative_gravity: The maximum amount of repulsion pull.
     perturbation: The default amount of noise for perturbation.
+    continuous_feature_perturbation_type: The type of perturbation to apply to
+      continuous features.
     categorical_perturbation_factor: A factor to apply on categorical params.
     pure_categorical_perturbation_factor: A factor on purely categorical space.
     prob_same_category_without_perturbation: Baseline probability of selecting
@@ -130,6 +143,9 @@ class EagleStrategyConfig:
   negative_gravity: float = 0.008
   # Perturbation
   perturbation: float = 0.16
+  continuous_feature_perturbation_type: ContinuousFeaturePerturbationType = (
+      ContinuousFeaturePerturbationType.ADDITIVE
+  )
   categorical_perturbation_factor: float = 1.0
   pure_categorical_perturbation_factor: float = 30
   prob_same_category_without_perturbation: float = 0.98
@@ -873,14 +889,28 @@ def _create_features(
     features_changes_continuous = jnp.matmul(
         scale, flat_features
     ) - flat_features_batch * jnp.sum(scale, axis=-1, keepdims=True)
-
-    new_features_continuous = (
-        features_batch.continuous
-        + jnp.reshape(
-            features_changes_continuous, features_batch.continuous.shape
-        )
-        + perturbations_batch.continuous
+    moved_features_continuous = features_batch.continuous + jnp.reshape(
+        features_changes_continuous, features_batch.continuous.shape
     )
+    if (
+        self.config.continuous_feature_perturbation_type
+        == ContinuousFeaturePerturbationType.ADDITIVE
+    ):
+      new_features_continuous = (
+          moved_features_continuous + perturbations_batch.continuous
+      )
+    elif (
+        self.config.continuous_feature_perturbation_type
+        == ContinuousFeaturePerturbationType.MULTIPLICATIVE
+    ):
+      new_features_continuous = moved_features_continuous * jnp.exp(
+          perturbations_batch.continuous
+      )
+    else:
+      raise ValueError(
+          "Unsupported continuous feature perturbation type:"
+          f" {self.config.continuous_feature_perturbation_type}"
+      )
     if self.max_categorical_size > 0:
       features_categorical_logits = (
           self._create_categorical_feature_logits(

vizier/_src/algorithms/optimizers/eagle_strategy_test.py

@@ -31,6 +31,9 @@
 from absl.testing import absltest
 
 tfd = tfp.distributions
+ContinuousFeaturePerturbationType = (
+    eagle_strategy.ContinuousFeaturePerturbationType
+)
 
 
 def _create_logits_vector_simple(
@@ -437,6 +440,60 @@ def test_optimize_with_eagle(self):
         score_fn=lambda x, _: -jnp.sum(x.continuous.padded_array, 1), count=1
     )
 
+  def test_continuous_feature_perturbation_type(self):
+    optimizer_additive_perturbation = vb.VectorizedOptimizerFactory(
+        strategy_factory=eagle_strategy.VectorizedEagleStrategyFactory(
+            eagle_config=eagle_strategy.EagleStrategyConfig(
+                continuous_feature_perturbation_type=(
+                    ContinuousFeaturePerturbationType.ADDITIVE
+                )
+            )
+        ),
+        max_evaluations=50,
+    )(self.converter)
+    expected_count = 4
+    new_features_additive_perturbation = optimizer_additive_perturbation(
+        score_fn=lambda x, _: -jnp.sum(x.continuous.padded_array, 1),
+        count=expected_count,
+    ).features
+    self.assertSequenceEqual(
+        new_features_additive_perturbation.continuous.shape,
+        (expected_count, 1, 2),
+    )
+    self.assertSequenceEqual(
+        new_features_additive_perturbation.categorical.shape,
+        (expected_count, 1, 2),
+    )
+    optimizer_mult_perturbation = vb.VectorizedOptimizerFactory(
+        strategy_factory=eagle_strategy.VectorizedEagleStrategyFactory(
+            eagle_config=eagle_strategy.EagleStrategyConfig(
+                continuous_feature_perturbation_type=(
+                    ContinuousFeaturePerturbationType.MULTIPLICATIVE
+                )
+            )
+        ),
+        max_evaluations=50,
+    )(self.converter)
+    new_features_mult_perturbation = optimizer_mult_perturbation(
+        score_fn=lambda x, _: -jnp.sum(x.continuous.padded_array, 1),
+        count=expected_count,
+    ).features
+    self.assertSequenceEqual(
+        new_features_mult_perturbation.continuous.shape,
+        (expected_count, 1, 2),
+    )
+    self.assertSequenceEqual(
+        new_features_mult_perturbation.categorical.shape,
+        (expected_count, 1, 2),
+    )
+    self.assertGreater(
+        np.abs(
+            new_features_additive_perturbation.continuous
+            - new_features_mult_perturbation.continuous
+        ).max(),
+        1e-1,
+    )
+
   def test_optimize_with_eagle_continuous_only(self):
     problem = vz.ProblemStatement()
     root = problem.search_space.select_root()