DiTok 相关工作回顾

扩散模型

从 DDPM 到 DDIM

扩散模型的理论根源可以追溯到非平衡热力学和随机过程。其核心思想是通过逐步添加噪声将数据转化为纯噪声，然后学习这个过程的逆向来生成新数据。

2015年，Sohl-Dickstein等人首次提出扩散概率模型的完整框架，但训练困难，生成质量不理想。

2020年，Ho等人发表《Denoising Diffusion Probabilistic Models》简化了训练目标，将其转化为去噪任务，首次在图像生成上展现出与GAN相当的质量。但存在致命缺陷：采样速度极慢（需要上千步迭代）。

2021 年，DDIM 提出非马尔可夫采样过程，将采样步数从1000步降到50步左右。并且，宋飏发表了Score-based Models,从分数匹配角度重新诠释扩散模型。

潜在扩散模型（LDM ）

为什么会进步到潜在扩散模型（Latent Diffusion Models, LDM）？

核心问题：

传统DDPM在像素空间直接操作存在严重瓶颈：

• 计算成本极高
  • 高分辨率图像（如512×512）意味着在262,144维空间中进行扩散
  • 每一步去噪都需要处理完整的像素数据
  • 训练和推理都需要巨大的计算资源
• 信息冗余
  • 像素空间包含大量低级细节（纹理、噪声）
  • 很多信息对语义理解贡献不大
• 难以扩展
  • 更高分辨率意味着成本呈平方增长
  • 限制了模型的实际应用

当然，这几个原因归根结底是一个问题： 计算成本高 。

Latent Diffusion Models 采用了两阶段策略。

阶段一：学习压缩表示。

使用预训练的VAE（变分自编码器）将图像压缩到低维潜在空间（如512×512→64×64），压缩比可达8倍或更高。

阶段2：在潜在空间进行扩散

扩散过程在压缩后的潜在表示上进行，计算量大幅降低（约降低64倍），保留了语义信息，丢弃了冗余细节。

流匹配

传统的 Diffusion Model 遇到了一个关键的问题：生成一张图需要模拟一整个基于复杂的深度模型的扩散过程，因此生成速度极慢。

流匹配模型就是想要解决这个问题。

无论是从噪声生成图片（图像重建），还是把人脸转换成猫脸（图像生成），都可以看作是从一个分布转化成另一个分布的过程。

传统的扩散模型采用的方法是去预测噪声，而 Rectified Flow 则是去预测噪声流动的方向。

我们有两个分布 \(\pi_0\) 与 \(\pi_1\) , 我们希望找到一个传输映射 T, 使用当 \(Z_0 \sim \pi_0\) , 我们有 \(Z_1 = T(Z_0)\sim \pi_1\)

要想从 $ Z_0$ 最快来到 \(Z_1\) ，最快的方式就是走直线。

因此，流匹配将分布 \(Z_0\) 与 \(Z_1\) 用直线相连。得到：

当然，我们不能从一开始就知道 \(Z_1\) 的分布形式。所以我们实际上采取的是差分方程：

这里的 \(\epsilon\) 是从 \(Z_0\) 向 \(Z_1\) 每一次前进的步长，而 \(v_\theta(Z_t, t)\) 就有赖于我们的模型进行计算。

单步扩散

有模型（One step diffusion via shortcut models）把步长也加进模型的输入参数中，从而有 \(v_\theta(Z_t, t, \epsilon)\)

因此我们有：

这一步是为了让模型学会在当前扩散时间步为 t 时应该进行多大幅度的跳跃，最终让扩散模型可以实现一步生成。

Mean Flow 是另一个单步扩散的工作，特点是引入了平均速度的概念。

DiT

LDM 使用的 backbone 一般是 UNet, 但是 DiT 将它换成了标准的 Transformer. 这提高了模型的 scalability.

同时， DiT 通过实验比较证明了，使用 AdaLN-Zero 传入条件信息是最优的选择。

在 DiT 中，传入的条件信息仅仅是 Timestep \(t\) 和 Label \(y\) 组成的条件。

这让我以为如果传递更复杂的信息，比如说参考图，则应该使用 Cross-Attention 这种更复杂的信息传入方法。

但是 DDT 证明了，图像 Token 也可以通过 AdaLN-Zero 当作参考信息传入。

DDT

DDT （Decoupled DiT）将模型的结构分为了 Condition Encoder 和 Velocity Decoder. Velocity Decoder 的 Condition 条件是通过 AdaLN-Zero 传入的 Condition Encoder 的输出结果。

