EM Algorithm

Start by defining the objective function in line with the previous discussions. We also drop the region subscript since it is implicit throughout.

\[ Q(\theta^\star, \theta) = \mathbb{E}_{\boldsymbol{s}} \left[ \log p(\boldsymbol{y}, \boldsymbol{s} | \theta) \middle| \boldsymbol{y}, \theta^\star \right] \tag{1}\]

The joint probability function, \(p(\boldsymbol{y}, \boldsymbol{s} | \theta)\), can be found recursively.

\[ p(\boldsymbol{y}, \boldsymbol{s} | \theta) = p(\boldsymbol{y} | \boldsymbol{s}, \theta) p(\boldsymbol{s} | \theta) = p(s_1 | \theta) \prod_{t=2}^T p(s_t | s_{t-1}, \theta) \prod_{t=1}^T p(y_t | s_t, \theta) \]

Substituting this into Equation 1 and taking advantage of the linearity of expectation leads to the following form for the objective function.

\[ Q(\theta^\star, \theta) = \mathbb{E}_{\boldsymbol{s}} \left[ \log p(s_1 | \theta) \middle| \boldsymbol{y}, \theta^\star \right] + \mathbb{E}_{\boldsymbol{s}} \left[ \sum_{t=2}^T \log p(s_t | s_{t-1}, \theta) \middle| \boldsymbol{y}, \theta^\star\right] + \mathbb{E}_{\boldsymbol{s}} \left[ \sum_{t=1}^T \log p(y_t | s_t, \theta) \middle| \boldsymbol{y}, \theta^\star \right] \]

This function can be further simplified by isolating each term. Beginning with the first,

\[\begin{align*} \mathbb{E}_{\boldsymbol{s}} \left[ \log p(s_1 | \theta) \middle| \boldsymbol{y}, \theta^\star \right] & = \sum_{s_1, \ldots, s_t} p(s_1, \ldots, s_t | \boldsymbol{y}, \theta^\star) \log p(s_1 | \theta) \\ & = \sum_{k=1}^K P(s_1 = k | \boldsymbol{y}, \theta^\star) \log \pi_k \end{align*}\]

followed by,

\[\begin{align*} \mathbb{E}_{\boldsymbol{s}} \left[ \sum_{t=2}^T \log p(s_t | s_{t-1}, \theta) \middle| \boldsymbol{y}, \theta^\star \right] & = \sum_{s_1, \ldots, s_t} p(s_1, \ldots, s_t | \boldsymbol{y}, \theta^\star) \sum_{t=2}^T \log p(s_t | s_{t-1}, \theta) \\ & = \sum_{t=2}^T \left( \sum_{s_1, \ldots, s_t} p(s_1, \ldots, s_t | \boldsymbol{y}, \theta^\star) \log p(s_t | s_{t-1}, \theta) \right) \\ & = \sum_{t=2}^T \left( \sum_{i=1}^K \sum_{j=1}^K P(s_{t-1} = i, s_t = j | \boldsymbol{y}, \theta^\star) \log A_{ij} \right) \end{align*}\]

and finally,

\[\begin{align*} \mathbb{E}_{\boldsymbol{s}} & = \left[ \sum_{t=1}^T \log p(y_t | s_t, \theta) \middle| \boldsymbol{y}, \theta^\star \right] \\ & = \sum_{s_1, \ldots, s_t} p(s_1, \ldots, s_t | \boldsymbol{y}, \theta^\star) \sum_{t=1}^T \log p(y_t | s_t, \theta) \\ & = \sum_{t=1}^T \sum_{k=1}^K P(s_t = k | \boldsymbol{y}, \theta^\star) \log f_k(y_t | \theta) \end{align*}\]

where \(f_k(\cdot)\) is the probability function for the observed distribution given latent state \(k\).

Letting \(\gamma_t(k) = p(s_t = k | \boldsymbol{y}, \theta^\star)\) and \(\phi_{t-1,t}(i, j) = p(s_{t-1} = i, s_{t} = j | \boldsymbol{y}, \theta^\star)\), we are left with the following.

\[ Q(\theta^\star, \theta) = \sum_{k=1}^K \gamma_1(k) \log \pi_k + \sum_{t=2}^T \sum_{i=1}^K \sum_{j=1}^K \phi_{t-1,t}(i, j) \log A_{ij} + \sum_{t=1}^T \sum_{k=1}^K \gamma_t(k) \log f_k(y_t | \theta) \tag{2}\]

E-Step

The expectation step consists of calculating \(\gamma_t(k)\) and \(\phi_{t-1,t}(i, j)\) for all \(k, i, j \in \{1, \ldots, K \}\) using the current iteration of \(\theta\).

The former can be expressed in terms of Bayes’ theorem and simplified.

