跳到主要内容

深度解析:从变分自编码器(VAE)到离散表示学习(VQ-VAE)

在这篇详细深入的文档中,我们将从图像生成的早期形态——自编码器(AutoEncoder)的探讨出发,一步步揭示为何需要引入带有概率视角的变分自编码器(VAE)。随后,我们将拆解 VAE 在生成高质量图像和控制力方面遇到的理论瓶颈。 最终,我们将引出生成模型历史上最具前瞻性的架构之一:向量量化变分自编码器(VQ-VAE),并拆解这一架构如何通过**直通估计器(STE)**解决数学计算壁垒,同时又是如何通过“离散密码本”(Codebook)机制将计算机视觉(CV)引入了像自然语言处理(NLP)一样的 Token 大时代。


第一部分:从表征学习谈起 —— 什么是自编码器 (AE) ?

1. 自编码器的初衷与哲学

神经网络除了能够进行“识别”之外,更为重要的一种能力叫作表示学习(Representation Learning)特征降维。 自编码器(AutoEncoder, AE)是最原始、最纯粹的表示学习架构。它的核心哲学非常简单:“我能否构造一个神经网络,它的输出目标就是尽可能完美地还原它自己的输入?”

它的架构可以粗略地分为两个对称的两侧:

  • 编码器(Encoder): 负责将高维度、冗余的输入数据 $x$ (比如一张 $1024 \times 1024$ 的高清照片)通过多层神经网络挤压、“降维”到一个极小的维度空间 $z$ 中。由于维度骤降,这部分也被称为瓶颈层。
  • 隐空间向量(Latent Vector): 这个 $z$ 是网络对原始数据最为精炼的抽象总结,也就是特征隐向量。
  • 解码器(Decoder): 负责接收这个狭小的隐向量 $z$,并且努力从中还原出完整的原始图像 $\hat x$。

2. 为什么普通的 AE 无法用于生成新的图像?

此时,我们面临一个思考:既然解码器可以把 $z$ 变成图像,那我们是不是可以在隐空间里面随机捏造一个隐向量 $z_{random}$,然后把它丢给解码器,直接生成一张现实世界上不存在的全新图像呢?

答案是:几乎不可能。普通的 AE 是无法用作生成模型 (Generative Model) 的。

原因在于 AE 缺乏对隐空间分布的约束。 假设我们用大量的美女图片训练了一个 AE。AE 会为了降低重构损失 (Reconstruction Loss),把每一张图片的特征极其勉强、稀疏地映射到一片没有边界的、参差不齐的隐空间里的某一个固定点。 这些特征点之间是散落的“孤岛”。如果我们在两个点之间取一个中间值作为 $z$,解码器由于在训练时从来没有见过这种分布的数据点,它解码出来的只会是一团模糊、扭曲的噪声。简而言之,AE 只是个压缩机,它学到的特征是过度拟合的,不是连续连贯的分布规律。

为了让这个隐空间变得平滑、连续,并具备真正的“生成能力”,变分自编码器 (VAE) 应运而生。


第二部分:生成模型的破冰 —— 变分自编码器 (VAE) 深度剖析

1. 为什么要加上“变分”的魔法?

变分自编码器 (Variational AutoEncoder, VAE, Kingma 等人, 2013) 彻底改写了编码器的任务。它规定: 编码器输出的不再是一个干瘪固定、冷冰冰的点 (Point Estimate) ,而应该输出一个概率分布 (Probabilistic Distribution)!

具体表现为: 对于每一张输入的图片 $x$ ,VAE 的编码器同时计算输出两个参数:

  • 均值 $\mu$
  • 方差的对数 $\log \sigma^2$

这两个参数规定了一个高斯分布:$N(\mu, \sigma^2)$。此时的隐特征从一个确定的位置变成了一个“带有一点模糊光晕、有半径范围的区域”。 接着,模型从这个分布中随机采样 (Sampling) 出一个具体的点 $z$,再把 $z$ 送给解码器还原。

这就是 VAE 最精彩的一笔:它强制要求编码器学到一种能够容忍一定“扰动和噪声”的映射。为了应对这种扰动,编码器必须把相似图像的特征分布范围彼此拉近、甚至稍微交叠,从而构建出一个平滑且拥有连贯过渡区的隐空间。

2. 数学基石:KL散度与重构平衡

如果只有均值和方差预测,网络很可能会投机取巧:它会把每一张图片的方差 $\sigma^2$ 压缩得非常非常小(接近于0),使得分布退化成原来的点,从而变成普通的 AE。

