【技术追踪】ADDP：通过交替去噪扩散过程学习用于图像识别和生成的通用表示（ICLR-2024）

  扩散模型交替去噪：助力图像识别与图像生成~

论文：ADDP: Learning General Representations for Image Recognition and Generation with Alternating Denoising Diffusion Process
代码：https://github.com/ChangyaoTian/ADDP

0、摘要

   图像识别（recognition）和生成（generation） 长期以来一直是独立发展的。（图像识别包含分类、分割、目标检测~）随着最近对通用表示学习的趋势，为识别和生成任务开发通用表示也得到了推动。然而，初步尝试主要关注生成性能，在识别任务上表现仍然较差。（研究背景与意义）

  这些方法是在矢量量化（vector-quantized，VQ）空间中建模的，而主要的识别方法使用像素作为输入。本文的关键见解主要有两点：（1）以像素作为输入对于识别任务至关重要；（2）将 VQ tokens 作为重建目标，对于生成任务大有裨益。基于这些观察，本文提出了交替去噪扩散过程（Alternating Denoising Diffusion Process，ADDP），旨在通过单一的表示学习框架，将这两种空间整合起来。

  在每个去噪步骤中，本文的方法首先从先前的 VQ tokens 中解码出像素，然后从这些解码后的像素中生成新的 VQ tokens。扩散过程逐渐屏蔽一部分 VQ tokens，以构建训练样本。这些学习到的表示可以用来生成各种各样的高保真图像，并且在识别任务上表现出优秀的转移性能。

  大量实验表明，本文的方法在无条件生成、ImageNet 分类、COCO 目标检测以及ADE20k 语义分割等任务上均取得了具有竞争力的性能。尤为重要的是，ADDP 是首个成功开发出适用于生成任务和密集识别任务的通用表示方法。

Figure 1 | 学习用于生成和识别的通用表示的统一方法的推理流程：以往的方法要么完全在原始像素空间中建模（例如 iGPT），要么完全在矢量量化（VQ）标记空间中建模（例如 ViT-VQGAN 和 MAGE）；相比之下，ADDP 同时利用了这两个空间，在识别和生成任务上均取得了具有竞争力的性能表现；

1、引言

1.1、研究意义与当前挑战

  （1）识别任务旨在感知和理解视觉世界，而生成任务则旨在为各种应用创造新的视觉数据。目前的识别算法已在许多基准测试中超越了人类的表现，而目前的生成模型能够合成多样化且高保真的图像；

  （2）现有研究试图通过特定的生成模型来学习识别和生成的一般表示范式，如 Masked Image Modeling (MIM)，如 图1 所示，在生成过程中，它们迭代地恢复部分掩盖区域的图像内容；

  （3）目前许多生成方法（ViT-VQGAN 和 MAGE 等）都在 VQ 空间，而最先进的表示学习方法，如 MAE 和 BEiT，则均以原始图像像素作为输入；

  （4）原始像素作为输入对于识别任务至关重要，像素比 VQ tokens 更能保留空间敏感信息这在密集识别任务中尤为有用。如 表1 所示，在典型的识别任务中，使用像素作为输入优于 VQ tokens；
  
Table 1 | ：在识别任务中，将像素或 tokens 作为输入的性能：采用预训练的 VQ tokenizer 来生成 tokens。在相同的训练计划下（详见附录），使用 VQ-tokens 作为输入在所有任务上的表现都不如使用像素作为输入，而在密集识别任务上，这种差距甚至更大；

  基于这些观察，一个自然的问题出现了：是否可以在一个单一的表示学习框架中将这两个空间联系起来，允许模型感知原始像素并生成 latent visual tokens ？

1.2、本文贡献

  （1）提出了一种通用表示学习框架，通过交替去噪扩散过程（ADDP）连接像素空间和 token 空间；
  （2）广泛的实验表明，ADDP 在图像生成和识别任务中表现出色；
  （3）ADDP 是首个开发出适用于生成和密集识别任务的通用表示的方法；

2、方法

