Principle:Pyro ppl Pyro Discrete Posterior Decoding

Metadata

Field	Value
Principle ID	Pyro_ppl_Pyro_Discrete_Posterior_Decoding
Title	Discrete Posterior Decoding
Project	Pyro (pyro-ppl/pyro)
Domains	Discrete_Inference, Bayesian_Inference
Implementation	Pyro_ppl_Pyro_Infer_Discrete
Repository	https://github.com/pyro-ppl/pyro

Summary

Discrete Posterior Decoding is the principle of inferring the posterior distribution of discrete latent variables after a model has been trained. Given a model with enumerated discrete sites and observed data, this principle enables two complementary modes of posterior inference: sampling from the exact discrete posterior (temperature=1) and MAP decoding to find the most likely discrete configuration (temperature=0, Viterbi-like).

Motivation

After training a probabilistic model that contains discrete latent variables -- such as a Hidden Markov Model (HMM), a mixture model, or a switching state-space model -- practitioners often need to answer the question: "What are the most likely discrete states given the observed data?" or "What is the distribution over discrete states?"

During training with TraceEnum_ELBO, discrete variables are marginalized out exactly. This is efficient for optimization but means the discrete states themselves are never explicitly sampled. Discrete posterior decoding bridges this gap by using the same enumeration machinery to compute the exact posterior over discrete variables and then either sample from it or find its mode.

Core Concepts

Forward-Filter Backward-Sample (temperature=1)

When temperature=1, the algorithm computes the exact posterior distribution over discrete latent variables using a forward-filter backward-sample procedure:

1. Forward pass (filtering): The model is run with enumeration enabled. Log-probabilities of all terms are collected and organized by their plate context (ordinal). The tensor variable elimination algorithm contracts these factors forward through the sequence or graph structure.
2. Backward pass (sampling): Starting from the contracted factors, the algorithm samples discrete values in reverse topological order, each conditioned on the values already sampled. This produces an exact sample from the joint posterior of all discrete variables.

This is the probabilistic analog of the forward-backward algorithm for HMMs, generalized to arbitrary plated factor graph structures.

MAP Decoding (temperature=0)

When temperature=0, the algorithm replaces the sum-product semiring with the max-product (or max-sum in log space) semiring:

1. Forward pass: Instead of summing over configurations (as in filtering), the algorithm takes the maximum. This is equivalent to the Viterbi algorithm for HMMs.
2. Backward pass: The algorithm traces back through the maxima to find the globally optimal discrete configuration -- the MAP estimate.

Relationship Between Modes

Temperature	Algorithm	Semiring	Output
1	Forward-filter backward-sample	Sum-product (SampleRing)	Posterior sample
0	Viterbi-like MAP decoding	Max-product (MapRing)	Most likely configuration

How It Works

The decoding procedure operates in several stages:

1. Enumerated trace construction: The model is run with EnumMessenger, which expands each discrete sample site to include all possible values along dedicated enumeration dimensions. The resulting trace contains tensors with extra dimensions corresponding to enumerated variables.
2. Log-probability collection: For each sample site, the log-probability is computed and tagged with its plate context (ordinal). These are organized into a tensor tree.
3. Tensor contraction: The contract_tensor_tree function processes the tensor tree using the appropriate ring (SampleRing or MapRing). This performs the forward pass of the message-passing algorithm.
4. Backward pass: Each query site's backward method is invoked to either sample (temperature=1) or argmax (temperature=0) the discrete values.
5. Trace replay: A collapsed trace is constructed with the decoded discrete values, and the model is replayed against this trace to produce the final output.

Example

import torch
import pyro
import pyro.distributions as dist
from pyro.infer import config_enumerate, infer_discrete

# Define an HMM model
@config_enumerate
def hmm_model(data, hidden_dim=10):
    transition = 0.3 / hidden_dim + 0.7 * torch.eye(hidden_dim)
    means = torch.arange(float(hidden_dim))
    states = [0]
    for t in pyro.markov(range(len(data))):
        states.append(
            pyro.sample("state_{}".format(t),
                        dist.Categorical(transition[states[-1]]))
        )
        pyro.sample("obs_{}".format(t),
                     dist.Normal(means[states[-1]], 1.0),
                     obs=data[t])
    return states

# MAP decoding (Viterbi)
viterbi_decoder = infer_discrete(
    hmm_model, first_available_dim=-1, temperature=0
)
map_states = viterbi_decoder(data)

# Posterior sampling
posterior_sampler = infer_discrete(
    hmm_model, first_available_dim=-1, temperature=1
)
sampled_states = posterior_sampler(data)

Relationship to Other Principles

• Pyro_ppl_Pyro_Enumeration_Configuration -- The discrete sites must be configured for enumeration (via @config_enumerate) before infer_discrete can decode them.
• Pyro_ppl_Pyro_Tensor_Variable_Elimination -- The forward pass of decoding uses the same tensor variable elimination algorithm as TraceEnum_ELBO.
• Pyro_ppl_Pyro_Markov_Dependency -- For sequential models, Markov annotations are essential to make the decoding tractable.

Related Pages

Implemented By

• Implementation:Pyro_ppl_Pyro_Infer_Discrete

References

• Pyro HMM tutorial: https://pyro.ai/examples/hmm.html
• Rabiner, L.R. "A tutorial on hidden Markov models and selected applications in speech recognition", 1989
• Viterbi, A. "Error bounds for convolutional codes and an asymptotically optimum decoding algorithm", 1967