Haar and DCT Time Reparameterization in NumPyro

Prepare Notebook

We need NumPyro 0.22 or later and numpyro_forecast 0.4 or later.

from collections.abc import Callable
from enum import StrEnum
from time import perf_counter
from typing import NamedTuple

import arviz as az
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import numpyro.distributions as dist
import polars as pl
import seaborn as sns
from jax import random
from jax.example_libraries.optimizers import exponential_decay
from jaxtyping import Array, Float, ScalarLike, UInt
from numpyro import handlers
from numpyro.distributions.transforms import (
    DiscreteCosineTransform,
    HaarTransform,
    Transform,
)
from numpyro.infer import (
    MCMC,
    NUTS,
    SVI,
    Predictive,
    Trace_ELBO,
    init_to_uniform,
    init_to_value,
)
from numpyro.infer.autoguide import AutoNormal
from numpyro.infer.reparam import DiscreteCosineReparam, HaarReparam, LocScaleReparam
from numpyro.infer.util import log_density
from numpyro.optim import Adam
from numpyro_forecast import (
    Horizon,
    draw_posterior,
    eval_crps,
    forecast,
    innovations,
    predict,
    predict_in_sample,
    time_reparam,
    to_datatree,
)
from numpyro_forecast.datasets import load_bart_weekly
from numpyro_forecast.features import fourier_features
from scipy.optimize import minimize

numpyro.set_host_device_count(n=4)
numpyro.enable_x64()


HDI_PROB = 0.94

az.style.use("arviz-darkgrid")
az.rcParams["stats.ci_kind"] = "hdi"
az.rcParams["stats.ci_prob"] = HDI_PROB
plt.rcParams["figure.figsize"] = [10, 6]
plt.rcParams["figure.dpi"] = 100
plt.rcParams["figure.facecolor"] = "white"

%load_ext jaxtyping
%jaxtyping.typechecker beartype.beartype
%config InlineBackend.figure_format = "retina"

We fix the random key of the notebook. Every sampling call below splits a new key from it.

seed: int = 42
rng_key = random.PRNGKey(seed=seed)

Local Level Model

We start with the simplest model that has the problem. Let \(T\) be the number of time steps and \(y_t\) the observation at time \(t = 1, \ldots, T\). We use the notation of the local level model in the UnobservedComponents class of statsmodels. The observation \(y_t\) is a noisy measurement of a latent level \(\mu_t\), and the level is a random walk with increments \(\eta_t\).

Our target is the posterior of the increments given the observations. We want to know how well a mean-field guide can approximate it.

Let \(\sigma_\varepsilon\) be the scale of the observation noise \(\varepsilon_t\) and \(\sigma_\eta\) the scale of the increments \(\eta_t\). The model is

\[\begin{align*} y_t & = \mu_t + \varepsilon_t, & \varepsilon_t & \sim \text{Normal}(0, \sigma_\varepsilon), \\ \mu_t & = \mu_{t-1} + \eta_t, & \eta_t & \sim \text{Normal}(0, \sigma_\eta), \end{align*}\]

with \(\mu_0 = 0\). Here the second argument of \(\text{Normal}\) is the standard deviation. The statsmodels documentation writes the variances \(\sigma_\varepsilon^2\) and \(\sigma_\eta^2\) instead.

The latent variables are the increments \(\eta = (\eta_1, \ldots, \eta_T)\), not the levels. This is how the Pyro forecasting tutorial and the innovations building block of numpyro_forecast write a random walk: sample white noise and take a cumulative sum, \(\mu_t = \sum_{s \le t} \eta_s\). In vector form the level is \(\mu = L \eta\), where \(L\) is the \(T \times T\) lower-triangular matrix of ones.

In this first part we keep the two scales \(\sigma_\eta\) and \(\sigma_\varepsilon\) fixed. Then the model is linear and Gaussian and we can compute everything analytically. We come back to unknown scales and to a non-Gaussian likelihood in the ridership example below.

First, let’s write the model in NumPyro. We use the names of the Pyro tutorial: the increments are the site drift, \(\sigma_\eta\) is drift_scale and \(\sigma_\varepsilon\) is sigma. The time axis is an event dimension of the drift site. We write a generative model without observations, and we condition it on data later.

T_TOY = 256
DRIFT_SCALE = 0.2
SIGMA = 1.0


def local_level_model(t_obs: int, drift_scale: float, sigma: float) -> None:
    "Generative local level model with the increments as latent variables."
    drift = numpyro.sample(
        "drift", dist.Normal(0.0, drift_scale).expand([t_obs]).to_event(1)
    )
    level = numpyro.deterministic("level", jnp.cumsum(drift))
    numpyro.sample("y", dist.Normal(level, sigma).to_event(1))


toy_args = (T_TOY, DRIFT_SCALE, SIGMA)

Next, we simulate one data set from the model itself. We use \(T = 256\) time steps, \(\sigma_\eta = 0.2\) and \(\sigma_\varepsilon = 1\).

rng_key, rng_subkey = random.split(rng_key)
prior_draw = Predictive(local_level_model, num_samples=1)(rng_subkey, *toy_args)
toy_drift = prior_draw["drift"][0]
toy_level = prior_draw["level"][0]
toy_y = prior_draw["y"][0]

Let’s plot the simulated observations together with the level that generated them.

toy_time = jnp.arange(1, T_TOY + 1)

fig, ax = plt.subplots()
ax.plot(
    toy_time,
    toy_y,
    color="black",
    alpha=0.7,
    label="observations $y_t$",
)
ax.plot(toy_time, toy_level, color="C3", linewidth=2, label="true level $\\mu_t$")
ax.set(xlabel="time step $t$", ylabel="$y_t$")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=2)
fig.suptitle("Simulated Local Level Data", fontsize=18, fontweight="bold");

