Categorical分布：深度学习中的离散建模利器

引言

对于深度学习研究者来说，概率分布是模型设计和优化的基石。在许多生成模型中，如变分自编码器（VAE）及其变种VQ-VAE（Vector Quantized Variational Autoencoder），我们需要对离散的潜在变量进行建模。这时，Categorical分布（分类分布）就成为一个自然而强大的工具。本文将深入探讨Categorical分布的数学定义、通俗解释，以及它在深度学习中的应用场景，特别是结合VQ-VAE的背景，剖析其为何在此类模型中扮演关键角色。

---

Categorical分布的数学定义

Categorical分布是离散概率分布的一种，用于描述一个随机变量在有限个互斥类别中取值的概率。假设我们有一个离散随机变量 ( $z$ )，它可以取 ( $K$ ) 个可能的值（类别），记为 ( { 1 , 2 , … , K } \{1, 2, \dots, K\} {1,2,…,K} )。Categorical分布的参数是一个概率向量 ( $\mathbf{p}=[p_1,p_2,\dots ,p_K]$ )，其中 ( $p_k$ ) 表示 ( $z$ ) 取值为 ($k$ ) 的概率，且满足以下约束：
$$\sum _{k=1}^K{p}_k=1,p_k\ge 0,\forall k$$

其概率质量函数（PMF）定义为：
p ( z = k ) = p k , k ∈ { 1 , 2 , … , K } p(z = k) = p_k, \quad k \in \{1, 2, \dots, K\} p(z=k)=pk​,k∈{1,2,…,K}

例如，若 ( $K=3$ )，( $\mathbf{p}=[0.2,0.5,0.3]$ )，则 ( $p(z=1)=0.2$ )，( $p(z=2)=0.5$ )，( $p(z=3)=0.3$ )。

Categorical分布是更广为人知的伯努利分布（( $K=2$ )）的多类别推广。当 ( $K=2$ ) 时，它退化为伯努利分布；当 ( $K>2$ ) 时，它描述了多于两种可能结果的情况。

---

通俗解释

想象你在一个有 ( $K$ ) 个选项的抽奖活动中，每个选项被抽中的概率由 ( $p_k$ ) 决定。比如一个三色转盘，红色占20%（( $p_1=0.2$ )），蓝色占50%（( $p_2=0.5$ )），绿色占30%（( $p_3=0.3$ )）。转一次转盘，结果只能是红、蓝、绿中的一种，而Categorical分布就是用来描述这种“单次选择”的概率模型。

在深度学习中，Categorical分布的“离散性”特别重要。当我们需要从一组离散的潜在表示（比如VQ-VAE中的码本向量）中挑选一个时，Categorical分布提供了一种简洁的方式来分配概率。

---

Categorical分布的作用：建模什么？

Categorical分布主要用于离散随机变量的建模，特别适合以下场景：

1. 分类任务的输出：在分类问题中，模型的最后一层（如softmax）通常输出一个Categorical分布，表示样本属于各个类别的概率。
2. 离散潜在变量：在生成模型中，如VQ-VAE，潜在空间被量化为离散的码本（codebook），Categorical分布用来描述从码本中选择某个向量的概率。
3. 序列建模中的决策：在自回归模型（如PixelCNN或WaveNet）中，每一步生成一个离散值时，可以用Categorical分布来建模输出。

简单来说，它是深度学习中处理“离散选择”的数学工具，尤其在需要从有限选项中采样时非常有用。

---

VQ-VAE中的Categorical分布

在VQ-VAE中，Categorical分布被用来建模离散潜在变量 ( $z$ ) 的先验分布 ( $p(z)$ )。让我们逐步分析它在此的应用：

VQ-VAE简介

VQ-VAE是一种生成模型，通过将编码器的输出量化为码本中的离散向量来构建离散潜在空间。具体来说：

• 编码器将输入 ( $x$ ) 映射到一个连续的特征图 ( $e$ )；
• 通过最近邻搜索，( $e$ ) 被量化为码本中的某个向量 ( $z_q$ )；
• 解码器从 ( $z_q$ ) 重建输入 ( $x$ )。

训练时，潜在变量 ( $z$ ) 是离散的，取值范围是码本的大小 ( $K$ )。因此，( $p(z)$ ) 需要一个离散分布来建模，而Categorical分布是自然选择。

为什么用Categorical分布？

1. 离散性匹配：VQ-VAE的核心创新是离散潜在空间，Categorical分布完美适配这种离散选择的需求。每个 ( $z$ ) 对应码本中的一个向量，( $p(z=k)$ ) 表示选择第 ( $k$ ) 个向量的概率。
2. 初始假设简单：在训练VQ-VAE时，先验 ( $p(z)$ ) 被设为均匀的Categorical分布（即 ( $p_k=1/K$ )），这简化了训练过程，避免了先验设计的复杂性。
3. 生成时的灵活性：训练完成后，可以通过拟合一个自回归分布（如PixelCNN或WaveNet）来替换均匀先验，使得 ( $p(z)$ ) 捕捉潜在变量间的依赖关系。这时，Categorical分布仍然是每一步采样的基础。

自回归扩展

文中提到，训练后可以用PixelCNN（对图像）或WaveNet（对音频）拟合一个自回归的 ( $p(z)$ )。这意味着 ( $p(z_i)$ ) 不再是独立的，而是依赖于之前的 ( $z_{<i}$ )：
$$p(z)=\prod _{i}p(z_i|z_{<i})$$
每一步的 ( $p(z_i|z_{<i})$ ) 仍然是一个Categorical分布，只是参数由自回归模型动态生成。这种方式通过祖先采样（ancestral sampling）生成 ( $z$ )，再由解码器生成 ( $x$ )，极大地提升了生成样本的质量。

为什么不联合训练？

文中（https://arxiv.org/pdf/1711.00937）指出，联合训练先验和VQ-VAE可能增强结果，但被留作未来研究。这是因为联合优化需要平衡编码器、解码器和先验的损失，计算复杂且容易不稳定。分开训练（先VQ-VAE，后先验）是更稳健的策略。

---

总结

Categorical分布是深度学习中处理离散变量的利器，其简单性（参数仅为概率向量）和普适性使其广泛应用于分类、生成模型等领域。在VQ-VAE中，它为离散潜在变量 ( $z$ ) 的先验 ( $p(z)$ ) 提供了数学基础，支持了从均匀分布到自回归分布的灵活建模，最终助力高质量样本生成。对于研究者而言，理解Categorical分布不仅有助于掌握VQ-VAE的原理，还能启发更多离散潜在空间的设计。