Add LogEI

src/decijax/__init__.py

@@ -4,6 +4,7 @@
     acquisition_functions,
     acquisition_maximizer,
     decision_maker,
+    maths,
     models,
     search_space,
     test_functions,
@@ -15,6 +16,7 @@
     "acquisition_functions",
     "acquisition_maximizer",
     "decision_maker",
+    "maths",
     "models",
     "search_space",
     "test_functions",

src/decijax/acquisition_functions/__init__.py

@@ -8,6 +8,7 @@
 )
 from decijax.acquisition_functions.expected_improvement import (
     ExpectedImprovement,
+    LogExpectedImprovement,
 )
 from decijax.acquisition_functions.probability_of_improvement import (
     LogProbabilityOfImprovement,
@@ -20,6 +21,7 @@
     "AbstractAcquisitionFunctionBuilder",
     "AbstractSinglePointAcquisitionFunctionBuilder",
     "ExpectedImprovement",
+    "LogExpectedImprovement",
     "LogProbabilityOfImprovement",
     "SinglePointAcquisitionFunction",
     "ThompsonSampling",

src/decijax/acquisition_functions/expected_improvement.py

@@ -3,6 +3,7 @@
 from collections.abc import Mapping
 
 import jax.numpy as jnp
+from jax.nn import logmeanexp
 from jax.scipy.stats import norm
 from jaxtyping import (
     Array,
@@ -13,6 +14,10 @@
     AbstractSinglePointAcquisitionFunctionBuilder,
     SinglePointAcquisitionFunction,
 )
+from decijax.maths import (
+    _log1mexp,
+    _log_abs_z_cdf_div_pdf,
+)
 from decijax.models import (
     ProbabilisticModel,
     SupportsGaussianPrediction,
@@ -24,12 +29,73 @@
 )
 
 
+def _log_ei_helper(z: Float[Array, "..."]) -> Float[Array, "..."]:
+    r"""Compute $`\log(\phi(z) + z \Phi(z))`$ in a numerically stable manner.
+
+    Evaluated directly this underflows to zero, gradient included, once
+    $`z \lesssim -6`$ in double precision. Implements the piecewise formulation of
+    [Ament et al., 2023](https://arxiv.org/abs/2310.20708), Eq. 9:
+
+    ```math
+    \begin{cases}
+    \log\left(\phi(z) + z \Phi(z)\right) & z > -1 \\
+    -\frac{z^2}{2} - c_1 + \text{log1mexp}\left(
+        \log\left(\text{erfcx}\left(\frac{-z}{\sqrt{2}}\right)|z|\right) + c_2
+    \right) & \tau < z \leq -1 \\
+    -\frac{z^2}{2} - c_1 - 2\log|z| & z \leq \tau
+    \end{cases}
+    ```
+
+    where $`c_1 = \frac{\log(2\pi)}{2}`$ and $`c_2 = \frac{\log(\pi / 2)}{2}`$. The
+    threshold $`\tau`$ between the last two branches is $`-10^6`$ in double
+    precision and $`-10^3`$ in single, following BoTorch. Each branch clamps on
+    *both* sides of its boundaries to keep gradients clean.
+
+    Args:
+        z: Scaled improvement.
+
+    Returns:
+        $`\log(\phi(z) + z \Phi(z))`$, elementwise.
+    """
+    if z.dtype == jnp.float64:
+        eps_bound = -1e6
+    elif z.dtype == jnp.float32:
+        eps_bound = -1e3
+    else:
+        raise NotImplementedError(
+            f"LogExpectedImprovement does not support dtype {z.dtype}."
+        )
+
+    log_sqrt_2pi = 0.5 * jnp.log(2.0 * jnp.pi)
+
+    # Branch 1 (z > -1): direct computation.
+    z_upper = jnp.maximum(z, -1.0)
+    direct = jnp.log(norm.pdf(z_upper) + z_upper * norm.cdf(z_upper))
+
+    # Branch 2 (eps_bound < z <= -1): stable computation using erfcx.
+    z_mid = jnp.clip(z, eps_bound, -1.0)
+    w = _log_abs_z_cdf_div_pdf(z_mid)
+    stable = -0.5 * z_mid**2 - log_sqrt_2pi + _log1mexp(w)
+
+    # Branch 3 (z <= eps_bound): asymptotic, where `_log1mexp` can no longer resolve w.
+    z_lower = jnp.minimum(z, eps_bound)
+    asymptotic = -0.5 * z_lower**2 - log_sqrt_2pi - 2.0 * jnp.log(-z_lower)
+
+    return jnp.where(
+        z > -1.0,
+        direct,
+        jnp.where(z > eps_bound, stable, asymptotic),
+    )
+
+
 class ExpectedImprovement(AbstractSinglePointAcquisitionFunctionBuilder):
     """Standard Expected Improvement acquisition function.
 
-    As introduced by [Močkus, 1974](https://link.springer.com/chapter/10.1007/3-540-07165-2_55). The "best" incumbent value is defined as the highest posterior
-    mean value evaluated at the previously observed points. This enables the
-    acquisition function to be utilised with noisy observations.
+    As introduced by
+    [Močkus, 1974](https://link.springer.com/chapter/10.1007/3-540-07165-2_55). The
+    "best" incumbent value is defined as the highest posterior mean value evaluated
+    at the previously observed points. This enables the acquisition function to be
+    utilised with noisy observations.
     """
 
     def build_acquisition_function(
@@ -41,9 +107,12 @@ def build_acquisition_function(
 
         This computes the expected improvement over the "best" of the previously
         observed points, utilising the posterior distribution of the surrogate model.
-        For posterior distribution $`f(\cdot)`$, and best incumbent value $`\eta`$, this is defined as:
+        For posterior distribution $`f(\cdot)`$, and best incumbent value $`\eta`$,
+        this is defined as:
+
         ```math
-        \alpha_{\text{EI}}(\mathbf{x}) = \mathbb{E}\left[\max(0, f(\mathbf{x}) - \eta)\right]
+        \alpha_{\text{EI}}(\mathbf{x})
+        = \mathbb{E}\left[\max(0, f(\mathbf{x}) - \eta)\right]
         ```
 
         For models carrying a leading sample axis (e.g. fully Bayesian GPs), the
@@ -85,3 +154,82 @@ def _expected_improvement(x: Float[Array, "N D"]) -> Float[Array, "N 1"]:
             return jnp.mean(ei, axis=0)[:, None]  # marginalise over S -> [N, 1]
 
         return _expected_improvement
+
+
+class LogExpectedImprovement(AbstractSinglePointAcquisitionFunctionBuilder):
+    r"""Numerically stable Log Expected Improvement acquisition function [1].
+
+    Given a predictive posterior distribution of the objective function $f$, the log
+    expected improvement at a test point $x$ is defined as:
+
+    $$\text{LogEI}(x) = \log \mathbb{E}\left[\max(0, f(x) - f(x^*))\right]$$
+
+    where $x^*$ is the maximiser of the posterior mean at previously observed values
+    (to handle noisy observations).
+
+    Being a strictly increasing transform of the expected improvement, this shares
+    its maximiser exactly, but is far better behaved as an optimisation target:
+    expected improvement vanishes to *exactly* zero, gradient included, once the
+    scaled improvement falls below roughly $-40$ in double precision, and those flat
+    regions come to dominate the search space in higher dimensions.
+
+    References:
+    ----------
+    [1] Ament, S., Daulton, S., Eriksson, D., Balandat, M., & Bakshy, E. (2023).
+    Unexpected improvement to expected improvement for Bayesian optimization.
+    Advances in Neural Information Processing Systems, 36.
+    """
+
+    def build_acquisition_function(
+        self,
+        models: Mapping[str, ProbabilisticModel],
+        key: KeyArray,
+    ) -> SinglePointAcquisitionFunction:
+        r"""Build the Log Expected Improvement acquisition function.
+
+        The expected improvement factorises as $`\sigma \cdot h(z)`$, for scaled
+        improvement $`z = \frac{\mu - \eta}{\sigma}`$ and
+        $`h(z) = \phi(z) + z\Phi(z)`$, so that:
+
+        ```math
+        \alpha_{\text{LogEI}}(\mathbf{x}) = \log \sigma(\mathbf{x}) + \log h(z)
+        ```
+
+        with the second term computed by `_log_ei_helper`. For models carrying a
+        leading sample axis (e.g. fully Bayesian GPs), it is computed per sample and
+        reduced with a log-mean-exp, the correct marginalisation
+        $`\log \mathbb{E}_\theta[\alpha_{\text{EI},\theta}(\mathbf{x})]`$.
+
+        Args:
+            models: Dictionary of models used to form the acquisition function. One
+                model must correspond to the `OBJECTIVE` key and support Gaussian
+                prediction, as we use the objective posterior to form the acquisition
+                function.
+            key: JAX PRNG key used for random number generation. Since the log
+                expected improvement is computed deterministically, the key is not
+                used.
+
+        Returns:
+            The Log Expected Improvement acquisition function to be *maximised* in
+            order to decide which point to query next.
+        """
+        self.check_objective_present(models)
+        objective_model = models[OBJECTIVE]
+
+        if not isinstance(objective_model, SupportsGaussianPrediction):
+            raise ValueError(
+                "Objective model must support Gaussian prediction to compute the "
+                "Log Expected Improvement."
+            )
+
+        eta = get_best_latent_observation_val(objective_model)  # [S, 1]
+
+        def _log_expected_improvement(x: Float[Array, "N D"]) -> Float[Array, "N 1"]:
+            latent_dist = objective_model.predict(x)
+            mean = latent_dist.mean  # [S, N]
+            std = latent_dist.stddev  # [S, N]
+            z = (mean - eta) / std
+            log_ei = _log_ei_helper(z) + jnp.log(std)  # [S, N]
+            return logmeanexp(log_ei, axis=0)[:, None]  # marginalise over S -> [N, 1]
+
+        return _log_expected_improvement

src/decijax/maths.py

@@ -0,0 +1,100 @@
+"""A collection of numerically stable mathematical functions."""
+
+import jax.numpy as jnp
+from jax.scipy.special import erfc
+from jaxtyping import (
+    Array,
+    Float,
+)
+
+
+def _log1mexp(x: Float[Array, "..."]) -> Float[Array, "..."]:
+    r"""Compute $`\log(1 - e^x)`$ for $`x < 0`$ without catastrophic cancellation.
+
+    Which form is accurate depends on how close $`e^x`$ is to one, hence the split
+    at $`-\log 2`$. See [Mächler,
+    2012](https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf).
+
+    Args:
+        x: Strictly negative array.
+
+    Returns:
+        $`\log(1 - e^x)`$, elementwise.
+    """
+    log_2 = jnp.log(2.0)
+    # Each branch is clamped to its own side of the split, so the discarded one
+    # cannot poison the gradient of the surviving one.
+    return jnp.where(
+        x > -log_2,
+        jnp.log(-jnp.expm1(jnp.maximum(x, -log_2))),
+        jnp.log1p(-jnp.exp(jnp.minimum(x, -log_2))),
+    )
+
+
+def _erfcx(x: Float[Array, "..."]) -> Float[Array, "..."]:
+    r"""Compute the scaled complementary error function $`e^{x^2}\text{erfc}(x)`$.
+
+    We need to define our own erfcx function, because `jax.scipy.special.erfcx` has
+    a bug where it currently returns zero on $`[26.54, 26.64]`$ (double) and
+    $`[9.19, 9.42]`$ (single). Once
+    [jax-ml/jax#38607](https://github.com/jax-ml/jax/issues/38607) is fixed, replace
+    this function with the native JAX implementation.
+
+    Args:
+        x: Strictly positive array.
+
+    Returns:
+        $`e^{x^2}\text{erfc}(x)`$, elementwise.
+    """
+    # Beyond this point the asymptotic expansion below has reached machine precision.
+    # It must stay *below* the point at which erfc underflows (26.54 in double, 9.19
+    # in single), which is exactly the bug JAX's own implementation has.
+    if x.dtype == jnp.float64:
+        asymptotic_bound = 13.7
+    elif x.dtype == jnp.float32:
+        asymptotic_bound = 6.9
+    else:
+        raise NotImplementedError(f"_erfcx does not support dtype {x.dtype}.")
+
+    # erfcx ~ (sqrt(pi) x)^-1 sum_k c_k x^-2k, with c_k = (-1)^k (2k-1)!! / 2^k, in
+    # descending degree for `jnp.polyval`.
+    coeffs = jnp.asarray(
+        [
+            7918.06640625,
+            -1055.7421875,
+            162.421875,
+            -29.53125,
+            6.5625,
+            -1.875,
+            0.75,
+            -0.5,
+            1.0,
+        ],
+        dtype=x.dtype,
+    )
+
+    large = x > asymptotic_bound
+    # Clamped so the discarded branch can neither overflow nor poison the gradient.
+    safe_x = jnp.where(large, jnp.ones_like(x), x)
+    direct = jnp.exp(jnp.square(safe_x)) * erfc(safe_x)
+    asymptotic = jnp.polyval(coeffs, 1.0 / jnp.square(x)) / (x * jnp.sqrt(jnp.pi))
+    return jnp.where(large, asymptotic, direct)
+
+
+def _log_abs_z_cdf_div_pdf(z: Float[Array, "..."]) -> Float[Array, "..."]:
+    r"""Compute $`\log(|z| \Phi(z) / \phi(z))`$ for $`z < 0`$.
+
+    Deep in the tail $`\Phi(z)`$ and $`\phi(z)`$ both underflow to zero even though
+    their ratio is perfectly ordinary: at $`z = -40`$ both are around $`10^{-349}`$,
+    while $`|z| \Phi(z) / \phi(z)`$ is $`0.999`$. `_erfcx` cancels the shared
+    $`e^{-z^2/2}`$ factor analytically, so neither tiny number is ever formed.
+
+    Args:
+        z: Strictly negative array.
+
+    Returns:
+        Elementwise; strictly negative, approaching zero from below.
+    """
+    neg_inv_sqrt_2 = -(2.0**-0.5)
+    log_sqrt_pi_div_2 = 0.5 * jnp.log(jnp.pi / 2.0)
+    return jnp.log(_erfcx(neg_inv_sqrt_2 * z) * jnp.abs(z)) + log_sqrt_pi_div_2

tests/test_acquisition_functions/test_acquisition_functions.py

@@ -5,6 +5,7 @@
 )
 from decijax.acquisition_functions.expected_improvement import (
     ExpectedImprovement,
+    LogExpectedImprovement,
 )
 from decijax.acquisition_functions.probability_of_improvement import (
     ProbabilityOfImprovement,
@@ -29,7 +30,12 @@
 )  # Sampling with tfp causes JAX to raise a UserWarning due to some internal logic around jnp.argsort
 @pytest.mark.parametrize(
     "acquisition_function_builder",
-    [ExpectedImprovement, ProbabilityOfImprovement, ThompsonSampling],
+    [
+        ExpectedImprovement,
+        LogExpectedImprovement,
+        ProbabilityOfImprovement,
+        ThompsonSampling,
+    ],
 )
 def test_acquisition_function_no_objective_model_raises_error(
     acquisition_function_builder: type[AbstractSinglePointAcquisitionFunctionBuilder],
@@ -49,7 +55,12 @@ def test_acquisition_function_no_objective_model_raises_error(
 )  # Sampling with tfp causes JAX to raise a UserWarning due to some internal logic around jnp.argsort
 @pytest.mark.parametrize(
     "acquisition_function_builder",
-    [ExpectedImprovement, ProbabilityOfImprovement, ThompsonSampling],
+    [
+        ExpectedImprovement,
+        LogExpectedImprovement,
+        ProbabilityOfImprovement,
+        ThompsonSampling,
+    ],
 )
 def test_model_without_required_capability_raises_error(
     acquisition_function_builder: type[AbstractSinglePointAcquisitionFunctionBuilder],
@@ -66,7 +77,12 @@ def test_model_without_required_capability_raises_error(
 
 @pytest.mark.parametrize(
     "acquisition_function_builder",
-    [ExpectedImprovement, ProbabilityOfImprovement, ThompsonSampling],
+    [
+        ExpectedImprovement,
+        LogExpectedImprovement,
+        ProbabilityOfImprovement,
+        ThompsonSampling,
+    ],
 )
 @pytest.mark.parametrize(
     "test_target_function",