2.1、像素和 VQ tokens 交替去噪

  以往的研究主要是在连续的原始像素空间或离散的 VQ token 空间中完整地执行去噪过程。然而，鉴于原始像素和 VQ tokens 分别对识别和生成任务至关重要，本文提出交替地对像素和 VQ tokens 进行去噪，如 图2 所示。
  在每一步中，首先从先前生成的 VQ tokens 中解码出像素，然后基于解码后的像素生成新的 VQ tokens。为了关联这两个空间并实现交替去噪过程，引入了从 token 到像素的解码（token-to-pixel decoding）以及从像素到 token 的生成（pixel-to-token generation）。
  
Figure 2 | 交替去噪过程：

2.1.1、Token-to-Pixel Decoding

  token 到像素解码广泛用于图像生成，以将生成的 VQ tokens 恢复为可视图像像素。现成的预训练 VQ tokenizers 中的 VQ 解码器子网络可以直接用于执行此类解码。然而，现有的 VQ 解码器只能解码完整的 VQ 序列中的图像，而无法将部分 VQ tokens 作为输入。

  相比之下，本文的去噪过程在每一步中仅生成一部分 VQ tokens。为了解决这种不一致性，并便于使用现成的 VQ 解码器，本文提出将这些可靠的 tokens 与一些不可靠的 tokens 配对。

  具体地，在步骤 $t$ 中，给定部分生成的 VQ tokens $z_t$，进一步采样额外的互补 VQ tokens ${\bar{z}}_t$，以便形成用于解码像素图像 $x_t$ 的完整的 VQ 序列（这个 ${\bar{z}}_t$ 是怎么采的呢）。为了区分 $z_t$ 和 ${\bar{z}}_t$，分别称它们为可靠 tokens 和不可靠 tokens。注意，${\bar{z}}_t$ 仅用于解码像素图像，不会保留到下一步。条件概率 $p(x_t|z_t,{\bar{z}}_t)$ 定义为：

  其中，$\delta$ 表示狄拉克分布，$\odot$ 是逐元素乘积。$m_t$ 是一个二进制掩码，指示从 $z_t$ 导出的不可靠区域。$m_t$ 与 $z_t$ 和 ${\bar{z}}_t$ 具有相同的空间大小，其中二进制值为 1 的区域是不可靠的。可靠的 VQ tokens $z_t$ 和不可靠的 VQ tokens ${\bar{z}}_t$ 都将由模型预测。

2.1.2、Pixel-to-Token Generation

  在最近的表示学习工作中，像素到 token 的生成已被证明对识别任务是有效的。本文的方法引入了一个可学习的编码器-解码器网络，用于从噪声图像中预测 VQ tokens，从而实现从像素中进行表示学习。
  如 图4 所示，以先前解码的噪声图像 $x_t$ 作为输入，编码器 $E$ 首先从 $x_t$ 中提取表示，然后将提取的表示 $E(x_t)$ 中的不可靠区域（即 $m_t=1$) 替换为可学习的 $<$MASK$>$ token嵌入，然后再输入到解码器 $D$ 中。
  解码器将根据这些输入预测下一步的 $z_{t-1}$ 和 ${\bar{z}}_{t-1}$。给定带噪声的图像 $x_t$，生成在下一步 $t-1$ 时可靠的 VQ tokens $z_{t-1}$ 和不可靠的 VQ tokens ${\bar{z}}_{t-1}$ 条件概率表示为：

  其中 $m_t$ 与等式(1)中的二进制掩码相同，表示第 $t$ 步不可靠的区域。$E$ 和 $D$ 分别是可学习的编码器和解码器子网络。$e_{mask}$ 是一个可学习的 $<$MASK$>$ 标记嵌入，$\theta$ 表示整个网络的参数。由于网络以完整图像作为输入，因此可以自然地适应各种图像骨干网络（例如，CNNs 和 ViTs) 作为编码器。

2.1.3、Alternating Denoising Process

  交替去噪过程如 图2 所示。从一个空序列开始，该序列包含全 1 不可靠掩码 $m_{T+1}$，通过将所有 $<$MASK$>$ 标记输入到解码器 $D$ 中，预测 $p_{\theta }(z_T,{\bar{z}}_T|\varnothing )$。
  之后，在每一步 $t$，根据 公式(1) 解码带噪声的图像 $x_t$，然后根据 公式(2) 生成新的可靠 tokens $z_t$ 和不可靠 tokens ${\bar{z}}_t$。最后，合成的无噪声图像即为第 0 步的精炼图像 $x_0$。交替去噪过程中所有变量的联合分布定义为：