这种模型架构实现了提取低频语义与解码高频信息之间的解耦。

我之前觉得使用 AdaLN-Zero 传入，没有经过注意力机制，效果应该一般。但仔细想来，假设输入的 c 的形状是 [B, L, d], 使用 AdaLN-Zero 传入信息，其实也会带来 [6, B, L, d] 的参数量，通过乘法与加法将其信息融入，并且中间也经过了一个多头注意力。与其增加注意力宽度，直接使用 AdaLN-Zero ，能用更低的算力消耗达到更好的效果。

RAE

LDM 这种把图像先降维到低维空间，再使用扩散模型进行处理的做法，非常受限于 VAE 的能力。

根据以下实验可以看出，SD-VAE 在提取图片的语义信息提取方面并不突出，甚至可以说，相比于最新的 Encoder 和 Decoder, 相当有局限性。甚至，VAE 在计算时，所需要的 GFlops 都比基于 ViT 的编码器要高。

然而，哪怕扩散模型的 backbone已经从 UNet 被换成了 DiT, 作为潜空间提供者的 VAE, 却至今仍占据主流。

而图像压缩领域，自然是一个很依赖于图像语义信息的领域，可是至今为止所有的基于扩散模型的图像压缩模型，都依然在使用着 VAE, 这使得很多语义信息在 VAE 压缩与解压缩的过程中，被。

为了将 VAE 消灭，RAE 试图将 VAE 也换成基于 ViT 的 RAE。 RAE 使用的编码器是一个预训练过的表征编码器，比如 DINO-V2, SigLIP2 或者是 MAE, 解码器则使用的是一个对编码器进行匹配的 ViT.

然而，如果直接使用 RAE 替换 VAE 的话，效果会比普通的 SD-VAE 更差。

RAE 证明了，要想让基于 ViT 的编码、解码器起效果，就必须要提高 DiT Token 的维度，要使用隐空间的维度大于 Token 的数量。

对于其原因， RAE 认为是因为模型加噪以后，图片不再是一个低维空间上的流形，而变成了一个满秩的矩阵。由于 VAE 的隐空间被编码为 \(N(\mu, \sigma^2I)\), 这使得它天然对噪声有一定的容忍度，可是 RAE 提取的特征则干净而纯粹，这种影响就体现了出来。

然而，如果我们要在 768 甚至更多 token 的情况下进行训练，这显然会提高计算量，这是不可接受的。

因此 RAE 在 DiT 的基础上，参考 DDT 加入了DDT Head ，用来增加模型 Token 的数量，将扩散过程与用于计算 v 的过程解耦。

像素扩散模型

RAE 试图将 VAE 换成 RAE, 也有人试图直接去掉编解码器，完全使用纯 ViT 完成扩散过程。

参考链接

给 JiT 加上像素级解码器：三篇近期 Pixel DiT 论文对比解读 | 周弈帆的博客

PixelDiT

PixelDiT 使用语义信息做指导，对像素信息进行注意力处理。

由于完全没有损失任何注意信息与细节信息， PixelDiT 理论上能完成完美的重建。在图像编辑任务上，也能保持不应被修改的地方像素的准确。

如下图所示：

当然，对于我们的图像压缩任务来说，不可能像 PixelDiT 一样，把所有的像素信息都保留。

DiP

DiP 模型要解决的问题是：

不能直接在像素域使用小 Patch 的 DiT, 因为那样计算量会很大，因此我们必须要使用大 Patch 的 DiT.

我们可以做一个计算：

对于一张 256x256 的图像，如果使用 32x32 的 patch 来进行切割，我们总共会得到 \(8^2=64\) 个 tokens. 计算注意力时就需要进行： \(O(64^2d)\) 次；如果使用 16x16 的 patch 进行切割，我们总共会得到 \(16^2=256\) 个 tokens. 计算注意力时就需要进行 \(O(256^2d)\) 次。计算量与 patch size竟然是一种四次反比的关系。

因此，我们如果想要在像素域上使用 DiT 进行扩散，必须要使用大 Patch.

可是，大 Patch 传递的信息不够细致，会导致一种归纳偏置。

DiP 模型使用了一个 Detailer Head 的结构来传递细节信息。Detailer Head 的细节如图所示。

DiP 直接通过 Detailer Head 整合了原图的细节信息与 DiT 处理得到的语义信息。

我们最终如果想实现一个细节与语义信息双分支的结构，也需要想办法整合好细节信息与语义信息。