为了打破这种“作弊行为”,VAE 设计了一套复合强力的损失函数体系(这个底层是源于证据下界 ELBO 的最优化推导):

  • 重建损失 (Reconstruction Loss):确保解码出的图片 $\hat x$ 与原图相似(多使用 MSE 或 BCE 损失)。
  • KL 散度正则化项 (Kullback-Leibler Divergence):这部分强行要求编码器输出的高斯分布,要尽可能地接近标准正态分布 $N(0, I)$。

此时,损失函数就像一个天平在左右互搏: 如果 KL 散度约束太强,所有图片的特征都被强行拉向原点,大家混在一团,解码器啥也还原不出来; 如果重建损失要求过高,方差又会收缩并造成孤岛。VAE 就是在这两者的博弈收敛中,得到了一个高度连续的高级特征空间。

3. 工程落地的命门:重参数化技巧 (Reparameterization Trick)

但在使用反向传播 (Backpropagation) 更新网络权重时,VAE 遇到了一个地狱级的障碍。 正如我们在上文描述的那样:从分布 $N(\mu, \sigma^2)$ 中进行随机采样得出 $z$ ,这是一个抛骰子式的随机操作。这种带有随机性的操作是**不可微分(Non-differentiable)**的,梯度走到这里就会断掉,无法往前回传给编码器。

这时候重参数化技巧 (Reparameterization Trick) 挺身而出。 数学家把随机性和包含模型参数的部分进行了巧妙的分离,使用如下公式: $$ z = \mu + \epsilon \times \sigma $$ 其中:

  • $\mu$ 和 $\sigma$ 均是来自于神经网络可导的正常神经元输出。
  • $\epsilon$ 是从独立的外部标准正态分布 $N(0, 1)$ 中凭空抽取的随机噪声样本。

通过这个分离,随机性(不可导的根源)仅仅附着在 $\epsilon$ 这一个常数输入节点上,这样就完美地搭建了一条让梯度能够通过 $\mu$ 和 $\sigma$ 原路返回前场的桥梁!网络顺利闭环,得以全自动优化。

4. VAE 体系的致命局限性分析

VAE 虽然在数学理论和架构上是极其优美且具有划时代开创性的生成模型,但它也有被后世模型诟病的致命先天缺陷:

第一:生成的高清图片总是偏于模糊。 MSE 类的重建损失,实际上是在促使预测结果倾向于所有可能正确的结果之平均数。例如如果一只猫的胡须可以长在向左一点,也可以向右一点的位置,使用 MSE 会使得网络干脆直接生成一团均值的灰色虚影敷衍了事。它无法如同 GAN(判别网络)一样去捕捉高频细节。

第二:后验坍塌现象 (Posterior Collapse)。 如果是极为深度的网络,有时候由于 KL 散度的过度惩罚,解码器会直接“无视”隐向量传入的信息(因为强迫它们都是正态分布噪声),而仅凭自己的网络深度“盲猜”还原一些基础大面貌,从而导致输入特征彻底无效坍塌。

这就呼唤着下一次革命的到来——那就是**向量量化(Vector Quantization)**的离散突破。


第三部分:VQ-VAE(向量量化变分自编码器)横空出世,离散的狂欢

随着时间进入 2017 年,DeepMind 团队发表了极其重磅的论文—— Neural Discrete Representation Learning,正式将 VQ-VAE 推上历史的宏伟舞台。

1. 从连续转向离散的思想变迁 (Continuous to Discrete)

我们一直执念于在 $N(0, I)$ 等连续空间中描绘世界的概率分布。然而自然界和现实数据往往是基于某种离散形式架构而构成的。

  • 语音信号的底层,其实是有限的音素(Phonemes)。
  • 语言文本,其实是庞大但有限数量的词汇(Words/Sub-words)。
  • 从图像的解构规律来看,其局部视觉形态也许就是特定有限组合的边缘或色域集合。

如果能够舍弃传统 VAE 的那个很难约束、容易崩坏的连续高斯隐空间分布,并且构建一个离散的隐变量集合去对输入的信息发起“分类式总结映射”,那我们该如何做呢?这就是 矢量量化 (Vector Quantization) 发挥作用的地方。

2. VQ-VAE 的三大核心模块与运算流程

VQ-VAE 的本质是抛弃对高斯分布方差的学习,维护一张能够离散表义的二维查表字典。 它的具体流向涉及三大结构单元:

