Implementation:Pyro ppl Pyro MaternKernel

Property	Value
Module	pyro.ops.ssm_gp
Source	pyro/ops/ssm_gp.py
Lines	169
Classes	MaternKernel
Dependencies	torch, pyro.nn (PyroModule, PyroParam)

Overview

This module provides the MaternKernel class, which represents univariate Gaussian Processes (GPs) with Matern kernels as state space models (SSMs). This reformulation enables O(T) inference for GPs (instead of O(T^3) with standard GP methods) by leveraging the equivalence between Matern-kernel GPs and linear stochastic differential equations.

The class supports Matern kernels of order 0.5, 1.5, and 2.5, with corresponding state dimensions of 1, 2, and 3. The state space formulation follows Hartikainen & Sarkka (2010) and Solin (2016), computing:

• Transition matrices: How the GP latent state evolves over a time interval.
• Stationary covariance: The long-run covariance of the state.
• Process covariance: The noise covariance for a given time step.

Code Reference

Class: MaternKernel (PyroModule)

Constructor parameters:

• nu (float): Matern kernel order (0.5, 1.5, or 2.5).
• num_gps (int): Number of independent GPs.
• length_scale_init (Tensor): Initial length scales, shape (num_gps,).
• kernel_scale_init (Tensor): Initial kernel scales, shape (num_gps,).

Learnable parameters:

• length_scale: Positive constrained, shape (num_gps,).
• kernel_scale: Positive constrained, shape (num_gps,).

Methods:

• transition_matrix(dt): Computes the exponentiated transition matrix A(dt) with layout (num_gps, state_dim, state_dim). This matrix multiplies states from the right.
  • nu=0.5: exp(-dt / rho) (scalar)
  • nu=1.5: 2x2 matrix involving exp(-sqrt(3) * dt / rho)
  • nu=2.5: 3x3 matrix involving exp(-sqrt(5) * dt / rho)

• stationary_covariance(): Computes the stationary covariance matrix P_inf of shape (num_gps, state_dim, state_dim).

• process_covariance(A): Given transition matrix A, computes process covariance Q = P_inf - A.T @ P_inf @ A.

• transition_matrix_and_covariance(dt): Convenience method returning both the transition matrix and process covariance for a time interval dt.

I/O Contract

Method	Input	Output
__init__	nu: float, num_gps: int, optional init tensors	MaternKernel (PyroModule)
transition_matrix(dt)	dt: float (time interval)	Tensor(num_gps, state_dim, state_dim)
stationary_covariance()	(none)	Tensor(num_gps, state_dim, state_dim)
process_covariance(A)	A: Tensor(..., state_dim, state_dim)	Tensor(num_gps, state_dim, state_dim)
transition_matrix_and_covariance(dt)	dt: float	Tuple of two Tensor(num_gps, state_dim, state_dim)

Usage Examples

import torch
from pyro.ops.ssm_gp import MaternKernel

# Create a Matern-3/2 kernel for 2 independent GPs
kernel = MaternKernel(
    nu=1.5,
    num_gps=2,
    length_scale_init=torch.tensor([1.0, 2.0]),
    kernel_scale_init=torch.tensor([1.0, 0.5]),
)

# Compute transition matrix for dt=0.1
dt = 0.1
A = kernel.transition_matrix(dt)
print(A.shape)  # torch.Size([2, 2, 2])

# Get both transition and process covariance
A, Q = kernel.transition_matrix_and_covariance(dt)

# Get stationary covariance
P_inf = kernel.stationary_covariance()
print(P_inf.shape)  # torch.Size([2, 2, 2])

Related Pages

• Pyro_ppl_Pyro_Gaussian -- Gaussian operations used for Kalman filtering with SSM-GPs