\[ \begin{align*} \gamma_t(k) & = p(s_t = k | \boldsymbol{y}, \theta^\star) \\ & = \cfrac{p(y_1, \ldots, y_t | s_t = k, \theta^\star) P(s_t = k| \theta^\star)}{p(y_1, \ldots, y_t | \theta^\star)} \\ & = \cfrac{p(y_1, \ldots, y_t | s_t = k, \theta^\star) p(y_{t+1}, \ldots, y_t | s_t = k, \theta^\star) P(s_t = k | \theta^\star)}{\sum_{s_1, \ldots, s_t} p(\boldsymbol{y}, \boldsymbol{s} | \theta^\star)}\\ & = \cfrac{p(y_1, \ldots, y_t, s_t = k | \theta^\star) p(y_{t+1}, \ldots, y_t | s_t = k, \theta^\star)}{\sum_{i=1}^K p(y_1, \ldots, y_t, s_t = i | \theta^\star) p(y_{t+1}, \ldots, y_t | s_t = i, \theta^\star)} \\ & = \cfrac{\alpha_t(k) \beta_t(k)}{\sum_{i=1}^K \alpha_t(i) \beta_t(i)} \end{align*} \tag{3}\]

Note, the denominator is formed by marginalizing the joint probability function.

\[ p(\boldsymbol{y}) = \sum_{\boldsymbol{s}} p(\boldsymbol{y}, \boldsymbol{s}) = \sum_{s_t} \left( \sum_{\boldsymbol{s} \setminus s_t} p(\boldsymbol{y}, \boldsymbol{s}) \right) = \sum_{s_t} p(\boldsymbol{y}, s_t) \]

In addition, the variables \(\alpha_t(k)\) and \(\beta_t(k)\) were substituted to respectively represent the forward probabilities, i.e. the joint probability of observing \(y_1, \ldots, y_t\) and being in state \(k\) at time \(t\), and the backward probabilities, i.e., the probability of observing \(y_{t+1}, \ldots, y_T\) given latent state \(k\) at time \(t\).

These forward and backward probabilities are also used to express \(\phi_{t-1,t}(i, j)\).

\[ \begin{align*} \phi_{t-1,t}(i, j) & = P(s_{t-1} = i, s_t = j | \boldsymbol{y}, \theta^\star) \\ & = \cfrac{p(y_1, \ldots, y_t | s_{t-1} = i, s_t = j, \theta^\star) P(s_{t-1} = i, s_{t} = j | \theta^\star)}{p(\boldsymbol{y} | \theta^\star)} \\ & = \cfrac{p(y_1, \ldots, y_{t-1} | s_{t-1} = i, \theta^\star) p(y_t | s_t = j, \theta^\star) p(y_{t+1}, \ldots, y_t | s_t = j, \theta^\star) P(s_t = j | s_{t-1} = i, \theta^\star) P(s_{t-1} = i | \theta^\star)}{p(\boldsymbol{y} | \theta^\star)} \\ & = \cfrac{\alpha_{t-1}(i) f_j(y_t | \theta^\star) \beta_{t}(j) P(s_t = j | s_{t-1} = i, \theta^\star)}{p(\boldsymbol{y} | \theta^\star)} \end{align*} \tag{4}\]

It is clear that we need to be able to find expressions for \(\alpha_t(k)\) and \(\beta_t(k)\) as part of the E-Step. Unsurprisingly, in accordance with today’s theme, this will be done using recursion. Starting with the forward probabilities,

\[ \begin{align*} \alpha_t(k) & = p(y_1, \ldots, y_t | s_t = k, \theta^\star) P(s_t = k, \theta^\star) \\ & = p(y_t | s_t = k ,\theta^\star) \sum_{s_{t-1}} p(y_1, \ldots, y_{t-1}, s_{t-1}, s_t | \theta^\star) \\ & = p(y_t | s_t = k, \theta^\star) \sum_{s_{t-1}} p(y_1, \ldots, y_{t-1} | s_{t-1}, \theta^\star) P(s_t = k | s_{t-1} = i, \theta^\star) \\ & = f_k(y_t | \theta^\star) \sum_{i=1}^K \alpha_{t-1}(i) P(s_t = k| s_{t-1} = i, \theta^\star) \end{align*} \tag{5}\]

with initial value, \(\alpha_1(k) = p(y_1 | s_1 = k, \theta^\star) P(s_1 = k | \theta^\star) = \pi_k f_k(y_1 | \theta^\star)\). Note, as part of this derivation recognize that \(y_1, \ldots, y_{t-1} {\perp\!\!\!\perp} s_t | s_{t-1}\). In other words, the observed data up until \(t-1\) are independent of the future latent states given the present state allowing us to express \(p(y_1, \ldots, y_{t-1} | s_t, s_{t-1}) = p(y_1, \ldots, y_{y-1} | s_{t-1})\).