The level moves slowly relative to the noise, because one increment has a scale of \(0.2\) and the noise a scale of \(1\). Over \(256\) time steps the increments add up, and the level moves over a range of several units. The observations therefore carry a clear signal about the path of the level, but not about any single increment.

Finally, we condition the generative model on the simulated observations with the effect handler handlers.condition. The resulting toy_model is the model that we study in the first part of the notebook.

toy_model = handlers.condition(local_level_model, data={"y": toy_y})

numpyro.render_model(
    toy_model,
    model_args=toy_args,
    render_params=True,
    render_distributions=True,
)

Posterior Geometry

What does the posterior of the increments look like? The contours of a Gaussian density are ellipsoids, and the eigenvectors of \(\Lambda\) are their axes.

A large eigenvalue is a stiff direction, where the data leave little freedom. A small eigenvalue is a loose direction.

Let’s compute the eigendecomposition of \(\Lambda\) numerically. We sort the eigenvalues from the largest to the smallest, so that direction \(k = 0\) is the stiffest one, and we fix the sign of each eigenvector so that its first entry is positive. We also check that direction \(k\) changes sign \(k\) times, which confirms that the directions are sorted from the smoothest to the most oscillating.

eigenvalues, eigenvectors = jnp.linalg.eigh(precision)
order = jnp.argsort(eigenvalues)[::-1]
eigenvalues = eigenvalues[order]
eigenvectors = eigenvectors[:, order]
eigenvectors = eigenvectors * jnp.sign(eigenvectors[0])

sign_changes = (jnp.diff(jnp.sign(eigenvectors), axis=0) != 0).sum(axis=0)
np.testing.assert_array_equal(sign_changes, jnp.arange(T_TOY))

The check passes. Now we plot the four stiffest directions (left) and, for every direction, the posterior variance relative to the prior variance (right).

Along direction \(k\) the posterior variance is \(1 / \text{eig}_k(\Lambda)\) and the prior variance is \(\sigma_\eta^2\). We call this ratio the shrinkage of direction \(k\). A ratio close to \(1\) means that the data say nothing about that direction, and a ratio close to \(0\) means that the data determine it.

shrinkage = 1 / (DRIFT_SCALE**2 * eigenvalues)

fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(15, 5))

for i in range(4):
    axes[0].plot(toy_time, eigenvectors[:, i], color=f"C{i}", label=f"$v_{i}$")

axes[0].set(
    xlabel="time step $t$",
    ylabel="weight on the increment $\\eta_t$",
    title="The four stiffest directions",
)
axes[0].legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=4)
axes[1].plot(toy_time, shrinkage, "o-", color="black", markersize=3)
axes[1].axhline(0.5, color="gray", linestyle="--", label="half of the prior variance")
axes[1].set(
    xscale="log",
    xlabel="direction $k + 1$ (log scale)",
    ylabel="posterior variance / prior variance",
    title="Shrinkage by direction",
    ylim=(0, 1.05),
)
axes[1].legend(loc="upper center", bbox_to_anchor=(0.5, -0.15))
fig.suptitle(
    "Principal Directions of the Posterior of the Increments",
    fontsize=18,
    fontweight="bold",
);

The left panel shows the weights of the four stiffest directions on the increments. The right panel shows how much of the prior variance is left along each direction.

The stiffest direction \(v_0\) moves all the increments in the same direction. This shifts the level at every later time step, and every observation constrains it. The next directions look like cosines of increasing frequency. The oscillating directions barely move the level, because consecutive increments cancel out, and the data do not constrain them. Only \(16\) of the \(256\) directions lose more than half of their prior variance.

The ellipsoid is therefore very elongated. The ratio between the largest and the smallest eigenvalue of \(\Lambda\), its condition number, is about \(1057\):

eigenvalues.max() / eigenvalues.min()

Array(1057.01655449, dtype=float64)

Here is why this matters for the reparameterization. In the basis of the eigenvectors the posterior precision is diagonal, so a mean-field guide is exact there.

The right panel of the figure above adds that the data only constrain the few smoothest directions. The other directions stay close to the prior, which is isotropic, so any orthonormal basis of those directions is as good as the eigenvectors, as long as it does not mix them with the stiff ones. A fixed basis therefore only has to separate the smooth directions from the oscillating ones, and this is what the two transforms of the next section do.

We cannot draw the ellipsoid in \(256\) dimensions, but we can draw two-dimensional slices through its center. A slice fixes all the other coordinates at their posterior mean. The conditional posterior of the two remaining coordinates is Gaussian, and its precision is the corresponding \(2 \times 2\) block of \(\Lambda\).

We draw the contour at one standard deviation, that is the points \(x\) with \(x^\top \Lambda_{\text{block}}\, x = 1\), for two kinds of slices:

• Stiffest and loosest directions. The plane of \(v_{T-1}\) and \(v_0\), where the block is diagonal.

• Adjacent increments. The plane of two increments \(\eta_t\) and \(\eta_{t+1}\), for \(t = 1\), \(128\) and \(255\).

The dashed circle is the same contour for the prior, with radius \(\sigma_\eta = 0.2\).

