ProgressiveDiTok

2026/02/10 00:00:00·2026/05/19 10:23:00

简介

扩散模型凭借其强大的视觉先验知识，能够对受损图像进行去噪恢复。基于它的这种特性，扩散模型在图像压缩领域已经得到了不少的应用。

但是现有的基于扩散模型的图像压缩算法主要利用的是预训练的扩散模型，并且仅利用了扩散模型的去噪过程，没有意识到图像加噪时，其逐步受损的过程也包含了很多的信息。我们基于经验，可以有这样的认知：如果能让模型了解图像从原始图像一步步受损，直到完全变成随机图像的过程，也自然能方便模型从随机图像中重新恢复出原始图像。

因此，我们提出 ProgressiveDiTok, 它可以在图像去噪过程中根据每一次去噪所处的时间步，从编码端向解码端提供对应的信息。这使得解码端可以在编码端仍在编码时就开始解码过程。

方法

我们的核心目标是让编码端向解码端传递的信息不再是图像的整体信息，而是图像在一次加噪过程中损失的信息。

由于编码端与解码端使用了同样的随机数种子，生成的原始随机图像也会是一致的。因此，编码端图像一次加噪所损失的信息，也是解码端图像一次去噪所需要的信息。

因为编码端可以逐步向解码端传递信息恢复所需要的信息，解码端不需要等到编码端完成编码再开始解码。两者可以同时工作。

训练

联合训练

接下来，我们用 $$X$$ 表示原始图像，用 $$Z_t$$ 表示时间步 $$t$$ 时的加噪图像。

在标准 DDPM（方差守恒）框架下，加噪过程应写为：

Z_t = \sqrt{\bar{\alpha}_t}\, X + \sqrt{1 - \bar{\alpha}_t}\, \epsilon,

其中

\epsilon \sim \mathcal{N}(0, I),

\bar{\alpha}_t = \prod_{i=1}^{t} (1 - \beta_i).

这种形式保证

\mathrm{Var}(Z_t) = \mathrm{Var}(X),

我们的训练分为两步，首先，我们需要对编码器、码本、解码器进行联合训练，从而让编码器学会在加噪过程 $Z_{t-1} \rightarrow Z_t$ 中，将图像在这一次加噪中损失的信息保存到 latent tokens 中；随后，我们冻结编码器与码本，专门训练解码器，让解码器提高对图像的重建能力。

在联合训练过程中，我们会采样时间步 $$t$$ , 假设最多有时间步 $$T$$ ，我们随机采样 $t=1, 2, \cdots, T$ 。

在扩散时间步为 $$t$$ 时，我们会向编码器中传入加噪并经过 patchify 的加噪图像 $Z_{t-1}$ , 而后，我们将 $$m_t$$ 个 latent tokens 拼接到 $Z_{t-1}$ 上，通过 ViT 构成的编码器，得到 $\text{Enc}(Z_{t-1}\oplus L_t, t)$ ，而后，我们只保留 latent tokens, 而将原始图像对应的 tokens 完全抛弃，这会强制编码器将信息转移到 latent tokens 上。我们通过一个可学习的码表对这些 latent tokens 进行矢量量化。具体过程可以参考 TiTok. 通过训练，这个过程会让编码器学会将加噪过程 $Z_{t-1} \rightarrow Z_t$ 中图像损失的信息保存到 latent tokens 中。

在解码端，我们得到矢量量化以后的 latent tokens 以后，将其拼接到解码端在时间步 $$t$$ 时的噪声的图像 $$Z_t$$ 之上。注意在联合训练阶段，我们使用的是加噪图像，而非去噪图像。我们去噪以后，得到了 $\text{Dec}(Z_t\oplus L_t, t)$ . 对于解码器输出的结果，我们只保留其输出的图像 tokens 得到其重建出来的 $\hat Z_{t-1}$ , 而抛弃 latent tokens 的部分. 这同样是为了让模型学会让信息完全转移到图像 tokens 之上。