Before we actually implement this in code, however, consider what happens to \(\alpha_t(k)\) as \(t \to \infty\). The forward probabilities will decay exponentially and risk underflowing. Therefore, to improve numeric stability we replace \(\alpha_t(k)\) with its normalized form, \(\hat{\alpha}_t(k) = \frac{\alpha_t(k)}{c_t}\), where \(c_t = \sum_{i=1}^K \alpha_t(k)\).².

"""
Calculate normalized forward probabilities.

# Arguments
- `π`: Initial state probabilities
- `A`: Transition matrix
- `B`: Emission likelihoods, ie ``N(y | μ_k, σ_k)``
"""
function forward(π, A, B)
    α = Array{Float64}(undef, size(B, 1), length(π))
    c = Vector{Float64}(undef, size(B, 1))

    # Initial forward probabilities
    α[begin, :] .= B[begin, :] .* π

    # Normalize by scaling factor
    c[begin] = sum(α[begin, :])
    α[begin, :] ./= c[begin]

    for t in Iterators.drop(axes(B, 1), 1)
        for j in eachindex(π)
            α[t, j] = B[t, j] * sum(α[t-1, :] .* A[:, j])
        end

        c[t] = sum(α[t, :])
        α[t, :] ./= c[t]
    end

    return α, c
end

The backward probabilities, \(\beta_t(k)\), can be obtained similarly.

\[\begin{align*} \beta_t(k) & = p(y_{t+1}, \ldots, y_t | s_t = k, \theta^\star) \\ & = \sum_{s_{t+1}} p(y_{t+1}, \ldots, y_t, s_{t+1} | s_t = k, \theta^\star) \\ & = \sum_{s_{t+1}} p(y_{t+1}, \ldots, y_t | s_{t+1}, \theta^\star) p(s_{t+1} | s_t = k) \\ & = \sum_{i=1}^K p(y_{t+2}, \ldots, y_t | s_{t+1}) p(y_{t+1} | s_{t+1} = i) p(s_{t+1} = i, s_t = k) \\ & = \sum_{i=1}^K \beta_{t+1}(i) f_i(y_{t+1} | \theta^\star) p(s_{t+1} = i | s_t = k) \end{align*}\]

We also normalize the backward probabilities with the same scaling factors such that \(\hat{\beta}_t(k) = \frac{\beta_t(k)}{c_{t+1}}\).

"""
Calculate normalized backward probabilities.

# Arguments
- `π`: Initial state probabilities
- `A`: Transition matrix
- `B`: Emission likelihoods, ie ``N(y | μ_k, σ_k)``
- `c`: Scaling factors
"""
function backward(π, A, B, c)
    β = Array{Float64}(undef, size(B, 1), length(π))

    β[end, :] .= 1
    for t in axes(B, 1) |> reverse |> (x -> Iterators.drop(x, 1))
        for j in eachindex(π)
            β[t, j] = sum(β[t+1, i] * B[t+1, i] * A[j, i] for i in 1:length(π))
        end

        β[t, :] ./= c[t+1]
    end

    return β
end

Finally, implement the E-Step to update \(\gamma_t(k)\) and \(\phi_{t-1,t}(i,j)\) using Equation 3 and Equation 4.³

using Distributions

"""
EM algorithm E-step.

# Arguments:
- `X`: Observed data
- `π`: Initial state probabilities
- `A`: Transition matrix
- `μ`: Means of the emission distributions
- `σ`: Standard deviations of the emission distributions
"""
function E_step(X, π, A, μ, σ)
    B = hcat([pdf(Normal(μ[i], σ[i]), X) for i in 1:length(π)]...)

    α, c = forward(π, A, B)
    β = backward(π, A, B, c)

    γ = α .* β
    γ ./= sum(γ, dims=2)

    ϕ = Array{Float64}(undef, length(X) - 1, length(π), length(π))
    for t in Iterators.take(eachindex(X), length(X) - 1)
        ϕ[t, :, :] = α[t, :] .* A .* (B[t+1, :]' .* β[t+1, :]')
        ϕ[t, :, :] ./= sum(ϕ[t, :, :])
    end

    return  sum(log.(c)), α, ϕ, γ
end

M-Step

Returning to Equation 2, we update \(\theta = \{\boldsymbol{A}, \pi, \mu, \sigma \}\) in the M-Step by maximizing the objective function, \(\theta = \text{arg max}_{\theta \in \Theta} Q(\theta^\star, \theta)\), while treating \(\gamma_t(k)\) and \(\phi_{t-1,t}(i, j)\) as constant.

For brevity I will skip the full derivations, but solving the maximization problems leads to the following updating equations:⁴