def contour_2d(
    block: Float[Array, "2 2"], num_points: int = 200
) -> Float[Array, "2 n"]:
    "Points x with x^T block x = 1, the one-sigma contour of a 2D Gaussian."
    angles = jnp.linspace(0, 2 * jnp.pi, num_points)
    circle = jnp.stack([jnp.cos(angles), jnp.sin(angles)])
    # If block = C C^T (Cholesky), then x = C^{-T} u satisfies x^T block x = u^T u = 1.
    cholesky = jnp.linalg.cholesky(block)
    return jnp.linalg.solve(cholesky.T, circle)


principal_block = jnp.diag(eigenvalues[jnp.array([-1, 0])])

fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(15, 7), sharex=True, sharey=True)
axes[0].plot(*contour_2d(principal_block), color="black", label="posterior")

for i, t in enumerate([1, 128, 255]):
    pair = jnp.array([t - 1, t])  # zero-based positions of eta_t and eta_{t+1}
    block = precision[jnp.ix_(pair, pair)]
    axes[1].plot(*contour_2d(block), color=f"C{i}", label=f"$t = {t}$")

for ax in axes:
    ax.plot(
        *contour_2d(jnp.eye(2) / DRIFT_SCALE**2),
        color="gray",
        linestyle="--",
        label="prior",
    )
    ax.set_aspect("equal")

axes[0].set(
    xlabel="coordinate along $v_{T-1}$ (loosest)",
    ylabel="coordinate along $v_0$ (stiffest)",
    title="Stiffest and loosest directions",
)
axes[1].set(
    xlabel="$\\eta_t - \\bar{\\eta}_t$",
    ylabel="$\\eta_{t+1} - \\bar{\\eta}_{t+1}$",
    title="Pairs of adjacent increments",
)
axes[0].legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=2)
axes[1].legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=4)
fig.suptitle("Slices of the Posterior Ellipsoid", fontsize=18, fontweight="bold");

The two panels show slices of the same posterior through different planes.

Along the principal directions the ellipse is a thin sliver. Along adjacent increments it is tilted, and the tilt fades toward the end of the series.

Observe that the axes of the tilted ellipses are diagonals and not coordinate axes. In the full posterior each stiff direction mixes all the increments, not only two of them. This is the cause of the problem: a mean-field guide works coordinate by coordinate.

Orthonormal Time Transforms

The idea of the reparameterization is to rotate the coordinates such that the axes of the ellipsoid become coordinate axes, at least approximately.

Let \(U\) be a \(T \times T\) orthonormal matrix, that is \(U U^\top = I\), and define the new latent variable \(z = U \eta\). We sample \(z\) and we recover the increments as \(\eta = U^\top z\). Three properties make this change of variables convenient:

• The density does not change. The density of \(z\) is the density of \(\eta\) evaluated at \(U^\top z\), times \(|\det U^\top| = 1\). There is no Jacobian term, and the model is unchanged.

• White noise stays white noise. If the increments \(\eta_t\) are independent \(\text{Normal}(0, \sigma_\eta)\) variables, then the coefficients \(z_k\) are independent \(\text{Normal}(0, \sigma_\eta)\) variables too. The prior keeps its simple form.

• The precision is rotated. The posterior precision of \(z\) is \(\Lambda_U = U \Lambda U^\top\).

A rotation does not change methods that treat all directions alike. For example, a full-rank Gaussian guide reaches the same ELBO in every orthonormal basis. The rotation matters for methods that work coordinate by coordinate, such as a mean-field guide.

The best choice of \(U\) is the matrix of eigenvectors of \(\Lambda\), which makes \(\Lambda_U\) diagonal. For the local level model these eigenvectors do not depend on the scales, because \(\Lambda = \sigma_\eta^{-2} I + \sigma_\varepsilon^{-2} L^\top L\) and a multiple of the identity does not change the eigenvectors. In a real model we do not know the eigenvectors: the likelihood is not Gaussian, and other parameters, such as an intercept or a regression, couple with the path.

For a random walk we do know that the stiff directions are smooth and that the loose ones oscillate. Therefore we use a fixed basis that sorts the coordinates from coarse to fine. Pyro offers two, and the first one resembles the stiffest directions of the previous section.

Mean-Field Gap

Now we quantify what a mean-field guide loses in each basis. First, let’s recall what SVI optimizes. For a detailed introduction see PyData Berlin 2025: Introduction to Stochastic Variational Inference with NumPyro.

Variational inference replaces the posterior \(\text{P}(z \mid y)\) with the closest member \(q\) of a simpler family of distributions, the guide. The discrepancy is the Kullback-Leibler (KL) divergence

\[\text{KL}\left(q \,\|\, \text{P}(z \mid y)\right) = \text{E}_{q}\left[\log q(z) - \log \text{P}(z \mid y)\right] ,\]

which is never negative and is zero only when \(q\) equals the posterior. With the natural logarithm we measure it in nats. The expectation is under \(q\), so the guide pays a large penalty where it puts mass and the posterior does not. A guide that cannot match the shape of the posterior therefore stays inside its high-probability region, and it underestimates the uncertainty of each coordinate. A sum of coordinates, such as the level, can still come out too wide, because the guide also drops the negative correlations. We see both effects below.

We cannot compute the KL divergence directly, because the posterior contains the unknown evidence \(\text{P}(y)\). SVI maximizes the evidence lower bound (ELBO) instead, which only needs the joint density \(\text{P}(y, z)\) of the model:

\[\text{ELBO}(q) = \text{E}_{q}\left[\log \text{P}(y, z) - \log q(z)\right] = \log \text{P}(y) - \text{KL}\left(q \,\|\, \text{P}(z \mid y)\right) .\]

The log evidence does not depend on \(q\), so maximizing the ELBO is the same as minimizing the KL divergence.

A mean-field Gaussian guide for \(z\) is a product of independent Normal distributions,

\[q(z) = \prod_{i = 1}^{T} \text{Normal}(z_i \mid m_i, s_i),\]

with one location \(m_i\) and one scale \(s_i\) per coordinate. This is the family that AutoNormal fits.

The log evidence does not depend on the basis either. Therefore the smallest KL divergence that a mean-field guide can reach in a basis sets its best possible ELBO there. We call that divergence the mean-field gap.

For a Gaussian posterior with precision \(\Lambda_U\) we can minimize in closed form. The optimal locations are the entries of the posterior mean and the optimal scales are \(s_i^2 = 1 / (\Lambda_U)_{ii}\). In other words, each factor of the guide gets the conditional variance of its coordinate, which is never larger than the marginal variance (see equation (10.12) in Section 10.1.2 of Bishop (2006)).

Let \(\tilde{\Lambda}_U\) be the precision \(\Lambda_U\) rescaled to have ones on the diagonal, which we call the normalized precision. The gap is

\[\text{gap}(U) = \frac{1}{2} \left( \sum_{i = 1}^{T} \log (\Lambda_U)_{ii} - \log \det \Lambda \right) = - \frac{1}{2} \log \det \tilde{\Lambda}_U .\]

Here is where the formula comes from. The guide and the posterior are both Gaussian, so their KL divergence has a closed form (see the Wikipedia article on the multivariate normal distribution), which we write out in the Mean-Field Check section below. In that expression the locations only appear in a quadratic term, which is zero at the posterior mean, and the variances separate into one term \((\Lambda_U)_{ii}\, s_i^2 - \log s_i^2\) per coordinate.

Each of these terms is smallest at \(s_i^2 = 1 / (\Lambda_U)_{ii}\), where it equals \(1 + \log (\Lambda_U)_{ii}\). Summing over the coordinates and putting the sum back into the KL formula gives \(\text{gap}(U)\).

We make two remarks:

• The gap is zero only for a diagonal precision. This is Hadamard’s inequality: the determinant of a positive definite matrix is at most the product of its diagonal entries, with equality only for a diagonal matrix.

• Only the first term depends on the basis. The determinant of \(\Lambda\) is invariant under rotations. A good basis is one where the product of the diagonal entries of \(\Lambda_U\) is small, that is, where the off-diagonal entries of the normalized precision are close to zero.

Let’s compute the gap for the three coordinate systems. We plot the normalized precision matrices (top row) and a zoom on their first \(16\) coordinates (bottom row). White is zero, which means no coupling, and the diagonal is \(1\) by construction.

def normalize(matrix: Float[Array, "t t"]) -> Float[Array, "t t"]:
    "Rescale a positive definite matrix to have ones on the diagonal."
    scale = jnp.sqrt(jnp.diag(matrix))
    return matrix / jnp.outer(scale, scale)


def mean_field_gap(
    precision: Float[Array, "t t"], rotation: Float[Array, "t t"]
) -> Float[Array, ""]:
    "KL divergence of the best mean-field Gaussian relative to the posterior, in nats."

    rotated = rotation @ precision @ rotation.T
    return 0.5 * (jnp.log(jnp.diag(rotated)).sum() - jnp.linalg.slogdet(precision)[1])


gaps = {name: mean_field_gap(precision, u) for name, u in rotations.items()}

# In the eigenbasis the precision is diagonal and the gap is zero.
np.testing.assert_allclose(mean_field_gap(precision, eigenvectors.T), 0, atol=1e-8)

fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(16, 10.5))

for column, (name, u) in enumerate(rotations.items()):
    normalized = normalize(u @ precision @ u.T)

    for row, size in enumerate([T_TOY, 16]):
        ax = axes[row, column]
        image = ax.imshow(normalized[:size, :size], cmap="RdBu_r", vmin=-1, vmax=1)
        ax.grid(visible=False)
        ax.set(xlabel="coordinate", ylabel="coordinate")

    axes[0, column].set(title=f"{name}: gap = {float(gaps[name]):.1f} nats")
    axes[1, column].set(title=f"{name}: first 16 coordinates")

fig.colorbar(image, ax=axes, shrink=0.6, label="normalized precision")
fig.suptitle(
    "Normalized Posterior Precision in the Three Coordinate Systems",
    fontsize=18,
    fontweight="bold",
);

Each column is the same posterior in one coordinate system, and the title gives the mean-field gap.

In the original coordinates every pair is coupled. The matrix is dense, and the gap is \(187.3\) nats, so a mean-field guide is very far from this posterior. The Haar basis leaves few entries, and the gap drops to \(7.4\) nats.

The DCT leaves even fewer, with a gap of \(1.4\) nats, which is less than \(1\%\) of the original gap. The eigenbasis has zero gap, and the DCT is \(1.4\) nats away from it without any knowledge of the scales.

The entries that remain are large. The zoom shows entries up to about \(0.8\) among the coarsest coefficients. For the DCT they are all in the first row: the coefficient \(z_0\), the average increment, is still coupled to the others. A small gap does not mean that every coordinate is well approximated. We see the consequence for the level below.

Next, we check that this result is not specific to the scale ratio \(\rho = \sigma_\eta / \sigma_\varepsilon = 0.2\) of our example. We compute the gap on a grid of scale ratios.

ratio_grid = jnp.logspace(-2, jnp.log10(3.0), 26)


def gap_by_ratio(rotation: Float[Array, "t t"]) -> Float[Array, " r"]:
    "Mean-field gap in a coordinate system, on the grid of scale ratios."
    return jax.vmap(
        lambda rho: mean_field_gap(posterior_precision(T_TOY, rho, 1.0), rotation)
    )(ratio_grid)


