Reworking Minibatching for ADVI

pymc_extras/inference/advi/__init__.py

@@ -5,11 +5,32 @@
     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,
+)
+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",
 ]

pymc_extras/inference/advi/compile.py

@@ -1,17 +1,142 @@
 from typing import Protocol
 
 import numpy as np
+import pytensor
 
 from pymc import Model, compile
+from pymc.pytensorf import rewrite_pregrad
+from pytensor import config
+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.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_training_fn(
+    model: Model,
+    guide: AutoGuideModel,
+    draws: int = 1,
+    path_derivative_gradient: bool = True,
+    logp_scalings: dict | None = None,
+    **compile_kwargs,
+) -> TrainingFn:
+    # draws is a compile-time constant: backends like JAX cannot handle inputs that
+    # determine the shapes of random variables
+    params = guide.params
+    inputs = list(params)
+
+    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)
+
+    negative_elbo_grads = pt.grad(rewrite_pregrad(negative_elbo), wrt=params)
+
+    compile_kwargs.setdefault("trust_input", True)
+
+    f_loss_dloss = compile(
+        inputs=inputs, outputs=[negative_elbo, *negative_elbo_grads], **compile_kwargs
+    )
+
+    return f_loss_dloss
+
+
+def compile_svi_step_fn(
+    model: Model,
+    guide: AutoGuideModel,
+    draws: int = 1,
+    path_derivative_gradient: bool = True,
+    logp_scalings: dict | None = None,
+    clip_norm: float | None = 10.0,
+    beta1: float = 0.9,
+    beta2: float = 0.999,
+    epsilon: float = 1e-8,
+    **compile_kwargs,
+) -> tuple[TrainingFn, dict[str, SharedVariable]]:
+    """Compile one full SVI step, with clipped-Adam updates applied in-graph.
+
+    The guide parameters and the optimizer state live in shared variables that the
+    compiled function updates in place. Its only input is the learning rate and its
+    only output the negative ELBO estimate, so no parameters or gradients round-trip
+    through Python during training.
+
+    Returns
+    -------
+    step_fn :
+        Compiled function ``step_fn(learning_rate) -> negative_elbo``.
+    shared_params : dict
+        Maps each guide parameter name to the shared variable holding its value.
+    """
+    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)
+
+    if clip_norm is not None:
+        global_norm = pt.sqrt(pt.sum([pt.sum(pt.square(g)) for g in grads]))
+        scale = pt.minimum(1.0, clip_norm / (global_norm + 1e-12))
+        grads = [g * scale for g in grads]
+
+    learning_rate = pt.scalar("learning_rate", dtype=config.floatX)
+    t = pytensor.shared(np.zeros((), dtype="int64"), name="adam_t")
+    t_new = t + 1
+    # The bias-correction powers `beta**t_new` mix a (weak) python float base with an int64
+    # exponent, which pytensor resolves to float64. Under floatX=float32 that upcasts param_new
+    # to float64 and the shared-variable update fails the dtype check. Compute the powers in
+    # floatX so the whole update stays in the parameter dtype.
+    t_new_float = t_new.astype(config.floatX)
+    updates = {t: t_new}
+    for shared_param, grad in zip(shared_params, grads):
+        value = shared_param.get_value(borrow=True)
+        m = pytensor.shared(np.zeros_like(value), name=f"adam_m_{shared_param.name}")
+        v = pytensor.shared(np.zeros_like(value), name=f"adam_v_{shared_param.name}")
+        m_new = beta1 * m + (1 - beta1) * grad
+        v_new = beta2 * v + (1 - beta2) * pt.square(grad)
+        m_hat = m_new / (1 - beta1**t_new_float)
+        v_hat = v_new / (1 - beta2**t_new_float)
+        param_new = shared_param - learning_rate * m_hat / (pt.sqrt(v_hat) + epsilon)
+        updates.update({m: m_new, v: v_new, shared_param: param_new})
+
+    compile_kwargs.setdefault("trust_input", True)
+
+    step_fn = compile(
+        inputs=[learning_rate], 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,97 @@
+from __future__ import annotations
+
+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,
+    convergence_window: int | None = 200,
+    relative_tolerance: float = 1e-3,
+    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. Training may stop
+        earlier, controlled by ``convergence_window`` and ``relative_tolerance``.
+    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,
+        clipped Adam on a :func:`linear_onecycle_schedule` peaking at 0.008 over
+        ``n_steps`` is compiled *into* the step function (fast path); passing an explicit
+        optimizer uses the Python-side update loop instead.
+    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.
+    convergence_window : int, optional
+        Number of steps per convergence window, by default 200. Set to None to always
+        run for ``n_steps``.
+    relative_tolerance : float, optional
+        Relative loss change between consecutive windows under which training stops,
+        by default 1e-3.
+    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,
+        convergence_window=convergence_window,
+        relative_tolerance=relative_tolerance,
+        model=model,
+        backend=backend,
+        compile_kwargs=compile_kwargs,
+        random_seed=init_seed,
+    )
+    state = trainer.fit(n_steps, random_seed=train_seed)
+    idata = trainer.sample_posterior(draws, 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