论文研读系列——“Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions”

Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

arxiv [Submitted on 10 Feb 2021 , last revised 22 Oct 2021 ]

链接：arXiv:2102.05379v3

代码：https://github.com/ehoogeboom/multinomial_diffusion

https://github.com/didriknielsen/argmax_flows

摘要

本文介绍了两种创新方法，即Argmax Flows和Multinomial Diffusion，它们专为从分类数据中学习而设计，分类数据可以分为特定组，例如单词或图像的一部分，而不仅仅是数字。Argmax Flows 将连续模型与选择最大值的步骤相结合，旨在通过学习 argmax 函数的概率逆向来有效地处理分类数据。多项式扩散引入了一种新方法，逐步添加基于类别的噪声并学会将其移除，这是在生成或建模适合类别的数据方面向前迈出的一步。所提出的方法在测试中表现优于先前的模型，在文本建模和图像分割等任务中显示出更好的结果，这表明该领域取得了显著进步。该论文表明，未来的研究可能侧重于通过解决将离散变量提升到连续空间和移除自回归分量所面临的挑战来提高性能。

1 Introduction（介绍）

许多高维数据源是分类的，例如语言和图像分割。尽管自然图像在很大程度上已经通过生成流和扩散模型进行了研究，但分类数据却没有得到同样广泛的处理。目前，它们主要由自回归模型建模，从中采样成本高昂。

归一化流很有吸引力，因为它们可以设计成在评估和采样方向上都很快。通常，归一化流对连续分布进行建模。因此，直接优化离散数据上的流可能会导致任意高的可能性。在文献中，这个问题是通过在离散值周围的单位间隔内添加噪声来解决的。但是，由于这些方法是为有序数据设计的，因此它们不能很好地处理分类数据。

其他有吸引力的生成模型是扩散模型，由于目标会随着时间的推移而分解，因此训练速度很快。扩散模型通常具有固定的扩散过程，该过程会逐渐增加噪声。该过程由可学习的生成过程补充，该过程对信号进行降噪。前人经验表明，扩散模型也可以设计用于快速采样。到目前为止，扩散模型主要用于学习有序数据分布，例如自然图像。

因此，在本文中，作者介绍了分类变量的流动和扩散模型的扩展（如图 1 所示）： 1） Argmax Flows使用argmax变换和argmax的相应概率逆系列来弥合分类数据和连续归一化流动之间的差距；2）作者引入了多项式扩散，这是一个直接在分类变量上定义的扩散模型。与归一化流动相反，直接定义离散变量的扩散不需要梯度近似，因为扩散轨迹是固定的。作为我们工作的结果，生成归一化流和扩散模型可以直接学习分类数据。

2 Background（背景）

生成流（Generative flows）和扩散模型（Diffusion models）主要用于序数数据，例如自然图像。这些模型通常设计为在评估和采样方向上都快速。正规化流通常模拟连续分布，因此直接在离散数据上优化流可能导致任意高的似然。为了解决这个问题，文献中通过在离散值周围添加单位间隔的噪声来解决这个问题。然而，这些方法是为了序数数据设计的，它们在分类数据上表现不佳。

扩散模型由于其目标随时间步分解而训练速度快。扩散模型通常有一个固定的扩散过程，逐渐添加噪声。这个过程由一个可学习的生成过程补充，该过程去噪信号。至今，扩散模型主要被训练来学习序数数据分布，如自然图像。

$$p_V(v)=p_Z(z)\cdot |\det \frac{\partial z}{\partial v}|,v=g(z),$$

为了学习序数离散数据（如自然图像）上的密度，通常会添加量化噪声。Nielsen等人将量化重新解释为一个在一方面确定性（$x=round(v)$）和在另一方面随机性的映射 $v\mapsto x$ 的流层。使用这种解释，量化可以被视为在潜在变量模型中四舍五入操作的概率右逆：

$$P(x)=\int P(x|v)p(v)dv,P(x|v)=\delta (x=round(v)),$$

其中四舍五入操作是逐元素应用的。在这种情况下，密度模型 $p(v)$ 使用一个正规化流来建模。学习过程通过引入变分分布 $q(v|x)$ 来建模四舍五入映射的概率右逆，并优化证据下界（ELBO）：

