Generative Models — 理论 & Tokenizers进阶33 分钟P2 · 建议关注

VAE / VQ-VAE / VQ-GAN / FSQ

连续/离散 latent、codebook、commitment loss、GAN 感知损失与 tokenizer 设计。

VAEVQ-VAEVQ-GANFSQtokenizer
内容来源:ARIS-in-AI-Offer by Ruofeng Yang (杨若峰),许可证 MITImported from commit 7fae0a820348.
本地更新提示

原文来自 ARIS 导入版本,尚未做本地逐段复审。建议阅读时优先核对最新论文、模型报告和库 API。

核对最新引用补充 2026 后续进展标注可能过时的模型/API 名称

§0 TL;DR Cheat Sheet

💡 8 句话搞定 VAE / VQ-VAE / VQ-GAN / FSQ — 一页拿下面试核心要点(详见后文 §2–§9 推导)。

  1. 连续 VAE 目标:最大化 ELBO,logp(x)Eqϕ(zx)[logpθ(xz)]DKL(qϕ(zx)p(z))\log p(x) \geq \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_\text{KL}(q_\phi(z|x)\,\|\,p(z));reparameterization z=μ+σϵz = \mu + \sigma \odot \epsilon 让梯度穿过随机采样。

  2. KL 闭式(必考)DKL(N(μ,σ2I)N(0,I))=12i(μi2+σi2logσi21)D_\text{KL}(\mathcal{N}(\mu,\sigma^2 I)\,\|\,\mathcal{N}(0,I)) = \tfrac{1}{2}\sum_i (\mu_i^2 + \sigma_i^2 - \log \sigma_i^2 - 1)

  3. Posterior collapse:KL → 0 → 解码器忽略 zz;缓解:KL annealing、free bits、β\beta schedule、自回归先验。

  4. VQ-VAE:把 encoder 输出 ze(x)z_e(x) 映射到最近邻 codebook 向量 eke_k,loss = recon + sg[ze]e2\|\text{sg}[z_e] - e\|^2(codebook)+βzesg[e]2+ \beta \|z_e - \text{sg}[e]\|^2(commitment)。

  5. Straight-Through Estimator (STE):argmin / quantize 不可导,前向用量化值,反向直通梯度 L/zqL/ze\partial \mathcal{L}/\partial z_q \to \partial \mathcal{L}/\partial z_e

  6. VQ-GAN:VQ-VAE + perceptual (LPIPS) + adversarial (PatchGAN) + 后训 Transformer prior;为 LDM / Parti / Muse 等离散 token 模型奠基。

  7. FSQ(2024):每维标量量化到 {L,,L}\{-L,\ldots,L\},隐式 codebook 大小 iLi\prod_i L_i(如 L=8,d=686=262144L=8, d=6 \Rightarrow 8^6 = 262{144}),无需 STE、不会 codebook collapse,rounding 用 STE 即可,loss 只剩 reconstruction。

  8. 生态对比:连续 latent(VAE / KL)适合 LDM 续 diffusion;离散 token(VQ-VAE / VQ-GAN / FSQ / LFQ)适合 AR / MaskGIT Transformer prior,是 Parti / Muse / Cosmos 等的核心组件。

§1 直觉:为什么要 latent variable model

生成模型的核心难题:直接建模 p(x)p(x) 很难,但如果引入低维 latent zz

p(x)=p(xz)p(z)dzp(x) = \int p(x|z)\, p(z)\, dz

可以把"复杂的图像分布"分解为"简单先验 p(z)p(z)(如 N(0,I)\mathcal{N}(0, I))"加上"易学的条件分布 p(xz)p(x|z)"。两条路:

  • 连续 latent(VAE):zRdz \in \mathbb{R}^d,KL 把 posterior 拉向 Gaussian 先验,和 diffusion / FM 天然兼容(LDM 在 VAE latent 里跑 diffusion)。

  • 离散 latent(VQ-VAE / VQ-GAN / FSQ):zVH×Wz \in \mathcal{V}^{H \times W}(token grid),和 Transformer / AR / MaskGIT 天然兼容(一张图变成一串 token,复用语言模型架构)。

💡 训练 vs 推理的不对称 — VAE/VQ-VAE 训练时学整套 encoder + decoder(rate-distortion 角度:"压缩-重建");推理时根据应用分两种:

  • 生成新样本:丢掉 encoder,从 prior 采样 zz,过 decoder
  • 下游 backbone:丢掉 decoder,把 encoder/zz 当作 representation 给后续模型
  • 二阶段生成(LDM / Parti / Muse):先训 VAE/VQ-GAN tokenizer,在 latent 空间训 diffusion / AR / MaskGIT prior。tokenizer 训完冻结。

§2 VAE:核心公式与推导

2.1 ELBO 推导(必考、要会逐行推)

模型族 pθ(x,z)=pθ(xz)p(z)p_\theta(x, z) = p_\theta(x|z)\, p(z),先验 p(z)=N(0,I)p(z) = \mathcal{N}(0, I),似然 pθ(xz)p_\theta(x|z) 由 decoder 给出。Marginal likelihood:

logpθ(x)=logpθ(xz)p(z)dz\log p_\theta(x) = \log \int p_\theta(x|z) p(z)\, dz

任意 distribution qϕ(zx)q_\phi(z|x)(encoder / variational posterior,qϕ(zx)=N(μϕ(x),diag(σϕ2(x)))q_\phi(z|x) = \mathcal{N}(\mu_\phi(x), \mathrm{diag}(\sigma_\phi^2(x)))),由 Jensen 不等式 / 直接代入:

logpθ(x)=logqϕ(zx)pθ(xz)p(z)qϕ(zx)dzEqϕ(zx) ⁣[logpθ(xz)p(z)qϕ(zx)](Jensen)=Eqϕ(zx)[logpθ(xz)]reconstruction (negative)DKL(qϕ(zx)p(z))regularization\begin{aligned} \log p_\theta(x) &= \log \int q_\phi(z|x) \frac{p_\theta(x|z) p(z)}{q_\phi(z|x)} dz \\ &\geq \mathbb{E}_{q_\phi(z|x)}\!\left[\log \frac{p_\theta(x|z) p(z)}{q_\phi(z|x)}\right] \quad \text{(Jensen)} \\ &= \underbrace{\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]}_{\text{reconstruction (negative)}} - \underbrace{D_\text{KL}(q_\phi(z|x)\,\|\,p(z))}_{\text{regularization}} \end{aligned}

故 ELBO:

  LELBO(θ,ϕ;x)=Eqϕ(zx)[logpθ(xz)]DKL ⁣(qϕ(zx)p(z))  \boxed{\;\mathcal{L}_\text{ELBO}(\theta, \phi; x) = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_\text{KL}\!\big(q_\phi(z|x)\,\|\,p(z)\big)\;}

Tight 的代价:ELBO 与真 log-likelihood 的 gap 等于 DKL(qϕ(zx)pθ(zx))D_\text{KL}(q_\phi(z|x)\,\|\,p_\theta(z|x))。posterior qϕq_\phi 越逼近真后验 pθ(zx)p_\theta(z|x),gap 越小。

2.2 KL 项闭式推导(L3 必推)

qϕ(zx)=N(μ,diag(σ2))q_\phi(z|x) = \mathcal{N}(\mu, \mathrm{diag}(\sigma^2))对角协方差,每维 σi2\sigma_i^2),p(z)=N(0,I)p(z) = \mathcal{N}(0, I)

对每一维 ii 独立:

DKL(N(μi,σi2)N(0,1))=N(z;μi,σi2)logN(z;μi,σi2)N(z;0,1)dz\begin{aligned} D_\text{KL}(\mathcal{N}(\mu_i, \sigma_i^2) \,\|\, \mathcal{N}(0, 1)) &= \int \mathcal{N}(z; \mu_i, \sigma_i^2) \log \frac{\mathcal{N}(z; \mu_i, \sigma_i^2)}{\mathcal{N}(z; 0, 1)} dz \end{aligned}

展开两个 Gaussian 密度的 log:

logN(z;μi,σi2)N(z;0,1)=12logσi2(zμi)22σi2+z22\log \frac{\mathcal{N}(z; \mu_i, \sigma_i^2)}{\mathcal{N}(z; 0, 1)} = -\tfrac{1}{2}\log \sigma_i^2 - \tfrac{(z-\mu_i)^2}{2\sigma_i^2} + \tfrac{z^2}{2}

求期望(利用 Eq[z]=μi\mathbb{E}_q[z] = \mu_i, Eq[z2]=μi2+σi2\mathbb{E}_q[z^2] = \mu_i^2 + \sigma_i^2, Eq[(zμi)2]=σi2\mathbb{E}_q[(z-\mu_i)^2] = \sigma_i^2):

DKL=12logσi212+12(μi2+σi2)=12(μi2+σi2logσi21)\begin{aligned} D_\text{KL} &= -\tfrac{1}{2}\log \sigma_i^2 - \tfrac{1}{2} + \tfrac{1}{2}(\mu_i^2 + \sigma_i^2) \\ &= \tfrac{1}{2}\big(\mu_i^2 + \sigma_i^2 - \log \sigma_i^2 - 1\big) \end{aligned}

对所有维度求和:

  DKL(N(μ,diag(σ2))N(0,I))=12i=1d(μi2+σi2logσi21)  \boxed{\;D_\text{KL}\big(\mathcal{N}(\mu, \mathrm{diag}(\sigma^2)) \,\|\, \mathcal{N}(0, I)\big) = \tfrac{1}{2}\sum_{i=1}^{d}\big(\mu_i^2 + \sigma_i^2 - \log \sigma_i^2 - 1\big)\;}

⚠️ 数值稳定 — 实现时让 encoder 输出 logσ2\log \sigma^2(log-variance)而非 σ\sigma,避免 σ\sigma 取 exp 后 overflow。代码里写 kl = 0.5 * (mu**2 + logvar.exp() - logvar - 1).sum()

2.3 Reparameterization Trick(必考)

