视觉 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
|
模型自己学习了一套离散表示,使得连续空间变成了有限的离散状态
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 逐步恢复数据
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
基本结构是标准 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