2.2、扩散过程

  根据之前的去噪扩散训练范例，本文同样为 ADDP 提出相应的扩散过程，如 图3 所示，给定图像 $x_0$，使用现成的预训练 VQ 编码器将 $x_0$ 映射到其对应的 VQ tokens ：$z_0=VQ$-$Encoder(x_0)$。随后，扩散过程通过马尔可夫链 $q(z_t|z_{t-1})$ 逐渐掩盖 $z_0$ 的某些区域。$z_{t-1}$ 对应于步骤 $t-1$ 时可靠的 VQ tokens，而步骤 $t$ 时的可靠 VQ tokens $z_t$ 是 $z_{t-1}$ 的一部分。整个过程总共包含 $T+1$ 个步骤，最终所有 tokens 都会被屏蔽。（这个 ${\bar{z}}_t$ 是这么来的）
  
Figure 3 | 扩散过程：

  对于不可靠的 VQ tokens ${\bar{z}}_{t-1}$，设计了一个正向过程 $q({\bar{z}}_{t-1}|z_t)$ 用于从 $z_t$ 生成 ${\bar{z}}_{t-1}$，使用可靠的 tokens $z_t$ 作为条件，而不是 $z_{t-1}$，表3 中的实验结果表明，使用 $q({\bar{z}}_{t-1}|z_t)$ 比使用 $q({\bar{z}}_{t-1}|z_{t-1})$ 更有优势。

  条件分布 $q({\bar{z}}_{t-1}|z_t)$ 是通过一个 token 预测器获得的，该预测器接收可靠的 tokens $z_t$ 作为输入，并预测不可靠的 tokens ${\bar{z}}_{t-1}$。由于我们的目标是用不可靠的 tokens 生成原始图像，$q({\bar{z}}_{t-1}|z_t)$ 的最优值应该是 $q(z_0|z_t)$。然而，$q(z_0|z_t)$ 通常难以处理，因此 token 预测器需要从数据样本中学习来近似 $q(z_0|z_t)$。

  为实现这一目标，模型仅从可靠的 tokens $z_t$ 中训练预测所有 tokens $z_0$。最后，由于我们的模型是训练来预测不可靠的 tokens，因此只需在训练过程中使用 token 预测器。整个扩散过程定义为：

2.3、学习去噪过程

  考虑到所提出的交替去噪扩散过程，ADDP 可以通过最小化 $log{p}_{\theta }(x_0)$的证据下界（ELBO）来优化，该证据下界由以下项组成：

  其中 $D_{KL}$ 表示 KL 散度。$L_{VQ}$ 对应于 VQ 分词器的训练，在本文的训练中被省略，因为使用了预训练的 VQ 分词器。$L_{1:T+1}$ 用于优化模型参数 $\theta$。

2.3.1、训练目标

  与之前的 MIM 方法一样，本文仅优化掩码标记上的损失函数 $L_t$（即 $m_t=1$ 的不可靠区域）。根据 VQ 扩散中的重新参数化技巧，$L_t$ 可以进一步简化为：

  这意味着 $q(z_0|z_t)$ 和 $q({\bar{z}}_{t-1}|z_t)$ 是模型的训练目标。由于 $q(z_0|z_t)$ 通常难以处理，使用 token 预测器 $q({\bar{z}}_{t-1}|z_t)$ 来估计 $q(z_0|z_t)$。因此，仅预测 $q({\bar{z}}_{t-1}|z_t)$ 是可行的。

2.3.2、训练过程

  如 图4 所示，给定图像 $x_0$，首先根据扩散过程计算可靠的 tokens $z_t$ 和不可靠的 tokens ${\bar{z}}_{t-1}$。然后将 $z_t$ 和 ${\bar{z}}_{t-1}$ 结合并输入到 VQ 解码器中生成合成图像 $x_t$。之后，像素到 token 生成网络以 $x_t$ 为输入，预测 ${\bar{z}}_{t-1}$ 的分布。训练目标 $q({\bar{z}}_{t-1}|z_t)$ 可以从 token 预测器中计算得出。（图4还是非常清楚明了的~）
  