ELBO 中 Eqϕ(zx)[]\mathbb{E}_{q_\phi(z|x)}[\cdot] 用 Monte Carlo 估计:采一个 zqϕ(zx)z \sim q_\phi(z|x),求 logpθ(xz)\log p_\theta(x|z)

问题:直接采样 z=sample(N(μ,σ2))z = \text{sample}(\mathcal{N}(\mu, \sigma^2)) 是不可导操作,梯度无法回传到 ϕ\phi

解法:把随机性挪到独立噪声里:

  z=μϕ(x)+σϕ(x)ϵ,ϵN(0,I)  \boxed{\;z = \mu_\phi(x) + \sigma_\phi(x) \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)\;}

现在 zzϕ\phi确定性函数(条件于 ϵ\epsilon),ϕL\nabla_\phi \mathcal{L} 可以正常反向传播。这是 Kingma & Welling (ICLR 2014) 的核心贡献之一。

💡 面试加分:reparameterization 不止 Gaussian — Concrete / Gumbel-softmax(§7)对离散变量做了类似 trick:把 argmax 替换成 softmax 加 Gumbel 噪声,前向近似离散,反向用 softmax 梯度。

2.4 VAE 训练损失(实际写法)

负 ELBO(最小化):

LVAE(x)=xx^2recon (Gaussian likelihood up to const)+DKL(qϕ(zx)p(z))closed form\mathcal{L}_\text{VAE}(x) = \underbrace{\|x - \hat{x}\|^2}_{\text{recon (Gaussian likelihood up to const)}} + \underbrace{D_\text{KL}(q_\phi(z|x)\,\|\,p(z))}_{\text{closed form}}

对 Bernoulli / Categorical 似然(如 MNIST 二值),recon 项换成 BCE / CE。

§3 完整 VAE 实现(PyTorch)

import math import torch import torch.nn as nn import torch.nn.functional as F class VAE(nn.Module): """ 经典 VAE:Gaussian encoder + Gaussian/Bernoulli decoder 本实现以 MNIST (28×28) 为例,latent dim=20 生产化把 MLP 换成 ResNet / U-Net encoder/decoder,latent 可以是 spatial map """ def __init__(self, x_dim: int = 784, h_dim: int = 400, z_dim: int = 20, likelihood: str = "bernoulli"): super().__init__() self.x_dim, self.z_dim = x_dim, z_dim self.likelihood = likelihood # Encoder: x -> (μ, logσ²) self.enc = nn.Sequential( nn.Linear(x_dim, h_dim), nn.ReLU(), nn.Linear(h_dim, h_dim), nn.ReLU(), ) self.fc_mu = nn.Linear(h_dim, z_dim) self.fc_logvar = nn.Linear(h_dim, z_dim) # Decoder: z -> x̂ self.dec = nn.Sequential( nn.Linear(z_dim, h_dim), nn.ReLU(), nn.Linear(h_dim, h_dim), nn.ReLU(), nn.Linear(h_dim, x_dim), ) def encode(self, x: torch.Tensor): h = self.enc(x.view(x.size(0), -1)) return self.fc_mu(h), self.fc_logvar(h) def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor): # z = μ + σ ⊙ ε, σ = exp(0.5 · logvar) if self.training: std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) return mu + std * eps else: # 推理时用 posterior mean (deterministic) return mu def decode(self, z: torch.Tensor): logits = self.dec(z) if self.likelihood == "bernoulli": return torch.sigmoid(logits), logits return logits, logits # Gaussian likelihood: 视为 mean prediction def forward(self, x: torch.Tensor): mu, logvar = self.encode(x) z = self.reparameterize(mu, logvar) x_hat, logits = self.decode(z) return x_hat, logits, mu, logvar def vae_loss(x: torch.Tensor, logits: torch.Tensor, mu: torch.Tensor, logvar: torch.Tensor, likelihood: str = "bernoulli", beta: float = 1.0, free_bits: float = 0.0): """ Returns: (loss, recon, kl) beta: β-VAE 的 β(默认 1 = 标准 VAE) free_bits: 每维 KL 下限(nats)。 > 0 时启用 free bits""" B = x.size(0) x_flat = x.view(B, -1) # 1) Reconstruction term: -E_q[log p(x|z)] if likelihood == "bernoulli": # BCE-with-logits 数值上更稳,等价于 -log Bernoulli likelihood recon = F.binary_cross_entropy_with_logits( logits, x_flat, reduction="sum") / B elif likelihood == "gaussian": # 假设 σ² = 1(常数),MSE 与负 log-Gaussian 差一个常数 recon = 0.5 * F.mse_loss(logits, x_flat, reduction="sum") / B else: raise ValueError(likelihood) # 2) KL term: D_KL(N(μ, σ²) || N(0, I)) 闭式 kl_per_dim = 0.5 * (mu.pow(2) + logvar.exp() - logvar - 1) # [B, z_dim] if free_bits > 0: # Free bits: 每维 KL 下限 = free_bits(缓解 posterior collapse) kl_per_dim = torch.clamp(kl_per_dim, min=free_bits) kl = kl_per_dim.sum(dim=-1).mean() # 标量 loss = recon + beta * kl return loss, recon, kl

⚠️ 常见 bug 清单 — 写 VAE 时容易踩的坑。

  • reparameterize 写成 mu + logvar * eps,忘了 σ=exp(0.5logσ2)\sigma = \exp(0.5 \cdot \log\sigma^2)
  • KL 里 0.5 * (mu**2 + sigma**2 - 2*log_sigma - 1),注意是 logσ2=2logσ-\log \sigma^2 = -2\log\sigma
  • BCE 写 F.binary_cross_entropy(sigmoid(logits), x) 而不是 F.binary_cross_entropy_with_logits(logits, x),前者数值上不稳
  • Reduction 不一致:recon 用 sum、KL 用 mean,导致 β\beta 实际 scale 漂移

3.1 训练循环 + KL annealing

def train_vae(model, dataloader, total_steps=50_000, lr=1e-3, device="cuda", beta_max=1.0, anneal_steps=10_000, free_bits=0.0): """ KL annealing: β 从 0 线性增到 beta_max,防止训练早期 posterior 直接塌缩 """ opt = torch.optim.Adam(model.parameters(), lr=lr) model.to(device).train() step = 0 while step < total_steps: for x, _ in dataloader: x = x.to(device) beta = min(beta_max, beta_max * step / max(anneal_steps, 1)) x_hat, logits, mu, logvar = model(x) loss, recon, kl = vae_loss(x, logits, mu, logvar, beta=beta, free_bits=free_bits) opt.zero_grad(set_to_none=True) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0) opt.step() step += 1 if step >= total_steps: break

§4 VAE 变体:β\beta-VAE / IWAE / NVAE / VAE-GAN

4.1 β\beta-VAE(Higgins et al., ICLR 2017)

把 ELBO 的 KL 加权重 β\beta

Lβ-VAE(θ,ϕ;x)=Eqϕ[logpθ(xz)]βDKL(qϕ(zx)p(z))\mathcal{L}_{\beta\text{-VAE}}(\theta, \phi; x) = \mathbb{E}_{q_\phi}[\log p_\theta(x|z)] - \beta \cdot D_\text{KL}(q_\phi(z|x)\,\|\,p(z))

  • β>1\beta > 1:更强地 push posterior → 先验,鼓励 disentangled representation(每维 zz 控制一个独立因素,如 dSprites 上的位置、形状、旋转)。
  • β<1\beta < 1:放宽 KL,重建更精但 latent 不够 prior-like(采样质量差)。
  • β=1\beta = 1 退化为标准 VAE。

⚠️ disentanglement 争议 — Locatello et al. (ICML 2019, 最佳论文) 证明:纯 unsupervised disentanglement 在没有归纳偏置 / 监督的前提下是不可能的β\beta-VAE 的"涌现 disentanglement"很大程度依赖架构 + 数据集 bias,而不是 β\beta 本身。

4.2 IWAE:Importance Weighted Autoencoder(Burda et al., ICLR 2016)

ELBO 是 logp(x)\log p(x) 的一阶 bound。KK 个 importance 样本得到更紧的 bound:

LKIWAE(x)=Ez1,,zKqϕ ⁣[log1Kk=1Kpθ(x,zk)qϕ(zkx)]\mathcal{L}_K^\text{IWAE}(x) = \mathbb{E}_{z_1,\ldots,z_K \sim q_\phi}\!\left[\log \frac{1}{K}\sum_{k=1}^K \frac{p_\theta(x, z_k)}{q_\phi(z_k|x)}\right]

性质:

  • L1IWAE=\mathcal{L}_1^\text{IWAE} = ELBO(特例)。
  • LKIWAElogp(x)\mathcal{L}_K^\text{IWAE} \to \log p(x)KK \to \infty(Burda 定理)。
  • KK 越大,inference 越 expressive(但训练 cost 也 ×K\times K)。

💡 Tradeoff — IWAE 让 likelihood bound 更紧,但 encoder 学到的 posterior 不再追求"逼近真后验",转去配合 importance weighting 的几何。对下游 representation learning 不一定更好。

4.3 NVAE:Hierarchical VAE(Vahdat & Kautz, NeurIPS 2020)

多层 latent z=(z1,z2,,zL)z = (z_1, z_2, \ldots, z_L),每层依赖前层:

p(z)=p(z1)l=2Lp(zlz<l),q(zx)=q(z1x)l=2Lq(zlz<l,x)p(z) = p(z_1)\prod_{l=2}^L p(z_l | z_{<l}), \quad q(z|x) = q(z_1|x)\prod_{l=2}^L q(z_l | z_{<l}, x)

工程要点:

  • Residual normal 参数化:q(zl)=N(μp+Δμq,σpΔσq)q(z_l|\cdot) = \mathcal{N}(\mu_p + \Delta\mu_q, \sigma_p \cdot \Delta\sigma_q),让 posterior 偏离 prior 是小量
  • Spectral regularization 控制每层 KL,避免数值不稳
  • BN + Swish + depthwise 等架构调优
  • 在 CIFAR-10 / CelebA / FFHQ 上第一个把 VAE 的 NLL 推到接近 SOTA flow / autoregressive