gap_grid = {name: gap_by_ratio(u) for name, u in rotations.items()}

fig, ax = plt.subplots()

for name, values in gap_grid.items():
    ax.plot(ratio_grid, values, color=COLORS[name], label=name)

ax.axvline(
    DRIFT_SCALE / SIGMA,
    color="black",
    linestyle=":",
    label="$\\rho = 0.2$ (our example)",
)
ax.set(
    xscale="log",
    yscale="log",
    xlabel="scale ratio $\\rho = \\sigma_\\eta / \\sigma_\\varepsilon$ (log scale)",
    ylabel="mean-field gap (nats, log scale)",
)
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=4)
fig.suptitle("Mean-Field Gap by Coordinate System", fontsize=18, fontweight="bold");

The figure shows the gap of the best mean-field guide as a function of the scale ratio, on logarithmic axes.

The ranking is the same for every \(\rho\): the original coordinates are the worst, and both rotations help. The gain grows with the scale ratio. For a very small \(\rho\) the posterior is close to the prior and every coordinate system is fine, while for \(\rho \ge 0.1\) the original gap is in the tens or hundreds of nats and the DCT gap stays below \(3\) nats.

The DCT has a smaller gap than Haar for this model, because its rows are closer to the stiffest directions of the posterior.

Local Level Experiments

We now check the derivation of the mean-field gap on the local level model with a fitted guide.

First, we need the model in the three coordinate systems. NumPyro implements the two rotations as the reparameterizers HaarReparam and DiscreteCosineReparam (see the NumPyro documentation). The effect handler handlers.reparam applies a reparameterizer to the sites we name, without changes to the model code.

The reparameterizer replaces the latent site drift with a new latent site, drift_haar or drift_dct, which holds the coefficients \(z = U \eta\). The site drift becomes a deterministic function of the new site. Let’s create the three versions of the conditioned model.

Model = Callable[..., None]

toy_models: dict[Coordinates, Model] = {
    Coordinates.NONE: toy_model,
    Coordinates.HAAR: handlers.reparam(
        toy_model, config={"drift": HaarReparam(dim=-1)}
    ),
    Coordinates.DCT: handlers.reparam(
        toy_model, config={"drift": DiscreteCosineReparam(dim=-1)}
    ),
}

Let’s check that the three versions define the same model: we evaluate the log joint density of the three models at the same point, written in each coordinate system.

log_joint, _ = log_density(toy_model, toy_args, {}, {"drift": toy_drift})

for name, transform in {
    Coordinates.HAAR: HaarTransform(dim=-1),
    Coordinates.DCT: DiscreteCosineTransform(dim=-1),
}.items():
    log_joint_rotated, _ = log_density(
        toy_models[name],
        toy_args,
        {},
        {f"drift_{name}": transform(toy_drift)},
    )
    np.testing.assert_allclose(log_joint_rotated, log_joint, rtol=1e-5)

The three log densities agree. We have three coordinate systems for one model.

Ridership Forecasting Model

Now we move to real data and to a model without a closed-form posterior. We use the example of the Pyro tutorial Forecasting I: univariate, heavy tailed, in the numpyro_forecast version of the univariate forecasting example.

Variational Inference

We fit an AutoNormal guide to each version of the model. We take Adam and the \(50000\) steps from the documentation example. The example uses a constant step size of \(0.005\). We let the step size decay from \(0.005\) to \(0.0005\), so that the parameters settle at the end of the run and do not jitter around the optimum. We collect these settings in the SVIConfig tuple below.

The result of one fit depends on the random key, which initializes the guide and draws the ELBO samples. We therefore fit every version of the model \(5\) times, once per random key, which gives \(15\) fits in total. The three coordinate systems share the same five keys, so that we can compare fits that only differ in the rotation.

We also initialize the location of bias at the mean of the training data. The data are close to \(14.6\) on the log scale, and the default initialization draws every location uniformly in \((-2, 2)\). We study the default initialization in the next subsection.

%%time

NUM_FITS = 5
NUM_STEPS = 50_000

# The step size decays from 0.005 to 0.0005 over the run.
step_schedule = exponential_decay(0.005, NUM_STEPS, 0.1)


class SVIConfig(NamedTuple):
    "Settings of one SVI run of the ridership model."

    init_loc_fn: Callable = init_to_uniform
    step_size: float | Callable = step_schedule
    num_steps: int = NUM_STEPS


class Fit(NamedTuple):
    "Result of one SVI fit."

    guide: AutoNormal
    params: dict[str, Array]
    losses: Float[Array, " steps"]
    seconds: float


def fit_guide(
    rng_key: UInt[Array, "2"],
    model: Model,
    config: SVIConfig,
    covariates: Float[Array, "t_obs features"],
    data: Float[Array, "t_obs 1"],
) -> Fit:
    "Fit a mean-field guide to a version of the ridership model."
    guide = AutoNormal(model, init_loc_fn=config.init_loc_fn)
    svi = SVI(model, guide, Adam(step_size=config.step_size), Trace_ELBO())
    start = perf_counter()
    result = svi.run(rng_key, config.num_steps, covariates, data, progress_bar=False)
    result.losses.block_until_ready()
    return Fit(guide, result.params, result.losses, perf_counter() - start)


def fit_all(config: SVIConfig) -> dict[Coordinates, list[Fit]]:
    "Fit every version of the model on the training data, once per random key."
    return {
        name: [
            fit_guide(key, model, config, covariates_train, y_train) for key in svi_keys
        ]
        for name, model in models.items()
    }