$$\log P(x)\ge {\mathbb{E}}_{v\sim q(v|x)}[\log P(x|v)+\log p(v)-\log q(v|x)]={\mathbb{E}}_{v\sim q(v|x)}[\log p(v)-\log q(v|x)].$$

扩散模型由预定义的变分分布 $q(x_t|x_{t-1})$ 组成，这些分布随时间逐步添加噪声。扩散轨迹定义为 $q(x_t|x_{t-1})$ 在 $x_{t-1}$ 周围添加少量噪声。这样，信息逐渐被破坏，以至于在最终时间步，$x_T$ 几乎没有关于 $x_0$ 的信息。它们的生成对应物由可学习的分布 $p(x_{t-1}|x_t)$组成，这些分布学习去噪数据。当扩散过程添加的噪声量足够小的时候，它足以定义去噪轨迹，使用在维度轴上分解（无相关性）的分布。分布 $p(x_T)$ 被选择为与扩散轨迹逼近的分布相似。扩散模型可以使用变分推断进行优化：

$$\log P(x_0)\ge {\mathbb{E}}_{x_1,\dots ,x_T\sim q}\left[\log p(x_T)+\sum _{t=1}^T\log \frac{p(x_{t-1}|x_t)}{q(x_t|x_{t-1})}\right].$$

扩散的一个重要见解是，通过条件 $x_0$ ，后验概率 $q(x_{t-1}|x_t,x_0)=q(x_t|x_{t-1})q(x_{t-1}\mid x_0)/q(x_t\mid x_0)$ 是可行的，并且可以直接计算，允许通过低方差的KL散度重述：

$$\log P(x_0)\ge {\mathbb{E}}_q\left[\log p(x_0|x_1)-KL\left(q(x_T|x_0)|p(x_T)\right)-\sum _{t=2}^{T}KL\left(q(x_{t-1}|x_t,x_0)|p(x_{t-1}|x_t)\right)\right]$$

3 Argmax Flows

Argmax 流使用模型根据连续数据创建离散类别，从而更容易处理诸如从众多数据中选择最佳选项之类的事情。从该模型中采样很容易，但是要弄清楚特定结果的可能性可能非常困难，需要一种称为变分推断的变通方法。为确保样本运行良好，特殊规则（argmax 约束）确保模型仅考虑与预期最佳选择相匹配的结果，从而简化计算。

3.1 Probabilistic Inverse

Argmax 层有助于从列表中选择最大值，可以看作是一种特殊的流层。它使用变分分布从许多可能的结果中进行选择，类似于从四舍五入的结果中猜出原始数字。将该层与去量化层进行比较，后者处理有序数据，而 argmax 适用于类别。我们可以选择任何变分分布，只要它遵循 argmax 规则。一种名为 Thresholding 的方法会选择一个值，并通过为其他方法设置上限来确保该值是最大的。另一种方法使用 Gumbel 分布，该分布的特性使其适合选择最大值。Gumbel 方法会根据特定规则调整值，因此所选值始终是最大的。Gumbel 阈值结合了两种方法以获得更好的结果，混合使用噪点和调整。这些方法的有效性取决于它们模仿原始数据分布的程度。随着时间的推移，所选方法在匹配数据方面会变得更好，可以平滑所有粗糙的边缘。

3.2 Cartesian Products of Argmax Flows（Argmax Flows的笛卡尔积）

Argmax Flows 现在可以通过将单个类别分解成更小的部分来处理更多的类问题，例如将一个大数字转换为较小的数字或比特。此方法增加了表示类的方法数量，但可能会增加额外的未使用类，但它仍然可以很好地实现 ELBO 的优化目标。

4 Multinomial Diffusion

在本节中，作者引入了一种针对分类数据的替代似然模型：多项式扩散（Multinomial Diffusion）。与前几节不同，这里的$x_t$将以独热编码格式表示 x t ∈ { 0 , 1 } K x_t ∈ \{0, 1 \}^K xt​∈{0,1}K 。具体来说，对于类别$k$，$x_k=1$而 $x_j=0$ 对于所有的 $j\ne k$ 注意，为了清晰起见，这里省略了维度轴，因为所有分布都是独立于维度轴的。