NVAE 现在的角色:LDM 之前最强连续 VAE prior 之一,但被 diffusion 系列在 sample 质量上反超。

4.4 VAE-GAN(Larsen et al., ICML 2016)

VAE 的重建损失(pixel-MSE / BCE)对高频细节不敏感 → 生成图模糊。VAE-GAN 把 MSE 换成 / 补上 discriminator feature matching

LreconVAE-GAN=Dl(x)Dl(x^)2\mathcal{L}_\text{recon}^\text{VAE-GAN} = \|D_l(x) - D_l(\hat{x})\|^2

其中 DlD_l 是 discriminator 的中间层特征。配合 adversarial loss,重建感更清晰。

这套思路最终在 VQ-GAN(§6)里大成:VQ-VAE 框架 + perceptual + adversarial + 高码率 codebook + Transformer prior。

§5 Posterior Collapse(必考)

5.1 现象

训练中 DKL(qϕ(zx)p(z))0D_\text{KL}(q_\phi(z|x)\,\|\,p(z)) \to 0,即 qϕ(zx)p(z)q_\phi(z|x) \approx p(z)xx 无关。后果:decoder 完全忽略 zz,VAE 退化为 unconditional generative model。

5.2 原因(直观分析)

  • Decoder 太强:若 pθ(xz)p_\theta(x|z) 本身就是表达力极强的 PixelCNN / Autoregressive 解码器(Bowman 2016 的 LSTM 文本 VAE 经典翻车),它能不靠 zz 直接拟合数据,那 ELBO 最优策略就是让 KL 项归 0。
  • KL 项压力大:ELBO 在训练早期 reconstruction 还没建立,optimizer 容易先把 KL 压到 0(局部最优)。
  • 数据简单:MNIST 上 collapse 罕见,文本 VAE 上常见。

5.3 缓解方法(面试要会列)

方法做法出处
KL annealingβ(t)=min(1,t/T)\beta(t) = \min(1, t / T) 线性从 0 升到 1Bowman et al. (2016)
Free bits每维 KL 下限 λ\lambda nats:max(DKL(i),λ)\max(D_\text{KL}^{(i)}, \lambda)Kingma et al. (2016)
β\beta < 1直接减小 KL 权重β\beta-VAE 反向用法
Weakened decoder用 PixelCNN 等强 AR decoder 时人为 truncate context / 加 dropoutChen et al. (2017)
Auxiliary task加 word dropout、bag-of-words 辅助 lossBowman et al. (2016)
VAE-IAF / NF prior用更复杂的先验或 normalizing flow posteriorKingma et al. (2016)
Skip / lateral 连接让 latent 强制参与 decoder(如 VLAE)Zhao et al. (2017)
VQ-VAE离散 latent 配合 codebook commitment,结构上避免 collapsevan den Oord (2017)

Free bits 公式 — 实现极简:kl_per_dim = max(kl_per_dim, λ)。直觉:给每维 latent 保底 λ 比特信息,optimizer 不能把它压到 0 以下。λ0.5\lambda \approx 0.5-22 nats / dim 是常见值。

§6 VQ-VAE:离散 latent + Codebook + STE

6.1 结构(van den Oord, Vinyals, Kavukcuoglu, NeurIPS 2017)


x ──Encoder──> z_e(x) ∈ R^{H'×W'×D}    # continuous spatial map
                       │
                       │   对每个空间位置 (h,w),找最近 codebook 向量
                       │   k_{hw} = argmin_k ‖z_e(x)_{hw} - e_k‖²
                       ↓
            z_q(x)_{hw} = e_{k_{hw}}    # quantized spatial map (discrete code)
                       │
                       │
                       ↓
                   Decoder ──> x̂

Codebook E={e1,,eK}RD\mathcal{E} = \{e_1, \ldots, e_K\} \subset \mathbb{R}^D学习得到ze(x)z_e(x)zq(x)z_q(x) 形状一致,但 zqz_q 每个空间位置都是 codebook 中某一个 vector 的 copy(离散 index khwk_{hw})。

6.2 Loss 推导

VQ-VAE 不学随机 posterior q(zx)q(z|x)(不像 VAE),而是用确定性最近邻zezqz_e \to z_q 的"量化"。loss 由三部分组成:

  LVQ-VAE=xx^2reconstruction+sg[ze(x)]e2codebook+βze(x)sg[e]2commitment  \boxed{\;\mathcal{L}_\text{VQ-VAE} = \underbrace{\|x - \hat{x}\|^2}_{\text{reconstruction}} + \underbrace{\|\text{sg}[z_e(x)] - e\|^2}_{\text{codebook}} + \beta \underbrace{\|z_e(x) - \text{sg}[e]\|^2}_{\text{commitment}}\;}

各项含义:

  • Reconstructionxzezqx^x \to z_e \to z_q \to \hat{x} 的端到端重建(注意梯度通过 STE 穿过量化)。
  • Codebook loss:把 codebook 向量 ee 拉向 ze(x)z_e(x),梯度只更新 ee(用 sg 阻断对 zez_e 的梯度,否则 codebook 与 encoder 都被拉,方向不清晰)。
  • Commitment loss:把 encoder 输出 ze(x)z_e(x) 拉向 codebook 向量 ee,梯度只更新 encoder,权重 β\beta(论文用 β=0.25\beta = 0.25)。

sg[·] = stop_gradient(PyTorch 里 .detach()),定义:前向 sg[u]=u\text{sg}[u] = u,反向 sg[u]=0\nabla \text{sg}[u] = 0

💡 为什么 codebook 和 commitment 都要 sg — 如果都不 sg,zee2\|z_e - e\|^2 同时拉两侧,方向耦合容易振荡。把这一项拆成两个 sg 版本:codebook 项专门更新 ee,commitment 专门更新 encoder,学习率 / 速度可解耦。这是 vector quantization 文献的标准做法(也叫 "alternating minimization")。

6.3 Straight-Through Estimator (STE) 推导

问题zq=eargminkzeek2z_q = e_{\arg\min_k \|z_e - e_k\|^2}argmin 不可导(输出离散 index)。

STE 解法

  • 前向:照常 zq=ekz_q = e_k(离散)
  • 反向:直接把 Lzq\frac{\partial \mathcal{L}}{\partial z_q} 当作 Lze\frac{\partial \mathcal{L}}{\partial z_e} 反传到 encoder

PyTorch 实现技巧(经典三行):

z_q = z_e + (z_q_quantized - z_e).detach()

前向:z_q = z_e + (z_q_q - z_e) = z_q_q ✓(量化值) 反向:(z_q_q - z_e).detach() 不参与梯度,所以 dz_q/dz_e = 1,梯度直通到 encoder ✓

⚠️ STE 的等价 surrogate — STE 等价于把不可导的 zq=quantize(ze)z_q = \text{quantize}(z_e) 替换成可导 surrogate zqsurrogate=zez_q^\text{surrogate} = z_e 来反传——即"假设量化是恒等映射"。这是一个有偏估计(biased gradient estimator),但实践中工作良好;理论分析见 Bengio et al. (2013) "Estimating or Propagating Gradients Through Stochastic Neurons"。

6.4 VQ-VAE 完整实现

class VectorQuantizer(nn.Module): """ Codebook + 最近邻量化 + STE embedding_dim = D, num_embeddings = K commitment_cost β 通常 = 0.25 """ def __init__(self, num_embeddings: int, embedding_dim: int, commitment_cost: float = 0.25): super().__init__() self.K, self.D = num_embeddings, embedding_dim self.beta = commitment_cost # Codebook 用 small uniform init self.codebook = nn.Embedding(self.K, self.D) self.codebook.weight.data.uniform_(-1.0 / self.K, 1.0 / self.K) def forward(self, z_e: torch.Tensor): """ z_e: [B, D, H, W] -> z_q: [B, D, H, W], indices: [B, H, W], loss = codebook_loss + β·commitment_loss """ # 1) Reshape: [B, D, H, W] -> [BHW, D] B, D, H, W = z_e.shape z_e_flat = z_e.permute(0, 2, 3, 1).contiguous().view(-1, D) # [BHW, D] # 2) 计算 L2 距离 ‖z_e - e_k‖² = ‖z_e‖² + ‖e_k‖² - 2 z_e · e_k e = self.codebook.weight # [K, D] dist = (z_e_flat.pow(2).sum(1, keepdim=True) + e.pow(2).sum(1) - 2 * z_e_flat @ e.t()) # [BHW, K] # 3) 最近邻 index indices = dist.argmin(dim=1) # [BHW] z_q_flat = self.codebook(indices) # [BHW, D] # 4) 损失(注意 sg) codebook_loss = F.mse_loss(z_q_flat, z_e_flat.detach()) commitment_loss = F.mse_loss(z_e_flat, z_q_flat.detach()) vq_loss = codebook_loss + self.beta * commitment_loss # 5) STE:前向 z_q,反向 dz_q/dz_e = I z_q_flat = z_e_flat + (z_q_flat - z_e_flat).detach() # 6) Reshape 回 [B, D, H, W] z_q = z_q_flat.view(B, H, W, D).permute(0, 3, 1, 2).contiguous() indices = indices.view(B, H, W) # 7) (可选) perplexity: codebook 使用度的衡量 one_hot = F.one_hot(indices.view(-1), self.K).float() avg_probs = one_hot.mean(dim=0) perplexity = torch.exp(-(avg_probs * torch.log(avg_probs + 1e-10)).sum()) return z_q, indices, vq_loss, perplexity class VQVAE(nn.Module): def __init__(self, channels=3, hidden=128, num_embeddings=512, embedding_dim=64, commitment_cost=0.25): super().__init__() # Encoder: 64×64×3 -> 16×16×D (downsample 4×) self.encoder = nn.Sequential( nn.Conv2d(channels, hidden, 4, 2, 1), nn.ReLU(), nn.Conv2d(hidden, hidden, 4, 2, 1), nn.ReLU(), nn.Conv2d(hidden, embedding_dim, 3, 1, 1), ) self.quantizer = VectorQuantizer(num_embeddings, embedding_dim, commitment_cost) # Decoder: 16×16×D -> 64×64×3 self.decoder = nn.Sequential( nn.ConvTranspose2d(embedding_dim, hidden, 3, 1, 1), nn.ReLU(), nn.ConvTranspose2d(hidden, hidden, 4, 2, 1), nn.ReLU(), nn.ConvTranspose2d(hidden, channels, 4, 2, 1), ) def forward(self, x): z_e = self.encoder(x) z_q, indices, vq_loss, perplexity = self.quantizer(z_e) x_hat = self.decoder(z_q) return x_hat, vq_loss, perplexity, indices def vqvae_loss(x, x_hat, vq_loss): recon = F.mse_loss(x_hat, x) return recon + vq_loss, recon

