Hidden Markov Model (HMM) (spec model)

This file defines an HMM with discrete observations:

• hidden states: nStates
• observations: nObservations (discrete symbols)

The model parameters are:

• initial distribution $\pi$,
• transition matrix $A$, and
• emission matrix $B$.

We represent a length-T observation sequence as a Tensor (Fin nObservations) [T]. The element type keeps observation symbols distinct from probabilities, while the tensor shape records the sequence length.

Notation and shapes

We use the conventional HMM notation:

• $\pi$: initial state distribution,
• $A$: transition matrix, with $A_{ij}=P(z_{t+1}=j\mid z_t=i)$, and
• $B$: emission matrix, with $B_{io}=P(x_t=o\mid z_t=i)$.

An observation sequence is $o_0,o_1,\ldots,o_{T-1}$, where each observation is represented by Fin nObservations.

References:

• Rabiner (1989), "A Tutorial on Hidden Markov Models and Selected Applications in Speech Recognition": https://ieeexplore.ieee.org/document/18626
• Baum and Petrie (1966), "Statistical Inference for Probabilistic Functions of Finite State Markov Chains": https://projecteuclid.org/journals/annals-of-mathematical-statistics/volume-37/issue-6/Statistical -Inference-for-Probabilistic-Functions-of-Finite-State-Markov-Chains/10.1214/aoms/1177699147.ful l

PyTorch analogy:

• emissions are categorical distributions (torch.distributions.Categorical),
• the forward algorithm corresponds to multiplying by $A$ and reweighting by the emission vector $B_{\mathord{:},o_t}$, then summing over previous states (often implemented in log-space in practice).

In practice, PyTorch users often reach for a dedicated HMM library (e.g. hmmlearn) or implement HMMs in log-space with logsumexp; TorchLean keeps the spec in a simple, explicit form that is good for reading and proofs.

Implementation status

No API builder implements this model. NN/Spec/Module/Hmm.lean wraps it as a Spec.Module, and NN/Tests/Runtime/Floats/TorchLeanOpsCheck.lean checks it numerically; no theorem is proved about it.

structure Spec.HMMSpec (α : Type) [TorchLean.Storage α] (nStates nObservations : ℕ) : Type

A discrete-observation HMM.

We do not enforce probabilistic validity (nonnegativity or rows summing to $1$) at the type level; that is a modeling assumption, similar to how PyTorch will happily store unconstrained tensors until you feed them to a distribution or a loss.

• initial : TorchLean.Tensor α [nStates]

  Initial distribution $\pi$.

• transition : TorchLean.Tensor α [nStates, nStates]

  Transition matrix $A$.

• emission : TorchLean.Tensor α [nStates, nObservations]

  Emission matrix $B$.

@[reducible, inline]

abbrev Spec.ObservationSeq (nObservations length : ℕ) : Type

A fixed-length sequence of symbols from the discrete observation alphabet.

Basic helpers

def Spec.getEmissionProbDiscrete {α : Type} [TorchLean.Storage α] {nStates nObservations : ℕ} (m : HMMSpec α nStates nObservations) (state : Fin nStates) (obs : Fin nObservations) : α

Get the emission probability $B_{\mathtt{state},\mathtt{obs}}$ for a discrete symbol.

Baum–Welch (EM) training

The forward-pass APIs above are enough to use a fixed HMM, but a “fully implemented” baseline should also include classical training. For discrete-observation HMMs, the standard training procedure is the Baum–Welch algorithm (an EM procedure):

• E-step: run forward–backward to compute expected state occupancies $\gamma$ and expected transition counts $\xi$.
• M-step: normalize those expected counts to update $\pi$, $A$, and $B$.

This implementation uses scaled forward–backward to reduce numerical underflow: each forward message $\alpha_t$ is normalized by a scalar $c_t$, and the backward messages divide by those same scalars. The sequence likelihood is then $\prod_t c_t$, so the log-likelihood is $\sum_t\log c_t$.

Concretely:

• forward recursion (unnormalized): $$ \widetilde{\alpha}_{t+1}(j) = B_{j,o_{t+1}}\sum_i \alpha_t(i)A_{ij}; $$
• scaling: $$ c_t=\sum_j\widetilde{\alpha}_t(j), \qquad \alpha_t=\frac{\widetilde{\alpha}_t}{c_t}, \qquad \sum_j\alpha_t(j)=1. $$