\[\begin{align*} \pi_k & = \cfrac{\gamma_1(k)}{\sum_{i=1}^K \gamma_1(i)} \\ a_{ij} & = \cfrac{\sum_{t=2}^T \phi_{t-1,t}(i, j)}{\sum_{k=1}^K \sum_{t=2}^T \phi_{t-1,t}(i, k)} \\ \mu_k & = \cfrac{\sum_{t=1}^T \gamma_t(k) y_t}{\sum_{t=1}^T \gamma_t(k)} \\ \sigma_k^2 & = \cfrac{\sum_{t=1}^T \gamma_t(k) (y_t - \mu_k)^2}{\sum_{t=1}^T \gamma_t(k)} \end{align*}\]

using LinearAlgebra

"""
EM algorithm M-Step.

# Arguments:
- `X`: Observed data
- `γ`: Probability of being in state `k` given `X`
- `ϕ`: Joint probability of being in state ``i`` and ``j`` at times ``t-1`` and ``t`` given `X`
"""
function M_step(X, γ, ϕ)
    π = γ[1, :]

    q = dropdims(sum(ϕ, dims=1), dims=1)
    A = q ./ sum(q, dims=2)

    μ = sum(γ .* X, dims=1) ./ sum(γ, dims=1)
    σ = sqrt.([dot(γ[:, i], (X .- μ[i]).^2) / sum(γ[:, i]) for i in axes(γ, 2)])

    return π, A, vec(μ), max.(σ, 1e-5)
end

EM Function

The entry point for the EM implementation sets the initial values for \(\theta\) according to the values selected by the original paper.

It then iterates between the E-Step and M-Step until convergence with the stopping criteria also in line with the paper. Specifically, the algorithm stops when the absolute difference in estimates between iterations drops below a threshold. Formally, with some slight abuse in notation, we express this condition as \(\forall i \in \{1, 2, \ldots, n \}, \; |\theta_i - \theta_i^\star| < \epsilon\) for \(n\) parameters and \(\epsilon > 0\).

"""
Entry point for EM algorithm.
"""
function EM(X; max_iter=1_000, tol=1e-8)
    π = fill(0.5, 2)
    A = [0.5 0.5; 0.5 0.5]

    μ = [0, 3]
    σ = [1, sqrt(3)]

    for i in 1:max_iter
        # E-step
        log_lik, α, ϕ, γ = E_step(X, π, A, μ, σ)

        @info "Iteration $i: log-likelihood = $log_lik: μ = $μ"

        # M-step
        π_new, A_new, μ_new, σ_new = M_step(X, γ, ϕ)

        # Stopping criteria
        delta = Iterators.flatten([π_new - π, A_new - A, μ_new - μ, σ_new - σ])
        if all(abs.(delta) .< tol)
            return A_new[1], A_new[4], μ_new, α
        end

        π, A, μ, σ = π_new, A_new, μ_new, σ_new
    end

    error("Model failed to converge")
end

Run it all

Finally, run each region specific model.

# Enable logging to print the log-likelihood for each iteration
using Logging
disable_logging(Logging.Info)

# Fit an HMM per region
groups = groupby(df, :region_name) |> collect
results = [EM(d.totaldeaths) for d in groups]

We can summarise the model fits by re-constructing Table 2 from the original paper. Most of the results are fairly in line, with the notable exception of Londonderry/Strabane which we estimate as having a significantly less degree of persistence between states.

tbl = map(((p, q, μ, _),) -> (μ[2], μ[1], q, p), results) |>
      DataFrame
rename!(tbl, ["Mean(conflict)", "Mean(peace)", "P(Conflict | Conflict)", "P(Peace | Peace)"])
tbl.Region = map(d -> d.region_name[1], groups)

using PrettyTables
pretty_table(select(tbl, :Region, :);
    formatters = ft_printf("%5.3f"),
    backend = Val(:markdown),
    show_subheader=false)

We can also recreate Figure 4 from the original paper in order to plot the posterior expectations for the amount of violence per yearly quarter across four different regions.⁵ Again, Londonderry/Strabane stands out as our fitted HMM is more likely to transition between conflict and peace than the original results. Otherwise, the remaining region predictions are fairly in line with Besley and Mueller (2012).

function pp(df, result)
    _, _, μ, α = result
    p = plot(df.time, df.totaldeaths, label="Observed deaths",
        legend=:topright, title=df.region_name[1], xrotation=45,
        background_color=:transparent)
    plot!(p, df.time, α * μ, label="Predicted deaths")

    return p
end

plot([pp(groups[i], results[i]) for i in [1, 5, 3, 9]]..., layout=[2,2])

There is plenty of room for further improvement. We can investigate more robust initialization strategies, run sensitivity analyses, and implement a number of easy optimizations to the code in order to improve performance. Have fun, the world is your oyster.