(A) 编码器 (Encoder)

与传统 VAE 相比,VQ-VAE 编码器不再负责输出 $\mu$ 和 $\sigma$,而是正常经过卷积层的挤压,去输出特征图(比如一幅 $256 \times 256$ 图像被卷缩到分辨率只有 $32 \times 32$ 但拥有长长通道向量空间的张量特征图)。 在这个降维小地图上的每一个通道上的连续向量,我们记作 $z_e(x)$。

(B) 离散密码本 (Codebook / Vector Quantisation 层)

模型专门划拨了一个可学习的参数矩阵,也就是密码本(Codebook 字典)。 设定我们的字典存在 $K$ 个单词(通常几千到几万),每个单词是一个长长的嵌入向量 (Embedding Vector)。我们可以用 $e_1, e_2, e_3 ... e_K$ 来表示这本字典。

接下来进行特征对齐转化

  • 对于编码器发出的 $z_e(x)$(由于它是个游移小数组成,我们可以将其视作一个自由坐标系的点)。
  • 网络强制执行:拿 $z_e(x)$ 与字典里包含的所有的 $K$ 个向量逐一去比赛,找出几何特征欧几里得距离最近的那一个(我们称为 $e_k$)。
  • 然后直接粗暴地:抹杀并抛弃 $z_e(x)$ ,用该字典向量 $e_k$ 直接取而代之! 这一替换的结果被命名为具有离散性质的 $z_q(x)$。

经过这一道工序后,不管你原图有多少稀奇古怪的光影细节,只要在这个格子位置,全都会被强行对应并限制在某个固定词上。整个图像从此变成了一连串包含几十乘几十数目的离散序号密码表

(C) 解码器 (Decoder)

当拿到了离散化的隐变量集特征 $z_q(x)$ 后,将其送入上采样解码网络进行卷积映射,最终尝试原原本本地绘制恢复出输入的清晰全彩原图 $\hat x$。

3. VQ-VAE 克服悬崖机制:直通估计器 (STE)

如果大家理解了前面的密码本运行机制,脑子敏锐的你一定又能闻到灾难的气味。 为了寻找最近邻词汇而执行的核心操作——求最小值索引(Argmin),这不仅是一个粗暴斩断函数的阶跃行为,更是彻底没有任何阶梯导数、绝对无法微分的不可导过程!由于查表操作打断了图的连接,解码网络往回倒推求计算反向传播 (Backpropagation) 梯度的道路到此灰飞烟灭中断,上游用于降维的编码器似乎永远也无法得到关于它的微小权重调整信号了。

然而,DeepMind 设计了名为 Straight-Through Estimator(STE,直通估计器) 的大胆绝妙设计。 STE 认为既然这一步斩断了,那咱们就蒙住“反向传播检测仪”的眼睛,强制让数据通行:

  • 既然前向计算时(Forward) 用的是量化好的新代表 $z_q(x)$ (即查到的字典向量)。
  • 那好,在发起反向计算时(Backward) 梯度会从解码网络抵达 $z_q(x)$。到了这位置,这串关于被重置代表信息的负梯度数值,直接一毫不差地通过“超时空平移复制手段”,跨过中间那个巨大的求最小值的深坑断层区,贴脸塞到了远端的真实连续特征源头——最初的 $z_e(x)$ 的怀里

正是靠这种假装量化没有存在、生搬硬套梯度数据的操作手段,成功激活了悬崖死锁网络前后段的所有神经链条。

4. VQ-VAE 的全新独占损失函数体系大解构

在具备直通手段后,该怎么能够保证:一是我们的 Encoder 输出靠谱特征,二是我们的密码本里的“词”也会因为长期学识积累不断进化和对齐?

VQ-VAE 配备了一个由三个大分块组成的组合损失公式设计: $$ L = ||x - \hat x||_2^2 + ||sg[z_e(x)] - e||_2^2 + \beta ||z_e(x) - sg[e]||_2^2 $$

:公式中的 $sg$ 代表停梯算子 (Stop Gradient),在诸如 PyTorch 等现代深度学习框架中即相当于使用方法: .detach() 截住求梯度链并将其仅视为一个数值常数不参与模型参数运算。

下面我们将此一分为三种损失来逐一详解:

(一)重构回归损失 (Reconstruction Loss) : 这第一项 $||x - \hat x||_2^2$ 非常直观。如同我们对旧时代普通 AE 的描述相仿。它的唯一目标是努力教导且让 Decoder 把量化特征映射出来的这幅赝品图像尽力复原出原来的 $x$。在这过程当中,Encoder 也因为直接承载了越境丢过来的微调梯度而被同步拉拽训练优化。