def final_loss(fit: Fit) -> float:
    "Average ELBO loss over the last 1000 steps."
    return float(fit.losses[-1_000:].mean())


init_bias = init_to_value(values={"bias": y_train.mean()})
svi_config = SVIConfig(init_loc_fn=init_bias)

rng_key, *svi_keys = random.split(rng_key, NUM_FITS + 1)
svi_fits = fit_all(svi_config)

CPU times: user 2min 46s, sys: 59.6 s, total: 3min 46s
Wall time: 1min 27s

We plot the ELBO loss (the negative ELBO, in nats, so lower is better) of the first fit in each coordinate system. The loss starts above \(10000\) nats and drops fast in the first \(1000\) steps, so the left panel starts at step \(1000\). The right panel shows the last \(20000\) steps. The legend gives the average loss over the last \(1000\) steps, which we call the final loss.

steps = jnp.arange(NUM_STEPS)
early = (steps >= 1_000) & (steps < 10_000)
late = steps >= NUM_STEPS - 20_000

fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(15, 5))

for name, fits in svi_fits.items():
    losses = fits[0].losses
    axes[0].plot(
        steps[early],
        losses[early],
        color=COLORS[name],
        linewidth=0.5,
        alpha=0.7,
        label=f"{name} (final loss: {final_loss(fits[0]):.0f})",
    )
    axes[1].plot(
        steps[late], losses[late], color=COLORS[name], linewidth=0.5, alpha=0.7
    )

axes[0].set(xlabel="SVI step", ylabel="ELBO loss (nats)", title="Steps 1000 to 10000")
axes[0].legend(loc="upper right")
axes[1].set(xlabel="SVI step", ylabel="ELBO loss (nats)", title="Last 20000 steps")
fig.suptitle("ELBO Loss by Coordinate System", fontsize=18, fontweight="bold");

The three curves are the same guide family fitted to the same model, in three coordinate systems.

The rotated fits reach a much better ELBO. The final loss is about \(-452\) in the original coordinates, \(-530\) with the Haar transform and \(-535\) with the DCT, so the gain is about \(78\) nats with Haar and \(83\) nats with the DCT. The rotated fits are also ahead from step \(1000\) on.

The log evidence is the same for the three models, so a better ELBO means a guide that is closer to the posterior 🚀. We look at what this means for the parameters and for the forecasts in the next section.

Observe that the band of the original fit stays wide in the last \(20000\) steps, although the step size decays from \(0.0013\) to \(0.0005\) over these steps. Each point is a single-sample Monte Carlo estimate of the ELBO, and the variance of that estimate is zero when the guide equals the posterior and large when the guide is far from it. The wide band is therefore one more symptom of the problem that the rotation addresses, and a smaller step size does not remove it.

Each step costs more in the rotated coordinates. Let’s time the fits. The first fit of each model includes the compilation, so we show the other four of each coordinate system. The bar is their mean and each black point is one fit.

fit_seconds = pl.DataFrame(
    [
        {"coordinates": name, "seconds": fit.seconds}
        for name, fits in svi_fits.items()
        for fit in fits[1:]
    ]
)

fig, ax = plt.subplots()
sns.barplot(
    data=fit_seconds,
    x="coordinates",
    y="seconds",
    hue="coordinates",
    palette=COLORS,
    alpha=0.7,
    errorbar=None,
    legend=False,
    ax=ax,
)
sns.stripplot(
    data=fit_seconds, x="coordinates", y="seconds", color="black", jitter=False, ax=ax
)
ax.plot([], [], "s", color="gray", alpha=0.7, markersize=12, label="mean over the fits")
ax.plot([], [], "o", color="black", label="one fit")
ax.legend(loc="upper left")
ax.set(xlabel="coordinate system", ylabel="seconds per fit (after compilation)")
fig.suptitle("Time of a Fit by Coordinate System", fontsize=18, fontweight="bold");

The spread of the points shows how much the time varies between runs.

A fit takes about \(4.4\) seconds in the original coordinates, \(5.9\) with the Haar transform and \(6.3\) with the DCT. A fit in the rotated coordinates takes about two seconds longer and ends about \(80\) nats lower.

Posterior and Forecast Comparison

We now judge the variational fits against a reference. For this model we do not have the exact posterior, so we use NUTS in the original coordinates.

We run \(4\) chains with \(1000\) draws each. The model is harder than the local level model, so we use \(2000\) warmup steps, as the comparison of inference methods does. We use the same initialization of bias as for the guides.

We convert the run to a DataTree with az.from_numpyro, which keeps the sampler statistics, and we check it with az.diagnose. It reports divergences, parameters with fewer than \(100\) effective draws per chain, and parameters with an \(\hat{R}\) above \(1.01\).

%%time

rng_key, rng_subkey = random.split(rng_key)
mcmc = MCMC(
    NUTS(ridership_model, init_strategy=init_bias),
    num_warmup=2_000,
    num_samples=1_000,
    num_chains=4,
    progress_bar=False,
)
mcmc.run(rng_subkey, covariates_train, y_train)
nuts_reference = mcmc.get_samples()
nuts_tree = az.from_numpyro(mcmc)

az.diagnose(nuts_tree);

Divergences
No divergent transitions found.

ESS
Effective sample size satisfactory for all parameters.

R-hat
R-hat values satisfactory for all parameters.

Processing complete, no problems detected.
CPU times: user 3min 27s, sys: 6min 15s, total: 9min 43s
Wall time: 1min 12s

The reference run converges, which is good! There are no divergences, and no parameter is flagged for its effective sample size or its \(\hat{R}\).

