Parametric Partial Pooling

Overview

A one-way random-effects model (Rubin 1981) whose group-level effects come from a parametric program template. The template school_effects declares the non-centered parametrization of a set of partially pooled effects once, typed over a scalar spread : Real and a plate K : FinSet; the two host programs instantiate it at a tight (0.6) and a loose (2.5) between-group spread. A score step adds a soft sum-to-zero factor on the effects (Morris et al. 2019), and a labeled return tuple names the program's outputs.

QVR Source

object School : FinSet 8

program school_effects(spread : Real, K : FinSet) : K -> K
    sample z : K <- Normal(0.0, 1.0)
    let effect = spread * z
    return effect

program pooled_tight : School -> School
    sample theta <- school_effects(0.6, School)
    sample sigma <- LogNormal(0.0, 0.5)
    let total_effect = sum(theta)
    score centering = -50.0 * total_effect * total_effect
    observe y : School <- Normal(theta, sigma)
    return (effects: theta, scale: sigma)

program pooled_loose : School -> School
    sample theta <- school_effects(2.5, School)
    sample sigma <- LogNormal(0.0, 0.5)
    let total_effect = sum(theta)
    score centering = -50.0 * total_effect * total_effect
    observe y : School <- Normal(theta, sigma)
    return (effects: theta, scale: sigma)

export pooled_tight

Walkthrough

school_effects is a parametric program: its parameter list carries typed parameters rather than data parameters, so the compiler stores it as a template instead of compiling it directly. The call site sample theta <- school_effects(0.6, School) substitutes 0.6 for spread and School for K, then inlines an alpha-renamed copy of the body into the host program: the local z becomes theta$z and the return variable effect becomes theta itself. Each call site contributes fresh latents, which is why pooled_tight and pooled_loose share structure but not parameters.

The non-centered form routes spread through the let arithmetic (effect = spread * z) rather than through a plate-draw scale, so the bound value shapes both forward simulation and the joint density. score centering = -50.0 * total_effect * total_effect adds \(-\tfrac{1}{2}(\sum_j \theta_j / 0.1)^2\) to the program's log-joint: a soft sum-to-zero constraint that pins the effects' grand mean near zero without a hard reparametrization.

return (effects: theta, scale: sigma) labels the output tuple. The compiled program's rsample returns a dict keyed by the labels (effects, scale) rather than by the internal variable names, and log_joint accepts intermediates keyed either way. export pooled_tight selects which of the two host programs the compiled module wraps; swapping the export line for export pooled_loose selects the weakly pooled variant with no other change.

DSL Features

  • Typed program parameters: (spread : Real, K : FinSet) declares a parametric template; admissible parameter kinds are Real, Nat, FinSet, Space, Object, and Mor[A, B].
  • Template instantiation: sample v <- template(args) inlines the substituted, alpha-renamed body at the call site.
  • Labeled return tuple: return (label: var, ...) names the program's outputs on the compiled output path.
  • score: adds an arbitrary tensor expression to the program's log-joint as an extra density factor.
  • export: selects the module's exported morphism among the declared candidates.

Try it

The SVI step counts and NUTS warmup, sample, and chain budgets in the snippets below are illustrative: each block is sized to run in tens of seconds and demonstrate the API surface. Production fits typically need 10x to 100x more SVI steps, longer NUTS warmup, and multiple chains to actually converge to the data-generating parameters.

Generating synthetic data

import torch
from quivers.dsl import load

torch.manual_seed(0)
prog = load("docs/examples/source/parametric_pooling.qvr")
model = prog.morphism

n_school = 8
true_spread = 0.6
true_sigma = 0.4
theta_true = true_spread * torch.randn(n_school)
theta_true = theta_true - theta_true.mean()
y = torch.distributions.Normal(theta_true, true_sigma).sample()

observations = {"y": y}
x_in = torch.zeros(n_school, 1)

draws = model.rsample(x_in)
print(sorted(draws.keys()))

SVI fit

from quivers.inference import AutoNormalGuide, ELBO, SVI

oracle_nll = float(
    -torch.distributions.Normal(theta_true, true_sigma).log_prob(y).mean()
)

torch.manual_seed(1)
guide = AutoNormalGuide(model, observed_names={"y"})
optim = torch.optim.Adam(
    list(model.parameters()) + list(guide.parameters()), lr=5e-2,
)
svi = SVI(model, guide, optim, ELBO(num_particles=1))

losses = []
for _ in range(200):
    losses.append(svi.step(x_in, observations))

print(f"initial loss: {losses[0]:.2f}")
print(f"final loss:   {losses[-1]:.2f}")
print(f"oracle NLL:   {oracle_nll:.2f}")

NUTS posterior

from quivers.inference import MCMC, NUTSKernel

torch.manual_seed(2)
kernel = NUTSKernel(step_size=0.05, max_tree_depth=3, target_accept=0.8)
mc = MCMC(kernel, num_warmup=20, num_samples=20, num_chains=1)
result = mc.run(model, x_in, observations)

print(f"acceptance:  {float(result.acceptance_rates.mean()):.2f}")
print(f"divergences: {int(result.divergence_counts.sum())}")

Categorical Perspective

A parametric program denotes a dependent kernel \(\Pi (p_1 : P_1) \ldots \Pi (p_n : P_n).\ \mathbf{Kern}(\mathrm{dom}(p), \mathrm{cod}(p))\) in the indexed family of Kleisli arrows of the Giry monad over the parameter category; each call site realizes a section of the family at the supplied arguments. The score step is multiplication by a density factor in the subprobability monad \(\mathcal{G}_{\le 1}\), and the labeled return tuple names the projections out of the program's product codomain.

References