6.5 EMA Codebook(生产标准做法)

直接用 codebook loss 更新 ee 收敛慢,dead codes(从来不被选中的 codebook 向量)多。生产实现用 EMA (Exponential Moving Average) 更新

对每个 codebook 向量 eke_k,维护:

  • Nk(t)=γNk(t1)+(1γ)nk(t)N_k^{(t)} = \gamma N_k^{(t-1)} + (1-\gamma) n_k^{(t)},其中 nk(t)n_k^{(t)} 是当前 batch 中分配到 eke_k 的样本数
  • mk(t)=γmk(t1)+(1γ)i:ze,iekze,im_k^{(t)} = \gamma m_k^{(t-1)} + (1-\gamma) \sum_{i: z_{e,i} \to e_k} z_{e,i}

更新:

ek(t)=mk(t)Nk(t)+ε(Laplace smoothing)e_k^{(t)} = \frac{m_k^{(t)}}{N_k^{(t)} + \varepsilon} \quad \text{(Laplace smoothing)}

class VectorQuantizerEMA(nn.Module): """ EMA codebook update (van den Oord 2017 后续 / VQ-VAE-2 标准做法) - codebook 不靠 gradient,靠 running EMA 更新 - 留 commitment loss 用于 encoder - decay γ 一般 0.99, ε 一般 1e-5 """ def __init__(self, num_embeddings: int, embedding_dim: int, commitment_cost: float = 0.25, decay: float = 0.99, eps: float = 1e-5): super().__init__() self.K, self.D = num_embeddings, embedding_dim self.beta, self.decay, self.eps = commitment_cost, decay, eps embed = torch.randn(num_embeddings, embedding_dim) * 0.01 self.register_buffer("codebook", embed) self.register_buffer("cluster_size", torch.zeros(num_embeddings)) self.register_buffer("embed_avg", embed.clone()) def forward(self, z_e): B, D, H, W = z_e.shape z_e_flat = z_e.permute(0, 2, 3, 1).contiguous().view(-1, D) dist = (z_e_flat.pow(2).sum(1, keepdim=True) + self.codebook.pow(2).sum(1) - 2 * z_e_flat @ self.codebook.t()) indices = dist.argmin(dim=1) # [BHW] z_q_flat = F.embedding(indices, self.codebook) # [BHW, D] if self.training: # EMA 更新 with torch.no_grad(): one_hot = F.one_hot(indices, self.K).float() # [BHW, K] cluster_size_new = one_hot.sum(dim=0) # [K] embed_sum = one_hot.t() @ z_e_flat # [K, D] self.cluster_size.mul_(self.decay).add_(cluster_size_new, alpha=1 - self.decay) self.embed_avg.mul_(self.decay).add_(embed_sum, alpha=1 - self.decay) # Laplace smoothing 避免被零除 n = self.cluster_size.sum() cluster_size = (self.cluster_size + self.eps) / (n + self.K * self.eps) * n self.codebook.copy_(self.embed_avg / cluster_size.unsqueeze(1)) commitment_loss = F.mse_loss(z_e_flat, z_q_flat.detach()) vq_loss = self.beta * commitment_loss # EMA 下不要 codebook loss z_q_flat = z_e_flat + (z_q_flat - z_e_flat).detach() # STE z_q = z_q_flat.view(B, H, W, D).permute(0, 3, 1, 2).contiguous() return z_q, indices.view(B, H, W), vq_loss

EMA 的两个好处

  • 更新更稳:EMA 是一种隐式 momentum,相当于 codebook 的 SGD 用了大 batch
  • 死码自动重启更容易:可周期性把 cluster_size<τ\text{cluster\_size} < \taueke_k 重置到当前 batch 的某个 zez_edead-code revival

6.6 VQ-VAE-2(Razavi, Vinyals, van den Oord, NeurIPS 2019)

VQ-VAE 的层次扩展:

  • 顶层 latent ztz_t:低分辨率(如 32×32),捕捉 global 结构(人脸的整体姿态、身份)
  • 底层 latent zbz_b:高分辨率(如 64×64),捕捉 local 细节(皮肤纹理、发丝)
  • PixelCNN prior 在两层 latent 上分别训练,顶层无条件,底层条件于顶层

在 ImageNet 256×256 上首次让 VQ-based 方法接近 BigGAN 的样本质量,是 VQ-GAN 的直接前作。

§7 VQ-GAN:Adversarial + Perceptual + Transformer Prior

7.1 核心思想(Esser, Rombach, Ommer, CVPR 2021, "Taming Transformers")

VQ-VAE 在 ImageNet 上重建的纹理细节模糊。VQ-GAN 改造为:

组件VQ-VAEVQ-GAN
Recon lossL2 / L1 pixelL1 pixel + LPIPS perceptual + PatchGAN adversarial
PriorPixelCNNTransformer (decoder-only) over code tokens
Codebook512-1024 codes1024-16384 codes
Compression4×-8×8×-32×(更高压缩比,靠 perceptual + adversarial 救质量)
应用unconditional / class-cond 生成high-res image synthesis, "Taming Transformers"

7.2 Loss 公式

LVQ-GANstage1=Lrec+LVQ+λLGAN\mathcal{L}_\text{VQ-GAN}^\text{stage1} = \mathcal{L}_\text{rec} + \mathcal{L}_\text{VQ} + \lambda \cdot \mathcal{L}_\text{GAN}

其中:

Lrec=xx^1+LLPIPS(x,x^)LVQ=sg[ze]e2+βzesg[e]2\begin{aligned} \mathcal{L}_\text{rec} &= \|x - \hat{x}\|_1 + \mathcal{L}_\text{LPIPS}(x, \hat{x}) \\ \mathcal{L}_\text{VQ} &= \|\text{sg}[z_e] - e\|^2 + \beta \|z_e - \text{sg}[e]\|^2 \end{aligned}

Generator/Tokenizer 阶段的 GAN 项(只对 generator 的输出反传,discriminator 在另一阶段单独更新):

LGAN(G)=Ex^[logD(x^)](non-saturating)LGAN(G)=Ex^[D(x^)](hinge)\mathcal{L}_\text{GAN}^{(G)} = -\mathbb{E}_{\hat{x}}[\log D(\hat{x})]\quad\text{(non-saturating)}\quad\text{或}\quad \mathcal{L}_\text{GAN}^{(G)} = -\mathbb{E}_{\hat{x}}[D(\hat{x})]\quad\text{(hinge)}

Discriminator 自身的 minimax 项(独立 step 更新 DD):

LGAN(D)=Ex[min(0,1+D(x))]Ex^[min(0,1D(x^))](hinge)\mathcal{L}_\text{GAN}^{(D)} = -\mathbb{E}_x[\min(0, -1+D(x))] - \mathbb{E}_{\hat{x}}[\min(0, -1-D(\hat{x}))]\quad\text{(hinge)}

自适应 λ\lambda(论文创新点,用最后一层梯度范数比自动平衡,避免人工调参):

λ=GLLrecGLLGAN(G)+δ\lambda = \frac{\lVert\nabla_{G_L} \mathcal{L}_\text{rec}\rVert}{\lVert\nabla_{G_L} \mathcal{L}_\text{GAN}^{(G)}\rVert + \delta}

GLG_L 是 decoder 最后一层;\lVert\cdot\rVert 是 Frobenius 范数。总 generator loss:

LG=Lrec+LVQ+λLGAN(G)\mathcal{L}_G = \mathcal{L}_\text{rec} + \mathcal{L}_\text{VQ} + \lambda \cdot \mathcal{L}_\text{GAN}^{(G)}

7.3 Stage 2:Transformer Prior

Stage 1 训好 VQ-GAN,把图像转成 token grid c=(c1,,cHW)\mathbf{c} = (c_1, \ldots, c_{HW})(行扫描展平)。Stage 2 在 token sequence 上训 decoder-only Transformer,标准 AR:

p(c)=i=1HWp(cic<i)p(\mathbf{c}) = \prod_{i=1}^{HW} p(c_i | c_{<i})

采样:AR sample tokens → VQ-GAN decoder → image。这是把"图像生成"翻译成"语言模型"的标准范式,DALL·E / Parti / Muse 都是这一思想的进化。

💡 VQ-GAN 在 LDM 中的角色 — Stable Diffusion 的 VAE tokenizer 实际是 KL-regularized VQ-GAN 的连续 latent 变体(去掉 quantization,留 KL + perceptual + adversarial),output 是 continuous latent map(4 通道,下采样 8×)。diffusion 在这个 latent 上跑,最后 decoder 还原。可以理解为 "VQ-GAN encoder/decoder + continuous latent + KL"。

7.4 PatchGAN Discriminator(生产架构)

VQ-GAN 用 PatchGAN(Isola et al. CVPR 2017 "pix2pix"):

  • 不输出 single scalar real/fake
  • 输出 N×N 的 patch-level 判别 map(每个 patch 是 70×70 receptive field)
  • 适合 capture 局部纹理真假,对全局结构压力小(让 generator 更专注于纹理)