Next, we draw \(4000\) samples from each fitted guide with draw_posterior. We pack the draws of the first fit of each coordinate system, and the draws of the reference, as ArviZ DataTree objects with to_datatree from numpyro_forecast.

%%time

# The same key for every fit, so that the draws are paired across coordinate systems.
rng_key, rng_subkey = random.split(rng_key)
svi_posteriors = {
    name: [draw_posterior(rng_subkey, fit.guide, fit.params, 4_000) for fit in fits]
    for name, fits in svi_fits.items()
}

FIT_COLORS = {"NUTS": "black"} | {
    f"SVI: {name}": color for name, color in COLORS.items()
}

# We keep one chain for every tree, so that the four trees share the same `chain` and
# `draw` coordinates.
fit_trees = {
    "NUTS": to_datatree(
        rng_subkey, ridership_model, nuts_reference, y_train, covariates_train
    )
} | {
    f"SVI: {name}": to_datatree(
        rng_subkey, models[name], posteriors[0], y_train, covariates_train
    )
    for name, posteriors in svi_posteriors.items()
}

CPU times: user 30.3 s, sys: 2.05 s, total: 32.3 s
Wall time: 6.83 s

First, we compare the marginal posteriors of the scalar parameters for the first fit. The variational posteriors of drift_scale are very narrow, so we plot cumulative distributions, which stay readable when the widths differ.

scalar_names = ["drift_scale", "sigma", "nu"]

pc = az.plot_dist(
    fit_trees,
    var_names=scalar_names,
    kind="ecdf",
    visuals={"point_estimate_text": False},
    color=list(FIT_COLORS.values()),
    figure_kwargs={"figsize": (15, 4.5)},
)
pc.add_legend("model", title="", loc="outside lower center", ncols=4)
fig = pc.viz["figure"].item()
fig.axes[0].set(ylabel="cumulative probability")
fig.suptitle(
    "Posterior of the Scalar Parameters: NUTS and SVI", fontsize=18, fontweight="bold"
);

Each curve is the empirical cumulative distribution of the draws of one fit. The point is its mean and the horizontal line is its \(94\%\) HDI. A steep curve is a narrow posterior. The black curve is the NUTS reference.

The NUTS reference puts the degrees of freedom nu around \(1.6\), far from a Gaussian likelihood. In the original coordinates the fit shrinks the level: the scale of the increments drift_scale is three times smaller than the NUTS mean, and the noise scale sigma is about twice as large. This is the pattern of the local level model with unknown scales.

The rotated fits are much closer to NUTS. Their estimates of drift_scale are inside the NUTS interval, and sigma and nu move about two thirds of the way toward the reference.

A difference remains: sigma and nu are still above the reference, and the variational posterior of drift_scale is about \(20\) times narrower than the NUTS one. A mean-field guide also ignores the dependence between the global parameters and the latent path, and the rotation cannot help there, because it only changes the coordinates of the path.

Next, we compare the posterior mean of the level, which is the quantity that the forecast extends.

level_draws = {"NUTS": jnp.asarray(nuts_reference["level"])[..., 0]} | {
    f"SVI: {name}": jnp.asarray(posteriors[0]["level"])[..., 0]
    for name, posteriors in svi_posteriors.items()
}

level_means = pl.DataFrame(
    [
        {"week": int(week), "fit": name, "mean": float(value)}
        for name, draws in level_draws.items()
        for week, value in zip(weeks_train, draws.mean(axis=0), strict=True)
    ]
)

fig, ax = plt.subplots(figsize=(12, 6))
sns.lineplot(data=level_means, x="week", y="mean", hue="fit", palette=FIT_COLORS, ax=ax)
sns.move_legend(ax, "upper center", bbox_to_anchor=(0.5, -0.12), ncol=4, title="")
ax.set(xlabel="week", ylabel="level")
fig.suptitle(
    "Posterior Mean of the Level: NUTS and SVI", fontsize=18, fontweight="bold"
);

Each line is the posterior mean of the level for one fit. The rotated fits recover the path of the level: their posterior means follow the NUTS reference, while the fit in the original coordinates gives a smoother path with a smaller amplitude that misses the turns. This is the shrinkage of the level that we saw in the scale estimates.

Finally, we look at the forecasts. We generate the in-sample predictions and the \(52\)-week forecast of every fit and of the NUTS reference.

We score them with the continuous ranked probability score (CRPS), where lower is better. For an introduction to this score see Intuition behind CRPS.

We compute the scores of every fit. Then we plot the mean over the fits (bars) and the score of every fit (points), in sample and out of sample. Thin gray lines connect the three SVI fits that share a random key.

def score(
    rng_key: UInt[Array, "2"], model: Model, posterior: dict[str, Array]
) -> tuple[float, float, Float[Array, "samples horizon 1"]]:
    "In-sample CRPS, out-of-sample CRPS and forecast draws of a posterior."
    key_train, key_test = random.split(rng_key)
    in_sample = predict_in_sample(key_train, model, posterior, covariates_train)
    forecast_draws = forecast(key_test, model, posterior, y_train, covariates)
    return (
        float(eval_crps(in_sample, y_train)),
        float(eval_crps(forecast_draws, y_test)),
        forecast_draws,
    )


rng_key, rng_subkey = random.split(rng_key)
scores = {"NUTS": [score(rng_subkey, models[Coordinates.NONE], nuts_reference)]} | {
    f"SVI: {name}": [
        score(rng_subkey, models[name], posterior) for posterior in posteriors
    ]
    for name, posteriors in svi_posteriors.items()
}