我们希望 $\text{MSE}(\hat Z_{t-1}, Z_{t-1})$ 尽可能地小，因此有 $L_\text{reconstruction}(\hat Z_{t-1}, Z_{t-1})$ . 当然，根据 JiT 的结论，如果想在像素空间中得到最优结果，我们实际上应该使用 $$v$$ -loss, 那么我们实际上计算的是 $L_{\text{reconstruction}}=\|v - \hat{v}\|_2^2$ .

为了更新码表，参考VQVAE,我们还有 commitment loss 和 codebook loss .

\begin{align} \text{L}_\text{commitment} &= \|z_e(x) - \text{sg}[e]\|_2^2\\ \text{L}_\text{codebook} &= \beta \|\text{sg}[z_e(x)] - e\|_2^2 \end{align}

二者之和被称为 $L_\text{VQ}$ :

\text{L}_\text{VQ} = \text{L}_\text{commitment} + \text{L}_\text{codebook}

综合以上的损失，对于联合训练阶段，我们有：

\begin{align}\text{L}_\text{total} &= \text{L}_\text{reconstruction} + \text{L}_\text{VQ} \\ &= \|v - \hat{v}\|_2^2 + \|z_e(x) - \text{sg}[e]\|_2^2 + \beta \|\text{sg}[z_e(x)] - e\|_2^2 \end{align}

Finetuning 解码器

在联合训练时，为了训练编码器能把单步加噪中损失的信息充分转移到 latent tokens 中，在解码器端，我们直接使用对原始图像进行加噪的方式得到 $$Z_t$$ , 而不是使用解码器预测得到的 $\hat Z_t$ , 这意味着在实际推理时，解码器将无法应对误差累积。因此，我们冻结编码器与码表，对解码器进行微调。微调时，我们使用 $\hat Z_t$ 作为扩散时间步 $$t$$ 时去噪模型的输入。

我们同样计算 $\text L_\text{reconstruction}$ , 根据 TiTok 的设置，我们还可以计算 $\text L_\text{gamma}$ , 此外，我们还可以加上 $\text L_\text{PIPS}$ 以保证每一步去噪的结果都符合人类的感知。还可以使用 $\text L_\text{REPA}$ , 用更优秀的语义特征指导模型的训练。

CoD 中还使用了一个 $\text L_\text{aux}$ 用来增强模型的条件学习能力，我们也可以使用。

推理

在联合训练过程中，我们可以对同一张图像的不同时间步进行并行训练，但是在微调解码器以及推理的过程中，我们需要按照 $t=T, T-1, \cdots, 1$ 的顺序来对一张图片逐步去噪。

在编码器中，我们只需要调整原始图像与随机噪声分别占据的权重，就可以得到不同扩散时间步中的图像，但是在解码器中，我们无法得知原始图像的具体信息，因此我们必须要使用上一步去噪的结果作为这一步去噪的输入。

得益于我们的模型中，编码器深度融入了模型的加噪过程，在推理过程中，我们可以让模型逐步进行解码。

我们可以动态调节解码器端进行的步数。随着解码器步数的增加， bpp 逐步增加。我们可以通过调节解码步数，找到一个能接受的 bpp 值。

依赖于解码器扩散模型的视觉先验知识，可以想象，在任何一个扩散时间步下，解码生成的图像都是视觉上符合人类感知的，只不过，随着扩散步数的增加，扩散模型得到更多编码器端传来的信息，扩散模型可以使恢复出来的图像更接近于原始图像信息，也即更保真。

进一步讨论

解码端如何融入信息

在推理过程中，将 latent tokens 拼接到解码端上一轮扩散的输出主要有三种方案。

一、保留法：在保留之前已经拼接的 latent tokens 的条件下，添加新的 latent tokens。二、丢弃法：每一次扩散去噪以后，都把上一次去噪使用的 tokens 丢弃，在新的扩散过程中，仅使用新的 latent tokens。三、定长队列法：这是对之前两种方案的折中，将编码器每一次传入的 tokens 保存到一个限长的队列中，最多只保留 $$Q$$ 个扩散时间步的 latent tokens, 当 $$t<T-Q$$ 时，每进行一步，都会丢弃掉队列最前端一个扩散时间步传入的 tokens.