作者使用具有以下参数的多项式分布定义多项式扩散过程：$q(x_t|x_{t-1})$，它具有${\beta }_t$的概率重新对类别进行均匀采样：

$$q(x_t|x_{t-1})=C(x_t|(1-{\beta }_t)x_{t-1}+{\beta }_t/K),$$

这里 $C$ 表示具有概率参数的多项式分布。进一步地，标量和向量的加法（和减法）是逐元素进行的。这个约定在本节中一直使用。由于这些分布形成了一个马尔可夫链，作者可以表示给定 $x_0$ 的任何 $x_t$ 的概率：

$$q(x_t|x_0)=C(x_t|{\overline{\alpha }}_t{x}_0+(1-{\overline{\alpha }}_t)/K)$$

这里 ${\alpha }_t=1-{\beta }_t$ 并且 ${\overline{\alpha }}_t={\prod }_{\tau =1}^t{\alpha }_{\tau }$ 。直观地说，对于每一个下一步时间，引入了少量的均匀噪声${\beta }_t$在K个类别上，并且以很大的概率$(1-{\beta }_t)$对$x_{t-1}$进行采样。我们可以封闭形式地计算多项式后验$q(x_{t-1}|x_t,x_0)$：

$$q(x_{t-1}|x_t,x_0)=C(x_{t-1}|{\theta }_{post}(x_t,x_0)),where{\theta }_{post}(x_t,x_0)=\overline{\theta }/\sum _{k=1}^K{\overline{\theta }}_{k}and\overline{\theta }=[{\alpha }_t{x}_t+(1-{\alpha }_t)/K]\odot [{\overline{\alpha }}_{t-1}{x}_0+(1-{\overline{\alpha }}_{t-1})/K].$$

图2：多项式扩散概览。一个生成模型 $p(x_{t-1}|x_t)$ 学习从左到右逐步去噪信号。一个推理扩散过程 $q(x_t|x_{t-1})$​​ 从右到左逐步添加噪声。

Ho等人（2020）的一个创新之处在于，他们不是直接预测生成轨迹的参数，而是预测噪声，使用后验方程$q$。尽管对离散数据预测噪声是困难的，我们预测从$x_t$到一个概率向量${\hat{x}}_0$，然后使用来自 $q(x_{t-1}|x_t,{\hat{x}}_0)$ 的概率向量参数化 $p(x_{t-1}|x_t)$ ，其中$x_0$是通过神经网络${\hat{x}}_0=\mu (x_t,t)$ 近似得到的。总结如下：

$$p(x_0|x_1)=C(x_0|{\hat{x}}_0)andp(x_{t-1}|x_t)=C(x_{t-1}|{\theta }_{post}(x_t,{\hat{x}}_0))where{\hat{x}}_0=\mu (x_t,t)$$

$KL$项可以通过枚举前两个方程中的概率，并计算离散分布的$KL$散度来简单计算，对于 $t\ge 2$ ：

$$KL(q(x_{t-1}|x_t,x_0)|p(x_{t-1}|x_t))=KL(C({\theta }_{post}(x_t,x_0))|C({\theta }_{post}(x_t,{\hat{x}}_0))),$$

5 Related Work（相关工作）

深度生成模型主要分为以下几类：自回归模型（ARMs）、变分自编码器（VAEs）、对抗网络（GANs）、正规化流（Normalizing Flows）、基于能量的模型（EBMs）和扩散模型（Diffusion Models）。

正规化流通常学习一个连续分布，并且需要量化（dequantization）来训练这些方法在序数数据如图像上的应用。大量工作致力于构建更具表现力的连续正规化流。为了使用正规化流学习序数离散分布，提出了在序数类别之间添加均匀噪声的方法，并在后续理论上进行了论证。基于变分推断提出了一种更强大的量化方法，并将其与自回归模型联系起来。还为二元变量提出了量化方法。Tran等人直接为分类变量提出了可逆变换。然而，这些方法可能因为梯度偏差而难以训练，并且在图像上的应用至今尚未展示。此外，还探索了针对序数离散数据（整数）的流。在其他工作中，VAEs已被适配为学习潜在空间的正规化流。但这些方法通常仍然使用argmax启发式进行采样，即使这不是训练期间指定的分布。

