Tutorial 2: Stochastic Relations

The FinStoch category has finite sets as objects and Markov kernels as morphisms. Their entries are conditional probabilities. This chapter constructs and composes kernels, conditions them on observations, and queries probabilities and expectations.

Concepts

  • Markov Algebra: the algebra where the tensor product is multiplication and the join is summation
  • StochasticMorphism: a morphism in FinStoch whose rows are probability distributions (each row sums to 1)
  • Discretized families: continuous distributions discretized into finite bins
  • Conditioning: updating a kernel given evidence
  • Probabilities and expectations: extracting entries and computing expected values

Setup

import torch
from quivers.core.objects import FinSet
from quivers.core.morphisms import identity, observed
from quivers.stochastic import (
    MARKOV,
    stochastic,
    DiscretizedNormal,
    DiscretizedBeta,
    condition,
    prob,
    marginal_prob,
    expectation,
)
from quivers.program import Program

Creating stochastic morphisms

The FinStoch category uses the Markov algebra: composition is Markov kernel composition (matrix multiplication), and the join is summation.

Create two finite sets:

X = FinSet(name="Latent", cardinality=3)
Y = FinSet(name="Observed", cardinality=4)

Create a latent (learnable) stochastic morphism from X to Y using stochastic:

kern = stochastic(X, Y)
print(kern.tensor.shape)  # [3, 4]

Verify it is row-stochastic (each row sums to 1):

row_sums = kern.tensor.sum(dim=1)
print(row_sums)  # tensor([1., 1., 1.])

The entries represent conditional probabilities:

print(kern.tensor)
# Each row: P(y | x) for a fixed x

The morphism is learnable; parameters are adjusted via softmax normalization internally.

Create an observed stochastic morphism from Y to a third set Z:

Z = FinSet(name="Output", cardinality=2)
data = torch.tensor([
    [0.7, 0.3],
    [0.4, 0.6],
    [0.5, 0.5],
    [0.2, 0.8],
])
kern_obs = observed(Y, Z, data, algebra=MARKOV)

Verify the data is valid:

print(kern_obs.tensor.sum(dim=1))  # all 1.0

Composition of Markov Kernels

Compose stochastic morphisms with >>:

composed = kern >> kern_obs
print(composed.tensor.shape)  # [3, 2]
print(composed.tensor.sum(dim=1))  # all 1.0

The composition is standard Markov kernel composition (matrix multiplication):

\[ (\kappa_2 \circ \kappa_1)(x, z) = \sum_y \kappa_1(x, y) \cdot \kappa_2(y, z) \]

Wrap in a Program for training:

program = Program(composed)
output = program()
print(output.shape)  # [3, 2]

Discretized distributions

For models mixing discrete and continuous variables, discretize continuous distributions into finite bins.

Create a DiscretizedNormal: a normal density evaluated at bin centers spanning the interval [low, high], normalized to a probability vector. The mean and log-scale are learnable parameters, one pair per domain element:

Z_bin = FinSet(name="DiscretizedValue", cardinality=10)
Unit = FinSet(name="Unit", cardinality=1)

disc_normal = DiscretizedNormal(Unit, Z_bin, low=0.0, high=1.0)
print(disc_normal.tensor.shape)        # [1, 10]
print(disc_normal.tensor.sum(dim=-1))  # ~[1.0]

This is a stochastic morphism from the terminal object to a discretized space. Each entry is the probability mass in a bin.

Similarly, create a DiscretizedBeta:

disc_beta = DiscretizedBeta(Unit, Z_bin)

These discretized morphisms compose with other stochastic morphisms. For instance, chain a discretized normal into a learnable kernel from Z_bin to a coarser output set:

Output = FinSet(name="Output", cardinality=3)
kern_to_output = stochastic(Z_bin, Output)

combined = disc_normal >> kern_to_output
print(combined.tensor.shape)        # [1, 3]
print(combined.tensor.sum(dim=-1))  # ~[1.0]

Observations and Conditioning

Given non-negative evidence over the codomain, condition reweights each row of the kernel and renormalizes:

\[ f|e(a, b) \;\propto\; f(a, b)\, e(b) \]

For instance, soft evidence over Y:

evidence = torch.tensor([0.1, 0.4, 0.3, 0.2])

conditioned = condition(kern, evidence)
print(conditioned.tensor.shape)        # [3, 4]
print(conditioned.tensor.sum(dim=-1))  # ~[1, 1, 1]

Hard evidence (only Y = 2 is allowed) is just a zero-one mask:

hard = torch.zeros(4)
hard[2] = 1.0
posterior_kernel = condition(kern, hard)

The result is still a kernel X -> Y (every row a distribution), now concentrated on the evidence support. To recover a posterior over X you push a prior through it and marginalize, as shown next.

Probabilities and Expectations

prob reads off entries of a kernel at specific (domain, codomain) index pairs:

x_idx = torch.tensor([0, 1, 2])
y_idx = torch.tensor([3, 1, 0])
print(prob(kern, x_idx, y_idx))  # shape [3]

marginal_prob marginalizes a kernel over its domain under a uniform prior, returning the codomain marginal at the requested indices:

y_idx = torch.tensor([0, 1, 2, 3])
print(marginal_prob(kern, y_idx))  # shape [4], sums to 1

expectation computes the conditional expectation of a real-valued function on the codomain, for each domain element:

values = torch.tensor([10.0, 5.0, 1.0, 0.5])
print(expectation(kern, values))  # shape [3]: E[v | x]

Multiple stages

Build a longer chain:

# X -> Y -> Z
kern1 = stochastic(X, Y)
kern2 = stochastic(Y, Z)
composed = kern1 >> kern2

print(composed.tensor.shape)        # [3, 2]
print(composed.tensor.sum(dim=-1))  # ~[1, 1, 1]

# Marginal of Z under uniform X
z_idx = torch.tensor([0, 1])
print(marginal_prob(composed, z_idx))  # shape [2], sums to 1

Summary

You have:

  • Created stochastic morphisms (Markov kernels)
  • Verified row-stochasticity
  • Composed kernels with >>
  • Used DiscretizedNormal and DiscretizedBeta to mix discrete and continuous randomness
  • Applied conditioning to update kernels given evidence
  • Computed entries, marginals, and expectations
  • Built multi-stage inference chains

Next, learn how to work with continuous spaces and probabilistic programs in Tutorial 3.