JiT

JiT 这篇文章对于像素扩散模型很重要。

本文通过以下的玩具数据集上的实验，检验了直接预测干净图像与预测噪声在高维空间中并无本质区别。而直接预测干净图像对模型来说会更轻松。

可以想象，对于一个高维空间来说，噪声所处的空间是几乎是全秩的，维数最高， v 所处的维数更低。 x 的维数最低。

根据以下的转化公式，我们可以在 \(x\), \(\epsilon\), \(v\) 之间进行相互转化。

看以下的实验结果可以发现，在高维大 patch 的情况下， \(x\)-pred 的效果更好，而使用 \(\epsilon\)-pred 与 \(v\)-pred 都遇到了灾难性失败。

但是在较低维的情况下， x-pred 和 x-loss 的效果就显得更差了。

因此，本文发现的是：在模型容量有限的情况下，使用 x-pred 更适合拟合高维数据。

另外，这三种预测方法在 U-Net 中区别不算特别大，在 github 官方仓库的 issue 里面，作者提到，可能是因为 UNet 中有使用跳跃连接等办法避免信息瓶颈。

pMF

是 JiT 的后续工作。在 JiT 的基础上，结果了 MeanFlow, 实现了在像素空间进行单步扩散。

可以看到， pMF 在用大 patch 处理高维图片的时候，并没有出现灾难性失败。

这篇工作里还有一些结论我们可以学习使用：

感知损失：

Muon 优化器：

这些东西都能提高效果。

基于 Token 的图像重构与生成

VAE 与 VQVAE

传统的潜在扩散模型使用 VAE 来给扩散模型提供一个稳定的潜空间。在这个潜空间中，维度更低，扩散模型需要的计算量也就不那么大，可以部分解决最高的扩散模型计算量大的问题。

VQ-VAE 的离散码本本质上是一个分类分布。分类分布天然支持多峰。它可以把特征空间切分成无数个不连续的子空间。特征只需要映射到最近的Codebook 向量，这是一种硬分配。它不需要在眼镜和无眼镜之间强行建立一条平滑过渡的连续路径。解码器只需要学会解码明确的code，而不需要处理那些模糊的中间态。它能够保留数据分布的非凸性。

为什么传统 VAE 会模糊，是因为VAE用单峰的连续高斯分布，去覆盖真实数据复杂的多峰分布，真实图像的潜特征分布是高度非凸、离散和多峰的，比如说眼镜和无眼镜是两个截然不同的状态，中间不存在半副眼镜这样的连续过渡，VAE 的 KL 散度约束强迫后验分布去匹配正态分布。为了满足这个单峰连续的约束，模型被迫将这些本该分离的峰挤压到一起，或者用一个宽广的高斯分布去同时覆盖它们。从这个被强行平滑的潜空间采样时，会采样到大量处于峰与峰之间的低概率区域。解码器面对这些处于眼镜和无眼镜之间的线性插值点，为了最小化全局误差，只能输出它们的统计平均值。在图像上，像素的平均就是模糊。

极端一点，如果后验分布完全满足标准正态分布，那就只能输出噪声了。所以在训练的时候，如果把KL散度损失权重放的太大，就非常容易出现平均脸甚至后验崩塌。其实现在的VAE训练已经是比较成熟了，不太会出现这种平均情况，因为相对于最原始的VAE加了非常多的tricks。

但是VQ-VAE这样也有明显坏处，它的离散化操作会导致不可导，使得它的训练过程就像被截断了一样，Encoder 实际上并没有收到真正指导它如何更好地离散化的梯度，而是收到了假设没有离散化时的梯度。这种梯度失配会导致训练过程不稳定，尤其是在训练初期。导致非常难训练，可能很多码本在训练过程中都没有使用到，解码器只能通过少数几个码本去映射整个图像空间，很容易就模式崩溃了。

在 Regularized Vector Quantization for Tokenized Image Generation 这篇论文中，就有体现这个问题。

现在模型用的VQ-VAE和最开始的VQ-VAE 已经有了很大的改进。

VQGAN

VQGAN 在 VQVAE 的基础上，融合了 GAN 和 GPT 的思想。很有一种博采百家之长的意思。

VQGAN 使用 GAN 来判别一张图片是重构生成的图片，还是原始图片。

此外，在经过了编码以及矢量量化以后，对于每一个 token ，得到的是一个编码。

由此，一张图片就可以表示成一个编码序列。

