VQ-VAE

VQ-VAE（van den Oord et al., 2017）把矢量量化搬进自编码器：潜在空间不再是连续高斯，而是一本离散码本。连续生成有模糊病，离散 token 才能接自回归——这一步让「图像当语言建模」成为可能，是 矢量量化（VQ）：从码本设计到率失真理论 在生成模型里最成功的化身。

VAE 的连续潜空间有两个先天缺陷：高斯先验下采样出的潜变量落在 decoder 没见过的空隙，生成图像发糊；潜变量是连续的，没法直接喂给自回归模型——而 GPT 范式只吃离散 token。VQ-VAE 的答案是把瓶颈层换成一次最近邻查表：encoder 输出 $z_e(x)$ ，在码本 $E=\{e_1,\dots ,e_K\}$ 里找最近码字替换，注意这里的码本 $E\in {\mathbb{R}}^{K\times d}$ 是一张 $K$ 行的可学习查找表。写成数学表达式：

$z_q(x)=e_k,k=\arg {\min }_i\Vert z_e(x)-e_i{\Vert }_2$

对图像 VQ-VAE，encoder 输出 $H^{\prime }\times W^{\prime }\times d$ 的特征图，每个空间位置独立查一次表——一张图变成 $H^{\prime }\times W^{\prime }$ 个码字索引。这个索引网格正是后续自回归模型要预测的 token 图（见下文 VQGAN）。查表这一步不可微，梯度怎么回传是整个架构的核心工程问题。

设计损失函数

先把需要更新参数的模块清点一遍：decoder 消费码字、输出重建，梯度正常反传，不添麻烦；encoder 的输出被 $\arg \min$ 换成常数码字，梯度在查表处断流；码本则天然够不到重建梯度——$\arg \min$ 挡路之外，一个 mini-batch 还只激活少数几个码字，信号稀疏。所以设计三项损失对着三个位置分而治之，一项管一方（ $\mathrm{sg}$ 为 stop-gradient，「这一项里把我当常数、别给我传梯度」）：

$\mathcal{L}=\underset{\text{重建}}{\underbrace{\log p(x\mid z_q)}}+\underset{\text{码本}}{\underbrace{\Vert \mathrm{sg}[z_e]-e{\Vert }_2^2}}+\underset{\text{commitment}}{\underbrace{\beta \Vert z_e-\mathrm{sg}[e]{\Vert }_2^2}}$

重建项 $\log p(x\mid z_q)$ 是 decoder 的负对数似然：给码字，还原输入——实践里就是 MSE 或感知损失（原文对图像用 discretized logistic 似然，不必纠结形式）。decoder 的全部参数靠它训练；encoder 也指望它——「往哪边改、重建会更好」正是 encoder 最想要的信号，只是这份梯度得先解决怎么过查表，留给下一节的 STE。唯独码本指望不上它：$\arg \min$ 切断了路径，退一步讲，即便梯度能流进去，一个 mini-batch 也只够到被选中的那几个码字，信号还要穿过整个 decoder——稀疏且微弱，靠它训码本，没被选中的码字永远原地踏步，利用率必然崩坏。码本需要一条不走查表的专属梯度通道，这就是第二项。

码本损失 $\Vert \mathrm{sg}[z_e]-e{\Vert }_2^2$ 只训码本。sg 拴住了 $z_e$ ，梯度只落在 $e$ 一侧：方向是把码字 $e_k$ 拉向当前选它的那些 $z_e$ 。这正是 LBG 质心条件的在线随机梯度版——胞腔内样本的平均位置就是新码字，理论篇里 LBG 是批量算质心，这里是每个 mini-batch 小步走。EMA 更新（下文）则干脆回到滑动平均版质心。

commitment 损失 $\beta \Vert z_e-\mathrm{sg}[e]{\Vert }_2^2$ 只训 encoder——sg 换到了 $e$ 头上，把 $z_e$ 拉向它选中的码字。没有它，encoder 输出可以在胞腔间随意游走，码本更新永远在追一个移动目标，训练剧烈抖动；有了它，encoder 被「承诺」在码字附近。 $\beta =0.25$ 是原文经验值，作者报告对具体取值相当不敏感，量级在 0.1–2.0 都能训稳。

两个 sg 的位置是这套设计的精髓：若不加 sg，第二、三项是同一根弹簧两头互拉—— $z_e$ 与 $e$ 在中途会合，谁也没被真正聚到位；加上 sg 后，划分（哪个 $z_e$ 归哪个码字）由 argmin 硬执行，质心（码本学习）与承诺（编码器稳定）被拆成两个各训一边的独立目标——K-means 的 E 步/M 步在这里被改写成了两个损失项加一次硬分配。

梯度怎么过查表：直通估计器

