深入理解DL-Workshop中的Dirichlet过程高斯混合聚类

引言

在机器学习和统计学中，聚类分析是一项基础而重要的任务。传统的聚类方法如K-means需要预先指定聚类数量，这在实际应用中往往是一个挑战。本文将深入探讨DL-Workshop项目中关于Dirichlet过程高斯混合模型(DP-GMM)的实现，这是一种非参数贝叶斯方法，能够自动确定最优的聚类数量。

问题背景

复杂的数据分布

我们首先考虑一个具有挑战性的数据分布问题：一个混合了多个高斯分布的数据集，但混合后的整体分布视觉上难以辨别真实的成分数量。

# 生成混合高斯数据
weights_true = np.array([2, 10, 1, 6])
locs_true = np.array([-2., -5., 3., 8.])
scale_true = np.array([1.1, 2, 1., 1.5])

base_n_draws = 1000
key = random.PRNGKey(42)
keys = random.split(key, 4)

draws = []
for i in range(4):
    shape = int(base_n_draws * weights_true[i]),
    draw = scale_true[i] * random.normal(keys[i], shape=shape) + locs_true[i]
    draws.append(draw)
data_mixture = np.concatenate(draws)

使用直方图可视化时，数据看起来可能只有2-3个主要成分，但实际上包含4个高斯分布。这展示了传统可视化方法的局限性。

Dirichlet过程高斯混合模型(DP-GMM)

模型概述

DP-GMM是一种非参数贝叶斯方法，它假设存在无限(或足够大)数量的潜在状态，每个状态对应一个高斯分布。与有限混合模型不同，DP-GMM不需要预先指定成分数量，而是通过Dirichlet过程自动确定。

数据生成过程

1. 从Dirichlet过程中生成状态概率分布
2. 根据状态概率选择对应的高斯分布
3. 从选定的高斯分布中生成观测数据

模型构建

组件权重对数似然

def component_probs_loglike(log_component_weights, log_concentration, num_components):
    """计算Dirichlet过程下组件权重的对数似然"""
    # 实现细节...

这部分计算在给定浓度参数下，组件权重向量的对数似然。

高斯混合对数似然

def mixture_loglike(log_component_weights, component_mus, log_component_scales, data):
    """计算高斯混合模型的对数似然"""
    # 实现细节...

这部分计算观测数据在高斯混合模型下的对数似然。

联合对数似然

def joint_loglike(log_component_weights, log_concentration, num_components,
                 component_mus, log_component_scales, data):
    """联合对数似然函数"""
    component_probs = np.exp(log_component_weights)
    probs_ll = component_probs_loglike(...)
    mix_ll = mixture_loglike(...)
    return probs_ll + mix_ll

联合似然将两部分结合起来，完整描述了模型的数据生成过程。

模型优化

损失函数

def make_joint_loss(num_components):
    def inner(params, data):
        # 解包参数
        ll = joint_loglike(...)
        return -ll  # 最大化对数似然等价于最小化负对数似然
    return inner

使用闭包模式固定组件数量，便于优化其他参数。

梯度计算

djoint_loss = grad(joint_loss)  # 自动微分计算梯度

训练过程

1. 参数初始化
2. 使用Adam优化器
3. 通过lax.scan高效执行训练循环

# 初始化参数
n_components = 50
log_component_weights_init = random.normal(...)
log_concentration_init = random.normal(...)
component_mus_init = random.normal(...)
log_component_scales_init = random.normal(...)

# 训练循环
adam_init, adam_update, adam_get_params = adam(0.05)
step_scannable = make_step_scannable(...)
step_scannable = jit(step_scannable)  # 即时编译加速

initial_state = adam_init(params_init)
final_state, state_history = lax.scan(step_scannable, initial_state, np.arange(N_STEPS))

结果分析

训练可视化

通过动画展示训练过程中高斯分布的变化：

animation = animate_training(params_for_plotting, int(N_STEPS / 200), data_mixture)

可以观察到模型如何逐步识别数据中的主要成分。

损失曲线

plt.plot(losses)
plt.yscale("log")

损失函数随着训练逐渐下降，表明模型在不断改进。

最终权重分布

plt.plot(normalize_weights(component_weights_opt), marker="o")

最终模型能够正确识别数据中的主要成分，并为次要成分分配较低的权重。

技术要点与注意事项

1. JAX的优势：充分利用vmap、lax.scan、grad和jit等特性，实现高效计算
2. 标签交换问题：在训练过程中可能出现成分标签交换，这在MCMC中是个问题，但在梯度优化中影响较小
3. 初始化敏感性：模型对初始参数可能敏感，需要适当调整学习率和训练步数
4. 组件数量选择：虽然DP-GMM理论上支持无限组件，实践中仍需选择足够大的初始组件数量

总结

DL-Workshop中的DP-GMM实现展示了如何将梯度下降应用于无监督学习问题。通过构建适当的对数似然函数并利用自动微分，我们能够:

1. 自动确定数据中的聚类数量
2. 估计每个聚类的参数(均值、方差)
3. 计算每个聚类的相对权重

这种方法不仅限于聚类问题，其核心思想——将问题表述为优化问题并使用梯度下降求解——可以推广到许多统计建模和机器学习任务中。DP-GMM为非参数贝叶斯建模提供了一个强大的工具，特别适用于聚类数量未知或可能变化的应用场景。