Reworking Minibatching for ADVI

pymc_extras/inference/advi/__init__.py

@@ -5,11 +5,38 @@
     AutoMultivariateNormal,
     get_value_shapes_and_dims,
 )
+from pymc_extras.inference.advi.fit import fit_advi
+from pymc_extras.inference.advi.optimizers import (
+    GradientTransformation,
+    adam,
+    apply_updates,
+    chain,
+    clip_by_global_norm,
+    clipped_adam,
+    linear_onecycle_schedule,
+    rmsprop,
+    scale_by_rmsprop,
+    sgd,
+)
+from pymc_extras.inference.advi.training import SVIState, Trainer
 
 __all__ = [
     "AutoDiagonalNormal",
     "AutoGuideModel",
     "AutoLowRankMultivariateNormal",
     "AutoMultivariateNormal",
+    "GradientTransformation",
+    "SVIState",
+    "Trainer",
+    "adam",
+    "apply_updates",
+    "chain",
+    "clip_by_global_norm",
+    "clipped_adam",
+    "fit_advi",
     "get_value_shapes_and_dims",
+    "linear_onecycle_schedule",
+    "rmsprop",
+    "scale_by_rmsprop",
+    "sgd",
 ]

pymc_extras/inference/advi/compile.py

@@ -1,17 +1,88 @@
 from typing import Protocol
 
 import numpy as np
+import pytensor
 
 from pymc import Model, compile
+from pymc.pytensorf import rewrite_pregrad
+from pytensor import tensor as pt
+from pytensor.compile.sharedvalue import SharedVariable
+from pytensor.graph.replace import graph_replace
 
 from pymc_extras.inference.advi.autoguide import AutoGuideModel
+from pymc_extras.inference.advi.objective import advi_objective, get_logp_logq
+from pymc_extras.inference.advi.optimizers import GradientTransformation
 from pymc_extras.inference.advi.pytensorf import vectorize_random_graph
 
 
 class SamplingFn(Protocol):
     def __call__(self, *params: np.ndarray) -> tuple[np.ndarray, ...]: ...
 
 
+class TrainingFn(Protocol):
+    def __call__(self, *params: np.ndarray) -> tuple[np.ndarray, ...]: ...
+
+
+def compile_svi_step_fn(
+    model: Model,
+    guide: AutoGuideModel,
+    optimizer: GradientTransformation,
+    draws: int = 1,
+    path_derivative_gradient: bool = True,
+    logp_scalings: dict | None = None,
+    **compile_kwargs,
+) -> tuple[TrainingFn, dict[str, SharedVariable]]:
+    """Compile one full SVI step, with optimizer updates applied in-graph.
+
+    The guide parameters and the optimizer state live in shared variables that the
+    compiled function updates in place.  It takes no inputs and returns only the
+    negative ELBO estimate, so no parameters or gradients round-trip through Python
+    during training.
+
+    Returns
+    -------
+    step_fn :
+        Compiled function ``step_fn() -> negative_elbo``.
+    shared_params : dict
+        Maps each guide parameter name to the shared variable holding its value.
+    """
+    if optimizer.pytensor is None:
+        raise ValueError(
+            f"The optimizer {optimizer} does not have a PyTensor implementation "
+            "and cannot be compiled into the step function."
+        )
+
+    logp, logq = get_logp_logq(
+        model,
+        guide,
+        path_derivative_gradient=path_derivative_gradient,
+        logp_scalings=logp_scalings,
+    )
+    scalar_negative_elbo = advi_objective(logp, logq)
+    [negative_elbo_draws] = vectorize_random_graph([scalar_negative_elbo], batch_draws=draws)
+    negative_elbo = negative_elbo_draws.mean(axis=0)
+
+    params_to_shared = {
+        param: pytensor.shared(np.asarray(value), name=param.name)
+        for param, value in guide.params_init_values.items()
+    }
+    [negative_elbo] = graph_replace([negative_elbo], replace=params_to_shared)
+    shared_params = list(params_to_shared.values())
+
+    grads = pt.grad(rewrite_pregrad(negative_elbo), wrt=shared_params)
+
+    new_grads, updates = optimizer.pytensor(grads, shared_params)
+
+    for param, grad in zip(shared_params, new_grads):
+        updates[param] = param + grad
+
+    compile_kwargs.setdefault("trust_input", True)
+
+    step_fn = compile(inputs=[], outputs=negative_elbo, updates=updates, **compile_kwargs)
+
+    return step_fn, {param.name: shared for param, shared in params_to_shared.items()}
+
+
 def compile_sampling_fn(
     model: Model, guide: AutoGuideModel, draws: int, **compile_kwargs
 ) -> SamplingFn:

pymc_extras/inference/advi/fit.py

