Expectation maximization (EM) algorithm: full analysis from theory to practice

[Live broadcast preview] How far do we have to go to replace Oracle?

This article provides an in-depth exploration of the principles, mathematical foundations and applications of the expectation maximization (EM) algorithm. Through detailed definitions and specific examples, the article explains the application of EM algorithm in Gaussian Mixture Model (GMM), and conducts practical demonstrations through Python and PyTorch code implementation.

Follow TechLead and share all-dimensional knowledge of AI. The author has 10+ years of Internet service architecture, AI product development experience, and team management experience. He holds a master's degree from Tongji University in Fudan University, a member of Fudan Robot Intelligence Laboratory, a senior architect certified by Alibaba Cloud, a project management professional, and research and development of AI products with revenue of hundreds of millions. principal.

file

I. Introduction

The Expectation-Maximization Algorithm (EM algorithm for short) is an iterative optimization algorithm mainly used to estimate the parameters of probability models containing latent variables. It has extensive applications in machine learning and statistics, including but not limited to Gaussian Mixture Model (GMM), Hidden Markov Model (HMM), and various clustering and classification problems.

Probabilistic models and latent variables

A probabilistic model is a mathematical representation of a data-generating process. In statistics and machine learning, a probabilistic model is often used to describe the relationship between observable data and latent structure.

Example : Suppose we have a data set containing the height and weight of a group of people. A simple probabilistic model might assume that both height and weight are normally distributed.

**Latent Variables** refer to variables that cannot be directly observed, but will affect the observed data. Parameter estimation is generally more difficult in probabilistic models that include latent variables.

Example : When inferring whether a group of people likes sports, we may be able to observe their height and weight, but the latent variable "whether they like sports" cannot be directly observed.

Maximum Likelihood Estimation (MLE)

**Maximum Likelihood Estimation (MLE)** is a method used to estimate the parameters of a probabilistic model. It seeks a set of parameters that maximize the likelihood of occurrence of given observed data (i.e., the likelihood function).

Example : In a coin tossing experiment, 10 heads and 15 tails are observed, MLE will look for a parameter (the probability of the coin landing on heads) that makes it most likely to observe such data.

Jensen's inequality

Jensen's inequality is a basic inequality in convex optimization theory and is often used to prove the convergence of the EM algorithm. Simply put, Jensen's inequality states that for a convex function, the value of the function on the convex combination will not be greater than the average of the values of the points in the convex combination.

file

2. Basic mathematical principles

Before understanding the working mechanism of the EM algorithm, we need to master some key mathematical concepts and principles. These principles not only form the mathematical basis of the EM algorithm, but also help us understand the convergence and efficiency of the algorithm.

Conditional probability and joint probability

file

likelihood function

file

Kullback-Leibler divergence

file

Bayesian inference

Bayesian inference is a parameter estimation and model selection method based on Bayes' theorem. It uses prior probabilities, likelihood functions, and evidence (or normalization factors) to calculate the posterior probabilities of parameters.

Example : In spam classification, Bayesian inference can be used to update the probability of a spam (or not spam) email every time a user flags a new email.

These mathematical principles provide us with the solid foundation needed to understand the EM algorithm. By understanding these concepts, we can explore more deeply how the EM algorithm performs parameter estimation, especially in complex models with hidden variables.

3. The core idea of EM algorithm

file

The main purpose of the EM algorithm is to find parameter estimates for probabilistic models containing latent variables. This goal is particularly important when direct application of maximum likelihood estimation (MLE) is difficult or infeasible. The EM algorithm achieves this goal by alternating two steps: the expectation (E) step and the maximization (M) step.

Expectation (E) step

The Expectation step involves calculating the conditional expectation of the latent variable given the observed data and the current parameter estimate. This is often used to construct a function, called the Q-function, that approximates the target function (usually the likelihood function).

Example : In a Gaussian mixture model, the expectation step involves calculating the conditional probability that each observed data point belongs to the respective Gaussian distribution. These probabilities are also called posterior probabilities.

Maximize (M) steps

The maximization step (Maximization step) is to find parameter values that maximize the Q function given the Q function.

Example : Continuing with the Gaussian mixture model example above, the maximization step involves adjusting the mean and variance of each Gaussian distribution to maximize the Q function resulting from the expectation step.