sample_labels = [
    f"in-sample (train, {t_train} weeks)",
    f"out-of-sample (test, {duration - t_train} weeks)",
]
crps_scores = pl.DataFrame(
    [
        {"fit": name, "key": key_index, "sample": sample_label, "CRPS": value}
        for name, runs in scores.items()
        for key_index, run in enumerate(runs)
        for sample_label, value in zip(sample_labels, run[:2], strict=True)
    ]
)
svi_fit_names = [name for name in scores if name != "NUTS"]

fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(15, 5.5), sharey=True)

for ax, sample_label in zip(axes, sample_labels, strict=True):
    sample_scores = crps_scores.filter(pl.col("sample") == sample_label)
    sns.barplot(
        data=sample_scores,
        x="fit",
        y="CRPS",
        hue="fit",
        palette=FIT_COLORS,
        alpha=0.7,
        errorbar=None,
        legend=False,
        ax=ax,
    )
    sns.stripplot(
        data=sample_scores, x="fit", y="CRPS", color="black", jitter=False, ax=ax
    )

    for key_index in range(len(svi_keys)):
        ax.plot(
            svi_fit_names,
            sample_scores.filter(
                pl.col("key") == key_index, pl.col("fit").is_in(svi_fit_names)
            )["CRPS"],
            color="gray",
            linewidth=0.8,
            alpha=0.7,
            label="SVI fits with the same random key" if key_index == 0 else None,
        )

    ax.plot(
        [], [], "s", color="gray", alpha=0.7, markersize=12, label="mean over the fits"
    )
    ax.plot([], [], "o", color="black", label="one fit")
    ax.set(xlabel="fit", ylabel="CRPS (lower is better)", title=sample_label)

fig.legend(*axes[0].get_legend_handles_labels(), loc="outside lower center", ncol=3)
fig.suptitle("CRPS by Fit", fontsize=18, fontweight="bold");

The left panel is the in-sample score, on the \(417\) training weeks, and the right panel is the out-of-sample score, on the held-out year. Each bar is the mean over the fits (five for SVI and one for NUTS), and each black point is one fit. Each thin gray line connects the three SVI fits (none, haar and dct) that use the same random key, so its slope shows the effect of the rotation for that key.

The rotated fits have a better in-sample score. Their in-sample CRPS is about \(0.025\), between the \(0.026\) of the original coordinates and the \(0.024\) of NUTS.

They also forecast better on the held-out year. The mean out-of-sample CRPS goes from about \(0.035\) in the original coordinates to about \(0.030\) with both rotations, which is the score of the NUTS reference. Every gray line ends lower at both rotated fits, so the gain holds for each of the five random keys.

Haar and DCT are close: the difference between their means is smaller than the spread over the random keys, which is itself below \(0.002\).

We end with the forecasts themselves. We use az.plot_lm to plot the forecast of the NUTS reference, of the first fit in the original coordinates, and of the first DCT fit.

forecast_fits = ["NUTS", "SVI: none", "SVI: dct"]

# `az.plot_lm` draws one panel per variable, and each variable needs its own x variable.
# The three fits have the same number of draws.
forecast_tree = az.from_dict(
    {
        "posterior_predictive": {
            name: np.asarray(scores[name][0][2][..., 0])[None] for name in forecast_fits
        },
        "observed_data": {name: np.asarray(y_test[:, 0]) for name in forecast_fits},
        "constant_data": {
            f"week ({name})": np.asarray(weeks_test, dtype=float)
            for name in forecast_fits
        },
    },
    coords={"time": np.asarray(weeks_test)},
    dims={name: ["time"] for name in forecast_fits}
    | {f"week ({name})": ["time"] for name in forecast_fits},
)

pc = az.plot_lm(
    forecast_tree,
    y=forecast_fits,
    x=[f"week ({name})" for name in forecast_fits],
    smooth=False,
    color=[FIT_COLORS[name] for name in forecast_fits],
    aes_by_visuals={"pe_line": ["color"]},
    visuals={
        "pe_line": {"label": "posterior mean"},
        "ci_band": {"alpha": 0.3, "label": "$94\\%$ HDI"},
        "observed_scatter": False,
    },
    col_wrap=1,
    figure_kwargs={"figsize": (12, 11), "sharex": True, "sharey": True},
)
fig = pc.viz["figure"].item()

for ax, name in zip(fig.axes, forecast_fits, strict=True):
    ax.plot(weeks_test, y_test[:, 0], color="C3", alpha=0.9, label="observed (test)")
    ax.set(xlabel="", ylabel="log(rides)", title=name)
    ax.legend(loc="lower left")

fig.axes[-1].set(xlabel="week")
fig.suptitle("Forecasts of the Held-Out Year", fontsize=18, fontweight="bold");

Each panel is the forecast of one fit for the \(52\) held-out weeks. The solid line is the posterior mean and the band is the \(94\%\) HDI, both in the color of the fit. The magenta line is the observed series.

The three forecasts follow the seasonal pattern, which the regression on the Fourier features captures in every fit. The DCT forecast is closer to the NUTS forecast, and the band of the fit in the original coordinates is wider, consistent with its larger noise scale sigma. The largest holiday drops are outside every band. The heavy-tailed likelihood treats them as outliers and keeps the band narrow instead of widening it to cover them.

Conclusion

We studied the Haar and DCT reparameterizations of time-indexed latent variables in NumPyro. We cnsidered two examples: a local level model with an exact posterior and on a ridership forecasting model. In both of them we saw how the rotation trick can help the mdoel fir better with SVI and AutoNormal guides (mean field approximation).