Implementation:Pyro ppl Pyro Trace MMD

Overview

The trace_mmd module (Template:Code) implements Trace_MMD, a variational inference objective that replaces the KL divergence in the standard ELBO with Maximum Mean Discrepancy (MMD) between the marginal variational posterior q(z) and the prior p(z). This yields the MMD-VAE loss function:

L(theta, phi) = -E_{p_data(x)} E_{q(z|x; phi)} log p(x|z; theta) + MMD(q(z; phi) || p(z))

where MMD between two distributions is:

MMD(q || p) = E_{p,p'} k(z,z') + E_{q,q'} k(z,z') - 2 E_{p,q'} k(z,z')

and k is a user-specified kernel function (typically from Template:Code).

This is based on the InfoVAE framework (Zhao et al.) which balances learning and inference. The MMD penalty encourages the aggregate posterior to match the prior without requiring per-sample KL divergence computation.

Important: The current implementation treats only the particle dimension as the batch dimension for MMD computation. All other dimensions are treated as event dimensions. This means large Template:Code values are needed for reasonable MMD variance, and Template:Code is recommended.

The kernel and MMD scaling factor can be specified either globally (applied to all latent variables) or per-variable via dictionaries.

Code Reference

File: Template:Code

Key Classes

Class	Parent	Description
Template:Code	Template:Code	ELBO variant with MMD replacing KL divergence between marginal posterior and prior.

Trace_MMD Methods

Method	Description
Template:Code	Initialize with kernel, MMD scale, and standard ELBO parameters.
Template:Code	Compute the differentiable MMD-VAE loss. Returns a tensor suitable for backpropagation.
Template:Code	Compute the MMD-VAE loss as a float.
Template:Code	Compute the loss, perform backward, and return the loss as a float.
Template:Code	Internal: computes the log-likelihood and MMD penalty separately.
Template:Code	Returns paired model/guide traces for importance tracing.

Properties

Property	Description
Template:Code	The kernel(s) used for MMD computation. Can be a single Template:Code or a dict mapping site names to kernels.
Template:Code	Scaling factor(s) for the MMD penalty. Can be a float or a dict mapping site names to floats.

Internal Helper

Function	Description
Template:Code	Computes the MMD between sample sets X and Z using the given kernel: E[k(X,X)] + E[k(Z,Z)] - 2*E[k(X,Z)].

I/O Contract

Constructor

Inputs:

• Template:Code -- A GP kernel instance (Template:Code) or a dict mapping latent variable names to kernel instances. Kernel parameters are frozen (Template:Code).
• Template:Code -- Float or dict of floats for per-variable MMD scaling. Default is 1.
• Template:Code -- Number of particles (default 10). Large values needed for good MMD estimates.
• Template:Code -- Whether to vectorize particles (default True, recommended).
• Other standard ELBO parameters (Template:Code, Template:Code, etc.).

differentiable_loss

Inputs:

• Template:Code -- A Pyro model. Latent distributions must be reparameterizable.
• Template:Code -- A Pyro guide. All sample sites must be reparameterizable.
• Template:Code -- Passed to model and guide.

Output:

• Template:Code -- Scalar loss: -loglikelihood + mmd_penalty.

Requirements

• All latent sample sites in both model and guide must be reparameterizable (Template:Code). Raises Template:Code otherwise.
• The model is traced twice per step: once with the guide (for log-likelihood) and once independently (for prior samples used in MMD).

Usage Examples

Basic MMD-VAE

import torch
import pyro
import pyro.distributions as dist
from pyro.infer import SVI, Trace_MMD
from pyro.optim import Adam
from pyro.contrib.gp.kernels import RBF

def model(data):
    z = pyro.sample("z", dist.Normal(torch.zeros(2), torch.ones(2)).to_event(1))
    loc = pyro.param("dec_loc", torch.zeros(5))
    pyro.sample("obs", dist.Normal(loc + z.sum(-1, keepdim=True), 0.1).to_event(1),
                obs=data)

def guide(data):
    loc = pyro.param("enc_loc", torch.zeros(2))
    scale = pyro.param("enc_scale", torch.ones(2),
                       constraint=dist.constraints.positive)
    pyro.sample("z", dist.Normal(loc, scale).to_event(1))

kernel = RBF(input_dim=2)
mmd_loss = Trace_MMD(kernel=kernel, mmd_scale=10.0, num_particles=50)

svi = SVI(model, guide, Adam({"lr": 0.01}), loss=mmd_loss)
for step in range(1000):
    loss = svi.step(data)

Per-Variable Kernels and Scales

from pyro.contrib.gp.kernels import RBF, Matern32

kernels = {
    "z1": RBF(input_dim=3),
    "z2": Matern32(input_dim=5)
}
scales = {
    "z1": 1.0,
    "z2": 5.0
}
mmd_loss = Trace_MMD(kernel=kernels, mmd_scale=scales, num_particles=100)

Related Pages

• Pyro_ppl_Pyro_EnergyDistance -- Another divergence-based inference objective
• Pyro_ppl_Pyro_TraceGraph_ELBO -- Standard graph-based ELBO for comparison
• Pyro_ppl_Pyro_RenyiELBO -- Renyi divergence ELBO alternative
• Pyro_ppl_Pyro_Infer_Utilities -- Validation and utility functions