CRF Training › The Forward-Backward Algorithm

Computing the partition function $Z(\mathbf{x})$ by enumerating all sequences is intractable. The forward-backward algorithm exploits the sequential structure of the model to compute it in polynomial time.

The key insight is that the CRF score decomposes into local terms. At each position $t$, the contribution depends only on the label at $t$ and the label at $t-1$. This Markov property enables dynamic programming: we can build up the sum over all sequences incrementally, position by position.

Think of the forward-backward algorithm as a two-pass bookkeeping system for computing partial sums efficiently. The forward pass sweeps left to right, and at each position it records the total score of all partial paths that end at that position with each possible label. The backward pass sweeps right to left and records the total score of all partial paths that extend from each position to the end. When you multiply the forward score at a node with its backward score, you get the total score of all complete paths that pass through that node. This product is proportional to the marginal probability of being in that state at that position.

The algorithm draws its name and structure from the forward-backward algorithm for Hidden Markov Models, which you may have encountered in the previous chapter on HMMs. The CRF version works the same way conceptually, but in log space to avoid numerical underflow when sequences grow long. In practice, working in log space converts multiplications into additions and uses the log-sum-exp trick to combine scores safely. The transition from probability space to log space is a consistent theme in modern sequence modeling, and it applies here with the same motivation: multiplying many small probabilities together quickly produces values too tiny for floating-point arithmetic to represent accurately.

Think of the forward variable as an accumulator. As you sweep through the sequence from left to right, the forward variable at each position records everything you need to know about paths that have arrived at each possible label state. You never need to remember which specific path got you there, only the total aggregated score of all such paths. This is the essence of dynamic programming: discard the details, keep the aggregate.

The forward variable $\alpha_t(j)$ represents the sum of scores for all partial label sequences that end with label $j$ at position $t$:

$$\alpha_t(j) = \sum_{\mathbf{y}_{1:t}: y_t = j} \exp\left(\sum_{t'=1}^{t} \sum_{k} \lambda_k f_k(y_{t'-1}, y_{t'}, \mathbf{x}, t')\right)$$

where:

• $\mathbf{y}_{1:t}$: all possible label sequences from position 1 to $t$
• $y_t = j$: constrained to end with label $j$
• The summation accumulates scores of all paths reaching state $j$ at time $t$

We compute forward variables recursively:

Base case (t = 1):

$$\alpha_1(j) = \exp\left(\sum_{k} \lambda_k f_k(y_0, j, \mathbf{x}, 1)\right)$$

where $y_0$ is a special start symbol.

Recursive case:

$$\alpha_t(j) = \sum_{i \in L} \alpha_{t-1}(i) \cdot \exp\left(\sum_{k} \lambda_k f_k(i, j, \mathbf{x}, t)\right)$$

This recursion says: to get the total score of paths ending in $j$ at position $t$, sum over all possible previous labels $i$, multiplying the forward score at $t-1$ (all paths ending in $i$) by the transition score from $i$ to $j$.

Notice that the recursion reuses the previous time step's forward variables rather than recomputing everything from scratch. Each $\alpha_t(j)$ requires only the $|L|$ values from $\alpha_{t-1}(\cdot)$, so the total storage is $O(T \cdot |L|)$ and the total computation is $O(T \cdot |L|^2)$: for each of $T$ positions and each of $|L|$ possible current labels, we sum over $|L|$ possible previous labels. This is a dramatic improvement over the $O(|L|^T)$ cost of brute-force enumeration.

The partition function is simply the sum of all forward variables at the final position:

$$Z(\mathbf{x}) = \sum_{j \in L} \alpha_T(j)$$

The backward variable is symmetric to the forward variable but runs in the opposite direction. While the forward variable answers "what is the total score of all paths that reach state $i$ at position $t$ coming from the left?", the backward variable answers "what is the total score of all paths that extend from state $i$ at position $t$ to the right end of the sequence?" Together, a forward value and a backward value bracket a specific state at a specific position, and their product gives the total score of all complete paths that pass through that state.

The backward variable $\beta_t(i)$ represents the sum of scores for all partial sequences starting from position $t+1$, given that position $t$ has label $i$:

$$\beta_t(i) = \sum_{\mathbf{y}_{t+1:T}} \exp\left(\sum_{t'=t+1}^{T} \sum_{k} \lambda_k f_k(y_{t'-1}, y_{t'}, \mathbf{x}, t')\right)$$

where $y_t = i$ is fixed.

Base case (t = T):

$$\beta_T(i) = 1 \quad \text{for all } i$$

Recursive case:

$$\beta_t(i) = \sum_{j \in L} \beta_{t+1}(j) \cdot \exp\left(\sum_{k} \lambda_k f_k(i, j, \mathbf{x}, t+1)\right)$$

The forward and backward variables together give us the probability of any label at any position. In practice, this is why we run both passes: the forward pass alone gives us the partition function, but both passes together give us the full set of marginal probabilities that we need for gradient computation. The posterior marginal probability of label $j$ at position $t$ is:

$$P(y_t = j | \mathbf{x}) = \frac{\alpha_t(j) \cdot \beta_t(j)}{Z(\mathbf{x})}$$

And the probability of a specific transition from $i$ to $j$ at positions $t-1$ to $t$:

$$P(y_{t-1} = i, y_t = j | \mathbf{x}) = \frac{\alpha_{t-1}(i) \cdot M_t(i, j) \cdot \beta_t(j)}{Z(\mathbf{x})}$$

where $M_t(i, j) = \exp\left(\sum_k \lambda_k f_k(i, j, \mathbf{x}, t)\right)$ is the transition score matrix at position $t$.

Let's implement the forward-backward algorithm: