视觉 Tokenizer

最基础的 Autoencoder (自编码器,AE) 可以写成

$$ x \xrightarrow{\text{Encoder}} z \xrightarrow{\text{Decoder}} \hat{x} $$
  • $x$ :输入图像
  • $z$:潜在表示(Latent Representation)
  • $\hat{x}$:重建得到的图像

训练目标是使重建结果尽可能接近原始输入

VAE

[1312.6114] Auto-Encoding Variational Bayes

Variational Autoencoder 变分自编码器

核心:与普通 Autoencoder 把输入编码为一个确定的 latent vector 不同,将输入编码为一个潜变量概率分布

1
2
3
4
5
6
7
8
9
10
11
12
13
x
↓
Encoder
↓
μ, σ
↓
Gaussian Sampling
↓
连续 z
↓
Decoder
↓
x̂

经过 Encoder 后得到一组分布参数,从该分布中采样

$$ z\sim \mathcal N(\mu,\sigma^2) $$
实际训练时通常写成
$$ z=\mu+\sigma\odot\epsilon,\qquad\epsilon\sim\mathcal N(0,I) $$
随机性被转移到 $ \epsilon$ 上,而 $ \mu$ 和 $\sigma$ 仍然可以通过梯度进行学习

VAE 的训练不仅要求 Decoder 能够重建输入,还会约束潜变量分布

$$ \mathcal L=\mathcal L_{\text{recon}}+D_{\mathrm{KL}}\left(q_\phi(z|x)\,\|\,p(z)\right) $$
  • Reconstruction Loss:保证 $z$ 中保留足够的信息来重建输入
  • KL 散度:约束使不同样本的潜在分布整体接近 $\mathcal N(0,I)$

最终得到一个连续且较规整的 Latent Space

VQ-VAE

[1711.00937] Neural Discrete Representation Learning

Vector Quantized Variational AutoEncoder 向量量化变分自编码器

核心:将 Encoder 得到的连续特征通过 Vector Quantization(向量量化) 映射到有限 Codebook 中,从而得到离散的潜在表示(Discrete Latent Representation)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
x
↓
Encoder
↓
连续特征 z_e(x)
↓
Vector Quantization
↓
离散 Code Index k
↓
Codebook Lookup
↓
z_q(x) = e_k
↓
Decoder
↓
x̂

Vector Quantization

VQ-VAE 维护一个可以学习的 embedding table

$$ E=\{e_1,e_2,\ldots,e_K\}, \qquad e_i\in\mathbb R^D $$
  • $K$:离散 latent space 的大小,也就是 Codebook Size
  • $D$:每个 embedding vector 的维度

对于 Encoder 输出的每一个向量 $z_e(x)$,在 Codebook 中寻找距离最近的 embedding

$$ k=\operatorname*{argmin}_{j}\left\|z_e(x)-e_j\right\|_2 $$
然后令:
$$ z_q(x) = e_k $$
定义成一个确定性的 one-hot categorical distribution
$$ q(z=k\mid x)=\begin{cases} 1, & k= \operatorname*{argmin}_j \|z_e(x)-e_j\|_2 \\ 0, & \text{otherwise} \end{cases} $$
本质上很像 NLP 中的
1
2
3
4
5
Token ID
↓
Embedding Table
↓
Token Embedding

而 VQ-VAE 在编码阶段相当于反过来

1
2
3
4
5
Encoder Feature
↓
寻找最近的 Embedding
↓
得到 Code Index

模型自己学习了一套离散表示,使得连续空间变成了有限的离散状态

image-20260922215138414
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
Image
[128, 128, 3]

↓ Encoder

z_e(x) [32, 32, D]

↓ Vector Quantization K=512

Code Index z
[32, 32,1]

↓ Codebook Lookup

z_q(x) [32, 32, D]

↓ Decoder

Reconstruction
[128, 128, 3]

Train

问题:Quantization 使用 argmin ,没有真正的梯度,所以无法正常反向传播

使用 Straight-Through Estimator,STE (直通估计器)的方法:

  • Forward :使用量化后的 $z_q (x) = e_k$
  • Backward :把 Decoder 输入 $z_q(x)$ 的梯度复制给 Encoder 输出 $z_e(x)$
$$ \frac{\partial z_q}{\partial z_e}\approx 1 $$

把 Quantization 在反向传播时近似当作 Identity

在Pytorch中:

1
z_q = z_e + (e_k - z_e).detach()

利用 .detach() 切断梯度,且前向传播中 $z_q = e_k$

反向传播时相当于 z_q = z_e + 常数

Loss

VQ-VAE 的训练目标包含三个部分:

$$ \mathcal L=\mathcal L_{\text{recon}}+\mathcal L_{\text{codebook}}+\beta\mathcal L_{\text{commit}} $$
**Reconstruction Loss**
$$ \mathcal L_{\text{recon}}=-\log p(x\mid z_q(x)) $$
衡量重建损失,具体展开取决于 Decoder 的输出分布

常见的就是MSE

$$ \mathcal L_{\text{recon}}=\|x-D(z_q(x))\|_2^2 $$
**Codebook Loss**
$$ \mathcal L_{\text{codebook}}=\left\|\operatorname{sg}[z_e]-e_k\right\|_2^2 $$
$z_e$ 被 Stop Gradient,所以只更新 $e_k$

目标是让 Codebook 中被选中的 embedding 向 Encoder 输出靠近

Commitment Loss

$$ \mathcal L_{\text{commit}}=\left\|z_e-\operatorname{sg}[e_k] \right\|_2^2 $$
$e_k$ 被 Stop Gradient,所以只更新 $z_e$

Encoder 去追 Codebook Embedding

训练目标

Loss Encoder Codebook Decoder
Reconstruction Loss ✓ × ✓
Codebook Loss × ✓ ×
Commitment Loss ✓ × ×
  • Encoder 学习如何把输入映射到有用的离散表示
  • Codebook 学习 embedding space
  • Decoder 学习如何根据量化后的 latent 重建输入

相比VAE,无需KL散度,因为潜空间不再连续

Prior

前面训练解决的是给定图片,如何得到离散 latent

但生成时并没有输入 $x$ ,因此还需要学习离散 latent 的概率分布 $p(z)$

VQ-VAE 采用 Two-Stage 的训练方式

1
2
3
4
5
6
7
8
9
10
11
12
Stage 1

训练 VQ-VAE
Image ↔ Discrete Latent

Stage 2

固定已经训练好的 VQ-VAE
↓
收集离散 latent z
↓
学习真正的 Prior p(z)

feature map 中不同位置的 latent 并不是独立的,因此可以使用自回归模型学习

$$ p(z)=\prod_{i=1}^{N}p(z_i\mid z_1,\ldots,z_{i-1}) $$
每一个 latent token 根据前面已经生成的 latent token 进行预测

论文中:

  • 图像使用 PixelCNN
  • Raw Audio 使用 WaveNet

在生成阶段其实就用不到Encoder了

$$ \text { PixelCNN prior } \rightarrow z_{1}, z_{2}, \ldots, z_{N} \rightarrow \text { Decoder } \rightarrow x $$

VQGAN

[2012.09841] Taming Transformers for High-Resolution Image Synthesis

核心:先使用 VQGAN 将图像压缩成由有限 Codebook 构成的离散视觉表示,再使用 Transformer 建模这些离散视觉 Token 的组合关系

Loss

继承 VQ-VAE 的量化训练方式

$$ \mathcal L_{\mathrm{VQ}}(E,G,Z) = \|x-\hat x\|_2^2 + \|\operatorname{sg}[E(x)]-z_q\|_2^2 + \|\operatorname{sg}[z_q]-E(x)\|_2^2 $$

VQ-VAE 的 Reconstruction Loss 更关注像素级重建,但如果压缩率进一步提高,仅使用 Pixel-level Reconstruction Loss 容易导致图像纹理和细节的感知质量下降

将原本的 $L_2$ Reconstruction Loss 替换成 Perceptual Loss(感知损失)


额外引入一个 Patch 二分类判别器 $D$,关注图像局部 Patch 的真实性

希望真实图 $D(x)$ 接近1,重建图 $D(\hat x)$ 接近0

$$ \mathcal L_{\mathrm{GAN}} (\{E,G,Z\},D) = \log D(x) + \log(1-D(\hat x)) $$
将 VQ 重建目标与 GAN 目标形成对抗

最后完整目标:

$$ Q^* = \operatorname*{argmin}_{E,G,Z} \operatorname*{max}_{D} \mathbb E_{x\sim p(x)} \left[ \mathcal L_{\mathrm{VQ}}(E,G,Z) + \lambda \mathcal L_{\mathrm{GAN}} (\{E,G,Z\},D) \right] $$
$\lambda$ 用于平衡两组损失,根据两个 Loss 对 Decoder 最后一层产生的梯度,自适应调整 GAN Loss 的权重
$$ \lambda=\frac{\nabla_{G_L}[\mathcal L_{\mathrm{rec}}]}{\nabla_{G_L}[\mathcal L_{\mathrm{GAN}}]+\delta} $$
VQ-VAE:Pixel Reconstruction → 重建容易偏平滑

VQGAN:Perceptual Loss + Patch-based GAN → 允许更强压缩

Prior

VQGAN 仍然采用和 VQ-VAE 类似的 Two-Stage 思路

但第二阶段不再使用 PixelCNN 建模离散 latent,而是使用 Transformer

$$ \mathcal L_{\text{Transformer}}=\mathbb E_{x\sim p(x)}[-\log p(s)] $$

扩散模型

U-Net

[1505.04597] U-Net: Convolutional Networks for Biomedical Image Segmentation

U-Net 是一种用于图像分割的 Encoder-Decoder 结构

论文将网络分为:

  • Contracting Path:收缩路径,对应 Encoder
  • Expansive Path:扩张路径,对应 Decoder

核心思想:收缩路径用于捕获上下文信息,扩张路径则用于恢复空间分辨率,从而实现精确的定位

整体结构:

1
2
3
4
5
6
7
8
9
10
11
12
13
Input
│
▼
[Encoder 1] ─────────────────────→ [Decoder 1]
│ ▲
▼ │
[Encoder 2] ───────────────→ [Decoder 2]
│ ▲
▼ │
[Encoder 3] ─────────→ [Decoder 3]
│ ▲
▼ │
[ Bottleneck ]

Encoder 和 Decoder 对应层之间通过 Skip Connection 连接

Encoder

逐渐降低空间分辨率,同时提取更高层的语义和上下文特征

1
2
3
4
5
6
7
8
9
10
11
12
3×3 Conv
│
ReLU
│
3×3 Conv
│
ReLU
│
2×2 Max Pool Stride = 2
│
▼
下一层

Decoder

将低分辨率的高级语义特征逐步恢复到高分辨率,并完成精确定位

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
  Upsampling
│
2×2 Up-Conv 通道数减半
│
▼
与对应 Encoder 特征 Concat
│
▼
3×3 Conv
│
ReLU
│
3×3 Conv
│
ReLU
│
▼
下一层

Skip Connection

如果只使用普通 Encoder-Decoder

在不断下采样过程中,精确的位置信息可能逐渐丢失

因此 U-Net 把 Encoder 中的高分辨率特征直接送给 Decoder

沿着 Channel 维拼接后通过 Conv 融合这些特征

输出

原始 U-Net 用于图像分割

1
2
3
4
5
6
7
Feature Map
│
▼
1×1 Conv
│
▼
Segmentation Map

将每一个像素位置上的特性向量映射到目标类别数

DDPM

[2006.11239] Denoising Diffusion Probabilistic Models

DDPM 包含两个方向相反的 Markov Process

1
2
3
4
5
Forward Process
x₀ → x₁ → x₂ → ... → x_T ≈ N(0,I)

Reverse Process
x_T → x_{T-1} → ... → x₁ → x₀

其中:

  • Forward Process:人为定义的逐步加噪过程
  • Reverse Process:由神经网络学习,从 Gaussian Noise 逐步恢复数据
image-20260923195554651

Forward Process

每一步向数据加入少量 Gaussian Noise

$$ q(x_t|x_{t-1})=\mathcal N\left(x_t;\sqrt{1-\beta_t}x_{t-1}, \beta_t I \right) $$
  • $q$:人为定义的 Forward Process
  • $\beta$:第 $t$ 步加入噪声的 variance schedule

定义:

$$ \alpha_t = 1-\beta_t \qquad \bar\alpha_t=\prod_{s=1}^{t}\alpha_s $$
由于 Gaussian Distribution 的性质,可以直接从 $x_0$ 得到任意 timestep 的 $x_t$
$$ q(x_t\mid x_0)=\mathcal N \left( \sqrt{\bar\alpha_t}x_0, (1-\bar\alpha_t)I \right) $$
通过重参数化:
$$ \boxed{x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon \qquad \varepsilon\sim\mathcal N(0,I)} $$
因此训练时可以直接随机采样到任一时间步的位置

训练阶段知道原始数据 $x_0$,因此 Forward Process 的一步反向可以解析求出,并且仍然是 Gaussian:

$$ q(x_{t-1}\mid x_t,x_0)=\mathcal N \left( x_{t-1}; \tilde\mu_t(x_t,x_0), \tilde\beta_t I \right) $$
但是生成时无法已知,所以需要学习 Reverse Process

Reverse Process

逆过程希望学习:

$$ p_\theta(x_{t-1}\mid x_t) = \mathcal N \left( x_{t-1}; \mu_\theta(x_t,t), \Sigma_\theta(x_t,t) \right) $$
原则上 Gaussian 有两个需要确定的部分
  • 均值:$\mu_\theta(x_t,t)$
  • 协方差:$\sum_\theta(x_t,t)$

但在论文里直接规定

$$ \Sigma_\theta(x_t,t)=\sigma_t^2I $$
明确 $\sigma_t^2$ 是随 $t$ 变化,但不参与训练的常数,因此问题简化为只需要学习 mean

Loss

在 variance 固定之后,对应的 KL term 可以写成:

$$ L_{t-1} = \mathbb E_q \left[ \frac{1}{2\sigma_t^2} \left\| \tilde\mu_t(x_t,x_0) - \mu_\theta(x_t,t) \right\|^2 \right] +C $$
最直观的训练方式就是:
$$ \mu_\theta(x_t,t) \approx \tilde\mu_t(x_t,x_0) $$
但是 DDPM 没有直接让网络预测 mean

通过将 mean 重新参数化为:

$$ \boxed{\mu_\theta(x_t,t) = \frac{1}{\sqrt{\alpha_t}} \left( x_t- \frac{\beta_t} {\sqrt{1-\bar\alpha_t}} \epsilon_\theta(x_t,t) \right)} $$
于是只需要预测噪声 $\varepsilon _\theta$
$$ \boxed{ \varepsilon_\theta \rightarrow \mu_\theta \rightarrow p_\theta(x_{t-1}\mid x_t) } $$
符号 来源 含义
$\varepsilon$ Forward Process 真正加到 $x_0$ 上的 Gaussian noise
$\varepsilon_\theta(x_t,t)$ U-Net 网络预测出来的噪声

因此可以直接训练

$$ \epsilon_\theta(x_t,t) \approx \epsilon $$
最终采用训练 Loss:
$$ \boxed{ L_{\text{simple}} = \mathbb E \left[ \|\varepsilon-\varepsilon_\theta(x_t,t)\|^2 \right] } $$
DDPM 表面上训练的是 Noise Predictor,本质上是在学习 Reverse Process

Timestep

不同的 $t $ 对应不同的 Noise Level,但所有 timestep 共用同一个 U-Net 参数

因此网络不仅需要输入 $x_t$ 还需要知道 $t$

用 Transformer 的位置编码方式来编码 timestep

LDM

[2112.10752] High-Resolution Image Synthesis with Latent Diffusion Models

Latent Diffusion Model,潜空间扩散模型

Pixel-space Diffusion Model 直接在高维 RGB 图像上进行反复去噪,训练和采样都需要大量计算

核心思想:使用 Autoencoder 将图像压缩到一个低维、感知上近似等价的 Latent Space,再在 Latent Space 中训练 Diffusion Model

感知压缩

目的是去掉高频、感知上不重要的信息,同时尽量保留图像的语义和视觉细节

训练一个 Autoencoder:

$$ z=E(x)\\ \tilde{x}=D(z)=D(E(x)) $$
但也不能压缩的太狠,论文实验中 LDM-4 和 LDM-8 表现出了较好的效率和质量平衡

LDM 的第一阶段与 VQGAN 的 Perceptual Compression 思想非常接近

Latent Diffusion

训练好 Encoder $E$ 和 Decoder $D$ 后,将它们固定

不再对 Pixel $x$ 做 Diffusion,而是在 Latent $z$ 中执行

1
z₀ → z₁ → z₂ → ... → z_T

与 DDPM 完全相同,可以写成

$$ \boxed{ L_{\text{LDM}} = \mathbb E_{E(x),\varepsilon,t} \left[ \left\| \varepsilon- \varepsilon_\theta(z_t,t) \right\|_2^2 \right] } $$
  • DDPM:预测 Pixel Space 中的噪声
  • LDM:预测 Latent Space 中的噪声