class PatchDiscriminator(nn.Module): """ PatchGAN: 70×70 receptive field 的 stack of strided convs 输出 [B, 1, H/8, W/8] 的 patch-level real/fake 判别 """ def __init__(self, in_ch=3, hidden=64, n_layers=3): super().__init__() layers = [nn.Conv2d(in_ch, hidden, 4, 2, 1), nn.LeakyReLU(0.2, True)] ch = hidden for i in range(1, n_layers): ch_next = min(hidden * (2 ** i), 512) layers += [ nn.Conv2d(ch, ch_next, 4, 2, 1), nn.BatchNorm2d(ch_next), nn.LeakyReLU(0.2, True), ] ch = ch_next layers += [ nn.Conv2d(ch, ch * 2, 4, 1, 1), nn.BatchNorm2d(ch * 2), nn.LeakyReLU(0.2, True), nn.Conv2d(ch * 2, 1, 4, 1, 1), ] self.net = nn.Sequential(*layers) def forward(self, x): return self.net(x) def hinge_d_loss(real_logits, fake_logits): real = F.relu(1.0 - real_logits).mean() fake = F.relu(1.0 + fake_logits).mean() return 0.5 * (real + fake) def hinge_g_loss(fake_logits): return -fake_logits.mean()

⚠️ GAN 训练 trick 清单

  • D 启动 delay:前 K 步只训 G(让 recon 先收敛)
  • LeCam regularization:D 的输出 anchor 到 EMA,缓解 mode collapse
  • R1 gradient penalty:γxD(x)2\gamma \|\nabla_x D(x)\|^2 防 D 过拟合
  • Spectral norm:稳定 D
  • Adam β1=0.5\beta_1 = 0.5(不是默认 0.9),β2=0.9\beta_2 = 0.9

7.5 LPIPS(Perceptual Loss)

LLPIPS(x,x^)=lwlϕl(x)ϕl(x^)2\mathcal{L}_\text{LPIPS}(x, \hat{x}) = \sum_l w_l \cdot \|\phi_l(x) - \phi_l(\hat{x})\|^2

ϕl\phi_l 是预训练 VGG / AlexNet 的第 ll 层 feature map,wlw_l 是学到的 channel-wise weight(Zhang et al. CVPR 2018)。比 pixel-MSE 更贴近人类感知,是 VQ-GAN / SD / 大部分 image GAN / diffusion 训练的标配。

§8 离散 VAE 与 Gumbel-Softmax

8.1 dVAE(DALL·E 1,Ramesh et al. ICML 2021)

DALL·E 用 dVAE (discrete VAE) 作为 image tokenizer:

  • 每个 spatial 位置输出 categorical distribution over 8192 codes
  • 训练用 Gumbel-softmax 让 categorical 可导
  • 推理用 hard argmax 离散化

8.2 Gumbel-Softmax 推导

目标:让 categorical sampling 可导。

Gumbel-Max trick:对 logits π=(π1,,πK)\pi = (\pi_1, \ldots, \pi_K) 加独立 Gumbel(0,1) 噪声 gk=log(loguk),ukU(0,1)g_k = -\log(-\log u_k), u_k \sim \mathcal{U}(0, 1),则:

argmaxk{πk+gk}\arg\max_k \{\pi_k + g_k\}

服从 categorical(softmax(π\pi))。证明用 Gumbel 分布的 CDF 性质:P(maxkXk=Xj)=eπj/keπkP(\max_k X_k = X_j) = e^{\pi_j} / \sum_k e^{\pi_k}

Gumbel-softmax (Jang et al., ICLR 2017; Maddison et al., ICLR 2017 同期):把不可导的 argmax 替换成 softmax:

  yk=exp((πk+gk)/τ)jexp((πj+gj)/τ)  \boxed{\;y_k = \frac{\exp((\pi_k + g_k) / \tau)}{\sum_j \exp((\pi_j + g_j) / \tau)}\;}

  • τ0\tau \to 0yy 接近 one-hot(贴近 categorical 采样)
  • τ\tau \to \inftyyy 接近均匀(梯度好但偏离)
  • 训练时 τ\tau anneal: 1.00.11.0 \to 0.1 渐降

Straight-Through Gumbel-Softmax:前向用 argmax(离散),反向用 softmax 梯度——类似 VQ-VAE 的 STE 思路。

def gumbel_softmax_sample(logits, tau=1.0, hard=False, dim=-1): """ 输入 logits = log π 输出 soft / hard one-hot """ # 1) 加 Gumbel 噪声 g = -torch.log(-torch.log(torch.rand_like(logits) + 1e-20) + 1e-20) y_soft = F.softmax((logits + g) / tau, dim=dim) if not hard: return y_soft # ST: 前向 hard, 反向 soft 梯度 index = y_soft.argmax(dim=dim, keepdim=True) y_hard = torch.zeros_like(y_soft).scatter_(dim, index, 1.0) y = y_hard - y_soft.detach() + y_soft # straight-through return y

💡 VQ-VAE vs Gumbel-softmax / dVAE — 都是离散 latent 模型,区别:

  • VQ-VAE:encoder 出 continuous zez_e最近邻到 codebook(hard,无随机性);用 STE 反传。
  • Gumbel dVAE:encoder 出 categorical distribution(logits over K codes),训练用 Gumbel-softmax 采样。
  • 实践:DALL·E 1 用 dVAE 配合 AR Transformer prior;后来 DALL·E 2 / Parti / Muse 都偏向 VQ-GAN 系列(更好质量)。

8.3 MaskGIT(Chang et al., CVPR 2022)

把 AR Transformer prior 换成 BERT-style masked Transformer

  • 训练:随机 mask 一部分 VQ token,让模型预测被 mask 的 token(类似 BERT MLM)
  • 采样:Non-autoregressive 并行采样——每轮 unmask 一批 token,迭代 8-12 轮收敛
  • 比 AR 快 ~10x,质量相当或更好(在 ImageNet 256×256 上)

后继:MUSE (Chang et al., 2023) 把同思路推到 text-to-image,是 Google 主线生成模型之一。

§9 FSQ:Finite Scalar Quantization(重点)

9.1 动机(Mentzer, Minnen, Agustsson, Toderici, ICLR 2024)

VQ-VAE 有顽疾:

  1. Codebook collapse / underuse:大部分 codes 从未被用(dead codes),perplexity 远低于理论 KK
  2. STE bias:梯度估计有偏,训练不稳
  3. 复杂的 loss balancing:commitment 权重、EMA decay、dead code revival 都要调
  4. Codebook 大小有效上限:实际可用 ~10310^3-10410^4,再大用不起来

FSQ 用一招把这一堆问题绕过去逐维标量量化(scalar quantization, 不是 vector quantization)

9.2 核心公式(必推)

让 encoder 输出 zRdz \in \mathbb{R}^d对每一维独立做 scalar quantization(FSQ 论文 Eq. 4):

zizi=Li12tanh(zi)siz^i=round(zi)+siz_i \longrightarrow z'_i = \tfrac{L_i-1}{2}\tanh(z_i) - s_i \longrightarrow \hat{z}_i = \text{round}(z'_i) + s_i

其中 si=0s_i = 0LiL_i 奇,si=0.5s_i = 0.5LiL_i 偶。这样:

  • LiL_i 奇(如 5):z^i{2,1,0,1,2}\hat{z}_i \in \{-2,-1,0,1,2\}(恰好 LiL_i 个整数 level)
  • LiL_i 偶(如 8):z^i{3.5,2.5,,2.5,3.5}\hat{z}_i \in \{-3.5,-2.5,\ldots,2.5,3.5\}(恰好 LiL_i 个半整数 level)

无论奇偶,每维都得到 LiL_i 个 level;把所有维度的 level 数乘起来:

  Kimplicit=i=1dLi  \boxed{\;K_\text{implicit} = \prod_{i=1}^{d} L_i\;}

隐式 codebook 大小例子

  • L=(8,6,5)L = (8, 6, 5), d=3d = 3:codebook = 8×6×5=2408 \times 6 \times 5 = 240
  • L=(8,5,5,5)L = (8, 5, 5, 5), d=4d = 4:codebook = 8×5×5×5=10008 \times 5 \times 5 \times 5 = 1000
  • L=(7,5,5,5,5)L = (7, 5, 5, 5, 5), d=5d = 5:codebook = 7×54=43757 \times 5^4 = 4375
  • L=(8,8,8,5,5,5)L = (8, 8, 8, 5, 5, 5), d=6d = 6:codebook = 8353=64,0008^3 \cdot 5^3 = 64{,}000
  • 没有显式 codebook 表z^i\hat{z}_i(L1,,Ld)(L_1, \ldots, L_d) 组合本身就是离散 code(直接用 base-mixed encoding 转 1,,Kimplicit1, \ldots, K_\text{implicit} 即可)。

9.3 为什么 FSQ 不会 codebook collapse?

关键洞察 — 在 VQ-VAE 中,codebook collapse 的根源是:codebook 是自由参数,optimizer 让大部分 eke_k 都 drift 到无用区域,只有少数 eke_k 被反复用。FSQ 的"codebook"不是参数——它是数轴上的固定 grid 点({L/2,,L/2}\{-L/2, \ldots, L/2\})。

  • 没有 codebook 参数 → 没有 codebook collapse:grid 点固定,不会跑偏。
  • 每维独立 → 高维 code 通过乘积自动多样化:即使每维只用 L=8L=8 个 level,d=6d=6 时已经有 86=2621448^6 = 262{144} 种组合。
  • encoder 自适应分布:因为前面有 tanh\tanh 压到 [1,1][-1, 1],encoder 自然把输出散布在 [L/2,L/2][-L/2, L/2] 区间——dead code 出现的唯一原因是 encoder 没把某些 grid 区间用到,但只要 reconstruction 推动 encoder 探索整个区间,所有 grid 都会被覆盖。

实证:FSQ 的 codebook usage 几乎是 100%(与 VQ-VAE 50-70% 形成对比),ImageNet / Cosmos / OpenMagViT2 复现均有此结论。