Q function and auxiliary function

The Q function is a core concept in the EM algorithm and is used to approximate the objective function (such as the likelihood function). The Q function usually depends on observed data, latent variables and model parameters.

Example : In the EM algorithm of Gaussian mixture models, the Q function is defined based on the observation data and the posterior probability of each Gaussian distribution.

**Auxiliary Function** is an important part of the EM algorithm and is used to ensure algorithm convergence. By maximizing the auxiliary function, we indirectly maximize the likelihood function.

Example : In some text classification problems, auxiliary functions can be constructed through the Lagrange multiplier method to simplify the maximization problem.

Convergence

In the EM algorithm, due to the use of Jensen's inequality and auxiliary functions, the algorithm is guaranteed to converge to the local maximum.

Example : After implementing the EM algorithm for Gaussian mixture models, you will find that each iteration causes the value of the likelihood function to increase (or remain the same) until a local maximum is reached.

By delving into these core concepts and steps, we can more fully understand how the EM algorithm works and why it is so effective when dealing with complex probabilistic models containing latent variables.

4. EM algorithm and Gaussian mixture model (GMM)

Gaussian Mixture Model (GMM) is a probability model built based on Gaussian probability density function (pdf). It is a typical example of the application of EM algorithm, especially when we want to cluster or density estimate data.

Definition of Gaussian Mixture Model

Gaussian mixture model is composed of multiple Gaussian distributions. Each Gaussian distribution is called a component , and each component has its own mean ((\mu)) and variance ((\sigma^2)).

Example : Suppose a data set exhibits two distinct clusters. A Gaussian mixture model might describe the two clusters with two Gaussian distributions, each with its own mean and variance.

Component weight

Each Gaussian component has a weight ((\pi_k)) in the model, and this weight describes the "importance" of the component to the entire data set.

Example : In a GMM consisting of two Gaussian distributions, if one distribution has a weight of 0.7 and the other is 0.3, it means that the first distribution has a greater influence on the entire model.

Application of E step in GMM

In the E step in GMM , we calculate the posterior probability of the data point for each Gaussian component , that is, given the data point, the probability that it comes from a specific component.

Example : Suppose a data point (x), in step E, we calculate the posterior probability that it comes from each Gaussian component in the GMM.

# 使用Python和PyTorch计算后验概率
import torch
from torch.distributions import MultivariateNormal

# 假设有两个分量
means = [torch.tensor([0.0]), torch.tensor([5.0])]
variances = [torch.tensor([1.0]), torch.tensor([2.0])]
weights = [0.6, 0.4]

# 数据点
x = torch.tensor([1.0])

# 计算后验概率
posterior_probabilities = []
for i in range(2):
    normal_distribution = MultivariateNormal(means[i], torch.eye(1) * variances[i])
    posterior_probabilities.append(weights[i] * torch.exp(normal_distribution.log_prob(x)))

# 归一化
sum_probs = sum(posterior_probabilities)
posterior_probabilities = [prob / sum_probs for prob in posterior_probabilities]

print("后验概率:", posterior_probabilities)

Application of M step in GMM

In the M step , we update the parameters (mean and variance) of each Gaussian component based on the posterior probability calculated in the E step.

Example : Assume that the posterior probabilities of the data points for two Gaussian components are obtained from the E step. We will use these posterior probabilities to update the mean and variance in a weighted manner.

By exploring the Gaussian mixture model in detail and its connection to the EM algorithm, we gain a deeper understanding of how this complex model works and what role the EM algorithm plays in it. This not only helps us understand the mathematical basis of the algorithm, but also provides practical insights for practical applications.

5. Practical cases

In a practical case, we will use Python and PyTorch to implement a simple Gaussian Mixture Model (GMM) to demonstrate the application of the EM algorithm.

Definition: goal

Our goal is to cluster one-dimensional data. We will use two Gaussian components (that is, K=2).

Example : Suppose we have a 1D dataset containing two clusters. We want to find the parameters (mean and variance) of these two clusters using a GMM model.

Definition: input and output

Input : one-dimensional data array
Output : Parameters (mean and variance) of the two Gaussian components and their weights.

Implementation steps