Figure 4 | ADDP 的训练流程：

  本文构建的模型以噪声合成图像 $x_t=VQ$-$Decoder(z_t,{\bar{z}}_t)$ 作为训练输入。不可靠 ${\bar{z}}_{t-1}$ 的不适当值可能会损害合成图像的质量，降低图像生成和图像识别任务的性能。为了解决这个问题，本文没有直接使用从 $q({\bar{z}}_t|z_{t+1})$ 中采样的 ${\bar{z}}_t$，而是引入了一个映射函数 $f(q({\bar{z}}_t|z_{t+1}))$，用于处理 等式(1) 中的不可靠部分。

  映射函数 $f(\cdot )$ 的候选设计为 {Sampling, ArgMax, WeightedSum}。Sampling 是根据 $q({\bar{z}}_t|z_{t+1})$ 对 tokens 进行采样的直观设计。ArgMax 则是选择在 $q({\bar{z}}_t|z_{t+1})$ 中概率最大的 token。WeightedSum 表示在将嵌入向量输入到 VQ 解码器网络之前，会根据条件概率分布 $q({\bar{z}}_t|z_{t+1})$ 对词元码本中的所有 VQ 嵌入向量进行加权求和。根据经验，本文发现使用 WeightedSum 映射可以产生高质量的图像，并有助于提高模型性能。

  在实践中，本文还发现将可靠的 tokens $z_t$ 作为额外的信息输入解码器 $D$ 对识别和生成任务都有好处，因为当掩码比例相对较高时，仅从具有相当噪声的图像预测 VQ tokens 是困难的。

2.3.3、应用到图像生成

  对于图像生成，遵循 图5 (a) 中的去噪过程。从一个空序列开始，模型根据纯 $<$MASK$>$ tokens 嵌入预测 $z_T$ 和 ${\bar{z}}_T$。每一步，VQ 解码器从当前步骤的 $z_{t+1}$ 和 ${\bar{z}}_{t+1}$ 生成 $x_{t+1}$。然后，像素 $x_{t+1}$ 被输入到编码器-解码器网络中，以预测下一步的 $z_t$ 和 ${\bar{z}}_t$。$x_0$ 是最终生成的无噪声图像。
  
Figure 5 | ADDP 图像生成与识别的推理流程：

2.3.4、应用到图像识别

  如 图5 (b) 所示，ADDP 在预训练后也可应用于各种识别任务。编码器直接接收原始像素作为输入，而输出表示则被传递到不同的任务特定头。在相应数据集上微调后，ADDP 表现出色。

3、实验结果

3.1、实施细节

   （1）网络： VQ tokenizer 采用了 MAGE 团队发布的现成的 VQGAN 模型。MAGE-Large 被用作 token predictor。需要注意的是，token predictor 仅用于训练，在推理时可以丢弃。ADDP 的编码器采用了 ViT-Large，解码器由 8 个 Transformer 模块组成。

   （2）训练： ADDP 在 ImageNet-1k 数据集上进行训练，去噪步骤总数为 $T=100$。在掩码策略方面，从 $1,2,...,T$ 的值中采样步骤 $t$，同时确保掩码比率分布接近 MAGE，以便进行公平比较。使用 AdamW 优化器，其峰值学习率为 $1.5\times 10^{-3}$。模型训练了 800 个 epoch，其中预热 epoch 为 40 个，batch size 大小为 4096。

   （3）图像识别： 分别在 ImageNet-1k 上评估图像分类的迁移性能，在 COCO 上评估目标检测的迁移性能，以及在 ADE20k 上评估语义分割的迁移性能。预训练的编码器被用作主干网络，并针对不同的任务附加了特定的任务头。

   （4）图像生成： 训练完成后，采用与 MaskGIT 相同的迭代解码方式，逐步填充被遮掩的 tokens 并生成图像。默认情况下，每一步恢复一个 token。

3.2、主要结果

  
Table 2 | ADDP与现有各种方法在视觉识别和生成任务上的比较：

3.3、消融研究

  
Table 3 | 不可靠 tokens 的消融：
Table 4 | 映射函数消融：
Table 5 | token 输入策略消融：
Table 6 | 预测目标与 masking 比例消融：
Figure 6 | 生成任务 masking 尺度消融：

  倒了两遍，有点意思( •̀ ω •́ )✧