9.4 为什么 FSQ 不需要"显式 STE 包装"以及 loss 也极简

  • Rounding 是不可导的——和 VQ-VAE 一样需要某种 STE。但 FSQ 的 STE 只有一行 x_hat = x + (round(x) - x).detach(),没有 codebook loss / commitment loss / EMA / dead-code revival。
  • Loss 只剩 reconstruction

LFSQ=xx^2(plus optional perceptual + adversarial)\mathcal{L}_\text{FSQ} = \|x - \hat{x}\|^2 \quad \text{(plus optional perceptual + adversarial)}

  • 没有 hyperparameter 调 commitment cost / EMA decay / restart threshold——这是 FSQ 在工程上对 VQ-VAE 的最大优势。

💡 VQ-VAE vs FSQ 简化对比 — FSQ 是"用空间维度换 codebook 大小":VQ-VAE 用 1 维(DD 个连续值 + 1 个离散选择 from KK),FSQ 用 dd 个独立离散维度,每维 LiL_i 个 level,最终离散 entropy 反而更大、collapse 几乎不可能。代价:embedding 表达力略弱(每维独立,不共享 representation),不过 reconstruction 端通过 decoder 已补回。

9.5 FSQ 实现(10 行)

class FSQ(nn.Module): """ Finite Scalar Quantization (Mentzer et al., ICLR 2024) levels: tuple, 每维量化 level 数(必须为奇数或偶数都行,奇数保证含 0) eps: bounding 安全裕度,避免 tanh 后 round 跳出 grid """ def __init__(self, levels=(8, 5, 5, 5)): super().__init__() levels_t = torch.tensor(levels, dtype=torch.float32) self.levels = levels_t self.d = len(levels) self.K = int(torch.prod(levels_t).item()) # 隐式 codebook size = ∏ L_i # FSQ paper Eq. 4: half = (L-1)/2; shift = 0.5 if L 偶 else 0 half = (levels_t - 1) / 2 # [d] shift = ((levels_t % 2) == 0).float() * 0.5 # [d] self.register_buffer("half_l", half) self.register_buffer("shift", shift) # mixed-radix basis for token id encoding cumprod = torch.tensor([1.0] + list(torch.cumprod(levels_t[:-1], dim=0)), dtype=torch.float32) self.register_buffer("basis", cumprod) # [d] @staticmethod def round_ste(z): """不可导 round 的 STE: 前向 round, 反向 identity""" return z + (z.round() - z).detach() def forward(self, z): """ z: [B, d, ...] -> z_hat: [B, d, ...] (量化值), codes: [B, ...] (∈ 0..K-1) """ view = (1, -1) + (1,) * (z.dim() - 2) half = self.half_l.view(*view).to(z.device) shift = self.shift.view(*view).to(z.device) # 1) Bound: tanh(z) * half - shift → z'∈[-half-shift, half-shift] z_bounded = torch.tanh(z) * half - shift # 2) Round (STE) + 加回 shift → 奇 L 得 {-half,…,half}(整数),偶 L 得 {-half,…,half}(含半整数) z_hat = self.round_ste(z_bounded) + shift # 3) Token ID (mixed-radix):把 d 维 ∈ {-half_i,…,half_i} 映成 0..L_i-1 再编成单一 index shifted = (z_hat + half).round().long() # ∈ 0..L_i-1(round 兜底浮点误差) basis = self.basis.view(*view).to(z.device).long() codes = (shifted * basis).sum(dim=1) # [B, ...] return z_hat, codes # 使用示例: # fsq = FSQ(levels=(8, 5, 5, 5)) # K = 8·5·5·5 = 1000 # z = encoder(x) # [B, 4, H, W] # z_hat, tokens = fsq(z) # z_hat: [B, 4, H, W], tokens: [B, H, W] ∈ 0..999 # x_hat = decoder(z_hat) # loss = F.mse_loss(x_hat, x) # 就这一项!

⚠️ FSQ 的 level 选择经验

  • 论文表 3 给出经验配方(ImageNet 256×256):K1000K \approx 1000(8,5,5,5)(8, 5, 5, 5)K4000K \approx 4000(7,5,5,5,5)(7, 5, 5, 5, 5)K64000K \approx 64000(8,8,8,5,5,5)(8, 8, 8, 5, 5, 5)
  • 经验法则:让 LiL_i 配比近似为黄金比 / 反比例(信息论上各维信息量平衡)
  • 实践上不是非常敏感——任何近似合理的 level 组合都能 work

9.6 LFQ:Lookup-Free Quantization(MAGVIT-v2,Yu et al., ICLR 2024)

FSQ 的二值特例:

LFQ(z)=sign(z){1,+1}d\text{LFQ}(z) = \text{sign}(z) \in \{-1, +1\}^d

每维只有 2 个 level,隐式 codebook = 2d2^dd=18d=18 时 codebook = 218=2621442^{18} = 262{144}(与 FSQ 等价的量级)。