扩散模型最初由Sohl-Dickstein等人引入，他们为高斯和伯努利分布开发了扩散。最近，Ho等人展示了通过架构改进和预测的重新参数化，去噪扩散模型能够生成高维图像。扩散模型相对快速训练，但采样速度慢，因为它们需要迭代多个时间步长。Song等人和Nichol和Dhariwal展示了在实践中可以使用更少的步骤生成样本。Nichol和Dhariwal证明了对目标组成部分进行重要性加权可以显著提高对数似然性能。在Song等人中提出了扩散模型的连续时间扩展。在本文最初发布后，作者发现Song等人同时描述了一个离散扩散的框架，但没有实证评估。

6 Experiments（实验）

在实验中，作者比较了他们的方法在语言建模任务和无条件学习图像分割图上的性能。

6.1 Language data（语言数据）

在本节中，作者比较了他们的方法在两个语言数据集text8和enwik8上的性能。text8包含27个类别（从‘a’到‘z’和空格），而enwik8直接对字节进行建模，结果有256个类别。

Model description（模型描述） 作者测试了两个版本的生成性Argmax流：使用自回归（AR）流和基于耦合的流来定义p(v)。在这些实验中，概率逆基于阈值方法。具体来说，他们训练了一个条件对角高斯分布q(u|x)并进行了阈值处理，从而得到了分布q(v|x)。Argmax流定义在二进制笛卡尔积上。这意味着对于K=27，他们使用了一个5维的二进制空间，而对于K=256，他们使用了一个8维的二进制空间。Argmax流与当前训练生成流直接在离散数据上的标准化方法进行了比较：量化。他们比较了均匀和变分量化，其中在分类数据的独热表示上添加了(0,1)区间的噪声。自回归密度模型基于Lippe和Gavves（2020）提出的模型。耦合密度模型由8个流层组成，每层包括一个1×1的卷积和混合逻辑变换，这是Ho等人（2019）提出的。在多项式文本扩散模型中，µ网络由一个12层的Transformer模型。

Comparison with Generative Flows（与生成流的比较）

作者首先比较了直接在语言数据上训练的生成流的性能。这些实验使用了相同的底层正规化流：基于耦合的流或自回归流。注意到Argmax流一致地优于均匀和变分量化。这表明，对于一个生成流来说，使用argmax流学习提升的连续分布更容易。可能解释这一差异的Argmax流的一个优点是，它们将变量提升到整个欧几里得空间，而传统的量化只在(0,1)区间引入概率密度，留下没有概率密度的间隙。当比较基于耦合的方法时，Argmax流的性能提升更为明显。同时注意到，耦合流的性能比自回归流差，而图像上的差异通常较小。这表明为文本设计更具表现力的耦合层是一个有趣的未来研究方向。

Comparison with other generative models（与其它生成模型的比较）

与文献中的模型性能进行比较，结果与作者的Argmax流和多项式扩散一起展示。包含自回归组件的潜在变量方法用(AR)标记。尽管自回归流与ARMs具有相同的缺点，但它们提供了一个视角，说明性能不足来自何处。作者发现，他们的自回归Argmax流比VAE方法表现更好，它们优于AF/AF (Ziegler and Rush, 2019)和CategoricalNF (Lippe and Gavves, 2020)。

当比较非自回归模型时，Argmax流也优于将分类空间提升到连续空间的方法：IAF / SCF (Ziegler and Rush, 2019)。有趣的是，多项式文本扩散是一个非自回归模型，其性能甚至优于argmax耦合流，但比自回归版本差。对于这个模型，可能不同的q的扩散轨迹会有更好的性能，因为在当前形式下，去噪模型必须对输入噪声非常鲁棒。这些实验还突出了标准ARMs和（自回归）连续密度模型在文本上之间仍然存在明显的性能差距，可能与量化差距有关 (Nielsen and Winther, 2020)。图3展示了在text8上训练的不同模型的样本。由于重现Discrete Flows的结果存在困难，本节没有包括对离散流的比较和分析。相反，它们在附录C中进行了广泛讨论。

