Add Renyi alpha divergence to variational inference and tests (#849)

* Add Renyi alpha divergence to variational inference and tests

* Add coverage for Gaussian VI objective helper

---------

Co-authored-by: Junpeng Lao <junpenglao@gmail.com>

Author: OleksiiBevza

blackjax/vi/_gaussian_vi.py

@@ -11,13 +11,86 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Shared ELBO optimization step for Gaussian VI variants (MFVI, FRVI)."""
-from typing import Callable
+"""Shared Gaussian VI optimization step for:
+   * mean field variational inference (MFVI)
+   * full rank variational inference (FRVI)"""
+from dataclasses import dataclass
+from typing import Callable, Union
 
 import jax
+import jax.numpy as jnp
+import jax.scipy as jsp
 from optax import GradientTransformation, OptState
 
 
+@dataclass(frozen=True)
+class KL:
+    """standard reverse-KL objective"""
+
+    pass
+
+
+@dataclass(frozen=True)
+class RenyiAlpha:
+    """Rényi alpha objective.
+
+    Notes
+    -----
+    A smooth interpolation from the evidence lower-bound to the
+    log (marginal) likelihood that is controlled by the value of alpha
+    that parametrises the divergence.
+    """
+
+    alpha: float
+
+
+Objective = Union[KL, RenyiAlpha]
+
+
+def _objective_value_from_log_ratio(
+    log_ratio: jax.Array,
+    objective: Objective,
+) -> jax.Array:
+    """Returns a scalar loss to minimize from the given log-ratio array and
+    supports two objective types.:
+
+    * KL: returns mean of the log-ratio, corresponding to KL divergence loss
+    * RenyiAlpha: returns negative Monte Carlo Rényi variational bound.
+      For alpha = 1.0 it recovers the reverse-KL objective.
+      For other alpha values, it computes:
+      (logsumexp((alpha - 1) * log_ratio) - log(N)) / (alpha - 1)
+      where N is the number of samples.
+
+     Parameters
+     ----------
+     log_ratio: A JAX array of log-ratio values (log q - log p)
+     objective: An instance of objective (KL or RenyiAlpha)
+
+     Returns
+     -------
+     A scalar JAX array representing the loss value to be minimized.
+
+    """
+    if isinstance(objective, KL):
+        return jnp.mean(log_ratio)
+
+    if isinstance(objective, RenyiAlpha):
+        alpha = objective.alpha
+
+        # for alpha = 1.0 it recovers the reverse-KL objective.
+        if alpha == 1.0:
+            return jnp.mean(log_ratio)
+
+        # negative Monte Carlo Renyi variational bound:
+        # -L_hat_alpha = (1 / (alpha - 1)) * log mean(exp((alpha - 1) * (logq - logp)))
+        scaled = (alpha - 1.0) * log_ratio
+        return (jsp.special.logsumexp(scaled) - jnp.log(log_ratio.shape[0])) / (
+            alpha - 1.0
+        )
+
+    raise TypeError(f"Unsupported objective type: {type(objective)!r}")
+
+
 def _elbo_step(
     rng_key,
     parameters: tuple,
@@ -27,13 +100,15 @@ def _elbo_step(
     sample_fn: Callable,
     logq_fn: Callable,
     num_samples: int,
-    stl_estimator: bool,
+    objective: Objective = KL(),
+    stl_estimator: bool = True,
 ) -> tuple[tuple, OptState, float]:
-    """Single ELBO optimization step shared by Gaussian VI variants.
+    """Single Gaussian VI optimization step shared by MFVI and FRVI.
 
-    Computes the KL divergence ``E_q[log q - log p]`` via Monte Carlo,
-    differentiates with respect to ``parameters``, and applies one optimizer
-    update.
+    Single step of variational optimisation (ELBO or Renyi bound)
+    shared by Gaussian VI variants. Computes a variational loss
+    (KL or Renyi) via Monte Carlo, differentiates with respect to
+    ``parameters``, and applies one optimizer update.
 
     Parameters
     ----------
@@ -55,6 +130,8 @@ def _elbo_step(
         function of the current approximation given its parameters.
     num_samples
         Number of Monte Carlo samples used to estimate the ELBO.
+    objective
+        The variational objective (KL or Rényi). Defaults to KL.
     stl_estimator
         If ``True``, apply ``stop_gradient`` to the parameters used in
         ``logq_fn`` (stick-the-landing estimator). Gradients still flow
@@ -66,21 +143,29 @@ def _elbo_step(
         Updated variational parameters after one optimizer step.
     new_opt_state
         Updated optimizer state.
-    elbo
-        Current ELBO estimate (scalar).
+    loss
+        Current estimate of the variational loss (scalar).
 
     """
 
-    def kl_divergence_fn(parameters):
+    if stl_estimator and isinstance(objective, RenyiAlpha) and objective.alpha != 1.0:
+        raise ValueError(
+            "stl_estimator is currently only supported with KL() or "
+            "RenyiAlpha(alpha=1.0). Use stl_estimator=False for "
+            "RenyiAlpha(alpha != 1.0)."
+        )
+
+    def objective_fn(parameters):
         z = sample_fn(rng_key, parameters, num_samples)
         logq_parameters = (
             jax.lax.stop_gradient(parameters) if stl_estimator else parameters
         )
         logq = jax.vmap(logq_fn(logq_parameters))(z)
         logp = jax.vmap(logdensity_fn)(z)
-        return (logq - logp).mean()
+        log_ratio = logq - logp
+        return _objective_value_from_log_ratio(log_ratio, objective)
 
-    elbo, elbo_grad = jax.value_and_grad(kl_divergence_fn)(parameters)
-    updates, new_opt_state = optimizer.update(elbo_grad, opt_state, parameters)
+    objective_value, objective_grad = jax.value_and_grad(objective_fn)(parameters)
+    updates, new_opt_state = optimizer.update(objective_grad, opt_state, parameters)
     new_parameters = jax.tree.map(lambda p, u: p + u, parameters, updates)
-    return new_parameters, new_opt_state, elbo
+    return new_parameters, new_opt_state, objective_value

blackjax/vi/fullrank_vi.py

@@ -21,9 +21,11 @@
 
 from blackjax.base import VIAlgorithm
 from blackjax.types import Array, ArrayLikeTree, ArrayTree, PRNGKey
-from blackjax.vi._gaussian_vi import _elbo_step
+from blackjax.vi._gaussian_vi import KL, Objective, RenyiAlpha, _elbo_step
 
 __all__ = [
+    "KL",
+    "RenyiAlpha",
     "FRVIState",
     "FRVIInfo",
     "sample",
@@ -88,6 +90,7 @@ def step(
     logdensity_fn: Callable,
     optimizer: GradientTransformation,
     num_samples: int = 5,
+    objective: Objective = KL(),
     stl_estimator: bool = True,
 ) -> tuple[FRVIState, FRVIInfo]:
     """Approximate the target density using the full-rank Gaussian approximation.
@@ -106,6 +109,9 @@ def step(
         The number of samples that are taken from the approximation
         at each step to compute the Kullback-Leibler divergence between
         the approximation and the target log-density.
+    objective:
+        The variational objective to minimize. `KL()` by default or
+        `RenyiAlpha(alpha)`. For alpha = 1, Renyi reduces to KL.
     stl_estimator
         Whether to use the stick-the-landing (STL) gradient estimator
         :cite:p:`roeder2017sticking`. Reduces gradient variance by removing
@@ -137,7 +143,8 @@ def logq_fn(parameters):
         sample_fn,
         logq_fn,
         num_samples,
-        stl_estimator,
+        objective=objective,
+        stl_estimator=stl_estimator,
     )
     new_state = FRVIState(new_parameters[0], new_parameters[1], new_opt_state)
     return new_state, FRVIInfo(elbo)
@@ -168,6 +175,8 @@ def as_top_level_api(
     logdensity_fn: Callable,
     optimizer: GradientTransformation,
     num_samples: int = 100,
+    objective: Objective = KL(),
+    stl_estimator: bool = True,
 ):
     """High-level implementation of Full-Rank Variational Inference.
 
@@ -180,6 +189,12 @@ def as_top_level_api(
         Optax optimizer to use to optimize the ELBO.
     num_samples
         Number of samples to take at each step to optimize the ELBO.
+    objective
+        The variational objective to minimize. `KL()` by default or
+        `RenyiAlpha(alpha)`. For alpha = 1, Renyi reduces to KL.
+    stl_estimator
+        Whether to use STL gradient estimator.
+        Only supported when `objective` is `KL()` or `RenyiAlpha(alpha=1.0)`.
 
     Returns
     -------
@@ -191,7 +206,15 @@ def init_fn(position: ArrayLikeTree):
         return init(position, optimizer)
 
     def step_fn(rng_key: PRNGKey, state: FRVIState) -> tuple[FRVIState, FRVIInfo]:
-        return step(rng_key, state, logdensity_fn, optimizer, num_samples)
+        return step(
+            rng_key,
+            state,
+            logdensity_fn,
+            optimizer,
+            num_samples,
+            objective=objective,
+            stl_estimator=stl_estimator,
+        )
 
     def sample_fn(rng_key: PRNGKey, state: FRVIState, num_samples: int):
         return sample(rng_key, state, num_samples)

blackjax/vi/meanfield_vi.py

@@ -20,9 +20,11 @@
 
 from blackjax.base import VIAlgorithm
 from blackjax.types import ArrayLikeTree, ArrayTree, PRNGKey
-from blackjax.vi._gaussian_vi import _elbo_step
+from blackjax.vi._gaussian_vi import KL, Objective, RenyiAlpha, _elbo_step
 
 __all__ = [
+    "KL",
+    "RenyiAlpha",
     "MFVIState",
     "MFVIInfo",
     "sample",
@@ -74,6 +76,7 @@ def step(
     logdensity_fn: Callable,
     optimizer: GradientTransformation,
     num_samples: int = 5,
+    objective: Objective = KL(),
     stl_estimator: bool = True,
 ) -> tuple[MFVIState, MFVIInfo]:
     """Approximate the target density using the mean-field approximation.
@@ -92,6 +95,9 @@ def step(
         The number of samples that are taken from the approximation
         at each step to compute the Kullback-Leibler divergence between
         the approximation and the target log-density.
+    objective
+        The variational objective to minimize. `KL()` by default or
+        `RenyiAlpha(alpha)`. For alpha = 1, Renyi reduces to KL.
     stl_estimator
         Whether to use the stick-the-landing (STL) gradient estimator
         :cite:p:`roeder2017sticking`. The STL estimator has lower gradient
@@ -120,7 +126,8 @@ def logq_fn(parameters):
         sample_fn,
         logq_fn,
         num_samples,
-        stl_estimator,
+        objective=objective,
+        stl_estimator=stl_estimator,
     )
     new_state = MFVIState(new_parameters[0], new_parameters[1], new_opt_state)
     return new_state, MFVIInfo(elbo)
@@ -140,7 +147,7 @@ def sample(rng_key: PRNGKey, state: MFVIState, num_samples: int = 1):
 
     Returns
     -------
-    A PyTree of samples with leading dimension ``num_samples``.
+    A PyTree of samples with leading dimension ``num_samples``
     """
     return _sample(rng_key, state.mu, state.rho, num_samples)
 
@@ -149,18 +156,26 @@ def as_top_level_api(
     logdensity_fn: Callable,
     optimizer: GradientTransformation,
     num_samples: int = 100,
+    objective: Objective = KL(),
+    stl_estimator: bool = True,
 ):
-    """High-level implementation of Mean-Field Variational Inference.
+    """High-level implementation of Mean-Field Variational Inference
 
-    Parameters
+     Parameters
     ----------
     logdensity_fn
         A function that represents the log-density function associated with
         the distribution we want to sample from.
     optimizer
-        Optax optimizer to use to optimize the ELBO.
+        Optax optimizer to use to optimize the variational objective.
     num_samples
         Number of samples to take at each step to optimize the ELBO.
+    objective
+        The variational objective to minimize. `KL()` by default or
+        `RenyiAlpha(alpha)`. For a = 1, Renyi reduces to KL.
+    stl_estimator
+        Whether to use the STL gradient estimator.
+        Only supported when `objective` is `KL()` or `RenyiAlpha(alpha=1.0)`.
 
     Returns
     -------
@@ -172,7 +187,15 @@ def init_fn(position: ArrayLikeTree):
         return init(position, optimizer)
 
     def step_fn(rng_key: PRNGKey, state: MFVIState) -> tuple[MFVIState, MFVIInfo]:
-        return step(rng_key, state, logdensity_fn, optimizer, num_samples)
+        return step(
+            rng_key,
+            state,
+            logdensity_fn,
+            optimizer,
+            num_samples,
+            objective=objective,
+            stl_estimator=stl_estimator,
+        )
 
     def sample_fn(rng_key: PRNGKey, state: MFVIState, num_samples: int):
         return sample(rng_key, state, num_samples)