VQ-VAE：学习离散字典

交互教程

如果编码器必须从有限字典中选择呢？VQ-VAE 将连续特征变为离散索引，再学习这些索引怎样共同出现。

之前已经能做什么

高斯 VAE 使用连续采样和可微重参数化。离散最近邻选择会切断该梯度路径。向量量化早已用于压缩，新的问题是如何共同训练编码器、字典和解码器。

编码器在每位置输出 z_e，选择最近的字典向量 e_k 交给解码器，并把 k 存为离散码。编码器用直通梯度估计，e_k 用字典损失更新，承诺损失则使 z_e 靠近选中的码。

阅读顺序从左到右。点击模块，将它与计算和代码对应。代码为教学伪代码；来源链接提供论文方法或已核查的实现。

每个空间向量使用共享字典独立量化。

z_e = encoder(x)

前向使用字典向量，反向使用直通估计。

ids = squared_distances(z_e, e).argmin(-1)
z_q = e[ids]
z_st = z_e + stopgrad(z_q - z_e)

重建训练 tokenizer；另一个先验模型预测索引。

x_hat = decoder(z_st)
L = recon(x_hat,x) + norm(stopgrad(z_e)-z_q)**2
L += beta * norm(z_e-stopgrad(z_q))**2

k=\arg\min_j\|z_e-e_j\|_2^2,\qquad z_q=e_k

演示中 z_e=0.8，四个码为 −3、−1、1、3，选中向量 1，平方误差 0.04。z_e 越过边界时索引突然变化，因此不能通过 argmin 的普通求导训练编码器。

改变参数，查看计算值，再把它与上方的真实模型说明对照。

解析教学模型 · 控件改变的是例子，不是训练后的研究模型。

编码值 z 0.8码本项数 K 4

橙色为编码结果，绿色为选中的码。

前向表达式 z_e + stopgrad(z_q−z_e) 等于 z_q，而对 z_e 的导数为 1。字典项阻止梯度进入编码器，承诺项阻止梯度进入码本。这是有偏梯度估计，不是最近邻选择的精确导数。

tokenizer 训练时，确定性后验和均匀分类先验产生每位置 log K 的常数 KL。之后用 PixelCNN 拟合编码。独立均匀采样 ID 不等价于从学习到的先验采样。

字典可能含有未使用项。增加项数提高潜在容量，也可能增加先验建模难度。原论文梯度更新和 EMA 替代方式不可无说明混用。VQ-VAE-2 接下来改变层次结构，而非最近码的基本思想。

展开公式推导

硬后验给单个 code 概率 1。在均匀 K 项先验下，KL 为 log K；除 code 切换外不依赖编码器参数。这并没有学习整图 codes 的联合分布。

D_{KL}(q(k|x)\Vert 1/K)=\log K

sg 表示停止梯度。中间项让码本靠近编码器，最后一项让编码器靠近码本；重建梯度另外通过直通路径传递。

L=L_{\mathrm{rec}}+\|\operatorname{sg}[z_e]-e_k\|^2+\beta\|z_e-\operatorname{sg}[e_k]\|^2

哪项 loss 直接移动选中的字典向量？

字典损失。重建直通路径更新编码器和解码器；承诺项约束编码器。

本页图形与伪代码为教学重绘。论文报告、固定实现观察和解析示例分别标记；未运行研究模型。