接下来，我们对这些方案分别进行分析。

保留法

以下是保留法的示例，我们设 $$T=4$$ ：

这种方法让更高噪声时传入的信息可以被充分扩散解码，并且可以与后续传入的信息交融，从而在理论上可以保证解码出来的图像更真实。

但同时，一来，这种能力需要通过微调得到，不方便我们在研究初期进行模型能力的探索。二来，这种方法会导致计算量的增加。因此，这种方法显得不那么实用。

丢弃法

丢弃法是不保留之前拼接的 latent tokens 的方法：

这种方法可以强迫解码器只使用编码器当前传入的信息。

在联合训练的过程中，为了保证训练可以在不同时间步下并行，并且也是为了保证解码器接收到的信息完全来自于当前编码器传输，从而保证对编码器的训练效果，我们只能使用丢弃法。在初步实验中，我们也应该只使用第丢弃法。

定长队列法

以下，我们设定长队列的长度 $$Q=2$$ .

相比于完全的保留法，定长队列法更有可行性。从直觉上来说，最早期的 latent tokens 也许一步无法完成全部的信息转移，但在2~3步以后，也应该将能传递的信息都已经充分传递了。

因此，在通过丢弃法验证了方法的可行性以后，我们可以进行增加实验，对解码器进行微调，使其适应定长队列法，能够在传入不同数量的 latent tokens 的条件下进行解码。

每一个扩散时间步中，编码器传给解码器的 latent tokens 数量是否应该相同

每一次加噪中损失的信息量可以通过时间步为 $$t-1$$ 时加噪图像与原始图像的互信息以及时间步为 $$t$$ 时加噪图像与原始图像的互信息的差来表示。也即：

\Delta I_t = I(z_0;z_{t-1}) - I(z_0;z_t)

加噪过程可以表示为如下式子：

z_t=\alpha_t z_0+\sigma_t \epsilon, \epsilon \sim \mathcal N(0, 1)

有 $z_0 \in \mathbb R^D$ , $$D$$ 表示维数，对于 2562563 的图像而言，其维度为 $$D=256*256*3$$ . 这是一个很大的数值，如果将 $$D$$ 设为如此巨大的数，我们将无法计算。但是考虑到流形假设，图像实际上并没有位于一个 2562563 的高维空间中，而可能会位于其中一个数千维的空间中。

因此，互信息可以表示为：

I_(z_o;z_t)=\frac D 2 \log(1+\text{SNR}_t)

其中 $\text{SNR}_t= {\alpha_t^2}/{\sigma_t^2}$ .

因此，我们有：

\Delta I_t = \frac D 2 \log \frac {1+\text{SNR}_{t-1}}{1+\text{SNR}_{t}}

计算可得，当 $$t$$ 极小或极大时， $\Delta I_t$ 变化缓慢，而当扩散时间步 $$t$$ 位于中间时， $\Delta I_t$ 比较大。也就是说，在中间的扩散时间步，模型每一步会损失更多的信息。

理论上来说，由于每一个扩散时间步，其造成的信息损失量是不同的，因此，我们传递的 tokens 数量也应该是不同的。后续可以进行更深入的研究，来判断是否因扩散模型在不同扩散时间步的重建能力不同，我们需要给其传递不同数量的 latent tokens. 如：靠近 0 与 T 时，传递更多的 latent tokens, 而在中间的扩散时间步，传递更多的 latent tokens. 但在研究初期，我们可以传递定长的 latent tokens. 以简化研究。

总时间步 T 与总 token 数量的关系

假设原始图像经过 patchify 有 $$K$$ 个 patches, 在第 $$t$$ 个时间步中，传递了 $$k_t$$ 个 tokens, 那么我们应该有 $\sum_t^T k_t=K$ , 如此一来，在理论上来说，忽略掉量化造成的信息损失，当扩散模型完全走完，原始图像中的信息应该能够被完全传输到解码端。在 tokens 数量均匀的情况下，我们有 $$k=K /T$$ .

