diff --git a/app/docs/ai/multimodal/RQVAE/index.assets/2.png b/app/docs/ai/multimodal/RQVAE/index.assets/2.png new file mode 100644 index 0000000..76f36ed Binary files /dev/null and b/app/docs/ai/multimodal/RQVAE/index.assets/2.png differ diff --git a/app/docs/ai/multimodal/RQVAE/index.mdx b/app/docs/ai/multimodal/RQVAE/index.mdx new file mode 100644 index 0000000..4d833bf --- /dev/null +++ b/app/docs/ai/multimodal/RQVAE/index.mdx @@ -0,0 +1,121 @@ +--- +title: RQ-VAE学习笔记 +description: 学习RQ-VAE时的心得 +--- + + +## RQ-VAE(残差量化变分自动编码器) + +### 背景与提出动机 + +**VQ-VAE 的局限**:向量量化变分自编码器(VQ-VAE)在高保真生成任务中引入离散潜变量,通过码本(codebook)将连续潜向量离散化。然而,当我们希望缩短离散编码序列(例如将高分辨率图像表示为更少的代码)时,传统 VQ-VAE 面临率失真权衡的挑战。具体而言,若减少量化后特征图的空间尺寸(即更少的位置代码),则需要指数级扩大码本大小以维持重建质量。过大的码本不仅模型参数剧增,还容易出现码本崩溃问题。 + +**残差量化的引入**:为了解决上述问题,研究者提出了残差量化变分自动编码器(RQ-VAE)。其核心思想是用多级残差向量量化代替单次量化:在给定固定码本大小的前提下,递归量化编码器输出的残差,从粗到细逐步逼近原始表示。通过这种 Residual Quantization(残差量化,缩写RQ),无需扩充码本也能精确近似高维特征。换言之,RQ-VAE 通过组合多层码字极大扩展了表示能力:若每层码本大小为 *K*、层数为 *L*,则组合空间相当于一个大小为 $K^L$ 的超大码本,但参数量却远低于直接训练如此巨大码本。这种方式既保留了重建质量又缩短了代码序列长度,为高分辨率图像等数据的生成提供了新方案。 + +**应用背景**:RQ-VAE 的提出最初旨在改进自回归图像生成的效率和质量。在高分辨率图像自回归模型中,表示图像的离散代码序列过长会导致生成过程缓慢且计算成本高。RQ-VAE 可以在保持图像细节的同时显著减少代码序列长度(例如将256×256图像压缩为仅8×8=64个位置的代码)。这一点对于加速自回归 Transformer 的建模非常关键。此外,残差量化思想也在音频编码和推荐检索等领域兴起,将连续信号离散化为多层语义代码以便高效处理。总之,RQ-VAE 的动机在于克服单级量化的瓶颈,实现更高压缩率与重建质量兼顾的离散表示学习。 + +### 模型结构 + +![](index.assets/2.png) + +#### 整体结构 + +RQ-VAE 延续了标准自编码器架构,由编码器、多级码本量化器和解码器三部分组成。 + +- 编码器将输入(如图像、音频片段或文本嵌入)映射到一个低维潜在表示 $z$。 +- 然后,潜在向量经过多级残差量化模块处理,被表示为一系列码字索引(code indices)的组合。 +- 解码器再根据这些离散码字重建出原始数据。 + +与传统 VAE 不同,RQ-VAE 的潜变量是离散且多层次的:每个输入对应多个码本索引,共同编码其信息。 + +#### 残差量化机制 + +在 RQ-VAE 中,潜在向量 $z$ 并非一次性用单个码字表示,而是通过递归残差量化逐级逼近。具体过程如下: + +- 初始残差向量 $r_0$ 定义为编码器输出 $z$ 本身。 +- 第一级量化时,模型在第0层码本 $C^0$ 中找到与 $r_0$ 最接近的码本向量 $v_{c_0}^0$,记录其索引 $c_0$ 作为第0级码字。然后计算残差 $r_1 = r_0 - v_{c_0}^0$,将未被码字解释的细节留存下来。 +- 接着进入下一级:在第1层码本 $C^1$ 中找到最接近 $r_1$ 的向量 $v_{c_1}^1$,得到索引 $c_1$,并更新残差 $r_2 = r_1 - v_{c_1}^1$。 +- 如此迭代,共进行 $L$ 级量化,得到码字索引序列 $(c_0, c_1, ..., c_{L-1})$。 + +最终的量化表征 $\tilde{z}$ 可表示为各级码本向量之和:$\tilde{z} = \sum_{d=0}^{L-1} v_{c_d}^d$,它近似地重构了原始潜在向量 $z$。解码器以 $\tilde{z}$ 为输入重新生成数据 $\hat{x}$,尽量逼近原始输入 $x$。 + +#### 多层码本设计 + +RQ-VAE 通常为每一量化级别配备一个独立的码本 $C^d$(大小为 $K$),而非使用单一巨大码本。这种层级码本设计的原因在于:随着量化深入,残差向量范数往往递减,每一层需要表示的细节信息规模不同。为此,各层码本可以针对不同尺度的信息进行不同粒度的量化。第一层码字通常捕获了输入最粗略、最主要的特征模式,后续层的码字逐步添加更精细的细节修正。这种逐层精炼使得 RQ-VAE 的离散表示本身带有语义分层结构:例如在推荐系统中,相似物品可能共享部分语义ID(码字序列的前缀),表示它们在高层语义上相近。值得注意的是,也有研究采用共享码本的策略,在所有层复用同一套码字向量。共享码本可以减少参数量并提高码字利用率——有实验观察到不同层之间存在大量码字复用现象。无论采用独立还是共享码本,RQ-VAE 都实现了一个层次离散表示:每个输入通过码字组合来表示,既包含总体语义又涵盖细节信息。 + +#### 编码器/解码器结构 + +编码器和解码器的具体网络结构可根据任务调整。例如在图像RQ-VAE中,编码器通常是卷积神经网络,将高维图像压缩成低分辨率的特征图,再铺平成一系列潜在向量供量化。解码器则是对称的卷积网络,将量化后的特征重建回图像。在音频或语言等一维序列任务中,编码器/解码器可以采用卷积堆栈或Transformer等结构,将输入信号转为潜在向量并复原。需要强调的是,RQ-VAE 的编码器-解码器与标准 VQ-VAE 相同点在于重建目标,不同点在于潜在代码表示的形式发生了变化——由单码字扩展为多码字组合。这一架构设计令 RQ-VAE 有能力在不牺牲太多重建质量的情况下,大幅提高压缩率和表示灵活性。 + +### 损失函数 + +RQ-VAE 的训练目标延续了 VQ-VAE 的基本思想,即在**重建误差**和**量化码本优化**之间取得平衡。总损失通常包含两部分: + +#### **重建损失** + +鼓励解码器输出 $\hat{x}$ 与原始输入 $x$ 尽可能接近。常用度量是重建误差的均方误差或感知损失等。对于图像等数据通常采用像素级 MSE 或感知特征误差,以确保多码字组合后仍能高质量重构输入。 + +#### 码本向量损失 + +这是针对量化层的特殊设计,用于学习和维持码本的有效性。典型做法是在量化过程中对编码器输出(或残差)和对应码本向量之间的差距施加约束。具体实现上,VQ-VAE 引入了“停止梯度”技术,将编码器输出视为目标常量,仅更新码本向量以最小化它们与编码器输出之间的距离。同时,添加一项“承诺损失”使编码器输出靠近选中的码本向量,避免编码器持续输出过大的残差不加收敛。在 RQ-VAE 的多级情境下,这类损失会对每一级量化都计算一次,鼓励各层码本向量能够贴近其负责逼近的残差信号。以第 $d$ 层为例,可对选中码字 $v_{c_d}^d$ 和对应残差 $r_d$ 施加 $| \text{sg}(r_d) - v_{c_d}^d |_2^2$ 的损失(sg表示停止梯度,仅优化码本),同时加上 $\beta | r_d - \text{sg}(v*{c_d}^d)|_2^2$ 约束编码器行为。各层损失求和后与重建损失加权合并,即构成完整训练目标。通过调整超参数权重 $\beta$ 等,可以平衡重建精度和码本学习稳定性。 + +### 模型训练 + +#### 联合训练 + +RQ-VAE 的编码器、解码器和多个码本通常采用端到端联合训练。每次迭代通过前向传播获得多级码字及重建结果,计算上述综合损失,再通过反向传播更新参数。由于量化操作本身不可导,通常使用**直通估计器**(straight-through estimator)对量化选取的码字传递梯度——将梯度直接传给对应的编码器输出,从而允许网络更新。码本向量则通过“停止梯度”技巧只根据码本损失更新,不干扰编码器输出梯度。实践中,这一联合训练过程需精心权衡各部分损失的系数,以防止重建误差压制码本更新或反之,导致模型陷入不收敛或码字未被充分利用。 + +#### 码本初始化与稳定性 + +为了缓解码本崩溃(大量输入映射到极少数码字)等训练不稳定现象,经验上常采用一些策略: + +**良好初始化码本**:例如对第一批训练数据进行 k-means 聚类,将聚类中心设置为初始码本向量。这样做可以提供一个较为分散、多样的初始码字分布,减少模型一开始陷入局部最优的风险。 + +**正则化与更新频率控制**:比如对很少被使用的码字添加微小扰动或周期性重新初始化未使用码字,以鼓励探索。 + +**分层训练**:如果码本层数较多,可考虑先训练少层的模型取得稳定重建,再逐步增加量化层数。 + +值得一提的是,Lee等人在提出 RQ-VAE 时观察到采用**共享码本**有助于提高码字利用率。共享码本意味着各层使用同一组向量,不存在某层码本闲置的情况,从而缓解了低频码字浪费的问题。不过,共享码本也可能降低每层针对不同残差范围优化的灵活性。总的来说,RQ-VAE 的训练在继承 VQ-VAE 经验的基础上,需要针对多级量化的特性作适当调整,以确保每一级码字都能有效学习并协同完成高质量重建。 + +### 模型优势与缺陷 + +**优势**: + +- **高效的离散表示**:RQ-VAE 能以较少的序列长度表示数据,同时保持细节信息。相比单码字方案,需要的码字总数大幅减少,方便下游自回归或序列建模。例如在图像任务中,RQ-VAE 可将256×256图像表示为8×8=64个位置,每个位置由多码字组合表示,极大降低了自回归模型的计算开销。 +- **高保真的重建质量**:通过残差递进量化,RQ-VAE 在固定码本容量下达到了比VQ-VAE更强的表示能力。实验表明,增加量化深度比扩大码本对提高重建质量更有效。RQ-VAE 可在比 VQ-VAE 更低的特征图分辨率下仍保持重建逼真,解决了单级量化降低分辨率即引发失真的难题。 +- **稳定的码本利用**:由于信息被分散到多层码字,单个码字无需承载所有信息,这在一定程度上减轻了码本塌陷的倾向。特别是共享码本策略下,不同量化层复用同一套码字,使得码字利用率更高,很多码字可以在粗层和细层均发挥作用。整体来看,RQ-VAE 的多级量化避免了出现“大部分码字闲置、少数码字过载”的极端情况。 +- **语义分层与组合灵活性**:RQ-VAE 的多码字组合形式天然形成层次化语义。第一层码字往往对应类别或粗略属性,后续码字添加细节(例如风格、纹理等),这种可解释性在某些应用中十分有益。同时,多码字离散表示也为组合泛化提供了可能——不同数据可能共享部分码字,从而在离散空间中体现相似度关系。 + +**缺陷**: + +- **自回归模型的局限**:RQ-VAE 常与自回归Transformer结合用于生成任务。虽然其在大数据集上效果突出,但在小规模数据上模型易过拟合,未能超越最先进的GAN(例如 StyleGAN2)。这是自回归范式的通病:在有限数据下难以学到像GAN那样的泛化表示。因此在小数据集上,RQ-VAE 并未解决自回归模型过拟合的问题,需要引入正则化或更强的先验。 +- **单向生成与速度**:自回归生成只能单向地逐步解码,相较扩散模型等可以并行采样的方法,其推理速度依然较慢。不过,RQ-VAE 减少序列长度在一定程度上加快了采样,但本质上仍需逐步生成序列。对于需要双向上下文(如图像编辑)的任务,自回归 + RQ-VAE 框架仍显局限。未来或需结合双向模型或其他生成范式来克服这一点。 +- **训练复杂度**:多级量化带来了更多的训练超参数(如量化层数、各层码本大小、损失权重等)。训练RQ-VAE往往比训练单级VQ-VAE更为复杂且耗时——每新增一层码本都会增加编码器计算和损失项,训练不当可能导致某些码字长期未被使用或重建质量下降。例如,研究者发现需要增加训练轮数才能充分发挥深层RQ-VAE的潜力,但这又提高了训练成本。因此,如何高效训练深层次的RQ-VAE成为一大挑战。 +- **模型参数与内存**:虽然RQ-VAE避免了单个超大码本,但如果采用独立码本方案且层数较多,整体参数量和存储开销也不可忽视。例如每层有K个码字、每个D维,如果L层独立码本则参数约为 $L \times K \times D$。当 L 很大时(如音频模型中可能几十层)、或 D 很高维时,总参数量仍然可观。不过在常用设置下(例如图像模型几层、每层码本几千维),这点开销通常在可接受范围。此外,多码字在推理时需要叠加求和,也略增加了解码计算量,但相对总计算而言影响不大。 + +综上,RQ-VAE 提供了优越的离散表示能力,但在特定方面仍存在改进空间,包括更快的生成、更稳健的训练以及结合其它生成架构等。 + +### 应用场景 + +RQ-VAE 作为离散表示学习的强大工具,被成功应用于多种模态的生成和建模任务,并取得优异表现: + +- **图像生成**:RQ-VAE 最初应用于高分辨率图像的生成。Lee等人提出将 RQ-VAE 编码的离散码作为 Transformer 的输入,构建两阶段模型(RQ-VAE + 自回归Transformer)用于图像生成。在ImageNet等数据集上,该方法在无条件和有条件生成任务中均超越了以往的自回归模型性能。尤其值得关注的是,在256×256图像上,RQ-VAE 将特征图缩减为8×8,提高了Transformer建模的效率,却仍能生成高质量图像,FID等指标显著改善。同时,相较于之前的像素CNN等逐像素AR模型,该框架的采样速度提高了数量级。不过在小规模数据集(如FFHQ人脸)上,其效果仍逊于当时最好的GAN模型。总体而言,RQ-VAE + Transformer 在图像生成方向证明了短码表示+自回归的可行性,为后来者提供了新思路。 +- **音频建模**:在神经音频编码和生成领域,残差向量量化已经成为标配技术之一。Google 提出的 SoundStream 模型将语音/音频信号经过卷积编码器并应用多层残差量化(高达数十层),实现了对宽带音频的有损压缩。SoundStream 在保持接近原始音质的同时,将比特率降至数 kbps 级别,其关键就在于多级码本提供了高精度的重建能力。这些离散音频码随后可用于音频生成模型:例如 Google 的 AudioLM、MusicLM 等方法先用SoundStream获取离散音频表示,再训练语言模型生成这些离散码,从而产生连贯的音乐或讲话。Facebook 的 EnCodec 也是类似思想的音频编码器。可以说,RQ-VAE(或称RVQ)的思路使得神经网络能学习端到端压缩音频,同时为后续的生成任务提供离散单位。因此在音频去噪、音乐生成、语音传输等方向,基于残差量化的VAE模型都表现出高保真、高压缩的优越性能。 +- **语言与推荐**:虽然文本天然是离散符号序列,但在语义表示方面,RQ-VAE 也开始发挥作用。例如在生成式推荐系统中,RQ-VAE 被用来将物品的内容特征(如商品描述文本)编码为语义ID。每个物品被表示为一组码字(语义ID元组),相似的物品会共享部分码字,从而在离散空间中保持语义上的邻近。像 Google 提出的 TIGER和国内快手的 OneSug 模型都使用 RQ-VAE 对文本描述或查询进行离散化,把用户历史序列映射成语义代码,再用序列到序列模型生成下一个推荐物品或查询。这些语义ID有效缓解了推荐场景中使用原始ID导致词表过大的问题,使得大语言模型也能直接处理物品序列。在语言建模方面,也有研究尝试将句子的连续语义表示离散化,以利于压缩和检索。例如生成式搜索(Generative Retrieval)任务中,RQ-VAE 提取查询的语义ID,再通过逐字解码生成最相关的文档ID。这些尝试表明,RQ-VAE 在语言相关领域可以作为语义嵌入离散化工具,提高跨模态或跨系统的对接效率。在多模态应用中,人们也设想将图像、音频通过各自的RQ-VAE离散化,和文本一起统一表示为离散序列,从而让一个通用Transformer同时理解多种信号。总之,在语言与推荐系统方向,RQ-VAE 提供了一种将高维语义向量转为可组合离散符号的方法,已展现出提升下游任务性能的潜力。 + +### 与VQ- VAE对比 + +#### VQ-VAE: + +经典 VQ-VAE使用单一级码本将连续潜在向量量化为一个离散码字,从而得到图像或序列的离散表示。它的优点是结构简单,离散码直接可用于训练自回归模型。然而,由于每个位置只有一个码字,VQ-VAE 为了表示足够细节常需要较高的空间分辨率或庞大的码本。如前所述,单码本方案在压缩比和保真度间存在权衡:若减少码字数目(降分辨率),必须增大码本大小来避免信息缺失。过大的码本又引发稳定性问题(如码本向量用不充分)。相比之下,RQ-VAE 利用多码字组合显著扩展了表达能力,无需极端依赖码本规模。在相同码本大小下,RQ-VAE 可以用更少的空间位置表示数据,同时保持高保真重建。可以说,RQ-VAE 在码本大小一定时突破了 VQ-VAE 在表示容量上的局限,是对 VQ-VAE 的扩展。 + +#### VQ-VAE-2 + +VQ-VAE-2 是对 VQ-VAE 的层次拓展,由 DeepMind 提出,用于更好地生成图像细节。它引入了两级(或多级)离散层次:例如高层码字捕获全局粗略结构,低层码字捕获局部精细信息。具体实现上,VQ-VAE-2 有两个(或多层)独立的VQ模块:先编码得到高层离散表示,再在解码时将高层表示作为条件去预测细粒度的低层离散表示。这与 RQ-VAE 的区别在于:VQ-VAE-2 的层次是在空间/语义上分离的,每层码字各自对应不同尺度的特征图;RQ-VAE 的层次是在同一潜在向量上的递进近似,多层码字共同构成一个位置的表示。直观地讲,VQ-VAE-2 好比“先概括后细化”地编码图像,不同层码字作用在不同抽象级别;RQ-VAE 则是在同一抽象级别通过残差叠加来细化编码。两者都体现了分层离散表示思想,但实现机制有所不同。VQ-VAE-2 能提高生成质量(特别是纹理细节)并提供多尺度语义控制,但它并未显著减少总的代码数量,反而因为引入多个层次可能需要更多码字参与重建。相反,RQ-VAE 的设计初衷之一就是减少离散序列长度,它把信息浓缩进较少的位置,通过每个位置的多码字叠加来确保细节不丢失。因此,在需要紧凑离散表示的场景下(如自回归建模需短序列),RQ-VAE 相对更有优势。 + +### 总结 + +RQ-VAE 将残差量化与变分自编码器思想相结合,成功实现了高压缩比且保真度高的离散表示学习。它通过多级码本的粗到细量化,克服了传统 VQ-VAE 在码本容量和序列长度上的两难困境,实现了离散潜代码的层次表示。这种方法在图像、音频、文本等多领域的应用中取得了显著成果,不仅提高了生成模型的效果和效率,也为跨模态表示、生成式检索等新兴任务提供了有力工具。 + + + diff --git a/app/docs/ai/multimodal/VQVAE/index.assets/1.png b/app/docs/ai/multimodal/VQVAE/index.assets/1.png new file mode 100644 index 0000000..41ab9a2 Binary files /dev/null and b/app/docs/ai/multimodal/VQVAE/index.assets/1.png differ diff --git a/app/docs/ai/multimodal/VQVAE/index.mdx b/app/docs/ai/multimodal/VQVAE/index.mdx new file mode 100644 index 0000000..c89853d --- /dev/null +++ b/app/docs/ai/multimodal/VQVAE/index.mdx @@ -0,0 +1,121 @@ +--- +title: VQ-VAE学习笔记 +description: 学习VQ-VAE时的心得 +--- + +## VQ-VAE(向量量化变分自编码器) + +### 概述 + +**VQ-VAE**(Vector Quantized VAE)是为了解决传统VAE存在的局限而提出的一种模型。传统VAE使用连续潜变量,可能会出现后验坍塌(posterior collapse)问题:当解码器很强大时,编码器的隐变量信息被忽略,模型几乎只依赖解码器来重构数据。VQ-VAE通过在编码器-解码器之间引入一个离散的代码簿(codebook)层,将连续隐空间离散化,使解码器被迫使用潜变量信息,从而避免了后验坍塌。同时,离散潜变量能够更有效地捕获数据中离散的结构特征(如语音中的音素、图像中的物体等),提升生成样本的质量。研究表明,采用离散潜变量的VQ-VAE在对数似然等指标上可以达到与连续潜变量模型相当的性能。此外,VQ-VAE为后续的生成模型奠定了基础,其离散表示已被用于诸如 OpenAI 的 DALL·E 等先进模型中,证明了在图像、语音等多种任务上的实用价值。 + +### 模型结构 + +![](index.assets/1.png) + +VQ-VAE 的结构可以看作是在自编码器(Autoencoder)基础上增加了一个向量量化层。包括三个主要部分:编码器(Encoder)、代码簿(Codebook)和解码器(Decoder)。模型工作过程如下: + +- 编码器网络将输入数据 $x$ 映射到隐空间得到连续隐表示 $z_e(x)$。 +- 然后,在隐空间插入向量量化操作:将 $z_e(x)$ 与代码簿中所有向量比较,找到距离最近的代码向量 $e_k$,并将其作为离散隐变量 $z_q(x)$ 传递给解码器。 +- 解码器接收选中的代码向量序列作为输入,尝试重构出原始数据 $\hat{x}$。 + +这个过程中,编码器输出经过一个非线性离散化步骤(选取最近邻的码簿向量),整体可以被视为带有离散瓶颈的自动编码器。需要注意的是,在图像等高维数据场景下,编码器通常不会只输出单个隐向量,而是输出一个如 $32\times32$ 大小的隐向量网格(每个位置对应局部区域特征)。该网格中的每个向量都会独立地进行量化,均从同一个代码簿中选取最近的嵌入向量,这样虽然代码簿大小固定,但因为输出多个code,解码器可以组合出指数级丰富的重构结果。例如,若代码簿大小为512,隐空间网格维度为$32\times32$,则解码器理论上可生成多达 $512^{32\times32}$ 种不同组合的图像。由此可见,VQ-VAE 的离散瓶颈既限制了每个隐变量只能取有限集合中的值,又通过多元组合提供了巨大的表示能力。 + +### 向量量化机制 + +VQ-VAE的关键在于**向量量化(Vector Quantization)**操作,即将连续的编码器表示映射为离散的索引。具体机制包括以下几点: + +#### 代码簿(Codebook): + +代码簿是一组可学习的嵌入向量集合,一般记为 $e = {e_1, e_2, \dots, e_K}$,其中 $K$ 是代码簿中向量的数量(离散latent空间的大小),每个向量维度为 $D$。这些向量可以看作原型向量或簇心,代表了离散潜在空间中可能取值的中心。模型的编码器输出将被强制映射到这些原型向量上,从而得到一个离散化的表示。 + +#### 最近邻查找 + +对于每一个编码器输出向量 $z_e(x)$,VQ-VAE 选择与其欧氏距离最近的代码簿向量 $e_k$ 来表示。用公式表示即为:$k = \arg\min_i | z_e(x) - e_i |_2$,然后令离散隐表示 $z_q(x) = e_k$。这个最近邻查找(nearest neighbor lookup)过程实现了连续向量到离散代码的量化,对应于选取一个 1-of-K 的离散编码。解码器随后接收 $z_q(x)$ 作为输入进行重构。由于每个 $z_e(x)$ 都被替换为最近的离散向量 $e_k$,该操作本质上相当于对隐空间进行了k-means量化(聚类),把隐表示压缩到有限的聚类中心上。 + +#### 直通估计(Straight-Through Estimator) + +由于最近邻的取整操作是非平滑、不可微的,直接反向传播梯度将无法传递到编码器。为了解决这一问题,VQ-VAE 使用了直通梯度估计器:在前向传播时使用离散的 $z_q(x)$,但在反向传播时忽略量化操作的离散性,将来自解码器的梯度直接传给编码器的输出 $z_e(x)$。简单而言,就是在反传过程中将量化操作视作恒等映射处理——解码器对 $z_q(x)$ 的梯度被直接拷贝给 $z_e(x)$,从而更新编码器参数。这种近似的梯度传播方法允许编码器仍然能够通过重构误差来更新,即使中间有一个非可导的选取最近邻操作。在实现中,这通常利用框架对张量的截断梯度(stop-gradient)操作来完成:量化步骤的前向输出用于后续计算,但对该输出梯度设为零,只让原始编码器输出承担梯度。这个技巧有效地绕过了量化的不连续性,实践中被证明可以稳定地训练VQ-VAE。 + +### 损失函数 + +VQ-VAE 的训练目标由三部分损失项组成,各自针对模型的不同组件进行优化: + +#### 重构损失(reconstruction loss) + +重构损失(reconstruction loss)是衡量模型重构输出 $\hat{x}$ 与原始输入 $x$ 之间差异的损失,用以训练编码器和解码器。根据任务不同,可采用均方误差(MSE)或对数似然损失(如像素级交叉熵)等形式表示。重构损失鼓励编码器-解码器能够尽可能准确地还原输入数据,是整个模型优化的主要驱动力。在公式中通常记为 $L_{\text{reconstruction}} = -\log p(x|z_q(x))$,表示对数据点在给定离散隐变量下的似然的负对数。 + +#### 代码簿损失(codebook loss) + +代码簿损失(codebook loss,又称嵌入对齐损失)用于训练代码簿向量的损失项。由于直通估计的使用,重构误差的梯度并不会更新代码簿(梯度已传给编码器而对代码簿截断),因此需要额外机制来学习代码簿。本项通过最小化编码器输出 $z_e(x)$ 与所选代码向量 $e_k$ 之间的平方误差来更新代码簿向量,使其朝向编码器输出的方向移动。具体形式为 $L_{\text{codebook}} = | \text{sg}[z_e(x)] - e_k |_2^2$,其中 $\text{sg}$ 表示停止梯度操作(其作用是让括号内的张量在反向传播时梯度为零)。由于 $z_e(x)$ 前面加了 stop-gradient,这部分损失只对代码簿参数 $e$ 产生梯度,迫使相应的 $e_k$ 靠近当前的编码器输出。直观来说,代码簿损失不断调整每个嵌入向量的位置,以更好地代表编码器实际输出的聚类中心,从而优化字典学习的效果。 + +#### 承诺损失(commitment loss) + +承诺损失(commitment loss,也称承诺代价)用于训练编码器输出的损失项。其目的在于约束编码器输出“承诺”于选定的离散向量,避免编码器输出无限增长或者频繁切换代码造成不稳定。形式为 $L_{\text{commitment}} = \beta | z_e(x) - \text{sg}[e_k] |_2^2$,其中系数 $\beta$ 为权重超参数,用于平衡重构误差和承诺代价的重要性这里对代码簿向量 $\text{sg}[e_k]$ 应用了停止梯度,因此此项损失只会使编码器参数更新(代码簿不受影响),推动 $z_e(x)$ 靠近已选的代码簇中心 $e_k$。这样可以防止编码器输出偏离代码簿过远,促使编码器尽可能产生稳定且贴合代码簇的表示。原论文指出,模型对 $\beta$ 的选择相对鲁棒,在 $0.1$到$2.0$ 范围内效果差别不大,一般可以取 $\beta \approx 0.25$。 + +#### 总损失 + +总的训练损失就是上述三部分之和: +$$ +L_\text{VQ-VAE}  =  L_\text{reconstruction}+L_\text{codebook}+L_\text{commitment} +$$ +其中重构损失影响编码器和解码器的更新,代码簿损失只更新代码簿参数,承诺损失则只作用于编码器输出。需要注意的是,在 VQ-VAE 中我们假设隐变量符合均匀分布的先验(即任何代码出现的概率相等),因此对应的 KL 散度项为常数,不随模型参数变化,在训练时通常被忽略不计。这意味着,与传统VAE不同,VQ-VAE的目标函数中没有显式的KL正则项,模型复杂度的约束主要由离散瓶颈和承诺损失来实现。 + +### 缺陷及改进 + +#### 代码簿崩溃(codebook collapse) + +在训练VQ-VAE时,需要注意一些潜在的问题和相应的对策。其中最重要的之一是**代码簿崩溃(codebook collapse)**现象。代码簿崩溃是指模型在训练过程中只使用了代码簿中很小一部分向量,大部分代码词(codeword)长期闲置未被使用,或者不同代码词收敛到几乎相同的取值,导致有效的离散表示容量大幅降低。这种情况会削弱模型的表示能力,相当于隐空间的多样性坍缩。造成崩溃的原因可能是编码器和解码器的更新不平衡、超参数设置不当(例如承诺损失权重过大使编码器不愿改变代码)等。为缓解和避免代码簿崩溃,实践中常用以下训练技巧。 + +#### 指数移动平均(EMA)更新 + +使用**EMA算法**来更新代码簿向量,而非直接反传梯度更新。这种方法在第二版VQ-VAE中得到采用,被发现可以使代码簿的更新更加平滑稳定,避免单个代码向量由于梯度波动出现剧烈变化或竞争失败。具体做法是每次前向传播时把每个隐向量指派给最近的代码向量,统计每个代码向量本轮对应的“样本和+样本数”,再用衰减系数把历史统计与本轮统计做加权融合,得到新的累计和与计数,最后用它们的比值更新码本(近似把簇心靠近最近一次观测到的数据分布)。EMA更新能使训练更平滑,避免梯度直接穿过量化带来的不稳定和码本崩溃,有效防止训练初期由于随机原因某些代码一直未得到更新而废弃的情况。 + +#### 代码重置(Codebook Reset) + +这是针对长期未被使用的代码词的一种处理策略。训练过程中持续监视每个代码向量被选中的频率,如果发现某些代码在较长一段迭代中从未或很少被使用,可以将这些“失效”代码向量重置——例如用当前mini-batch中随机某些编码器输出向量来替换它们。通过重新初始化闲置的代码词,使其有机会迁移到数据分布较密集的区域参与表示,确保代码簿的覆盖面不至于过偏。Code reset 的操作类似于 k-means 聚类中对空簇中心的处理,可以提高代码簿利用率,对抗崩溃。 + +#### 增加多样性正则 + +在损失函数中显式引入鼓励代码使用丰富度的正则项。例如,可以加入一个熵正则化或 Code usage 多样性损失,鼓励编码器的离散输出分布接近均匀。具体做法包括最大化离散后验 $q(z|x)$ 的熵,或在一段时间内统计代码使用频率分布,与均匀分布计算KL散度作为惩罚。这样的项在数学上推动模型充分利用更多的代码词,而不是集中在少数几个上。研究和实操表明,适度的熵/多样性正则能有效缓解代码簿崩溃,提高码本利用率。不过需要注意权衡,正则项过强可能反而干扰重构性能,因此通常将其作为辅助监控指标或微弱正则。 + +#### 超参数与训练策略 + +合理选择并调整一些关键超参数对训练稳定也很重要。例如,承诺损失权重 $\beta$ 需要平衡编码器输出移动和代码簿更新的速度:$\beta$ 过大会让编码器过于死板地黏着当前代码(减少代码使用的探索),过小则可能出现编码器输出变化太快代码簿跟不上等情况。通常在0.25左右取得较好平衡,但也需视具体任务调节。另外,代码簿大小应与数据复杂度匹配;过大的码本在数据不足时更易出现大量代码闲置而崩溃,必要时可以减少码本规模或采用分层码本来提升利用率。训练时使用较大的batch size、充分的训练步数也有助于编码器更稳定地使用各种代码词。一些实现中还采取分阶段训练策略:先固定解码器训练编码器和码本,使其学到初步离散表示,再释放解码器一起优化,以防止解码器过早主导重构而忽略潜变量。 + +总之,针对VQ-VAE特有的训练挑战,需要结合上述方法细心调校,监控代码簿的使用情况。幸运的是,这些技巧已经在很多实践中验证有效,使我们能够较稳定地训练出高质量的VQ-VAE模型而不出现严重的崩溃或失效问题。 + +### 应用场景 + +VQ-VAE 提供了一种将数据压缩到离散表示的强大手段,因而在多个领域获得了成功应用: + +- **图像生成与压缩**:VQ-VAE 最初在图像生成任务上展示了其威力。通过先训练VQ-VAE将图像压缩成离散代码,再训练一个像素级的自回归模型(如 PixelCNN 或 Transformer)作为代码序列的先验分布,模型可以生成高质量且多样的图像样本。这种两阶段方法中,VQ-VAE负责学习图像的离散表征,将复杂图像简化为易于建模的离散token序列;随后强大的生成模型在这个离散空间中学习数据分布。实验证明,该策略能够在无监督情况下生成逼真的图像。例如,VQ-VAE 在 ImageNet 上训练后结合 PixelCNN prior,可以生成清晰的自然图像。后续的VQ-VAE-2 引入分层多尺度的离散表示和更强的先验,进一步提高了生成图像的一致性和保真度,能够产生接近原始分辨率的高质量图像。此外,在图像压缩领域,VQ-VAE的离线码本量化与自回归建模也被用于无损/有损压缩算法中,以在保证感知质量的同时大幅降低比特率。 +- **语音和音频建模**:VQ-VAE 在语音生成与音频表示学习方面的应用也非常成功。原始论文中,将VQ-VAE应用于语音数据,模型在无监督情况下自动学出了类似**音素**的离散单元——即代码簿的向量开始对应语音中的基础声音单位。这说明VQ-VAE能够从原始音频中提取出有意义的离散特征,而这些特征对于下游任务(如语音合成、语者识别)非常有用。进一步地,通过在解码器端输入不同的说话人ID等条件,VQ-VAE还能实现说话人转换等应用:即编码内容不变但用另一人的声音重构,同一段代码可以由解码器转换成不同音色的语音。OpenAI 的 Jukebox 项目则将 VQ-VAE 扩展到音乐领域,采用分级VQ编码将原始音乐波形压缩为多层级的离散码,再用 Transformer 模型在这些代码序列上进行建模,成功地生成了长时段且结构丰富的歌曲音乐。这些应用展示了VQ-VAE 在音频领域捕捉长期结构和离散语义单元的能力,使得高保真语音和音乐的生成成为可能。 +- **自然语言与序列建模**:虽然自然语言本身天生是离散符号序列,但VQ-VAE思路在文本序列的潜在离散表示方面也有所探索。在一些文本生成或机器翻译任务中,研究者引入VQ-VAE作为中间瓶颈,将连续的隐含语义向量离散化,以获取更可控和解释的表示。例如,有工作将VQ-VAE用于句子级别的生成模型,利用离散潜变量来表示整个句子的高阶语义结构,避免了传统序列VAE中经常出现的posterior collapse问题。这些离散语义码可以看作模型自动学习出的“句子代码”,为生成器提供了更明确的信号。另一些研究则将VQ-VAE引入到Transformer语言模型中,获取隐层的离散概念单元,以改善生成的多样性和语义一致性。总体而言,在语言模型中使用离散表示可以增强模型对全局结构的把控,提供一种调节语义的“开关”。例如最新的T5VQVAE模型通过在Transformer VAE中嵌入离散空间,提升了生成语义控制能力,在诸多文本生成指标上优于以往方法。虽然这方面应用还在探索中,但初步结果显示离散潜变量有助于学习更抽象的语言结构,对于长文本或多样风格的生成具有潜力。 +- **强化学习与其他领域**:VQ-VAE 的离散表示思想同样被尝试应用在一些其他领域,例如强化学习(RL)中的状态表示。研究者将高维感知输入(如游戏画面)通过VQ-VAE压缩成离散代码,发现这些代码往往对应环境中的\关键物体或地形,从而为下游的决策提供了更简洁有用的状态特征。这种方式可以在不提供额外标注的情况下,让智能体自行学习离散的“概念”,从而改进长时策略的学习。在机器人控制、动作生成等序列决策任务中,也有尝试使用VQ-VAE提取离散的运动primitive或姿态代码,以简化连续控制的问题难度。总体来说,凡是需要从高维连续信号中提炼出抽象符号的场景,都可能看到VQ-VAE的身影。 + +### 与VAE的对比 + +**潜在空间类型**:VAE 的隐空间是连续的,多采用高斯等连续分布来建模潜变量;而 VQ-VAE 的隐空间由离散的代码组成,通过代码簿的有限集合来表示。换言之,VAE 的隐向量每一维可以是任意实数,而 VQ-VAE 的每个潜变量只能取有限个离散值(由代码簿索引决定)。 + +**编码器输出与推断方式**:VAE 的编码器输出的是潜在分布的参数(如高斯分布的均值和方差),训练时需通过重参数化技巧从该分布中采样连续隐向量 $z$,这会在训练目标中引入一个 KL 散度项用于正则化;相比之下,VQ-VAE 的编码器直接输出确定的隐向量,并通过最近邻量化得到离散的 $z_q$,因而后验分布是退化为 one-hot 的确定性分布($q(z=k|x)=1$ 对应选中的代码,其它为0)。VQ-VAE 训练时不需要随机采样步骤,也没有复杂的后验近似,编码器更像是被“硬编码”成一个近邻查询函数,这简化了训练过程且避免了VAE中因随机采样导致的梯度高方差问题。 + +**损失函数组成**:传统 VAE 的目标由重构误差和KL正则项组成,即在保证重构数据的同时逼近隐变量先验(通常为标准正态);而 VQ-VAE 没有显式的KL项,取而代之的是代码簿损失和承诺损失来确保编码器输出与离散代码簇相匹配,并维持隐空间稳定。也就是说,VAE 通过KL散度鼓励连续隐向量分布接近先验,而VQ-VAE通过额外的MSE损失来拉近编码器输出和嵌入向量,从另一个角度实现对隐空间的约束。 + +**先验分布处理**:在 VAE 中,隐变量先验 $p(z)$ 一般预先固定为标准正态分布或其他简单分布,训练时通过最小化KL散度将编码器后验拉向该先验;VQ-VAE 则是训练时假定先验均匀(不显式参与优化),模型训练完成后再对离散隐变量学习一个新的先验分布。通常做法是训练一个自回归模型(如 PixelCNN、Transformer)作为离散代码序列的先验,实现对隐空间复杂分布的建模。这样VQ-VAE的先验不局限于简单分布,而是可以由强大的模型来拟合数据的隐结构。 + +**后验坍塌问题**:VAE 常出现后验坍塌,尤其在解码器为强自回归模型时,优化倾向于让编码器输出接近先验、忽略潜变量,从而KL项很小但重构主要靠解码器记忆,导致潜变量的语义无效。相比之下,VQ-VAE 由于隐空间被量化为有限的离散值,解码器必须依赖这些离散代码才能重构输入,因此极大程度缓解了后验坍塌。实际上,VQ-VAE 的设计初衷之一就是通过离散瓶颈强制使用潜变量信息,使其不会被强力解码器所忽略。这使得VQ-VAE能在不引入KL惩罚的情况下,依然学到有用的隐表示。 + +**潜在表示的表达能力**:连续隐空间的VAE在表示某些离散的细节时可能效率不高,例如图像中的清晰边缘、重复纹理,或语音中的符号化特征,容易因为取均值而变得模糊。而VQ-VAE的离散隐变量更擅长捕获这类信息,使生成结果更清晰、多样。实验证明,原始VQ-VAE在图像生成质量上优于典型VAE,可生成分辨率更高、细节更锐利的样本。这是因为离散代码不会发生连续隐变量那种“插值模糊”效应,且配合强先验可以确保抽样结果接近真实数据分布。 + +**训练难点与稳定性**:VAE 的训练相对稳定,因为优化目标(ELBO)对参数是平滑可微的,但可能面临KL权重选择、posterior collapse等问题;VQ-VAE 的训练需要处理非连续的量化操作,尽管通过straight-through技巧解决了梯度传播,但仍可能出现码本利用不充分等问题。此外,VAE训练中编码器梯度直接来自重构和KL项,而VQ-VAE中部分梯度被截断,需要额外的嵌入更新机制,这使得调参有所不同。然而,总体而言VQ-VAE在避免后验坍塌、提高表示离散性方面的收益,使其在许多任务上成为优于标准VAE的选择。 + +### VQ-VAE 的发展 + +#### VQ-VAE-2 + +VQ-VAE-2 是对原始VQ-VAE的改进版本,由DeepMind团队在2019年提出。VQ-VAE-2 的核心思想是在模型中引入多层次的离散隐空间:它采用分层编码器-解码器结构,先提取高等级的粗粒度离散表示,再提取低等级的细粒度离散表示,从上到下逐级生成图像。这样做的好处是顶层的离散code可以捕捉全局概况(如图像的大致布局),底层code捕捉局部细节,从而提高生成图像的一致性和细节质量。此外,VQ-VAE-2 引入了一种基于 PixelCNN 的强先验模型:它在最高层离散代码上训练PixelCNN(或PixelSnail)来建模代码分布,并在下层生成时逐级条件采样。凭借分层表征和强先验,VQ-VAE-2 能生成高保真且多样的图像样本,显著优于一层结构的VQ-VAE。一些公开实验表明,VQ-VAE-2 可以在Imagenet这样的数据集上生成接近真实分辨率($256\times256$或更高)的图像,其感知质量可与GAN等方法相媲美。这确立了离散VAE在大型图像生成任务中的地位。同时,VQ-VAE-2 在训练技术上也做了改进,例如使用EMA更新码本(上文提及)来替代代码簿损失,进一步缓解了训练不稳定和code collapse问题。可以说,VQ-VAE-2 的出现证明了离散潜变量模型在复杂生成任务上的可行性和优越性。 + +#### 结合 Transformer 等自回归模型 + +近几年,随着Transformer在生成建模上的成功,将 VQ-VAE 与 Transformer 结合成为了热门趋势。基本思路是利用 VQ-VAE 将高维数据(如图像、音频)转换成离散 token 序列,然后用 Transformer(如GPT)在这些 token 上学习分布,用于生成新样本。这种“两段式”模型中,VQ-VAE负责学习离散表示,Transformer负责在离散空间里做序列预测。典型的例子就是 OpenAI 的 DALL·E:它先训练一个VQ-VAE将图像映射为离散码,再训练一个大规模Transformer以文本描述为条件来自回归地生成图像代码序列,最后通过VQ-VAE的解码器得到图像。DALL·E 的成功展示了该思路在多模态生成上的强大威力。同样,前述的 Jukebox 使用了层级VQ-VAE压缩音乐、Transformer生成音乐codes,实现了长时音乐的生成。将Transformer引入离散空间具有几个显著优点:Transformer强大的长程依赖建模能力可以充分发挥,离散token使得输出空间有限且利于概率估计,从而生成质量和样本多样性都显著提升。实践证明,VQ-VAE 与 Transformer 是天然契合的组合:前者提供表示能力,后者提供建模能力,两者结合几乎可以扩展生成模型到任意复杂的数据分布。只要有足够的算力和数据支撑,这种方法能够在图像、音频、文本等领域生成极其高质量的内容。目前,许多最前沿的生成模型都采用了这种范式(例如把 VQ-VAE 或其变种作为编码器,Transformer作为主生成器)。可以预见,在未来相当长一段时间内,VQ-VAE+Transformer 的架构仍将是生成建模领域的重要工具,离散表示学习也会持续为多模态AI的发展提供支持