Unsupervised spell-checking

文本扩散模型的一个有趣的副产品是它可以用来进行拼写检查，仅使用一个前向传播。为了演示这一点，作者从测试数据中取出一个句子，并通过改变一些字符使其损坏。这个损坏的序列作为x1输入到生成去噪模型，该模型在步骤0时接近数据。然后去噪模型预测p(x0|x1)，并且最有可能的x0可以被建议。注意，这个模型只适用于字符级别的损坏，不适用于插入。一个示例如图5所示。由于模型选择最匹配的单词，更大的损坏最终将导致单词变化。

6.2 Segmentation maps（分割图）

对于图像类型的数据，作者引入了一个分类图像数据集：重新利用cityscapes数据集进行无条件的图像分割学习。与标准设置不同，需要学习的分割目标的分布不能依赖于照片。为了降低计算成本，作者使用最近邻插值将cityscapes的分割图重新采样为32×64像素的图像。他们使用全局类别作为预测目标，结果是一个8类问题。

Model description（模型描述）

Argmax流直接定义在K=8的分类空间上。密度模型p(v)使用由DenseNets参数化的仿射耦合层定义。对于概率逆，作者学习了一个条件流q(u|x)，它也基于仿射耦合结构。根据方法的不同，应用了softplus或Gumbel阈值化来获得v。回想一下，对于他们的第一个Gumbel方法，相当于将q(u|x)设置为单位均匀分布，而对于Gumbel阈值化，q(u|x)是学习的。作者比较了文献中现有的量化策略：均匀(Uria et al., 2013)和变分量化(Ho et al., 2019)，它们应用于独热表示。所有模型都使用相同的底层流架构，因此参数数量大致相同。例外的是均匀量化和Gumbel分布，因为不需要额外的变分流分布。

Comparison（比较）

本实验的结果如表4所示，以ELBO和如果可用的IWBO（重要性加权界限）表示，以每像素的比特数测量，越低越好。与语言实验一致，传统的量化方法（均匀/变分）被Argmax流超越。有趣的是，尽管使用softplus阈值化的argmax流实现了最佳的ELBO，但使用Gumbel阈值化方法的argmax流实现了更好的IWBO。多项式扩散模型在测试上表现稍差，为0.37 bpp，而训练时为0.33 bpp。有趣的是，这是唯一一个出现过拟合问题的模型，需要数据增强，这可能解释了这部分性能差异。对于所有其他模型，训练性能与测试和验证性能相当。在cityscapes上训练的不同模型的样本如图4所示。另一个有趣的点是，耦合流在生成一致的文本样本方面有困难（见图3），但在更像图像的cityscapes数据上没有这个问题。由于耦合层最初是为图像设计的（Dinh et al., 2015），它们可能需要调整以在文本上增加表现力。

7 Social Impact and Conclusion（社会影响和结论）

Social Impact（社会影响）

本文描述的方法可以用来学习分类分布。因此，它们可能被用来比迭代方法更快地生成高维分类数据，例如文本或图像分割图。可能的负面影响是生成假媒体形式的文本，或者对于客户服务来说非常无用的自动化聊天机器人。此工作可能对社会产生积极影响，例如通过新的文本生成方法或改进的自动驾驶汽车的分割。此外，作者的方法也可能被用于异常检测，以标记假内容。同时，作者相信，目前形式的方法还远离上述直接应用。

Conclusion（结论）

在本文中，作者提出了正规化流和扩散模型的两种扩展，以学习分类数据：Argmax流和多项式扩散。他们的实验表明，他们的方法在负对数似然方面优于可比模型。此外，作者的实验突出了该领域的明显性能差距：标准ARMs、连续自回归模型和非自回归连续模型之间。这表明未来的工作可以集中在两个降低性能的来源：1) 当离散变量被提升到连续空间时，以及2) 当移除自回归组件时。

最后感谢你看到这里，以上观点均为本人对原论文的个人理解，仅作个人学习使用，如有错误或侵权，麻烦联系我，本人定修改或删除。

祝你天天开心，多笑笑。