Initialization parameters : Set initial values for mean, variance and weights.
Step E : Calculate the posterior probability that the data point belongs to each component.
M step : Update the mean, variance and weights using posterior probabilities.
Convergence check : Check whether the parameters converge. If not, return to step 2.

# Python和PyTorch代码实现
import torch
from torch.distributions import Normal

# 初始化参数
means = torch.tensor([0.0, 5.0])
variances = torch.tensor([1.0, 1.0])
weights = torch.tensor([0.5, 0.5])

# 假设的一维数据集
data = torch.cat((torch.randn(100) * 1.5, torch.randn(100) * 0.5 + 5))

# EM算法实现
for iteration in range(100):
    # E步骤
    posterior_probabilities = []
    for i in range(2):
        normal_distribution = Normal(means[i], torch.sqrt(variances[i]))
        posterior_probabilities.append(weights[i] * torch.exp(normal_distribution.log_prob(data)))
        
    # 归一化
    sum_probs = torch.stack(posterior_probabilities).sum(0)
    posterior_probabilities = [prob / sum_probs for prob in posterior_probabilities]

    # M步骤
    for i in range(2):
        responsibility = posterior_probabilities[i]
        means[i] = torch.sum(responsibility * data) / torch.sum(responsibility)
        variances[i] = torch.sum(responsibility * (data - means[i])**2) / torch.sum(responsibility)
        weights[i] = torch.mean(responsibility)

    # 输出当前参数
    print(f"Iteration {iteration+1}: Means = {means}, Variances = {variances}, Weights = {weights}")

Interpretation of results

After running the above code, you will see that the parameters for mean, variance, and weights are updated after each iteration. When these parameters no longer change significantly, we can consider the algorithm to have converged.

Input : A one-dimensional dataset containing two clusters.
Output : mean, variance and weight after each iteration.

Through this practical case, we not only demonstrated how to implement the EM algorithm in PyTorch, but also deeply understood each step of the algorithm through specific code examples. This content arrangement is designed to satisfy your need for content that is conceptually rich, full of detail, and well-defined.

6. Summary

After detailed theoretical analysis and practical examples, we have a more comprehensive understanding of the expectation maximization (EM) algorithm. From basic mathematical principles to specific implementation and applications, the EM algorithm demonstrates its powerful ability in statistical model parameter estimation, especially when we face missing or hidden data.

Selection of probabilistic model : Although we use Gaussian Mixture Model (GMM) in actual combat, the EM algorithm is not limited to this. In fact, it can be applied to any probabilistic model that meets certain conditions, which is especially important when studying and applying more complex data structures.
Importance of initialization : This article mentions the initial selection of parameters, but you should be more careful in actual applications. Poor initialization can cause the algorithm to fall into a local optimum, thus affecting model performance.
Convergence and efficiency : Although EM algorithms generally guarantee convergence, convergence speed can be an issue, especially in high-dimensional data and complex models. This may lead us to find more efficient optimization algorithms or use distributed computing.
Trade-off between model interpretability and complexity : The EM algorithm is capable of estimating the parameters of complex models, but this complexity may lead to reduced model interpretability. In practical applications, we need to carefully consider this trade-off.
Generalization ability of the algorithm : The EM algorithm is not only used for clustering problems, but is also widely used in many fields such as natural language processing and computational biology. Understanding its core ideas and working mechanisms can provide powerful tools for dealing with different types of data problems.

By deeply exploring these technical insights, we not only deepen our understanding of the core concepts and working mechanisms of the EM algorithm, but also better apply this algorithm to various practical problems. I hope this article furthers your understanding of complex probabilistic models and expectation maximization algorithms, and that you can find practical applications for this information in your own projects or research.

Follow TechLead and share all-dimensional knowledge of AI. The author has 10+ years of Internet service architecture, AI product development experience, and team management experience. He holds a master's degree from Tongji University in Fudan University, a member of Fudan Robot Intelligence Laboratory, a senior architect certified by Alibaba Cloud, a project management professional, and research and development of AI products with revenue of hundreds of millions. principal. If it helps, please pay more attention to TeahLead KrisChang, 10+ years of experience in the Internet and artificial intelligence industry, 10+ years of experience in technical and business team management, bachelor's degree in software engineering from Tongji, master's degree in engineering management from Fudan, Alibaba Cloud certified senior architect of cloud services, Head of AI product business with revenue of over 100 million.