CRF Inference

Chain-structured Gaussian Conditional Random Fields (CRFs) for efficient discrete-time probabilistic inference with O(log T) complexity on parallel hardware.

When to Use

• Discrete-time state-space model inference
• Computing marginal distributions at each timestep
• Sampling from joint distributions over sequences
• Kalman filtering and smoothing
• Any chain-structured Gaussian graphical model

Key Concepts

CRF Structure

A chain CRF has the joint distribution:

p(x1, ..., xT) ∝ ∏ φt(xt) × ∏ ψt(xt, xt+1)

Where:

• φt(xt) are node potentials (observations/priors at each timestep)
• ψt(xt, xt+1) are transition potentials (dynamics between timesteps)

Message Passing

The CRF performs message passing to compute:

• Forward messages (filtering)
• Backward messages (smoothing preparation)
• Marginals at each timestep
• Samples from the joint distribution

Parallel Scan

For long sequences, linsdex uses parallel scan for O(log T) complexity instead of O(T) sequential message passing.

Code Examples

Building a CRF

import jax.numpy as jnp
import jax.random as random
import equinox as eqx
from linsdex import CRF, NaturalGaussian, GaussianTransition, DiagonalMatrix

dim = 2
seq_len = 10

# Create node potentials (one per timestep)
J = DiagonalMatrix.eye(dim)
h = jnp.zeros(dim)
single_potential = NaturalGaussian(J, h)
node_potentials = eqx.filter_vmap(lambda _: single_potential)(jnp.arange(seq_len))

# Create transitions between adjacent nodes
A = DiagonalMatrix.eye(dim)  # Identity transition matrix
u = jnp.zeros(dim)
Sigma = DiagonalMatrix.eye(dim) * 0.1
single_transition = GaussianTransition(A, u, Sigma)
transitions = eqx.filter_vmap(lambda _: single_transition)(jnp.arange(seq_len - 1))

# Create the CRF
crf = CRF(node_potentials, transitions)

Computing Marginals

# Get marginal distributions at each timestep
marginals = crf.get_marginals()

# marginals is a batch of Gaussian distributions
# Access mean and covariance at each timestep
means = marginals.mu  # Shape: (seq_len, dim)

Forward and Backward Messages

# Forward messages (filtering)
fwd_messages = crf.get_forward_messages()

# Backward messages (smoothing preparation)
bwd_messages = crf.get_backward_messages()

Sampling from the CRF

key = random.PRNGKey(0)

# Sample a sequence from the joint distribution
sample = crf.sample(key)  # Shape: (seq_len, dim)

# Sample multiple sequences
keys = random.split(key, 100)
samples = jax.vmap(crf.sample)(keys)  # Shape: (100, seq_len, dim)

Evaluating Log Probability

# Compute log probability of a sequence
log_prob = crf.log_prob(sample)

Parallel vs Sequential Inference

# By default, uses parallel scan for O(log T) complexity
marginals = crf.get_marginals()

# Can also use sequential scan (useful for debugging or small sequences)
marginals_seq = crf.get_marginals(parallel=False)

CRF with Varying Observation Precision

# Create node potentials with varying observation certainty
obs_means = jnp.sin(jnp.linspace(0, 4*jnp.pi, seq_len))
obs_precision = 10.0  # Higher = more certain observations

def make_node_potential(obs_mean):
    J = DiagonalMatrix.eye(dim) * obs_precision
    h = jnp.zeros(dim).at[0].set(obs_mean * obs_precision)
    return NaturalGaussian(J, h)

node_potentials = eqx.filter_vmap(make_node_potential)(obs_means)

# Create transitions with decay
def make_transition(_):
    A = DiagonalMatrix(jnp.ones(dim) * 0.95)  # Decay factor
    u = jnp.zeros(dim)
    Sigma = DiagonalMatrix.eye(dim) * 0.1
    return GaussianTransition(A, u, Sigma)

transitions = eqx.filter_vmap(make_transition)(jnp.arange(seq_len - 1))

crf = CRF(node_potentials, transitions)
marginals = crf.get_marginals()

Using NaturalGaussian for Combining Observations

from linsdex import NaturalGaussian

# Natural parameterization is efficient for combining observations
# Adding natural Gaussians multiplies densities

dim = 3
J1 = DiagonalMatrix.eye(dim)
h1 = jnp.array([1.0, 0.0, 0.0])
obs1 = NaturalGaussian(J1, h1)

J2 = DiagonalMatrix.eye(dim) * 2.0
h2 = jnp.array([0.0, 2.0, 0.0])
obs2 = NaturalGaussian(J2, h2)

# Combine observations (product of Gaussians)
combined = obs1 + obs2  # J = J1 + J2, h = h1 + h2

Key Classes

• CRF(node_potentials, transitions) - Chain-structured CRF
• NaturalGaussian(J, h) - Node potentials in natural parameterization
• GaussianTransition(A, u, Sigma) - Linear-Gaussian transitions p(xt+1 | xt)

GaussianTransition Details

A GaussianTransition represents p(y | x) = N(y; Ax + u, Σ):

• A - Linear transformation matrix
• u - Offset/bias term
• Sigma - Noise covariance

transition = GaussianTransition(A, u, Sigma)

# Apply to a point to get conditional distribution
x = jnp.ones(dim)
p_y_given_x = transition.condition_on_x(x)  # Returns StandardGaussian

# Chain transitions: p(z|x) from p(y|x) and p(z|y)
transition2 = GaussianTransition(A2, u2, Sigma2)
chained = transition.chain(transition2)  # p(z|x)

Tips

• Use NaturalGaussian for node potentials because addition corresponds to multiplying densities
• The CRF automatically uses parallel scan on GPU for sequences longer than a threshold
• For very long sequences, the O(log T) complexity provides significant speedups
• Use DiagonalMatrix for covariances when dimensions are independent to save computation
• The equinox.filter_vmap pattern is useful for creating batched potentials and transitions