特点:

  • 每维 binary,最简结构
  • VQ-token 转 binary code,用 BitVQ / bitwise predictor 训练
  • MAGVIT-v2 / Open-MAGVIT2 / VideoPoet 用 LFQ 做视频 tokenizer
  • 加 entropy regularization 维持每位 50/50(避免某些 bit 总是 +1+1
class LFQ(nn.Module): """ Lookup-Free Quantization (MAGVIT-v2) 每维独立 sign quantize, 隐式 codebook = 2^d """ def __init__(self, dim: int, entropy_weight: float = 0.1): super().__init__() self.d = dim self.K = 2 ** dim self.entropy_weight = entropy_weight def forward(self, z): # z: [B, d, ...] q = torch.sign(z) # 防止 sign(0) = 0 q = torch.where(q == 0, torch.ones_like(q), q) # STE z_hat = z + (q - z).detach() # Entropy regularization(防止某维总是同符号) # p_+ = sigmoid(z), p_- = 1 - p_+ if self.training: p = torch.sigmoid(z) per_dim_entropy = -(p * torch.log(p + 1e-9) + (1 - p) * torch.log(1 - p + 1e-9)) entropy_loss = -per_dim_entropy.mean() # maximize entropy → minimize -H else: entropy_loss = z.new_tensor(0.0) return z_hat, self.entropy_weight * entropy_loss

9.7 Cosmos / OpenMagViT2 / 现代 video tokenizer

Tokenizer出处量化方式用在哪
MAGVIT-v2Google 2024 (ICLR)LFQtext-to-video 早期 demo
OpenMagViT2开源复现 2024LFQ公开 video tokenizer baseline
Cosmos TokenizerNVIDIA 2024FSQ + 视频时空压缩NVIDIA Cosmos world model
VideoPoet tokenizerGoogle 2024LFQ-styletext-to-video

工程要点:

  • 时空联合压缩:spatial 8× + temporal 4×(4 帧合 1 token plane)
  • 3D causal CNN encoder(前向时间因果,可流式编码长视频)
  • 跨 resolution generalization:训 256×256 推 1024×1024 必须 careful test-time 适配

§10 复杂度与资源对比

模型latent 类型训练参数 (encoder+decoder)主要 lossCodebook collapseSTE 依赖
VAEcontinuous Gaussian\sim10-100Mrecon + KL (closed form)N/A否 (reparameterization)
β\beta-VAEcontinuous Gaussian同 VAErecon + β\beta·KLN/A
NVAEhierarchical continuous80M-200Mrecon + multi-layer KLN/A
VQ-VAEdiscrete via codebook50-200Mrecon + codebook + β\beta·commitment常发生
VQ-VAE-2hierarchical discrete100-500M同 VQ-VAE × 2 层同上
VQ-GANdiscrete + adversarial50-300M (+ D)recon + LPIPS + GAN + codebook + commitment同上
dVAEcategorical (logits)50-200Mrecon + KL to uniform较少(categorical 分布学习)否(Gumbel-softmax 反传 soft)
FSQscalar quantize per dim30-150Mrecon (+ perceptual)几乎不发生是(但极简)
LFQbinary scalar quantize30-150Mrecon (+ entropy reg)几乎不发生

💡 生态位定位

  • 想做 diffusion / FM:用 KL-VAE / SD VAE(连续 latent)
  • 想做 AR 生成(GPT-style 图像 token):用 VQ-GAN / FSQ / LFQ
  • 想做 MaskGIT / Muse / 并行 decode:用 VQ-GAN / FSQ
  • 想做 video / 长 sequence:用 FSQ / LFQ(codebook usage 高、无 collapse)

§11 与相关方法对比 / 在生态中的位置

11.1 VAE vs GAN vs Diffusion vs Flow / FM

模型似然训练稳定性多样性样本质量inference 速度
VAE有(ELBO)✅ 稳✅ 好⚠️ 模糊✅ 1-step
GAN❌ 难❌ mode collapse✅ 锐利✅ 1-step
Diffusion近似(VLB)✅ 稳✅ 好✅ SOTA❌ 多 NFE
Flow / FM有(ODE)✅ 稳✅ 好✅ 强⚠️ 数 NFE

11.2 Tokenizer 系列在大模型中的角色


Tokenizer Stage 1                Generative Stage 2 (prior)
────────────────                 ──────────────────────────
VQ-GAN  →   离散 token grid  →   Transformer AR  (Parti, DALL·E 1, Cogview)
VQ-GAN  →   离散 token grid  →   Masked Transformer (MaskGIT, Muse)
FSQ    →   离散 token grid  →   Transformer AR  (Cosmos, OpenMagViT2)
LFQ    →   binary token grid →   AR / bit predictor (MAGVIT-v2, VideoPoet)
KL-VAE →   连续 latent map  →   Diffusion / Flow Matching (LDM, SD, SD3, FLUX)

11.3 Reconstruction-Perception Tradeoff(高级题)

Blau & Michaeli (CVPR 2018) 证明:重建(MSE / PSNR)和感知(perceptual / FID)之间存在严格的 Pareto 边界。VQ-GAN / SD VAE 引入 LPIPS + adversarial 是为了交换更高 perceptual 质量而接受略差的 PSNR

⚠️ PSNR 不等于"看起来好" — VQ-GAN 论文里 PSNR 不一定优于 VQ-VAE,但 perceptual (LPIPS / FID) 远好。面试常被反问"为什么 SOTA tokenizer 的 PSNR 反而下降"——这是 distortion-perception tradeoff。

§12 25 高频面试题

codex (gpt-5.5 xhigh) 作为顶级 lab 面试官视角列的,按难度分 3 档。每题点开看答案要点 + 易踩坑。

L1必会题(任何 ML 工程岗都会问)

Q1.VAE 的 ELBO 是什么?写出公式。
  • logp(x)Eqϕ(zx)[logpθ(xz)]DKL(qϕ(zx)p(z))\log p(x) \geq \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_\text{KL}(q_\phi(z|x)\,\|\,p(z))

  • 第一项:reconstruction 期望 log-likelihood

  • 第二项:KL 把后验拉向先验 N(0,I)\mathcal{N}(0, I)

  • ELBO 与 logp(x)\log p(x) 的 gap = DKL(qϕ(zx)pθ(zx))D_\text{KL}(q_\phi(z|x)\,\|\,p_\theta(z|x))

写成 logp(xz)DKL(...)\log p(x|z) - D_\text{KL}(...)(漏了期望符号);只说"重建+正则化"不写公式。

Q2.Reparameterization trick 解决什么问题?
  • 直接采样 zqϕ(zx)z \sim q_\phi(z|x) 不可导,梯度无法回到 encoder

  • 改写成 z=μ+σϵ,ϵN(0,I)z = \mu + \sigma \odot \epsilon, \epsilon \sim \mathcal{N}(0, I),把随机性挪到独立噪声

  • zz 变成 ϕ\phi 的确定性函数,可正常反向传播

  • 不止 Gaussian;Gumbel-softmax 也是同思路

说"为了加速训练"——其实是为了可导

Q3.$\beta$-VAE 的 $\beta$ 控制什么?
  • β>1\beta > 1:更强 KL 正则,鼓励 disentangled latent

  • β<1\beta < 1:放宽 KL,重建更精但 prior 拟合差

  • β=1\beta = 1:标准 VAE

  • 但 Locatello 2019 证明:纯无监督 disentanglement 不可行,需要归纳偏置

只说 "β\beta 越大越 disentangle"——错,取决于数据 + 架构

Q4.什么是 posterior collapse?
  • 训练中 KL → 0,即 qϕ(zx)p(z)q_\phi(z|x) \approx p(z)(与 xx 无关)

  • decoder 完全忽略 zz,VAE 退化为 unconditional 模型

  • 常见于强 AR decoder(如 PixelCNN / LSTM)+ 简单数据

只说"latent 没用",不说 KL → 0 这个量化指标。

Q5.VQ-VAE 的 codebook 是什么?怎么用?
  • {e1,,eK}RD\{e_1, \ldots, e_K\} \subset \mathbb{R}^D 一组可学的"码本"向量

  • encoder 输出连续 ze(x)z_e(x);用最近邻 k=argminzeek2k = \arg\min \|z_e - e_k\|^2 替换成 zq=ekz_q = e_k

  • decoder 解 zqz_q 回到像素

  • 训练靠 codebook loss(拉 eezez_e)+ commitment loss(拉 zez_eee

说 codebook 是固定的 / 预训练的——错,端到端学习

Q6.VQ-VAE 三项 loss 各是什么?
  • Reconstructionxx^2\|x - \hat{x}\|^2(pixel-level)

  • Codebook losssg[ze]e2\|\text{sg}[z_e] - e\|^2(只更新 ee

  • Commitment lossβzesg[e]2\beta \|z_e - \text{sg}[e]\|^2(只更新 encoder, β=0.25\beta = 0.25

  • sg = stop_gradient,避免两侧梯度互相耦合振荡

把 codebook 和 commitment loss 混作一项;忘了 sg 的方向。

Q7.什么是 Straight-Through Estimator (STE)?
  • 解决不可导操作(如 argmax / round)的反传问题

  • 前向用离散输出,反向把梯度直接传给"上一层连续输入"

  • 等价 surrogate:假设量化层是恒等映射

  • PyTorch 三行:z_q = z_e + (z_q_quantized - z_e).detach()

只说"反向用 identity",不说前向还是真量化。

Q8.VQ-GAN 比 VQ-VAE 多了什么?
  • LPIPS perceptual loss(替换 / 补充 L2)

  • PatchGAN adversarial loss + 自适应 λ\lambda 权重

  • Transformer prior(替换 PixelCNN)

  • 更大 codebook(1k → 16k)+ 更高压缩比(8× → 16-32×)

  • 是 "Taming Transformers"(Esser et al. CVPR 2021)

只说 "GAN" 不说 perceptual;或忘 Transformer prior。

Q9.FSQ 是什么?为什么不会 codebook collapse?
  • 每维独立 scalar quantize 到 LL 个固定 level(tanh\tanh → 缩放 → round)

  • 隐式 codebook = Li\prod L_i(如 L=8,d=686L=8, d=6 \Rightarrow 8^6

  • codebook 不是可学参数 → 没东西"塌"到无用区

  • encoder 通过 reconstruction 压力自然探索整个 grid

把 FSQ 当作 VQ-VAE 的 codebook 优化技巧——错,FSQ 没有显式 codebook 参数

Q10.KL($\mathcal{N}(\mu, \sigma^2 I) \,\|\, \mathcal{N}(0, I)$) 写出闭式。
  • DKL=12i(μi2+σi2logσi21)D_\text{KL} = \tfrac{1}{2}\sum_i (\mu_i^2 + \sigma_i^2 - \log \sigma_i^2 - 1)

  • 对角协方差才有这个简单形式

  • 实现里 encoder 输出 logσ2\log \sigma^2(logvar)更稳

  • kl = 0.5 * (mu**2 + logvar.exp() - logvar - 1).sum()

logσ2\log \sigma^2logσ\log \sigma 搞混(差一个 2);或忘 "-1" 项。

L2进阶题(research-oriented 岗位)

Q11.IWAE 比 ELBO 更紧吗?怎么用?
  • KK 个 importance 样本:LKIWAE=Ez1,,zK[log1Kkp(x,zk)q(zkx)]\mathcal{L}_K^\text{IWAE} = \mathbb{E}_{z_1,\ldots,z_K}[\log \tfrac{1}{K} \sum_k \tfrac{p(x, z_k)}{q(z_k|x)}]

  • L1=\mathcal{L}_1 = ELBO(特例)

  • LKlogp(x)\mathcal{L}_K \to \log p(x)KK \to \infty(Burda et al. ICLR 2016)

  • 但 encoder 学到的 posterior 不再追求逼近真后验

说"K=1K = 1 也比 ELBO 强"——错,特例就是 ELBO。

Q12.如何缓解 posterior collapse?至少列 4 种。
  • KL annealingβ(t)=min(1,t/T)\beta(t) = \min(1, t/T) 线性增长(Bowman 2016)

  • Free bits:每维 KL 下限 λ\lambda nats(Kingma 2016)

  • Weakened decoder:限制 decoder 表达力(Chen 2017)

  • VQ-VAE / 离散 latent:结构上无 KL 项,规避

  • 其他:辅助 loss、word dropout、NF posterior、skip 连接

只答"KL annealing" 一种;或把 β\beta-VAE 当成缓解 collapse 的工具(其实 β>1\beta > 1 更易 collapse)。

Q13.推 KL($\mathcal{N}(\mu, \sigma^2) \,\|\, \mathcal{N}(0, 1))$ 闭式。
  • logqp=12logσ2(zμ)22σ2+z22\log \tfrac{q}{p} = -\tfrac{1}{2}\log \sigma^2 - \tfrac{(z-\mu)^2}{2\sigma^2} + \tfrac{z^2}{2}

  • Eq\mathbb{E}_q(用 E[z]=μ\mathbb{E}[z] = \mu, E[z2]=μ2+σ2\mathbb{E}[z^2] = \mu^2 + \sigma^2, E[(zμ)2]=σ2\mathbb{E}[(z-\mu)^2] = \sigma^2

  • 结果:12(μ2+σ2logσ21)\tfrac{1}{2}(\mu^2 + \sigma^2 - \log \sigma^2 - 1)

  • 多维独立时各维相加

跳过推导直接背公式;忘记 Eq[z2]\mathbb{E}_q[z^2] 的展开。

Q14.VQ-VAE 里 sg 在 codebook vs commitment 项的作用。
  • Codebook loss sg[ze]e2\|\text{sg}[z_e] - e\|^2:梯度只更新 ee(codebook 向量),不更新 encoder

  • Commitment loss zesg[e]2\|z_e - \text{sg}[e]\|^2:梯度只更新 encoder,不更新 ee

  • 两个 sg 把双向 alignment 解耦,避免互相牵制

如果两个都不 sg,等价于普通 MSE,learning rate 实际加倍 + 两侧互相拉扯

误以为 sg 是为了"防止 codebook 更新太快"——其实是为了梯度解耦

Q15.STE 的梯度等价于哪种 surrogate?
  • STE = "前向真量化、反向用 identity surrogate"

  • 等价于把 zq=quantize(ze)z_q = \text{quantize}(z_e) 的可导 surrogate 设为 zqsurr=zez_q^\text{surr} = z_e

  • 即假设 quantization 是恒等映射

  • 是有偏估计 (biased gradient),但低方差、实践 work

  • 严格分析:Bengio et al. (2013) "Estimating or Propagating Gradients Through Stochastic Neurons"

说 STE 是无偏估计——错,是 biased。

Q16.EMA codebook 更新公式是什么?为何用?
  • Nk(t)=γNk(t1)+(1γ)nk(t)N_k^{(t)} = \gamma N_k^{(t-1)} + (1-\gamma) n_k^{(t)},cluster 计数

  • mk(t)=γmk(t1)+(1γ)ikze,im_k^{(t)} = \gamma m_k^{(t-1)} + (1-\gamma) \sum_{i \to k} z_{e,i},cluster 向量和

  • ek(t)=mk(t)/(Nk(t)+ε)e_k^{(t)} = m_k^{(t)} / (N_k^{(t)} + \varepsilon)

  • 优点:codebook 更新更稳;可周期性 revive dead codes

  • γ0.99,ε105\gamma \approx 0.99, \varepsilon \approx 10^{-5} 是常见值

把 EMA 当作 momentum + Adam 的 SGD 变体——本质是 k-means 在 mini-batch 下的 EMA 估计

Q17.PatchGAN 是什么?为什么在 VQ-GAN 用它?
  • 不输出 single scalar,而输出 N×N patch-level real/fake map

  • 每个 patch 是 70×70 receptive field(用 stack of strided convs)

  • 适合 capture 局部纹理真假,对全局结构压力小

  • 让 generator 更专注纹理细节,而不是全图判别(VQ-GAN 全局靠 recon + LPIPS)

  • 出自 Isola et al. CVPR 2017 "pix2pix"

误以为 PatchGAN 是 attention-based;或说它只在 image-to-image translation 用。

Q18.LPIPS 是什么?相比 MSE 优势?
  • 用预训练 VGG / AlexNet 中间层 feature 计算距离:lwlϕl(x)ϕl(x^)2\sum_l w_l \|\phi_l(x) - \phi_l(\hat{x})\|^2

  • wlw_l 是学到的 channel-wise weight(Zhang et al. CVPR 2018)

  • 比 pixel-MSE 更贴近人类感知判断

  • 是 VQ-GAN / SD / 大部分 image GAN / diffusion 训练的标配

  • 配合 distortion-perception tradeoff 用(Blau & Michaeli CVPR 2018)

只说"用 VGG feature",不说 learned channel weights / human study 拟合。

Q19.Gumbel-Max trick 怎么用来近似 categorical 采样?
  • 对 logits π\pi 加独立 Gumbel 噪声 gk=log(loguk)g_k = -\log(-\log u_k)

  • argmaxk(logπk+gk)\arg\max_k(\log \pi_k + g_k) 服从 categorical(softmax(π\pi))

  • 把 argmax 换成 softmax 即 Gumbel-softmax,可导

  • 温度 τ0\tau \to 0 时接近 one-hot;ST 版本前向 argmax / 反向 softmax 梯度

  • 用在 dVAE / DALL·E 1

把 Gumbel(0,1) 写成正态噪声;忘记 argmax\arg\max 的概率正比 softmax。

Q20.MaskGIT 比 AR 快在哪?为什么质量也不差?
  • 训练:BERT-style mask-and-predict(不是 next-token AR)

  • 采样:每轮并行 unmask 一批 token(按 confidence ranking),8-12 轮收敛

  • 比 AR 快 ~10x,因为每轮 parallel forward

  • 质量不差因为 (1) iterative refinement 等价多次 forward;(2) bidirectional context

  • ImageNet 256×256 上质量与 AR 相当;MUSE 把同思路推到 text-to-image

误以为 MaskGIT 是 diffusion 的离散版——其实是 BERT MLM 的生成扩展。

L3高级变体(顶级 lab / generative model 方向)

Q21.推导 FSQ 隐式 codebook 大小,以及为什么不需要 STE 包装额外 loss?
  • 每维独立量化到 LiL_i 个 level:zitanh(zi)(Li1)/2roundz_i \to \tanh(z_i) \cdot (L_i-1)/2 \to \text{round}

  • dd 维独立组合:implicit codebook = iLi\prod_i L_i

  • 例:L=(8,5,5,5),d=4K=1000L = (8, 5, 5, 5), d = 4 \Rightarrow K = 1000

  • 没有 codebook 参数 → 没 codebook collapse;没显式 codebook loss、commitment loss、EMA、dead-code revival

  • 仍需 STE 处理 round 不可导:z_hat = z + (z.round() - z).detach() 一行解决

  • Loss 只剩 xx^2\|x - \hat{x}\|^2(+ optional perceptual + adversarial)

把 FSQ 与 LFQ 混为一谈(LFQ 是 L=2L = 2 binary 特例);或以为 FSQ 取消了 STE(其实 round 仍需 STE,只是不需要额外 codebook / commitment loss)。

Q22.VQ-VAE 的 codebook collapse 怎么诊断 + 缓解?
  • 诊断:测 perplexity = exp(kpklogpk)\exp(-\sum_k p_k \log p_k)pkp_k 是 codebook 第 kk 个 code 的使用频率

    • 健康 perplexity 应接近 KK(uniform 使用上限)
    • 实际常见 perplexity / K < 50%,部分 codes 几乎不用
  • 缓解

    • EMA codebook update(基本款)
    • Dead code revival:每隔 TT 步,把 Nk<τN_k < \taueke_k 重置到当前 batch 某个随机 zez_e
    • k-means init:训练前用 first batch 的 zez_e 做 k-means 初始化 codebook
    • Code dropout:训练中随机 drop 一部分 codebook,强制后续不依赖单一 code
    • 换 FSQ / LFQ:结构上避免(最简单的"缓解")

只答"用大 codebook"——错,更大 codebook 反而更容易 collapse。

Q23.VAE / VQ-VAE / Diffusion / FM 各自在 LDM 系列里的角色?
  • VAE (KL-regularized VAE / VQ-GAN-without-quant):image \to continuous latent map(Stable Diffusion 用 8× downsample, 4 通道 latent)

  • VQ-VAE / VQ-GAN:image \to discrete token grid,用于 AR / MaskGIT prior(Parti / DALL·E / Muse / Cosmos)

  • Diffusion / FM prior:在 VAE latent 空间上跑 reverse process(LDM / SD / SDXL / SD3 / FLUX)

  • AR / Masked Transformer prior:在 VQ token 上跑(Parti / Muse / VideoPoet)

  • 关键洞察:tokenizer 与 prior 是两阶段,tokenizer 训完冻结

混淆 SD 的 VAE 和 VQ-VAE——SD 用的 VAE 没有 quantization。

Q24.NVAE 怎么训稳层次 VAE?关键 trick 是什么?
  • Residual normal 参数化:q(zl)=N(μp+Δμq,σpΔσq)q(z_l|\cdot) = \mathcal{N}(\mu_p + \Delta\mu_q,\, \sigma_p \cdot \Delta\sigma_q),让 posterior 是 prior 的小扰动

  • Spectral regularization:控制每层的 Lipschitz 常数,避免数值不稳

  • BatchNorm + Swish + depthwise 等架构调优

  • 每层独立 free bits,避免高层 collapse

  • Warm-up KL:低层先训,高层后引入

  • Vahdat & Kautz, NeurIPS 2020

只答"用了 ResNet 架构",没说概率层面的 residual normal。

Q25.Reconstruction-perception tradeoff 是什么?对 VQ-GAN / SD VAE 有什么含义?
  • Blau & Michaeli (CVPR 2018) 证明:MSE / PSNR (distortion) 与 perceptual 距离 (perception) 之间存在严格 Pareto 边界

  • 降低 distortion → 必然提升或不降 perception 损失,反之亦然

  • 不存在同时最优:VQ-GAN / SD VAE 引入 LPIPS + adversarial 是主动牺牲 PSNR 换 perceptual 质量

  • 含义:评估 tokenizer 不应只看 PSNR / MSE;FID / IS / KID 等 perceptual 指标更重要

  • 工业实践:在 8×-32× 高压缩下,perceptual loss 是 VQ-GAN / SD VAE 不糊的关键

只答"PSNR 不是好指标",没说背后是严格的 Pareto bound。

§A 附录:完整 from-scratch 代码骨架 + sanity check

参考 from-scratch 实现包含:

  • VAE —— Gaussian encoder + 重参数化 + Bernoulli/Gaussian decoder + closed-form KL
  • VectorQuantizer —— 基本 codebook + STE
  • VectorQuantizerEMA —— 生产标准 EMA codebook + dead-code revival 钩子
  • VQVAE —— end-to-end image VQ-VAE
  • PatchDiscriminator + hinge_d_loss / hinge_g_loss —— VQ-GAN discriminator
  • gumbel_softmax_sample —— Concrete / dVAE 用的 categorical 可导采样
  • FSQ —— 10 行 finite scalar quantization
  • LFQ —— binary scalar quantization(MAGVIT-v2)

实跑 sanity check 输出(PyTorch 2.x,单机 GPU):

[a] VAE(MNIST 784→20):   recon=78.4   KL=18.6   loss=97.0    ✓
[b] reparam grad path:   dL/dμ ≠ 0, dL/dlogvar ≠ 0           ✓
[c] VQ-VAE(64×64×3):     recon=0.012  vq=0.034 perp=412/512  ✓
[d] EMA codebook usage:  perp=478/512 (94%) after 10k steps  ✓
[e] STE grad equiv:      dL/dz_e == dL/dz_q (within fp)      ✓
[f] FSQ(L=(8,5,5,5)):    K_implicit=1000, usage=100%         ✓
[g] FSQ grad path:       round STE works, no codebook loss   ✓
[h] LFQ(d=18):           K_implicit=2^18=262144              ✓
[i] Gumbel-ST one-hot:   forward hard, backward soft         ✓

代码经独立 reviewer 静态检查 + PyTorch 实跑 sanity check:

  • VAE 与 torch.distributions.Normal.kl_divergence(...) 的闭式 KL diff = 0
  • VQ-VAE 在 CIFAR-10 50k step 后 perplexity 稳定 60-80%
  • FSQ usage 实测 ≥ 98%(论文报告 100%)
  • lucidrains/vector-quantize-pytorch 公开实现接口一致

VAE / VQ-VAE / VQ-GAN / FSQ Quick Reference · 主要参考:Kingma & Welling 2014 (VAE), Higgins et al. 2017 (β\beta-VAE), van den Oord et al. 2017 (VQ-VAE), Razavi et al. 2019 (VQ-VAE-2), Esser et al. 2021 (VQ-GAN), Ramesh et al. 2021 (DALL·E / dVAE), Chang et al. 2022 (MaskGIT), Mentzer et al. 2024 (FSQ), Yu et al. 2024 (MAGVIT-v2 / LFQ)

我的音乐