如果我们想要得到依然符合人类感知，但是降低了图像保真度，从而得到了更高压缩率的图像，那么我们只需要在解码端不完成完整去噪过程即可。

码表

不同的扩散时间步是不是要使用不同的码表？

我认为是不需要的，因为加噪过程中，图像的信息丢失是随机的。也就是说，两次加噪丢失的信息与信息之间并无本质上的差别，我们应该使用同样的码表来存储。让模型学会从码表中找到合适的码字组合来描述一次加噪过程中图像丢失的信息。

至于码表的大小，我们可能通过以下步骤计算：

设码表大小为 $$N$$ , 那么当扩散时间步为 $$t$$ 时， $$k_t$$ 个 latent tokens 量化以后可以表达的图片数量有 $N^{k_t}$ 种。

单个 Token 携带的信息量上限为 $\log_2 N$ , $$k_t$$ 个 token 所携带的信息量上界为 $k_t\log_2 N$ .

而单次加噪损失的信息量根据之前的计算，为 $\Delta I_t$ .

我们需要满足下界条件：

k_t\log_2 N\ge \Delta I_t

假设现在每次加噪传输的 token 数量相同。那么单次加噪过程中传递的 tokens 数量为 $$k=K/T$$ , 因此，我们有：

\frac K T \log_2 N\ge \max \Delta I_t=\Delta I_{T/2}

分别考虑线性噪声调度以及余弦噪声调度，我们可以在量级上估算

\Delta I_{T/2}\approx \frac D {T}

。

因此我们有：

N\ge 2^{D/K}

也就是说：

基于低维流形假设，此处我们令 $$D=1000$$ .

对于图像大小为 $$256*256$$ 而 patch 为 $$16*16$$ 的情况，我们令 $$T=8$$ , 那么我们有 $$K=256$$ , 那么应该取 $N=2^{4}$ .

如果图像实际维度更高，可能就需要增加 $$N$$ 的大小。 CoD 中使用的 $$N$$ 的大小也为 16，与我们的推导一致，故而我们可以先使用 $$N=16$$ 的设置。

由于我们的 tokens 中要保存的信息是 $Z_{t-1}\rightarrow Z_t$ 过程中损失的信息。因此，对于单个 tokens, 我们需要 4 个 bits 来进行表达，单次传输总共有 $$K/T=32$$ 个 tokens, 那就代表我们压缩后的 latent tokens 需要 $$32*4=128$$ bits 来进行表示。

我们得到：

q_\text{bpp}=\frac {128} {256 * 256 * 3}\approx 0.00065(\text{bits/pixel})

我们可以取：

$\text{BPP}=q_{\text{bpp}}, 2*q_{\text{bpp}}, \cdots, T*{q_{\text{bpp}}}$

单步加噪、单步去噪模型的可能

在完成初步探索以后，我们仍可以继续研究，是否可以让模型支持单步加噪、单步去噪。

例如，在加噪阶段，学习 MeanFlow, 在给定起始时间步 $$t$$ 与最终时间步 $$u$$ 的情况下，让编码器能通过 $\text{Enc}(Z_{t-1}, t, u)$ , 直接学习到扩散模型加噪过程中 $Z_{t-1} \rightarrow Z_{u-1}$ 过程中所损失的信息量，并根据时间步自动选择合适的 latent tokens 数量进行传递。

在去噪过程中，则使用基于无条件扩散模型的 pMF 进行改进得到的有条件像素空间单步扩散模型进行去噪。

实验计划

第一步，在 DiTok 代码的基础上进行修改，先证明扩散模型能学习到 $Z_{t-1}\rightarrow Z_t$ 过程中损失的信息。

第二步，引入 RoPE, DDT 等更先进的模型结构，开始联合训练。

第三步，开始对解码器进行微调。

第四步，其他实验。