Implementation:Pyro ppl Pyro GaussianScaleMixture

Knowledge Sources	Pyro
Domains	Probability_Distributions
Last Updated	2026-02-09 09:00 GMT

Overview

Description

GaussianScaleMixture is a probability distribution class in Pyro that implements a mixture of zero-mean Normal distributions with diagonal covariance matrices. The distribution is a mixture with K components, where each component is a D-dimensional Normal distribution with zero mean. The covariance structure is factored so that each component's diagonal covariance is determined by shared coordinate scales and per-component scale factors:

Sigma_ii = (component_scale_k * coord_scale_i) ** 2   # for i = 1, ..., D

The mixture weights are specified via softmax logits. A key feature of this distribution is that it implements pathwise derivatives (reparameterized gradients) for samples, using the method described in "Pathwise Derivatives for Multivariate Distributions" by Jankowiak and Karaletsos (arXiv:1806.01856). The backward pass is implemented via a custom autograd Function (_GSMSample) that computes analytic gradients with respect to all three parameters: coord_scale, component_logits, and component_scale.

The distribution supports both even and odd dimensionality (D >= 2), though even dimensions yield higher numerical precision since they avoid the use of error functions (erf) in the backward pass. Batched parameters are not currently supported.

Usage

GaussianScaleMixture is useful in variational inference and deep generative models where a flexible, heavy-tailed prior or variational family is desired. The mixture of Gaussians with shared coordinate scales provides a structured yet expressive distribution, and pathwise gradient support makes it compatible with stochastic variational inference methods that require low-variance gradient estimators.

Code Reference

Source Location

• File: pyro/distributions/gaussian_scale_mixture.py
• Repository: pyro-ppl/pyro

Signature

class GaussianScaleMixture(TorchDistribution):
    def __init__(self, coord_scale, component_logits, component_scale)

Import

from pyro.distributions import GaussianScaleMixture

I/O Contract

Inputs

Parameter	Type	Description
coord_scale	torch.Tensor	A 1-D tensor of shape (D,) containing positive scale parameters shared across all K mixture components. Each entry controls the scale along one coordinate dimension.
component_logits	torch.Tensor	A 1-D tensor of shape (K,) containing real-valued logits that determine the mixture weights via softmax normalization.
component_scale	torch.Tensor	A 1-D tensor of shape (K,) containing positive scale multipliers, one per mixture component. These multiply the shared coord_scale to define per-component covariance.

Outputs

Method	Return Type	Description
rsample(sample_shape)	torch.Tensor	Returns a reparameterized sample of shape sample_shape + (D,) from the mixture distribution with pathwise gradients.
log_prob(value)	torch.Tensor	Returns the scalar log probability of a 1-D input tensor value of shape (D,).

Usage Examples

import torch
from pyro.distributions import GaussianScaleMixture

# Define a mixture of 3 components in 4 dimensions
D = 4
K = 3

coord_scale = torch.ones(D)
component_logits = torch.zeros(K)          # uniform mixture weights
component_scale = torch.tensor([0.5, 1.0, 2.0])

dist = GaussianScaleMixture(coord_scale, component_logits, component_scale)

# Draw a reparameterized sample
sample = dist.rsample()
print(sample.shape)  # torch.Size([4])

# Compute log probability
log_p = dist.log_prob(sample)
print(log_p.shape)   # torch.Size([])

import pyro
import pyro.distributions as dist
import torch

# Using GaussianScaleMixture as a prior in a Pyro model
def model(data):
    coord_scale = pyro.param("coord_scale", torch.ones(2), constraint=dist.constraints.positive)
    component_logits = pyro.param("component_logits", torch.zeros(3))
    component_scale = pyro.param("component_scale", torch.ones(3), constraint=dist.constraints.positive)

    with pyro.plate("data", len(data)):
        z = pyro.sample("z", dist.GaussianScaleMixture(coord_scale, component_logits, component_scale))

Related Pages

• Pyro_ppl_Pyro_OMTMultivariateNormal - Multivariate Normal with Optimal Transport gradients, another distribution with custom pathwise derivatives
• Pyro_ppl_Pyro_MultivariateStudentT - Multivariate Student's t-distribution, another heavy-tailed multivariate distribution
• Pyro_ppl_Pyro_Rejector - Rejection sampling distribution with custom gradient support