@@ -0,0 +1,81 @@
+import numpy as np
+import xarray as xr
+
+from pymc import Model, modelcontext
+from xarray import DataTree
+
+from pymc_extras.inference.advi.optimizers import GradientTransformation
+from pymc_extras.inference.advi.training import Trainer
+
+
+def fit_advi(
+    model: Model | None = None,
+    *,
+    n_steps: int = 10_000,
+    n_particles: int = 1,
+    draws: int = 1_000,
+    optimizer: GradientTransformation | None = None,
+    path_derivative_gradient: bool = True,
+    random_seed=None,
+    backend: str | None = None,
+    compile_kwargs: dict | None = None,
+) -> DataTree:
+    """Fit a model with automatic differentiation variational inference (ADVI).
+
+    Fits a mean-field normal approximation to the model posterior in the unconstrained
+    space, then returns posterior draws from the fitted guide. A thin wrapper around
+    :class:`~pymc_extras.inference.advi.training.Trainer` with its default guide.
+
+    Parameters
+    ----------
+    model : Model, optional
+        The PyMC model to fit. If None, the model is inferred from context.
+    n_steps : int, optional
+        Maximum number of optimization steps, by default 10_000.
+    n_particles : int, optional
+        Number of guide draws per step used to estimate the ELBO gradient, by default 1.
+    draws : int, optional
+        Number of posterior draws to sample from the fitted guide, by default 1_000.
+    optimizer : GradientTransformation, optional
+        An optax-like optimizer (actual optax optimizers are compatible). By default,
+        :func:`clipped_adam` is used.
+    path_derivative_gradient : bool, optional
+        Whether to use the lower-variance path-derivative ("sticking the landing")
+        gradient estimator, by default True. It is an unbiased variance reduction (it changes
+        only the gradient, not the ELBO); numpyro's ``Trace_ELBO`` does not offer it.
+    random_seed : optional
+        Seed for the guide initialization, the training draws, and the posterior draws.
+    backend : str, optional
+        PyTensor backend to compile the training and sampling functions with
+        (e.g. "numba", "jax", "c"). Mutually exclusive with ``compile_kwargs["mode"]``.
+    compile_kwargs : dict, optional
+        Additional kwargs passed to pytensor compilation.
+
+    Returns
+    -------
+    DataTree
+        Posterior draws from the fitted guide, with the negative loss history in the
+        ``fit`` group (as ``elbo``).
+    """
+    model = modelcontext(model)
+
+    if random_seed is not None:
+        rng = np.random.default_rng(random_seed)
+        init_seed, train_seed, sampling_seed = (int(s) for s in rng.integers(2**30, size=3))
+    else:
+        init_seed = train_seed = sampling_seed = None
+
+    trainer = Trainer(
+        optimizer=optimizer,
+        n_particles=n_particles,
+        path_derivative_gradient=path_derivative_gradient,
+        backend=backend,
+        compile_kwargs=compile_kwargs,
+        random_seed=init_seed,
+    )
+    state = trainer.fit(n_steps, model=model, random_seed=train_seed)
+    idata = trainer.sample_posterior(draws, model=model, random_seed=sampling_seed)
+    idata["fit"] = DataTree(
+        dataset=xr.Dataset({"elbo": ("step", -np.asarray(state.loss_history, dtype=float))})
+    )
+    return idata

pymc_extras/inference/advi/objective.py

@@ -1,13 +1,20 @@
 from __future__ import annotations
 
+import pytensor.tensor as pt
+
 from pymc import Model
 from pytensor.graph.replace import graph_replace
 from pytensor.tensor import TensorVariable
 
 from pymc_extras.inference.advi.autoguide import AutoGuideModel
 
 
-def get_logp_logq(model: Model, guide: AutoGuideModel, path_derivative_gradient: bool = True):
+def get_logp_logq(
+    model: Model,
+    guide: AutoGuideModel,
+    path_derivative_gradient: bool = True,
+    logp_scalings: dict[TensorVariable, float] | None = None,
+):
     """
     Compute the log probability of the model and the guide, evaluated under draws from the guide.
 
@@ -23,6 +30,10 @@ def get_logp_logq(model: Model, guide: AutoGuideModel, path_derivative_gradient:
         change the value of logq, only its gradients: the score-function term, which has zero
         expectation, is dropped, yielding the lower-variance path-derivative gradient
         estimator of _[1] (also known as "sticking the landing").
+    logp_scalings : dict, optional
+        Maps model variables to a factor multiplying their logp term, e.g. the
+        ``total_size / batch_size`` rescaling that makes a minibatch log-likelihood an
+        unbiased estimate of the full-data one.
 
     Returns
     -------
@@ -43,7 +54,16 @@ def get_logp_logq(model: Model, guide: AutoGuideModel, path_derivative_gradient:
         if rv not in model.observed_RVs
     }
 
-    logp = graph_replace(model.logp(), inputs_to_guide_rvs)
+    if logp_scalings:
+        logps = model.logp(sum=False)
+        all_vars = model.free_RVs + model.observed_RVs + model.potentials
+        scales = pt.constant([logp_scalings.get(var, 1.0) for var in all_vars])
+        summed_logps = pt.stack([pt.sum(logp) for logp in logps])
+        model_logp = pt.dot(scales, summed_logps)
+    else:
+        model_logp = model.logp()
+
+    logp = graph_replace(model_logp, inputs_to_guide_rvs)
     logq = guide.stochastic_logq(path_derivative_gradient=path_derivative_gradient)
 
     return logp, logq