(二)词典集更新损失 (Codebook Loss) : 第二项 $||sg[z_e(x)] - e||_2^2$ 最为巧妙玄机。因为我们要确保能够挑选出来的那个词本里面的特征词 $e$ 不能因为长期不更新从而导致被输入淘汰,所以增加这部分强制词语 $e$ 会被推送到距离其附近选出的那些连续编码目标输出 $z_e(x)$ 更为逼近相连的位置。加入带有括号拦截的 $sg$ 在左侧项目是绝对关键点,这强制保障这部分数学优化仅修改和打磨能够决定密码集本身特征位置的权重!

(三)承诺紧靠损失 (Commitment Loss) : 如果上述手段会导致词库向中心发散的话,那么我们的最终三条公式体系中还需要防止原本前段产出的特征飘忽不定,所以需要一条公式: $\beta ||z_e(x) - sg[e]||_2^2$ ;其通过停梯拦截锁定住了字典数据不动,但是它逼着前段 Encoder 所放出的所有的原始特征图必须要去对其被赋予对应位置的字典项去做出“承诺(Commit)”——别太狂奔偏离那些离散空间太远,牢牢守住在附近!超参数 $\beta$ 通常保持在较小的数值(例如 0.25)用来避免强压效果引起坍塌扭曲。

5. 字典平滑变体:指数移动平均 (EMA) 技巧方案

值得顺带引申介绍的是在高规格的大规模架构工业训练下(比如后来用来做百亿参数多模态的时候),即便是前文的词库损失项,有时候如果面对极端特异性的批次数据传入(Batch Input),仍然会引发少许不稳定的动荡导致模型词频崩溃丧失(Index Collapse)。

所以有许多后续复现变体并不会依靠直接优化器去猛烈冲击字典层梯度变量以实现学习更新,而是广泛采纳使用了一种更加润物细无声、稳定的做法架构 —— 基于频率与指数移动平滑结合技术体系 (EMA Updates) 来不断缓和微调其字典层。EMA 字典更新法更像是一个非常巨大数据流级别的流动态 K-Means 中心平滑聚合调整手段。


第四部分:划时代的丰碑 —— VQ-VAE 作为大模型的奠基意义与启示录

当你看完了这一切原理架构和机制优化时,自然也会疑问它的生成能力能够超越同时代的其他所有技术吗?从后视镜的角度再看,历史的车轮印记说明了其实它给世界铺平了一条彻底变革的 AIGC 的新世界地图。

1. 突破 GAN 分解界限,拉近跨模态的平行距离 (Tokenizer 对图像处理的胜利)

由于把图像从一整个密集的“千万像素大阵列信息包”,全部通过量化器打包塞嵌成了类似于“几百一千个特定小方块密码索引词汇流列表(Token Sequence)”。由此这就证明:原本只会被语言研究工作者用在把一篇宏大的文章转化进小批量 Word Piece 的那套思路机制终于被验证可以彻底搬用到图片甚至是多模态声音的建模上去了!

VQ-VAE 确立了自己作为“计算机视觉世界首台完美无损通用语言翻译官 / 生成领域的 Tokenizer”的终极无上地位!

2. DALL-E 初代 与 自回归先兆

正是有了 VQ-VAE 体系把原本长篇复杂的连续世界,变成离散排列组合;随后在 2021 年发表的那个引爆全世界文生图大地震话题神坛初作 DALL-E 1 ,其本身架构最为庞大而难以克服的工作环节第一步,便是极其精准复刻调用了一枚加强型改进版的离散词组解码变换网络:dVAE(实质上为架构极其相似的体系)。

将视觉变成了短长序列数组之后,研发工作者把用户在框内写出的诸如“生成一只在地球上行走的鳄梨沙发图片”的大段英文文字序列语句经过 GPT (Transformer 自回归机制结构处理)同样地排成了前面的半串数据 Token,跟随后缀的这些画面代表着 Token 的大片长串数组,直接被极其简单蛮暴地“相拼合并衔接”并作为长句强喂给千亿万亿巨量计算级别堆砌下的 Transformer 进行一步步往下“单词接龙预测”(Autoregressive Predicting)。最后通过 VQ 的对应密码表层去倒序回馈拼回完整光鲜的艺术原图来!自回归艺术生图绘画被揭幕,就此展开狂欢新天地。