LDM 的 Diffusion Backbone 仍然采用 Time-Conditional U-Net

DiT

[2212.09748] Scalable Diffusion Models with Transformers

核心:用 Transformer 替换 Diffusion Model 中传统的 U-Net Backbone

Latent Diffusion

DiT 并不是直接在高分辨率 Pixel Space 上运行 Transformer,而是建立在 LDM 框架之上

先使用预训练 Autoencoder

1
2
3
4
5
6
7
Image
256 × 256 × 3

↓ Encoder

Latent
32 × 32 × 4

整个 Diffusion Process 都发生在这个 Latent Space

可以理解为:

  • LDM:Latent Diffusion + U-Net
  • DiT:Latent Diffusion + Transformer

Diffusion 数学过程基本没有发生变化

Patchify

Latent $z_t$ 仍然是二维 Feature Map:

Transformer 不能直接按照二维 Feature Map 的形式处理,所以与 ViT 类似,DiT 首先对 Latent 进行 Patchify

假设 Patch Size 为

$$ p \times p $$
那么 Token 数量为:
$$ \boxed{ T=\left(\frac{I}{p}\right)^2 } $$
每个 Patch 经过 Linear Projection,映射成一个 $d$ 维 token

对于位于二维网格坐标 $(x,y)$ 的 Patch,分别对 $x$ 和纵坐标 $y$ 使用不同频率的 sine / cosine 函数进行编码

单个坐标的位置编码形式与 sinusoidal positional encoding 类似

$$ PE_x(x)\in\mathbb{R}^{D/2}, \qquad PE_y(y)\in\mathbb{R}^{D/2} $$
拼接得到最终位置编码
$$ PE(x,y)\in\mathbb{R}^{D} $$
2D sine-cosine positional embedding 由固定数学公式生成,通常不作为模型参数参与训练

DiT Block

Snipaste_2026-09-26_22-12-13

基本结构是标准 Transformer,所以 DiT 没有 U-Net 的 Downsample → Bottleneck → Upsample,尺寸始终保持一致

Diffusion Model 除了输入 Noisy Latent 还需要知道

  • timestep $t$
  • class label $y$

论文比较了多种方法

1
2
3
4
5
6
7
8
9
Condition
│
├── In-Context Conditioning
│
├── Cross-Attention
│
├── adaLN
│
└── adaLN-Zero ← 最终效果最好

In-Context Conditioning

分别得到:

$$ e_t=\operatorname{Embed}(t),\qquad e_y=\operatorname{Embed}(y) $$
然后直接作为两个 Token
1
2
3
4
5
6
7
8
Image Tokens:
[x₁, x₂, ..., xₙ]

Condition Tokens:
[e_t, e_y]

拼接:
[x₁, x₂, ..., xₙ, e_t, e_y]

最后一个 Transformer Block 后,再把Condition Tokens 删除

Cross-Attention

在 Self-Attention 后加入一个 Cross-Attention Layer

1
2
3
4
5
Image Tokens
↓
Self-Attention
↓
Cross-Attention ← [e_t, e_y]

这种方式与 LDM 中的思路比较类似

但是额外的 Cross-Attention 会增加计算量

adaLN

Adaptive Layer Normalization

不把 $t$、$y$ 当作普通 token,而是利用它们去控制 LayerNorm

$$ c=e_t+e_y $$
然后由 Conditioning $c $ 预测 LayerNorm 的 Dimension-wise Scale 和 Shift
$$ \gamma(c),\beta(c) $$
LayerNorm 的输出经过 scale 和 shift
$$ \operatorname{adaLN}(x,c)=\gamma(c)\odot \operatorname{LN}(x)+ \beta(c) $$
Condition 不作为 Token 参与 Attention,而是动态改变 Transformer 中归一化后的 Feature

相比 Cross-Attention,adaLN 引入的额外 Gflops(十亿次浮点运算) 最少

adaLN-Zero

论文使用的不是普通 adaLN

在 adaLN 的基础上继续加入一个可学习的 Scale

初始化时:

$$ \alpha=0 $$
因此
$$ \operatorname{DiTBlock}(x)\approx x $$
每一个 DiT Block 初始化时近似 Identity Function

借鉴了 ResNet 和 Diffusion U-Net 中将 Residual Branch 初始化为零的思想,这种设计有利于深层网络训练

后续工作基本全部采用 adaLN-Zero