VQGAN 基于此，拿这些编码序列去训练 GPT 模型，而后使用 GPT 来预测编码，从而生成图像。

MaskGIT

MaskGIT 使用了 Bidirectional Transformer, 使用并行渐进式解码完成图像的生成。

每一次运行，以后，只保留 confidential 比较高的那些 tokens, confidential 比较低的那些 token 则保持被遮蔽的状态。

这种方法与 Mased AutoEncoder 是一脉相承的。

TiTok

TiTok 的一个核心观点是：现有的 token-based 的方法，在 token 与 token 之间保留了不必要的空间依赖，也就是说，每一个 token 内部，包含的不过是一个 patch 内部的信息。但是这就导致这种 token 编码方式没有充分利用好 token 与 token 之间的关系。

事实上，不同的 patch 之间在空间上的关系甚至能延伸到全图，而不仅仅是图片一个 patch 附近的几个patches.比如说，如果这是一个人脸照片，那么人脸上眼睛的位置，就会束缚嘴巴的位置。

只要充分利用好这种位置信息，我们能让模型哪怕去掉其中的某一个 patches，也可以通过其他的patches的信息来预测出这种信息。

Masked AutoEncoder 的成功已经证明了这一点。

但我们既然是想要做图像压缩，我们就希望能够完全消除这种空间上的依赖。

TiTok 是怎么做到这一点的呢？

他们在图像编码的过程中，给图像上添加了几个随机噪声 latent tokens, 而后在经过 ViT 处理以后，将原本的 patches 完全丢弃不用。后续仅使用 latent tokens 来恢复全图。这样在反向传播的作用下， ViT 就会学习将所有的信息都学习到这些 tokens 之上，并且这些 tokens 对于 image patches 来说，是没有什么区分的，因此信息可以均匀地分配到所有的 tokens 上。

同时，这种设置可以方便我们选择 token 的数量。

另一个角度而言，这种方式也使得单个 token 的语义指向性更弱了，转而更依赖与多个 tokens 组成的整体。我做一个不恰当的比方：VQGAN 的码本存的是一个又一个的词，单个词是有明确的意思的。因此， VQGAN 的码本的项之间依然不能随意组合。比如你不能说： Kill Money of. 这是不成立的东西。

而 TiTok 的码本更像是一个又一个的字母，单个字母并没有明确的意思，必须要把字母拼成一个整体，才能形成一个可表意的单元。

这代表着什么呢？这代表着 VQGAN 虽然 Token 空间很大，但是依然不是每一种组合都是有效的，真正有效的组合依然分布在少量的密码流形上。而 TiTok 的压缩表示虽然可能也位于流形上，但这个流形占了空间中更大的位置。这意味着，信息的冗余更少了。

如果能做到完全没有信息冗余的话，那就可能会有这种情况： 从 TiTok 的码本中随机选择32个 Tokens, 把它们拼接起来，使用 TiTok Decoder 进行解码，解码器的输出将会是一张图片。

需要注意的是，这些 latent tokens 使用的是学习式的位置编码。如果说改用 RoPE 编码，不知道会不会导致不同的 tokens 之间有了区别，降低信息分配的均匀性。

在解码的过程中， TiTok 采用的方法是拼接上与 image patches 数量相同的 mask tokens. mask tokens 是把同样的一个 tokens 复制了 n**2 个。

在生成图像中， TiTok 使用的是 MaskGIT 框架。由 TiTok 解码，而后不断解开 mask tokens.

如果我们要改用扩散模型，这里就不应该再使用 mask tokens, 而应该换成随机噪声图像。

TiTok 实际上已经使用了渐进式的方法来进行解码，所以改让 TiTok 的 Decoder 进行像素扩散过程不一定能让它的效果比使用 MaskGIT 的效果更好，这一点需要有心理预期。

DLF

TiTok 在语义上保留得比较好，但是它的一个问题是在细节上会与原图有比较大的误差。

为了解决这个问题， DLF 采取了一种双流的结构。用 TiTok 来提取语义信息，而使用另外的窗口自注意力分支来提取细节。

在窗口自注意力与细节分支之间，有 IT Module 进行信息交换（其实就是一个自注意力，让细节分支与语义分支同时参与了进去）。

对过去模型的总结

通过使用像素扩散模型，我们可以避免 VAE 丢失图像语义信息。

而使用 x-pred 在大的 patch size 下，归纳偏置也比较小, 这使得我们可以放心地采用较大的 patch, 而不用担心信息丢失。