3. VQGAN 及其感知超频与对新派巨头的孕育

仅仅依靠重建误差去压缩的 VQ-VAE 后来面临极度精细入微的极限光效挑战中会略显黯淡吃亏。于是 2020 年左右来自德国海德堡团队大神们通过引入对抗神经网络中的那个核心组件——(从重构重建误差,引申到了引入判别器的感知计算对抗特征优化技术:Adversarial/Perceptual Loss 环节)组成了闻名天下的王座神作:VQGAN 。不仅极其可怕地颠覆再创和超越复现了超精密真实的物理相片面貌,更为随后真正改变全人工内容开发世界的 Stable Diffusion 首版底层隐空间(Latent Space)编码解构埋下了无尽的技术底座指路灯!

结语与总结

在这篇近五大章节详实的报告中我们追随和穿梭着生成网络世界的理论演进步伐: 我们看到研究者不满足于 AE 的死板而试图利用带方差的重参数化技巧微调建立并尝试赋予数据带概率分布思维光芒,从而实现平缓插值的变分技术(VAE 时代);最终经过不可导断绝的难题考验,进而以看似简单但充满数学巧思的直通手段架构起了新的纪元体系(VQ-VAE / VQGAN 的离散词表时代),从而带领人们见识到了被划为不同数据模态(文本与图画)所隔绝了数百年的宏大沟通壁垒究竟是如何借此能够以 Token 为唯一的统御中轴而土崩瓦解的壮丽历史时刻。


附录:核心逻辑伪代码 (STE & 量化)

在文章的最后部分我们提供一份简化提炼版本的用来执行最为难以理解的核心技术流程直通特征映射 (Straight-Through Vector Quantization Flow) 伪代码实现:以方便大家在学习深度理论公式以外能够通过程序架构思维去从计算机视角完成印证代码落位。

import torch
import torch.nn as nn
import torch.nn.functional as F

class VectorQuantizer(nn.Module):
def __init__(self, num_embeddings, embedding_dim, commitment_cost=0.25):
super().__init__()
self.num_embeddings = num_embeddings # 词典包含的向量单词总数量 (K)
self.embedding_dim = embedding_dim # 每个单词持有的频道维度 (D)
self.commitment_cost = commitment_cost

# 定义核心武器密码本 (Codebook Vocabulary)
self.codebook = nn.Embedding(self.num_embeddings, self.embedding_dim)
# 初始化分配字典内分布权重
self.codebook.weight.data.uniform_(-1/self.num_embeddings, 1/self.num_embeddings)

def forward(self, inputs):
# inputs是前端传来的 z_e(x) 我们需要将其拉平以便和字典里的维度比较距离
flat_inputs = inputs.view(-1, self.embedding_dim)

# 计算每一组特征位置与全尺寸字典间极其庞大的距离数据阵列表网络
# 欧式距离平方展开公式推导化成了:(a-b)^2 = a^2 + b^2 - 2ab
distances = (torch.sum(flat_inputs**2, dim=1, keepdim=True)
+ torch.sum(self.codebook.weight**2, dim=1)
- 2 * torch.matmul(flat_inputs, self.codebook.weight.t()))

# 提取极值所在位置索引标签
encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
# 把序号通过独热矩阵和相乘来抽提出该词的完整属性数据表示
encodings = torch.zeros(encoding_indices.shape[0], self.num_embeddings, device=inputs.device)
encodings.scatter_(1, encoding_indices, 1)
quantized = torch.matmul(encodings, self.codebook.weight).view(inputs.shape)

# 损失体系核心组件拼凑(Stop Gradient 由此使用 detach 分隔开来实现拦截不互串导数)
# 密码字典调整推进
e_latent_loss = F.mse_loss(quantized.detach(), inputs)
# 约束固定偏移量
q_latent_loss = F.mse_loss(quantized, inputs.detach())
# 我们这里合并归纳上级传导的承诺系数调节参数量组合输出
loss = q_latent_loss + self.commitment_cost * e_latent_loss

# 神级操作:STE (直通估计技巧)
# 这里把原本已经断掉网络关联通过利用减加一个常数操作但是附带导数的梯度流直接搭了过来!
# quantized.detach() - inputs.detach() 作为常数,加上 inputs 就保留了 inputs 的梯度
quantized = inputs + (quantized - inputs).detach()

# 我们最后返回携带好具有隐形穿越特征但是却具备词典绝对固化表面性质的数据返回即可
return quantized, loss