Categorical分布：深度学习中的离散建模利器
引言
对于深度学习研究者来说，概率分布是模型设计和优化的基石。在许多生成模型中，如变分自编码器（VAE）及其变种VQ-VAE（Vector Quantized Variational Autoencoder），我们需要对离散的潜在变量进行建模。这时，Categorical分布（分类分布）就成为一个自然而强大的工具。本文将深入探讨Categorical分布的数学定义、通俗解释，以及它在深度学习中的应用场景，特别是结合VQ-VAE的背景，剖析其为何在此类模型中扮演关键角色。
Categorical分布的数学定义
Categorical分布是离散概率分布的一种，用于描述一个随机变量在有限个互斥类别中取值的概率。假设我们有一个离散随机变量 ( z z z )，它可以取 ( K K K ) 个可能的值（类别），记为 ( { 1 , 2 , ... , K } \{1, 2, \dots, K\} {1,2,...,K} )。Categorical分布的参数是一个概率向量 ( p = [ p 1 , p 2 , ... , p K ] \mathbf{p} = [p_1, p_2, \dots, p_K] p=[p1,p2,...,pK] )，其中 ( p k p_k pk ) 表示 ( z z z ) 取值为 ( k k k ) 的概率，且满足以下约束：
∑ k = 1 K p k = 1 , p k ≥ 0 , ∀ k \sum_{k=1}^K p_k = 1, \quad p_k \geq 0, \quad \forall k k=1∑Kpk=1,pk≥0,∀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 K = 3 K=3 )，( p = [ 0.2 , 0.5 , 0.3 ] \mathbf{p} = [0.2, 0.5, 0.3] p=[0.2,0.5,0.3] )，则 ( p ( z = 1 ) = 0.2 p(z=1) = 0.2 p(z=1)=0.2 )，( p ( z = 2 ) = 0.5 p(z=2) = 0.5 p(z=2)=0.5 )，( p ( z = 3 ) = 0.3 p(z=3) = 0.3 p(z=3)=0.3 )。
Categorical分布是更广为人知的伯努利分布（( K = 2 K=2 K=2 )）的多类别推广。当 ( K = 2 K=2 K=2 ) 时，它退化为伯努利分布；当 ( K > 2 K>2 K>2 ) 时，它描述了多于两种可能结果的情况。
通俗解释
想象你在一个有 ( K K K ) 个选项的抽奖活动中，每个选项被抽中的概率由 ( p k p_k pk ) 决定。比如一个三色转盘，红色占20%（( p 1 = 0.2 p_1 = 0.2 p1=0.2 )），蓝色占50%（( p 2 = 0.5 p_2 = 0.5 p2=0.5 )），绿色占30%（( p 3 = 0.3 p_3 = 0.3 p3=0.3 )）。转一次转盘，结果只能是红、蓝、绿中的一种，而Categorical分布就是用来描述这种"单次选择"的概率模型。
在深度学习中，Categorical分布的"离散性"特别重要。当我们需要从一组离散的潜在表示（比如VQ-VAE中的码本向量）中挑选一个时，Categorical分布提供了一种简洁的方式来分配概率。
Categorical分布的作用：建模什么？
Categorical分布主要用于离散随机变量的建模，特别适合以下场景：
- 分类任务的输出：在分类问题中，模型的最后一层（如softmax）通常输出一个Categorical分布，表示样本属于各个类别的概率。
- 离散潜在变量：在生成模型中，如VQ-VAE，潜在空间被量化为离散的码本（codebook），Categorical分布用来描述从码本中选择某个向量的概率。
- 序列建模中的决策：在自回归模型（如PixelCNN或WaveNet）中，每一步生成一个离散值时，可以用Categorical分布来建模输出。
简单来说，它是深度学习中处理"离散选择"的数学工具，尤其在需要从有限选项中采样时非常有用。
VQ-VAE中的Categorical分布
在VQ-VAE中，Categorical分布被用来建模离散潜在变量 ( z z z ) 的先验分布 ( p ( z ) p(z) p(z) )。让我们逐步分析它在此的应用：
VQ-VAE简介
VQ-VAE是一种生成模型，通过将编码器的输出量化为码本中的离散向量来构建离散潜在空间。具体来说：
- 编码器将输入 ( x x x ) 映射到一个连续的特征图 ( e e e )；
- 通过最近邻搜索，( e e e ) 被量化为码本中的某个向量 ( z q z_q zq )；
- 解码器从 ( z q z_q zq ) 重建输入 ( x x x )。
训练时，潜在变量 ( z z z ) 是离散的，取值范围是码本的大小 ( K K K )。因此，( p ( z ) p(z) p(z) ) 需要一个离散分布来建模，而Categorical分布是自然选择。
为什么用Categorical分布？
- 离散性匹配 ：VQ-VAE的核心创新是离散潜在空间，Categorical分布完美适配这种离散选择的需求。每个 ( z z z ) 对应码本中的一个向量，( p ( z = k ) p(z=k) p(z=k) ) 表示选择第 ( k k k ) 个向量的概率。
- 初始假设简单 ：在训练VQ-VAE时，先验 ( p ( z ) p(z) p(z) ) 被设为均匀的Categorical分布（即 ( p k = 1 / K p_k = 1/K pk=1/K )），这简化了训练过程，避免了先验设计的复杂性。
- 生成时的灵活性 ：训练完成后，可以通过拟合一个自回归分布（如PixelCNN或WaveNet）来替换均匀先验，使得 ( p ( z ) p(z) p(z) ) 捕捉潜在变量间的依赖关系。这时，Categorical分布仍然是每一步采样的基础。
自回归扩展
文中提到，训练后可以用PixelCNN（对图像）或WaveNet（对音频）拟合一个自回归的 ( p ( z ) p(z) p(z) )。这意味着 ( p ( z i ) p(z_i) p(zi) ) 不再是独立的，而是依赖于之前的 ( z < i z_{<i} z<i )：
p ( z ) = ∏ i p ( z i ∣ z < i ) p(z) = \prod_{i} p(z_i | z_{<i}) p(z)=i∏p(zi∣z<i)
每一步的 ( p ( z i ∣ z < i ) p(z_i | z_{<i}) p(zi∣z<i) ) 仍然是一个Categorical分布，只是参数由自回归模型动态生成。这种方式通过祖先采样（ancestral sampling）生成 ( z z z )，再由解码器生成 ( x x x )，极大地提升了生成样本的质量。
为什么不联合训练？
总结
Categorical分布是深度学习中处理离散变量的利器，其简单性（参数仅为概率向量）和普适性使其广泛应用于分类、生成模型等领域。在VQ-VAE中，它为离散潜在变量 ( z z z ) 的先验 ( p ( z ) p(z) p(z) ) 提供了数学基础，支持了从均匀分布到自回归分布的灵活建模，最终助力高质量样本生成。对于研究者而言，理解Categorical分布不仅有助于掌握VQ-VAE的原理，还能启发更多离散潜在空间的设计。