我们可以把这种思想迁移到 TiTok 的解码上来。

TiTok 在完成了编码以后，直接使用基于 x-pred 进行训练的 ViT 进行解码。

DLF 模型使用注意力机制完成了细节与语义信息的融合以后，使用 VQGAN 的 Decoder 来将融合了细节与语义信息的内容解码。这一点与 TiTok 使用 MaskGIT 的渐进式解码的方式是不同的。

而且 DLF 这个地方的处理很奇怪：TiTok 的输出结果本身就是像素结果，只不过效果比较差，所以使用 MaskGIT 渐进式解码而已，DLF 却给它加了一个 PixelGenerator.

如果二者都能进行解码的话，没道理我们使用 x-pred 的扩散模型直接解码会失败。

基于可变 Token 的图像处理方法

可变 Latent Tokens

我们继续思考：是不是可以设计出这样一种机制： token越多，图像越准确；token越少，图像越失真。而我们只需要训练一个模型，通过调整 Token 的数量来调节输出结果的质量。

但是 token 的数量不同，必然会造成模型需要改变。

为了解决这个问题，我们是不是可以学习OSCAR?

OSCAR 中的特殊之处就在于这里的 Hyper Encoder, 通过给每一个率失真函数对应训练一个单独的 HyperEncoder，确保了不同的压缩率不需要有不同的压缩模型，而只需要对 HyperEncoder 进行训练。

不过我们也不需要照搬他们的思路，我们可以给不同的 Token 使用不同的码表，训练模型使得有的码表中存储的信息，更倾向于模型整体的语义信息，有的码表中存储的信息，更倾向于模型的细节信息，让它们之中存储的信息有足够的差别。

比如说总共有4个码表，我们把128个 tokens 拆成32 + 32 + 32 + 32 的形式。

这就表示， ViT Block 的宽度会来到 M + 128 ，其中 M 是最大能接收的图像 Token 数量的大小。

参考这个思路，我们列出以下的图。

现在的问题在于：如何让码表中所存在的信息能有足够的差距，还是说模型自动就会让码表中的信息拉开差距？

我倾向于会自动让码表信息拉开差距。

我们在训练的时候，就需要让训练数据激活不同数量的分支。

对于每一张用于训练的图像，我们使用只激活一个分支、激活两个分支、激活三个分支、激活四个分支这四种模式分别生成图片，而后我们要求，其中生成的每一张图片与原始图片都应该要足够像，但是图片与图片之间却要足够地不像。

这样就能强迫不同的码表里面，存储着不同层次的信息。

Token 1-32 固定使用码表A，Token 33-64 固定使用码表B，Token 65-96固定使用码表C，Token 87-128 固定使用码表 D.

有两种可能的训练方法：

1. 从低分辨率到高分辨率，码表 A-D 依次启用。
2. 每一个码表都可能以任意方式调用。

我比较偏向于码表依次启用的做法。 要想实现这一点，有一个问题，就是 ViT 似乎没有任意增加 token 数量的先例。

不过这个问题不打紧。在NLP领域，大语言模型在处理数据的时候，不可能每一次计算遇到的 token 数量都是一样的，但是大语言模型也依然可以正常处理数据，说明 transformer在遇到token数量少于max context size的时候，也是可以正常运行的，只需要不去计算它们的注意力就行。对应的，需要的计算也会变少。通过 Norm 层以后，无论有 32 个 Token 还是 128 个 Token, 它们之间受到的注意力大小会自动分配到一个合理的值。

可变 image tokens

这一点，我不知道能否实现，可能需要验证一下： TiTok 也许能同时对不同像素大小的图像进行图像压缩。

首先， TiTok 的 latent tokens 与图片原本的 tokens 都没有关系，并且 TiTok 的 latent tokens 能与所有的 tokens 进行注意力运算，这代表着它能了解全局的信息。既然有了全局的信息，而 TiTok 的能力就是把全局信息总结成一串 Tokens, 那么 TiTok 理论上来说就能处理不同大小的图片，只不过在处理不同大小的图片的时候，需要激活的参数数量不一样而已。

我们可以只需要扩大 ViT 的宽度，然后同时用高清图像与低清图像进行训练（在训练低清图像时，有些地方被标为 [NULL], 出现 [NULL] 就不进行即可，既然大语言模型既然能处理短于 max context size 长度的文字，视觉领域没理由不可以）。

这样我们可以让模型学习从不同的尺度下提取信息。

---