This is the same basic idea used in many practical HMM implementations (sometimes also expressed as log-space forward–backward).

This is deterministic and written for clarity; it is not intended to be a high-performance HMM trainer.

def Spec.normalizeVec {α : Type} [TorchLean.Storage α] [Context α] {n : ℕ} (v : TorchLean.Tensor α [n]) : TorchLean.Tensor α [n] × α

Normalize a nonnegative vector $v$ to sum to $1$, returning $(v/\sum_i v_i,\sum_i v_i)$.

If the sum is $0$, the normalized message is totalized to a uniform vector, while the returned scale remains $0$. The zero scale is essential: it records that the observation prefix has probability zero.

def Spec.emissionVec {α : Type} [TorchLean.Storage α] {nStates nObservations : ℕ} (m : HMMSpec α nStates nObservations) (obs : Fin nObservations) : TorchLean.Tensor α [nStates]

Emission probabilities $B_{\mathord{:},\mathtt{obs}}$ as a vector over states.

structure Spec.HMMForwardStep (α : Type) [TorchLean.Storage α] (nStates nObservations : ℕ) : Type

One timestep of a scaled HMM forward trace.

• observation : Fin nObservations

  Observation consumed at this timestep.

• message : TorchLean.Tensor α [nStates]

  Normalized forward message $\alpha_t$.

• scale : α

  Normalization constant $c_t$.

def Spec.hmmForwardScaled {α : Type} [TorchLean.Storage α] [Context α] {nStates nObservations length : ℕ} (m : HMMSpec α nStates nObservations) (observations : ObservationSeq nObservations length) : TorchLean.Tensor (HMMForwardStep α nStates nObservations) [length]

Scaled forward pass, returning the observation, $\alpha_t$, and $c_t$ at each timestep.

• Each $\alpha_t$ is normalized to sum to $1$.
• Each $c_t$ is the normalization constant used at step $t$.

If you need the total likelihood, multiply the scales: $p(o_{0:T-1})=\prod_t c_t$.

def Spec.baumWelchStepSpec {α : Type} [TorchLean.Storage α] [Context α] {nStates nObservations length : ℕ} [DecidableEq (Fin nObservations)] (m : HMMSpec α nStates nObservations) (observations : ObservationSeq nObservations length) : HMMSpec α nStates nObservations × Option α

One Baum–Welch (EM) step on a single sequence.

def Spec.baumWelchEpochSpec {α : Type} [TorchLean.Storage α] [Context α] {nStates nObservations batch length : ℕ} [DecidableEq (Fin nObservations)] (m : HMMSpec α nStates nObservations) (dataset : TorchLean.Tensor (Fin nObservations) [batch, length]) : HMMSpec α nStates nObservations × Option α

One Baum–Welch epoch over a dataset of observation sequences (sums expected counts).

Forward / likelihood

def Spec.hmmForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {nStates nObservations length : ℕ} (m : HMMSpec α nStates nObservations) (observations : ObservationSeq nObservations length) : α

Forward algorithm (scaled) returning the total sequence likelihood.

Implementation note: we compute the likelihood from the per-timestep scaling factors produced by hmmForwardScaled. This avoids the worst underflow behavior of multiplying many small probabilities directly.

def Spec.hmmBatchedForwardSpec {α : Type} [TorchLean.Storage α] [Context α] {nStates nObservations batch length : ℕ} (m : HMMSpec α nStates nObservations) (observations : TorchLean.Tensor (Fin nObservations) [batch, length]) : TorchLean.Tensor α [batch]

Batched forward pass with statically matched batch and sequence dimensions.

def Spec.hmmInitSpec {α : Type} [TorchLean.Storage α] [Context α] {nStates nObservations : ℕ} : HMMSpec α nStates nObservations

Initialize an HMM with uniform (uninformative) parameters.

This is a deterministic uniform initializer (useful for examples/tests); it is not intended as a statistically meaningful random initialization.

def Spec.hmmLogLikelihoodSpec {α : Type} [TorchLean.Storage α] [Context α] {nStates nObservations : ℕ} (m : HMMSpec α nStates nObservations) {length : ℕ} (observations : ObservationSeq nObservations length) : Option α

Log-likelihood of an observation sequence.

We compute this from the same scaling factors used in the EM implementation: $$ \log p(x_{0:T-1})=\sum_t\log c_t. $$