三项损失清点下来还剩最后一个缺口：encoder 的重建梯度。decoder 收到的 $z_q$ 是查表后的常数码字，重建误差想对 $z_e$ 求导，先撞上 $\arg \min$ ——离散决策没有导数；即便选出了 $k$ ，查表替换也已经把 $z_e$ 抹掉，decoder 前向根本不消费它。STE（Straight-Through Estimator）的对策一句话说尽：前向照常量化，反向假装没量化——把 $\frac{\partial \mathcal{L}}{\partial z_q}$ 原封不动当作 $\frac{\partial \mathcal{L}}{\partial z_e}$ 传回 encoder，相当于只在反向传播时令 $z_q\approx z_e$ 。

代码里就一行：z_q = z + (z_q - z).detach()。前向求值时 detach 不改数值， $z+(z_q-z)=z_q$ ，输出仍是码字；反向时 detach 项导数为零、 $z$ 的导数为一，于是 $\frac{\partial z_q}{\partial z_e}=1$ ，梯度「直通」。论文的等价记法是 $z_q=z_e+\mathrm{sg}[e_k-z_e]$ ，与这行代码一一对应。

直觉上，STE 把「吸附到最近码字」当成对 encoder 输出的一次局部扰动：decoder 的梯度本想说「 $z_q$ 往这边挪一点，重建会更好」，STE 把这句话原样转交给 $z_e$ ——因为 $z_e$ 就在被选码字附近，邻域性质相近，方向大体正确、幅度略有偏差。它是有偏估计，但方差小、实现一行，VQ 系模型基本都靠它；Gumbel-Softmax 走的是另一条路——把离散选择松弛成带温度的连续分类采样、逐步退火逼近硬选择，理论更优雅，VQ-VAE 场景下收益有限。

import torch
import torch.nn.functional as F
 
class VectorQuantizer(torch.nn.Module):
    def __init__(self, k=1024, d=256, beta=0.25):
        super().__init__()
        self.codebook = torch.nn.Embedding(k, d)
        self.beta = beta
 
    def forward(self, z):            # z: (B, d)
        d = (z ** 2).sum(1, keepdim=True) \
            - 2 * z @ self.codebook.weight.T \
            + (self.codebook.weight ** 2).sum(1)
        idx = d.argmin(1)            # 最近邻查表（不可微）
        z_q = self.codebook(idx)
        # 码本损失 + commitment 损失
        loss = F.mse_loss(z_q, z.detach()) \
             + self.beta * F.mse_loss(z, z_q.detach())
        # STE：前向用 z_q，梯度直通回 encoder
        z_q = z + (z_q - z).detach()
        return z_q, idx, loss

梯度通道虽已接通，码本靠梯度下降更新在实践中仍不够稳——码本损失走的是 mini-batch 小步随机质心，步子碎、容易抖。实用变体是 EMA：不用梯度，按胞腔内样本的滑动平均更新码字（等效于 LBG 的质心步骤），训练稳得多，是大码本（几万级）的默认选择。

码本坍缩（codebook collapse）

大量码字长期无人问津变成死码，有效码本缩水、重建质量骤降。缓解手段：EMA 更新、死码重置（周期性把闲置码字重新播种到高频编码输出）、降低码本利用率敏感的初始化。根治方案是 2023 年的 FSQ（Finite Scalar Quantization）——干脆扔掉码本，每维做有限级标量量化，无坍缩、无查表，效果与 VQ 打平，新项目值得先试。

从 VQ-VAE 到自回归生成

训练完的这一套——encoder、码本、decoder——本质是一台有损压缩器：进去一张图，出来一张索引网格。它自己不会生成：训练目标里只有重建，从头到尾没有任何组件学过「索引网格本身服从什么分布」。要生成，缺的是一个先验——在 token 图上训练一个自回归模型，从它采样出索引网格，再查表、decode 还原成图像。VQ-VAE 原文拿 PixelCNN 当这个先验；生成能力的完整形态是 VQGAN（Esser et al., 2021）：tokenizer 一侧，重建损失换成感知损失 + patch 判别器，码本学到感知等价的 token；先验一侧，Transformer 配合滑动窗口注意力，256×256 图像压成 16×16 token 序列。此后「tokenizer + 自回归 Transformer」成为图像生成的一大流派——DALL·E 的 dVAE、Parti 的 ViT-VQGAN 都在这条线上。

一个流传很广的错误需要纠正：Stable Diffusion 并不用 VQ。它的潜空间是 KL 正则化的连续 VAE——训练稳定、推理一步出潜变量；离散 token 路线哺育的是自回归系（DALL·E、Parti、LlamaGen），两条路线在「潜空间离散还是连续」上分道扬镳。

离散化的收益不止生成：BEiT 拿 dVAE token 当掩码预测目标，wav2vec 2.0 用量化目标学语音表征——自监督 + 量化的交叉地带见 深入解析SSL VQ训练：矢量量化赋能的自监督学习范式全解。音频侧的 SoundStream、EnCodec 则用级联残差量化（乘积量化（PQ）：高维向量检索的工程基石 里的 RQ）把波形压成 token，撑起了音乐生成的前端。

回到选型直觉：要接自回归、要跨模态共享词表、要显式离散语义单元，选 VQ 潜空间；要训练稳、采样快、走扩散范式，连续潜空间更省心。离散与连续之争没有终局，只有范式匹配。