生成扩散模型漫谈（二十九）：用DDPM来离散编码
原文链接 | 2025-02-14 | 苏剑林
笔者前两天在arXiv刷到了一篇新论文《Compressed Image Generation with Denoising Diffusion Codebook Models》，实在为作者的天马行空所叹服，忍不住来跟大家分享一番。
如本文标题所述，作者提出了一个叫DDCM（Denoising Diffusion Codebook Models）的脑洞，它把DDPM的噪声采样限制在一个有限的集合上，然后就可以实现一些很奇妙的效果，比如像VQVAE一样将样本编码为离散的ID序列并重构回来。注意这些操作都是在预训练好的DDPM上进行的，无需额外的训练。
有限集合#
由于DDCM只需要用到一个预训练好的DDPM模型来执行采样，所以这里我们就不重复介绍DDPM的模型细节了，对DDPM还不大了解的读者可以回顾我们《生成扩散模型漫谈》系列的（一）、（二）、（三）篇。
我们知道，DDPM的生成采样是从\(\boldsymbol{x}_T\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\)出发，由下式迭代到\(\boldsymbol{x}_0\)：
\[\begin{equation}\boldsymbol{x}_{t-1} = \boldsymbol{\mu}(\boldsymbol{x}_t) + \sigma_t \boldsymbol{\varepsilon}_t,\quad \boldsymbol{\varepsilon}_t\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\label{eq:ddpm-g}\end{equation}\]
对于DDPM来说，每一步的迭代都需要采样一个噪声，加上\(\boldsymbol{x}_T\)本身也是采样的噪声，所以通过\(T\)步迭代生成\(\boldsymbol{x}_0\)的过程中，共传入了\(T+1\)个噪声向量，假设\(\boldsymbol{x}_t\in\mathbb{R}^d\)，那么DDPM的采样过程实际上是一个\(\mathbb{R}^{d\times (T+1)}\mapsto \mathbb{R}^d\)的映射。
DDCM的第一个奇思妙想是将噪声的采样空间换成有限集合（Codebook）：
\[\begin{equation}\boldsymbol{x}_{t-1} = \boldsymbol{\mu}(\boldsymbol{x}_t) + \sigma_t \boldsymbol{\varepsilon}_t,\quad \boldsymbol{\varepsilon}_t\sim\mathcal{C}_t\label{eq:ddcm-g}\end{equation}\]
以及\(\boldsymbol{x}_T\sim \mathcal{C}_{T+1}\)，其中\(\mathcal{C}_t\)是\(K\)个从\(\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\)预采样好的噪声集合。换句话说，采样之前就从\(\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\)采样好\(K\)个向量并固定不变，后面每次采样都从这\(K\)个向量中均匀采样。这样一来，生成过程变成了\(\{1,2,\cdots,K\}^{T+1}\mapsto \mathbb{R}^d\)的映射。
这样做对生成效果有什么损失呢？DDCM做了实验，发现损失很小：
可以看到，\(K = 64\)时已经基本追平FID，仔细观察就会发现其实\(K=2\)时也没损失多少，这表明采样空间有限化可以保持DDPM的生成能力。注意这里有个细节，每一步的\(\mathcal{C}_t\)是独立的，即所有Codebook加起来共有\((T+1)K\)个噪声，而非共享\(K\)个噪声。笔者做了简单的复现，发现共享的话需要\(K\geq 8192\)才能保持相近效果。
离散编码#
现在我们考虑一个逆问题：寻找给定\(\boldsymbol{x}_0\)的离散编码，即寻找相应的序列\(\boldsymbol{\varepsilon}_T\in \mathcal{C}_T,\cdots,\boldsymbol{\varepsilon}_1\in \mathcal{C}_1\)，使得迭代\(\boldsymbol{x}_{t-1} = \boldsymbol{\mu}(\boldsymbol{x}_t) + \sigma_t \boldsymbol{\varepsilon}_t\)生成的\(\boldsymbol{x}_0\)跟给定的样本尽可能接近。
既然将采样空间有限化后FID不会有明显变化，因此我们可以认为同分布的所有样本都可以通过式\(\eqref{eq:ddcm-g}\)来生成，所以上述逆问题的解理论上是存在的。可怎么把它找出来呢？直观的想法是从\(\boldsymbol{x}_0\)开始倒推，但细思之下我们会发现难以操作，比如第一步\(\boldsymbol{x}_0 = \boldsymbol{\mu}(\boldsymbol{x}_1) + \sigma_1 \boldsymbol{\varepsilon}_1\)，这里\(\boldsymbol{x}_1\)和\(\boldsymbol{\varepsilon}_1\)都是未知的，难以同时把它们定下来。
这时候DDCM的第二个奇思妙想登场了，它将\(\boldsymbol{x}_0\)的离散编码视为一个条件控制生成问题！具体来说，我们从固定的\(\boldsymbol{x}_T\)，通过如下方式来选择\(\boldsymbol{\varepsilon}_t\)：
\[\begin{equation}\boldsymbol{x}_{t-1} = \boldsymbol{\mu}(\boldsymbol{x}_t) + \sigma_t \boldsymbol{\varepsilon}_t,\quad \boldsymbol{\varepsilon}_t = \mathop{\text{argmax}}_{\boldsymbol{\varepsilon}\in\mathcal{C}_t} \boldsymbol{\varepsilon}\cdot(\boldsymbol{x}_0-\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t))\label{eq:ddcm-eg}\end{equation}\]
这里的\(\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)是用\(\boldsymbol{x}_t\)来预测\(\boldsymbol{x}_0\)的模型，它跟\(\boldsymbol{\mu}(\boldsymbol{x}_t)\)的关系是：
\[\begin{equation}\boldsymbol{\mu}(\boldsymbol{x}_t) = \frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\end{equation}\]
如果读者忘记了这部分内容，可以到《生成扩散模型漫谈（三）：DDPM = 贝叶斯 + 去噪》复习一下。
详细的推导我们下一节再讨论，现在先来观摩一下式\(\eqref{eq:ddcm-eg}\)，它跟随机生成的式\(\eqref{eq:ddcm-g}\)的唯一区别通过\(\text{argmax}\)来选择最优\(\boldsymbol{\varepsilon}_t\)，指标是\(\boldsymbol{\varepsilon}\)与残差\(\boldsymbol{x}_0-\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)的内积相似度，直观理解就是让\(\boldsymbol{\varepsilon}\)尽可能补偿当前\(\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)与目标\(\boldsymbol{x}_0\)的差距。通过迭代式\(\eqref{eq:ddcm-eg}\)，图片等价地被转化为\(T-1\)个整数（\(\sigma_1\)通常设为零）。
规则看上去很简单，那实际重构效果如何呢？下面截取了原论文的一个图，可以看到还是比较惊艳的，原论文还有更多效果图。笔者也在自己的模型尝试了一下，发现基本上能复现相近的效果，所以方法还是比较可靠的。
条件生成#
刚才我们说了，DDCM将编码过程视为一个条件控制生成过程，怎么理解这句话呢？我们从DDPM的式\(\eqref{eq:ddpm-g}\)出发，它可以等价地写成
\[\begin{equation}p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) = \mathcal{N}(\boldsymbol{x}_{t-1};\boldsymbol{\mu}(\boldsymbol{x}_t),\sigma_t^2\boldsymbol{I})\end{equation}\]
现在我们要做的事情是已知\(\boldsymbol{x}_0\)的前提下调控生成过程，所以我们往上述分布中多加一个条件\(\boldsymbol{x}_0\)，即改为\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\)。事实上在DDPM中，\(p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\)是有解析解的，我们在《生成扩散模型漫谈（三）：DDPM = 贝叶斯 + 去噪》已经求出过：
\[\begin{equation}p(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t, \boldsymbol{x}_0) = \mathcal{N}\left(\boldsymbol{x}_{t-1};\frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\boldsymbol{x}_0,\frac{\bar{\beta}_{t-1}^2\beta_t^2}{\bar{\beta}_t^2} \boldsymbol{I}\right)\end{equation}\]
或者写成
\[\begin{equation}\begin{aligned} \boldsymbol{x}_{t-1} =&\, \frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\boldsymbol{x}_0 + \frac{\bar{\beta}_{t-1}\beta_t}{\bar{\beta}_t}\boldsymbol{\varepsilon}_t \\ =&\, \underbrace{\frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)}_{\boldsymbol{\mu}(\boldsymbol{x}_t)} + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)) + \underbrace{\frac{\bar{\beta}_{t-1}\beta_t}{\bar{\beta}_t}}_{\sigma_t}\boldsymbol{\varepsilon}_t \end{aligned}\label{eq:ddcm-eg0}\end{equation}\]
其中\(\boldsymbol{\varepsilon}_t\sim\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\)。相比式\(\eqref{eq:ddpm-g}\)，上式多了\(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)一项，它用来引导生成结果向\(\boldsymbol{x}_0\)靠近。但别忘了我们的任务是离散编码\(\boldsymbol{x}_0\)，所以生成过程不能有\(\boldsymbol{x}_0\)显式参与，否则就因果倒置了。为此，我们寄望于\(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)这一项能从\(\boldsymbol{\varepsilon}_t\)中得到补偿，所以我们调整\(\boldsymbol{\varepsilon}_t\)的选择规则为
\[\begin{equation}\boldsymbol{\varepsilon}_t = \mathop{\text{argmin}}_{\boldsymbol{\varepsilon}\in\mathcal{C}_t} \left\Vert\frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)) - \frac{\bar{\beta}_{t-1}\beta_t}{\bar{\beta}_t}\boldsymbol{\varepsilon}\right\Vert\label{eq:ddcm-eps0}\end{equation}\]
由于\(\mathcal{C}_t\)的向量都是从\(\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\)预采样好的，所以跟《让人惊叹的Johnson-Lindenstrauss引理：理论篇》里面的“单位模引理”类似，我们可以认为\(\mathcal{C}_t\)的向量模长大致相同，在这个假设之下，上式也等价于
\[\begin{equation}\boldsymbol{\varepsilon}_t = \mathop{\text{argmax}}_{\boldsymbol{\varepsilon}\in\mathcal{C}_t} \boldsymbol{\varepsilon}\cdot(\boldsymbol{x}_0-\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t))\end{equation}\]
这就得到了DDCM的式\(\eqref{eq:ddcm-eg}\)。
重要采样#
在上述推导中，我们用到了“寄望于\(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)\)这一项能从\(\boldsymbol{\varepsilon}_t\)中得到补偿，所以将\(\boldsymbol{\varepsilon}_t\)的选择规则改为式\(\eqref{eq:ddcm-eps0}\)”的说法，这样虽然看起来比较直观，但从数学上来说是不严谨的，甚至严格来说是错误的。这一节我们来将这部分内容严格化。
再次回顾式\(\eqref{eq:ddcm-eg0}\)，前面的推导直到式\(\eqref{eq:ddcm-eg0}\)都是严谨的，不严谨的地方在于建立将式\(\eqref{eq:ddcm-eg0}\)与\(\mathcal{C}_t\)联系起来的方式。试想一下，按照\(\eqref{eq:ddcm-eps0}\)来选\(\boldsymbol{\varepsilon}_t\)，当\(K\to\infty\)时会发生什么？此时我们可以认为\(\mathcal{C}_t\)已经覆盖了整个\(\mathbb{R}^d\)，所以\(\frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)) = \frac{\bar{\beta}_{t-1}\beta_t}{\bar{\beta}_t}\boldsymbol{\varepsilon}_t\)，也就是说式\(\eqref{eq:ddcm-eg0}\)变成了确定性的变换：
\[\begin{equation}\boldsymbol{x}_{t-1} = \frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\bar{\boldsymbol{\mu}}(\boldsymbol{x}_t) + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)) = \frac{\alpha_t\bar{\beta}_{t-1}^2}{\bar{\beta}_t^2}\boldsymbol{x}_t + \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2}\boldsymbol{x}_0\end{equation}\]
换言之\(K\to\infty\)时无法恢复原本的随机采样轨迹，这其实并不是一件科学的事情。我们认为采样空间有限化应该是连续采样空间的一个近似，当\(K\to\infty\)时应该还原连续型的采样轨迹，这样离散化的每一步才更有理论保证，或者用一种更数学化的表达方式，就是我们认为必要条件是
\[\begin{equation}\lim_{K\to\infty} \text{DDCM} = \text{DDPM}\end{equation}\]
从式\(\eqref{eq:ddpm-g}\)变到式\(\eqref{eq:ddcm-eg0}\)，我们也可以认为是噪声分布从\(\mathcal{N}(\boldsymbol{0},\boldsymbol{I})\)换成了\(\mathcal{N}(\frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2\sigma_t}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)),\boldsymbol{I})\)，但出于编码的需求，我们不能直接从后者采样，而只能从有限集合\(\mathcal{C}_t\)中采样。为了满足必要条件，即为了使采样结果更贴近从\(\mathcal{N}(\frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2\sigma_t}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t)),\boldsymbol{I})\)采样，我们可以用它的概率密度函数对\(\boldsymbol{\varepsilon}\in\mathcal{C}_t\)进行加权：
\[\begin{equation}p(\boldsymbol{\varepsilon})\propto \exp\left(-\frac{1}{2}\left\Vert\boldsymbol{\varepsilon} - \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2\sigma_t}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t))\right\Vert^2\right),\quad \boldsymbol{\varepsilon}\in\mathcal{C}_t\label{eq:ddcm-p}\end{equation}\]
也就是说，最合理的方式应该是对\(-\frac{1}{2}\left\Vert\boldsymbol{\varepsilon} - \frac{\bar{\alpha}_{t-1}\beta_t^2}{\bar{\beta}_t^2\sigma_t}(\boldsymbol{x}_0 - \bar{\boldsymbol{\mu}}(\boldsymbol{x}_t))\right\Vert^2\)做Softmax后按概率做重要性采样，而不是直接取\(\text{argmax}\)。只不过当\(K\)比较小时，Softmax之后的分布会接近one hot分布，所以按概率采样约等于\(\text{argmax}\)了。
一般形式#
我们还可以将上述结果推广到Classifier-Guidance生成。根据《生成扩散模型漫谈（九）：条件控制生成结果》的推导，为式\(\eqref{eq:ddpm-g}\)加入Classifier-Guidance后的结果是
\[\begin{equation}\boldsymbol{x}_{t-1} = \boldsymbol{\mu}(\boldsymbol{x}_t) + \sigma_t^2 \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{y}|\boldsymbol{x}_t) + \sigma_t\boldsymbol{\varepsilon}_t,\quad \boldsymbol{\varepsilon}_t\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})\end{equation}\]
即新增了\(\sigma_t^2 \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{y}|\boldsymbol{x}_t)\)，其中\(p(\boldsymbol{y}|\boldsymbol{x}_t)\)是带噪样本分类器，如果我们只有无噪样本分类器\(p_o(\boldsymbol{y}|\boldsymbol{x})\)，那么我们也可以让\(p(\boldsymbol{y}|\boldsymbol{x}_t) = p_{o}(\boldsymbol{y}|\boldsymbol{\mu}(\boldsymbol{x}_t))\)。
在同样的推导之下，我们可以得到类似式\(\eqref{eq:ddcm-eps0}\)的选择规则
\[\begin{equation}\boldsymbol{\varepsilon}_t = \mathop{\text{argmin}}_{\boldsymbol{\varepsilon}\in\mathcal{C}_t} \left\Vert\sigma_t^2 \nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{y}|\boldsymbol{x}_t) - \sigma_t\boldsymbol{\varepsilon}\right\Vert\end{equation}\]
或者近似地
\[\begin{equation}\boldsymbol{\varepsilon}_t = \mathop{\text{argmax}}_{\boldsymbol{\varepsilon}\in\mathcal{C}_t} \boldsymbol{\varepsilon}\cdot\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{y}|\boldsymbol{x}_t)\end{equation}\]
当然也可以按照式\(\eqref{eq:ddcm-p}\)来构造采样分布
\[\begin{equation}p(\boldsymbol{\varepsilon})\propto \exp\left(-\frac{1}{2}\left\Vert\boldsymbol{\varepsilon} - \sigma_t\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{y}|\boldsymbol{x}_t)\right\Vert^2\right),\quad \boldsymbol{\varepsilon}\in\mathcal{C}_t\end{equation}\]
这些都只是前述结果的简单推广。
个人评析#
至此，我们对DDCM的介绍基本完毕，更多的细节大家请自行看原论文就好。作者还没有开源代码，这里笔者给出自己的参考实现：
如果大家对扩散模型已有一些了解并且手头上有现成的扩散模型，强烈推荐亲自试试。事实上DDCM的原理很好懂，代码也不难写，但只有亲自尝试过，才能体验到那种让人拍案叫绝的惊艳感。上一次给笔者相同感觉的是《生成扩散模型漫谈（二十三）：信噪比与大图生成（下）》介绍的一种免训练生成大图的技巧Upsample Guidance，同样体现了作者别出心裁的构思。
但从长远影响力来说，个人认为Upsample Guidance还是不如DDCM的。因为对图片做离散化编码是多模态LLM的主流路线之一，它充当着“图片Tokenizer”的角色，是相当关键的一环，而DDCM可以说是开辟了VQ、FSQ以外的全新路线，因此可能会有更深远的潜在影响。在原论文中，DDCM只将自己定义为Compression方法，反而有点“自视甚低”了。
作为一个离散编码模型，DDCM还有一个非常突出的优点是它出来的离散编码天然就是1D的，而不像VQ、FSQ等方案一样编码结果通常保留了图像的2D特性（TiTok等用了Q-Former思想去转1D的模型除外），这意味着我们在用这些编码做自回归生成时就不用再考虑“排序”这个问题了（参考《“闭门造车”之多模态思路浅谈（二）：自回归》），这一点会显得非常省心。
当然，目前看来它还有一些改进空间，比如目前的编码跟生成是同时进行的，这意味着DDPM的采样速度有多慢，DDCM的编码速度就有多慢，这在目前来说还是不大能接受的。偏偏我们还不能随便上加速采样的技巧，因为加速采样意味着减少了\(T\)，而减少\(T\)意味着缩短了编码长度，即增大了压缩率，那么会明显增加重构损失。
总的来说，个人认为DDCM是一种非常有意思的、潜力亟待挖掘、同时也亟待进一步优化的离散编码方法。
文章小结#
本文介绍了扩散模型的一个新脑洞，它将DDPM生成过程中的噪声限制在一个有限的集合上，并结合条件生成的思路，将DDPM免训练地变成一个类似VQ-VAE的离散自编码器。