On the Expectation-Maximization (EM) algorithm and regression models
Before the main content: I am creating an R Community on Google Groups. You can join the group using this form.
Linear Regression
The Expectation-Maximization (EM) algorithm is an iterative optimization framework used to find maximum likelihood estimates of parameters when a model depends on unobserved, latent variables.
For linear regression, the EM algorithm is unnecessary because Ordinary Least Squares (OLS) yields a closed-form analytical solution, which is \(\hat{\beta} = (X^T X)^{-1} (X^T y)\). However, framing linear regression through the EM algorithm is a quite clarifying exercise before jumping to Poisson (or binomial/logit) regression, models where data is missing (e.g., Tobit), or where the dataset is generated by a Mixture of Linear Regressions (multiple hidden lines).
To understand how EM applies, consider a dataset with \(n\) observations and \(p < n\) variables (or features). Each observation \(i\) corresponds to a vector \(x_i = (x_{i1}, x_{i2}, \ldots, x_{ip})^T \in \mathbb{R}^p\), a scalar response \(y_i\), and the model is
\[ y_i = \sum_{j = 1}^p \beta_j x_{ij} + e_i,\quad e_i \sim N(0, \sigma^2), \]
where \(\beta = (\beta_1, \beta_2, \ldots, \beta_p)^T \in \mathbb{R}^p\) are the weights to be estimated.
Stacking all \(n\) observations gives the compact matrix form \(y = X\beta + e\), where
\[ X = \begin{pmatrix} x_{11} & x_{12} & \ldots & x_{1p} \\ x_{21} & x_{22} & \ldots & x_{2p} \\ \vdots & \ddots & & \\ x_{n1} & x_{n2} & \ldots & x_{np} \end{pmatrix} \in \mathbb{R}^{n \times p}, \quad y = \begin{pmatrix} y_1 \\ y_2 \\ \vdots \\ y_n \end{pmatrix} \in \mathbb{R}^n, \quad e = \begin{pmatrix} e_1 \\ e_2 \\ \vdots \\ e_n \end{pmatrix} \sim N(0, \sigma^2 I_n). \]
In the EM framing we treat the unobserved (latent) variables \(z_i = \sum_{j = 1}^p \beta_j x_{ij}\) as the complete data, even though we can compute their distribution exactly. The point is to establish the machinery that generalises to other models.
Setup: complete-data log-likelihood
If we observed both \(y_i\) and \(z_i\), the complete-data log-likelihood for \(\beta\) and \(\sigma^2\) would be
\[ \ell_c(\beta, \sigma^2) = -\frac{n}{2}\log(2\pi\sigma^2) - \frac{1}{2\sigma^2}\sum_{i=1}^n (y_i - z_i)^2. \]
because \(y_i \mid z_i \sim N(z_i, \sigma^2)\) (the noise model), and \(z_i\) is deterministic given \(\beta\).
Expectation step
Given current parameter estimates \(\beta^{(t)}\) and \(\sigma^{2(t)}\), the E-step computes the expected complete-data log-likelihood with respect to the conditional distribution \(p(z \mid y, \beta^{(t)}, \sigma^{2(t)})\).
For linear regression the latent variable \(z_i\) is fully determined by \(\beta\). There is no uncertainty once \(\beta\) is fixed, so the conditional expectation is just the current fitted value
\[ \mathbb{E}\left[z_i \mid y_i, \beta^{(t)}\right] = \sum_{j = 1}^p \beta_j^{(t)} x_{ij}. \]
In matrix notation this is the \(i\)-th entry of \(X \beta^{(t)}\).
Substituting into the complete-data log-likelihood yields the Q-function
\[ Q(\beta, \sigma^2 \mid \beta^{(t)}, \sigma^{2(t)}) = -\frac{n}{2}\log(2\pi\sigma^2) - \frac{1}{2\sigma^2}\sum_{i=1}^n \left(y_i - \sum_{j = 1}^p \beta_j^{(t)} x_{ij}\right)^2. \]
Or equivalently, using vector notation for the residual vector \(y - X \beta^{(t)}\),
\[ Q(\beta, \sigma^2 \mid \beta^{(t)}, \sigma^{2(t)}) = -\frac{n}{2}\log(2\pi\sigma^2) - \frac{1}{2\sigma^2} \| y - X \beta^{(t)} \|^2. \]
The E-step collapses to plugging in the current fitted values. No integration is needed.
Maximization step
The M-step updates the parameters by maximising \(Q\) with respect to \(\beta\) (and \(\sigma^2\)).
Updating \(\beta\). The only term in \(Q\) that depends on \(\beta\) is the sum of squared residuals. Applying the chain rule to \(Q\) with respect to \(\beta_j\) gives
\[ \frac{\partial Q}{\partial \beta_j} = -\frac{1}{2\sigma^2} \sum_{i=1}^n 2\!\left(y_i - \sum_{k=1}^p \beta_k x_{ik}\right)(-x_{ij}) = \frac{1}{\sigma^2} \sum_{i=1}^n x_{ij}\!\left(y_i - \sum_{k=1}^p \beta_k x_{ik}\right). \]
Setting this to zero and multiplying through by \(\sigma^2\):
\[ \sum_{i=1}^n x_{ij} y_i = \sum_{i=1}^n x_{ij} \sum_{k=1}^p \beta_k x_{ik} = \sum_{k=1}^p \beta_k \underbrace{\sum_{i=1}^n x_{ij} x_{ik}}_{(X^T X)_{jk}}, \quad j = 1, \ldots, p. \]
The left-hand side is the \(j\)-th entry of \(X^T y\), since \((X^T y)_j = \sum_i x_{ij} y_i\). The right-hand side is the \(j\)-th entry of \(X^T X \beta\), since the \((j,k)\) entry of \(X^T X\) is exactly \(\sum_i x_{ij} x_{ik}\). Stacking all \(p\) equations (\(j = 1, \ldots, p\)) into a single matrix equation:
\[ X^T X \beta = X^T y. \]
where \(X^T X \in \mathbb{R}^{p \times p}\) is a symmetric positive-definite matrix (assuming the columns of \(X\) are linearly independent) and \(X^T y \in \mathbb{R}^p\) is a vector of inner products between each feature and the response. For illustration with \(p = 3\), these look like:
\[ \underbrace{\begin{pmatrix} \sum x_{i1}^2 & \sum x_{i1}x_{i2} & \sum x_{i1}x_{i3} \\ \sum x_{i1}x_{i2} & \sum x_{i2}^2 & \sum x_{i2}x_{i3} \\ \sum x_{i1}x_{i3} & \sum x_{i2}x_{i3} & \sum x_{i3}^2 \end{pmatrix}}_{X^T X} \begin{pmatrix} \beta_1 \\ \beta_2 \\ \beta_3 \end{pmatrix} = \underbrace{\begin{pmatrix} \sum x_{i1} y_i \\ \sum x_{i2} y_i \\ \sum x_{i3} y_i \end{pmatrix}}_{X^T y}. \]
Inverting \(X^T X\) gives the OLS formula
\[ \beta^{(t+1)} = (X^T X)^{-1} X^T y. \]
Note that this solution is independent of the current iterate \(\beta^{(t)}\), so the algorithm converges in a single M-step regardless of the initialisation. This is consistent with the fact that OLS has a closed-form solution.
Updating \(\sigma^2\). Setting \(\partial Q / \partial \sigma^2 = 0\):
\[ \sigma^{2(t+1)} = \frac{1}{n} \sum_{i=1}^n \left(y_i - \sum_{j=1}^p \beta_j^{(t+1)} x_{ij}\right)^2 = \frac{1}{n}\|y - X \beta^{(t+1)}\|^2. \]
which is the mean squared residual at the updated weights.
Convergence
Because the M-step yields the global maximum of the Q-function in closed form and that maximum does not depend on \(\beta^{(t)}\), the sequence \(\{\beta^{(t)}\}\) reaches \(\hat{\beta} = (X^T X)^{-1} X^T y\) after exactly one iteration. This is the standard EM convergence property specialised to the case where the complete-data problem is a convex problem with a unique global maximum.
Poisson Regression
Poisson regression models count data. As an aside, international trade models rely on Poisson pseudo maximum likelihood (PPML) and a continuous variable such as exports (or imports) does not follow a discrete Poisson distribution, which is why the PPML and not PML name. The PPML estimator is consistent if the conditional mean of the variate of interest is correctly specified. More on that on The Log of Gravity page.
The response \(y_i \in \{0, 1, 2, \ldots\}\) is assumed to follow a Poisson distribution whose mean depends on the covariates through a log-link:
\[ y_i \sim \text{Poisson}(\mu_i), \quad \mu_i = \exp\!\left(\sum_{j=1}^p \beta_j x_{ij}\right) = \exp(x_i^T \beta). \]
Unlike linear regression there is no closed-form solution for \(\beta\), so we need an iterative method. The EM algorithm provides one by introducing latent variables that make the complete-data problem tractable.
Latent-variable construction
Write the Poisson mean as \(\mu_i = \exp(x_i^T \beta)\) and introduce \(m\) latent binary indicators \(z_{i1}, \ldots, z_{im}\). This is one for each of \(m\) hypothetical sub-processes that together generate \(y_i\). Specifically, partition \(\mu_i\) into \(m\) equal parts \(\lambda = \mu_i / m\) and let
\[ z_{il} \sim \text{Bernoulli}(\lambda / (1 + \lambda)), \quad l = 1, \ldots, m, \]
so that \(y_i = \sum_{l=1}^m z_{il}\) in the limit \(m \to \infty\).
In practice the standard EM formulation for Poisson regression avoids this explicit construction and instead treats the complete data as the pair \((y_i, \eta_i)\), where \(\eta_i = x_i^T \beta\) is the linear predictor, and exploits the exponential-family properties that apply to Poisson log-likelihood.
For more on the exponential family and statistical sufficiency, you can check Casella and Berger. It is one of my favourite books (unlike others that put elegance over clarity).
Setup: complete-data log-likelihood
The Poisson log-likelihood for a single observation is
\[ \log p(y_i \mid \beta) = y_i \log \mu_i - \mu_i - \log(y_i!) = y_i (x_i^T \beta) - \exp(x_i^T \beta) - \log(y_i!). \]
Summing over all \(n\) observations gives the complete-data log-likelihood (dropping the constant \(\sum_i \log(y_i!)\)):
\[ \ell(\beta) \propto \sum_{i=1}^n \left[ y_i (x_i^T \beta) - \exp(x_i^T \beta) \right] = y^T X \beta - \mathbf{1}^T \exp(X\beta), \]
where \(\exp(X\beta)\) denotes element-wise exponentiation and \(\mathbf{1}\) is a vector of ones.
Expectation step
Unlike linear regression, the Poisson log-likelihood is not quadratic in \(\beta\), so the E-step does not collapse trivially. The standard approach is to construct a working quadratic surrogate (the Q-function) at the current iterate \(\beta^{(t)}\) using a second-order Taylor expansion of \(\exp(x_i^T \beta)\) around \(\eta_i^{(t)} = x_i^T \beta^{(t)}\):
\[ \exp(x_i^T \beta) \approx \exp(\eta_i^{(t)}) + \exp(\eta_i^{(t)})(x_i^T \beta - \eta_i^{(t)}) + \frac{1}{2}\exp(\eta_i^{(t)})(x_i^T \beta - \eta_i^{(t)})^2. \]
As an aside, Taylor expansions provide the foundation for Newton’s Method. Both are used heavily in industry, and a famous example is Quake’s III Fast Inverse Square Root.
Substituting into \(\ell(\beta)\) and keeping only terms that depend on \(\beta\) gives the Q-function
\[ Q(\beta \mid \beta^{(t)}) \propto -\frac{1}{2} \sum_{i=1}^n \mu_i^{(t)} \!\left(x_i^T \beta - \eta_i^{(t)} - \frac{y_i - \mu_i^{(t)}}{\mu_i^{(t)}}\right)^{\!2}, \]
where \(\mu_i^{(t)} = \exp(\eta_i^{(t)})\).
Defining the working response
\[ \tilde{y}_i^{(t)} = \eta_i^{(t)} + \frac{y_i - \mu_i^{(t)}}{\mu_i^{(t)}} \]
and the weight \(\mu_i^{(t)}\), the Q-function becomes
\[ Q(\beta \mid \beta^{(t)}) \propto -\frac{1}{2} \sum_{i=1}^n \mu_i^{(t)} \!\left(\tilde{y}_i^{(t)} - x_i^T \beta\right)^2. \]
This is exactly a weighted least-squares objective. It is the same structure as the linear regression log-likelihood with complete data, but with observation-specific weights \(\mu_i^{(t)}\).
Maximization step
Maximising the weighted-least-squares Q-function with respect to \(\beta\) is the same calculation as in the linear regression M-step, but with a diagonal weight matrix \(W^{(t)} = \text{diag}(\mu_1^{(t)}, \ldots, \mu_n^{(t)}) \in \mathbb{R}^{n \times n}\).
Updating \(\beta\). The weighted normal equations are obtained by the same entry-wise gradient argument as before. For each \(j = 1, \ldots, p\):
\[ \frac{\partial Q}{\partial \beta_j} = \sum_{i=1}^n \mu_i^{(t)} x_{ij}\!\left(\tilde{y}_i^{(t)} - x_i^T \beta\right) = 0, \]
which in matrix form is
\[ X^T W^{(t)} X \beta = X^T W^{(t)} \tilde{y}^{(t)}. \]
Here \(X^T W^{(t)} X\) is symmetric positive-definite (it is the weighted Gram matrix), and \(X^T W^{(t)} \tilde{y}^{(t)} \in \mathbb{R}^p\). Solving gives
\[ \beta^{(t+1)} = \left(X^T W^{(t)} X\right)^{-1} X^T W^{(t)} \tilde{y}^{(t)}. \]
This is the Iteratively Reweighted Least Squares (IRLS) update, which is the standard algorithm for fitting Poisson (and other GLM) models. Each M-step is a weighted OLS problem with the same structure as the linear-regression result, but the weights \(W^{(t)}\) and working responses \(\tilde{y}^{(t)}\) change at every iteration as \(\mu_i^{(t)}\) is updated.
Convergence
Because the Poisson log-likelihood \(\ell(\beta)\) is strictly concave in \(\beta\) (the Hessian \(-X^T W X\) is negative-definite if \(X\) has full column rank), the sequence \(\{\beta^{(t)}\}\) converges to the unique maximum-likelihood estimate \(\hat{\beta}\). Unlike linear regression, convergence requires multiple iterations. The algorithm terminates when \(\|\beta^{(t+1)} - \beta^{(t)}\| < \varepsilon\) for a chosen tolerance \(\varepsilon\).
More about Poisson and other models
Check McCullagh and Nelder. It covers Poisson, Logit, and Generalized Linear Models with a very detailed treatment. This is another of my favourite books.