重建与生成不再对立 | GigaTok借语义正则化统一视觉分词器,3B参数完胜VQ-GAN,刷新ImageNet纪录

大模型向量数据库数据中台

点击下方卡片,关注「集智书童」公众号

点击加入👉「集智书童」交流群

picture.image

picture.image

picture.image

picture.image

picture.image

picture.image

精简阅读版本

本文主要解决了什么问题

重建与生成的矛盾 :扩展视觉分词器可以提高图像重建质量,但通常会导致下游生成质量下降。作者研究了这一矛盾,并提出了解决方案。

潜在空间复杂性增加 :随着分词器规模的扩大,潜在空间复杂性增加,导致AR模型更难以学习token依赖关系。作者通过语义正则化缓解了这一问题。

分词器扩展的最佳实践 :作者探索了如何有效扩展分词器以同时提升重建和生成性能。

本文的核心创新是什么

语义正则化 :通过将分词器特征与预训练视觉编码器提供的语义一致特征对齐,约束潜在空间的复杂度,从而解决重建与生成之间的矛盾。

非对称扩展策略 :优先扩展解码器而非编码器,同时保持解码器比编码器更大,以实现最佳性能。

一维分词器的优势 :发现一维分词器相较于二维分词器具有更好的可扩展性,减少归纳偏差。

熵损失的应用 :在训练十亿级参数的分词器时,引入熵损失以稳定训练过程并提高codebook利用率。

结果相较于以前的方法有哪些提升

重建性能 :GigaTok在ImageNet上实现了最先进的重建性能(rFID)。

下游AR生成性能 :使用2.9B GigaTok分词器的1.4B AR模型在LLM风格的自动回归下一词预测模型中实现了最先进的图像生成性能(gFID)。

表示学习质量 :通过视觉生成预训练,1.4B AR模型达到了最高的线性检测精度,表明GigaTok有助于下游模型学习更好的表示。

局限性总结

计算成本 :尽管一维分词器更具可扩展性,但其训练和推理仍需要较高的计算资源。

特定任务适用性 :虽然GigaTok在统一多模态建模方面表现出色,但在某些视觉特定任务(如VAR的下一尺度预测)中可能不如专门设计的方法。

语义正则化的权衡 :语义正则化层的选择和权重调整需要在重建质量、生成性能和表示学习之间进行权衡,可能不适用于所有场景。

深入阅读版本

导读

在自回归(AR)图像生成中,视觉分词器将图像压缩成紧凑的离散潜在 Token ,从而通过下一 Token 预测高效地训练下游自回归模型进行视觉生成。虽然扩展视觉分词器可以提高图像重建质量,但它通常会导致下游生成质量下降——这是一个现有文献未能充分解决的问题。

为此,作者引入GigaTok,这是首个在扩展视觉分词器时同时提升图像重建、生成和表示学习的方法。作者识别出潜在空间复杂性的日益增长是重建与生成困境背后的关键因素。为缓解这一问题,作者提出了语义正则化,该正则化方法使分词器特征与预训练视觉编码器提供的语义一致特征对齐。这种约束在扩展过程中防止潜在空间过度复杂化,从而在重建和下游自回归生成两方面都实现了持续改进。

基于语义正则化,作者探索了扩展分词器的三种关键实践:

    1. 使用一维分词器以获得更好的可扩展性
    1. 在扩展编码器和解码器时优先扩展解码器
    1. 采用熵损失来稳定百亿规模分词器的训练

通过扩展到30亿参数,GigaTok在重建、下游AR生成和下游AR表示质量方面均达到了当前最佳性能。

  1. 引言

自回归语言模型(AR LM)已成为视觉生成领域的一种有前景的方法[15, 50, 65, 68],其优势在于已证明的可扩展性[2, 5, 14, 19, 37, 51, 52, 54, 55]以及统一多模态建模的潜力[12, 45, 61]。AR图像生成框架由一个视觉分词器和一个下游AR生成器组成。分词器将图像编码为离散的符号,并使用图像重建监督进行训练,而AR生成器通过下一符号预测来建模这些离散符号的分布。图像分词器在AR视觉生成中发挥着关键作用,它提供了一个紧凑且表达力强的潜在空间,使下游AR模型能够进行有效的生成建模。

尽管视觉分词器在关键作用,但其在文献中的扩展研究却鲜有探讨。事实上,与那些可扩展性已被广泛验证的下游AR模型[12, 30, 59, 61]不同,扩展视觉分词器面临重大挑战。具体而言,存在重建与生成之间的矛盾,即扩展分词器能提高重建保真度,但会降低下游生成质量,如图1所示。这一矛盾在先前工作中[13, 21]也有体现。在本工作中,作者旨在克服这一局限性,并探索有效扩展分词器的策略,以提升重建和生成性能。

picture.image

为探究这一困境的根本原因,作者提出一种AR检测方案,该方案训练一个轻量级下游生成式AR模型来监控分词器的训练过程。令人惊讶的是,作者发现随着分词器规模的扩大,下游AR模型学习所得的token分布变得更加困难,这从AR生成损失的不断增加中得到了证实。这表明,更大的分词器产生了更复杂的token空间,使得AR模型的有效学习变得更加困难。

为应对这一挑战,作者引入预训练的视觉表征模型(例如DINOv2)来规范分词器。具体而言,作者在分词器训练过程中利用语义正则化损失,促使分词器特征与预训练模型特征之间具有高度相似性。这种正则化有助于约束潜在空间的复杂度,防止分词器学习过于复杂的潜在 Token 依赖关系,从而阻碍下游AR生成建模。

此外,作者设计了一种基于混合CNNTransformer架构的矢量量化(VQ)分词器作为主干,适用于1D和2D分词器,并探索了分词器扩展的最佳实践:

    1. 与2D分词器相比,1D分词器表现出更好的可扩展性;
    1. 非对称模型扩展,优先考虑解码器扩展而非编码器扩展,被证明是有效的;
    1. 在训练具有十亿级参数的分词器时,熵损失[68]对收敛至关重要。

通过作者的语义正则化和三种关键扩展策略,作者有效地将GigaTok扩展到30亿参数,克服了重建与生成之间的困境。

作者总结作者的贡献如下:

  • • 作者识别出,在分词器扩展中,重建与生成之间的困境源于较大分词器中潜在空间复杂性的增加。为解决此问题,作者提出了语义正则化,有效缓解了这一困境,并实现了分词器扩展。
  • • 作者探讨了扩展分词器的最佳实践,包括混合CNN-Transformer架构的一维分词器、非对称编码器-解码器扩展,以及用于十亿规模分词器的熵损失。
  • • 作者的GigaTok是首个扩展到30亿规模的分词器,在ImageNet上实现了最先进的重建、下游AR生成和下游AR表征效果。
  1. 相关工作

图像分词器。图像分词器将图像输入映射为离散或连续的符号,这些符号可以被下游生成模型进行建模。连续分词器基于变分自编码器(VAE)[28, 29]构建,而离散分词器如VQ-GAN [15]在训练过程中将视觉特征量化为离散的视觉符号。向量量化(VQ)[15, 56, 65]被主要采用以离散化符号,而其他量化方法[49, 68, 74, 75]则专注于调整码本大小以获得更好的分词器。然而,如何在现有文献中适当扩展分词器模型的研究却很少。一项并行工作ViTok [75]尝试扩展基于连续VAE的分词器以用于下游扩散模型,但最终建议优先考虑分词器扩展,因为其对下游模型的影响难以预测。与以往工作不同,作者详细分析了扩展分词器时的重建与生成困境及其解决方案。

自回归视觉生成。自回归视觉生成模型[15, 33, 38, 40, 49, 50, 56, 58, 59, 65]遵循大语言模型(LLMs)的下一个词预测方法,使其能够利用LLMs的进展并简化统一多模态生成的路径。其他使用离散 Token 器的方法结合了视觉特定的范式,如 Mask 图像建模[8, 60, 68, 69]和下一尺度预测[36, 53],以提升性能。作者专注于自回归范式,并揭示扩大 Token 器规模有助于使AR模型与当前最佳视觉生成模型相媲美。

视觉模型语义引导。近期研究探索了利用视觉基础模型的引导[7, 23, 43, 46, 71]来增强压缩视觉潜空间,赋予其更丰富的语义[9, 10, 18, 36, 62-64, 67, 72, 75, 76],以及提升视觉生成模型中的表征[70]。与这些方法不同,作者分析了分词器扩展中的困境,并强调了语义正则化在有效扩展中的关键作用。

  1. 初步研究

作者首先将AR Probing作为 Agent 方法,用于有效监控分词器在下游生成任务中的有效性(第3.1节),随后进行一项初步实验,探究在盲目扩展视觉分词器时,重建与生成任务所面临的挑战(第3.2节)。

3.1. 评估Tokenizer的AR检测

在自回归视觉生成中,分词器的训练和下游自回归模型的训练在不同的阶段进行。在第一阶段,一个视觉分词器被训练以将图像压缩成离散的符号,该训练过程通过重建目标进行优化。在第二阶段,基于预训练分词器生成的离散符号来训练下游生成模型。然而,在第一阶段表现出色(在重建保真度方面)的分词器不一定能带来下游生成模型更好的性能。因此,在评估训练好的分词器的重建质量的同时,对其在下游生成中的有效性进行评估至关重要。

尽管其重要性,评估分词器如何影响下游生成模型可能计算成本高昂。例如,在64块V100 GPU上充分训练一个343M参数的下游AR生成器需要170小时。为应对这一挑战,作者引入了AR Probing,灵感来源于表征学习文献中的线性检测[11, 23]。其核心思想是利用小规模AR模型的性能作为 Agent ,反映大规模AR模型的性能趋势。

具体而言,作者使用分词器训练一个小的Llamastyle模型[50, 54](1.11亿参数),训练50个epoch,并评估其gFID[24]、验证损失和线性检测准确率[11, 23],以实现不同分词器之间的公平比较。使用所提出的AR检测模型评估分词器的效率比训练原始的343M下游AR模型高10倍。作者在第5.1节(图6)中的实验表明,AR检测观察到的趋势与大规模AR模型在充分训练后的性能一致。

picture.image

gFID. 生成FID [24] 是AR检测中两阶段框架整体图像生成性能的指标。它反映了分词器的重建保真度,以及下游AR检测模型学习视觉 Token 依赖关系的能力(即 Token 分布的可学习性)。

验证损失。作者使用AR检测模型的验证损失来衡量潜在 Token 作为解耦因子的可学习性。验证损失是在ImageNet [48] 50k验证集上,按照下一个 Token 预测范式计算的 Token 级交叉熵损失的均值。在相同的词汇量、相同的视觉 Token 数量和结构以及相同的AR检测模型下,较大的验证损失表明AR模型更难学习该潜在空间。因此,作者使用验证损失来反映AR模型的潜在空间复杂性和可学习性。

线性检测精度。除了视觉生成质量之外,作者还研究了扩展分词器是否会导致AR模型学习到更好的视觉表征,这可能为未来使用AR模型进行统一多模态理解和生成的研究提供启示。为了评估AR模型的表征质量,作者采用了标准做法[11, 23],使用AR检测模型中间Transformer层的特征来评估线性检测精度。

3.2. 简单扩展分词器并不奏效

为研究在简单扩展视觉分词器时面临的挑战,作者在ImageNet [48]数据集上以

分辨率训练了三个不同规模的矢量量化分词器。如图3所示,随着分词器规模的增加,虽然重建质量(rFID)持续提升,但生成性能(gFID)显著下降。这凸显了分词器扩展中的重建与生成困境。

picture.image

此外,作者观察到AR检测的验证损失随着分词器规模的扩大而持续增加,表明更大的分词器会导致更复杂的分词依赖关系,而AR模型更难以学习。这一观察促使作者在第4.2节设计了语义正则化方法,以约束分词器的潜在空间复杂度,从而打破重建与生成困境。

  1. GigaTok

在本节中,作者介绍了作者可扩展视觉分词器GigaTok的模型结构和训练策略。在4.1节中,作者提出了一种支持1D和2D分词结构的分词器主干,并讨论了编码器和解码器的非对称缩放策略。在4.2节中,作者引入了语义正则化,通过用预训练的视觉表示正则化潜在空间的复杂度来打破重建与生成之间的困境。在4.3节中,作者展示了熵损失[68]如何促进十亿规模分词器的收敛。

4.1. 架构

在当前文献中,CNN [32] 架构是图像分词器 [15, 40, 68, 75] 的主流选择,因为它们在捕捉细粒度局部细节方面非常有效。然而,Transformer 是更具可扩展性的架构,且具有较少的归纳偏差。因此,作者设计了一种混合架构的向量量化(VQ)分词器主干,该架构结合了CNN [15, 32] 和Transformer [6, 13, 57] 用于编码器和解码器(图4)。

picture.image

具体而言,作者的编码器由一系列CNN模块组成,这些模块逐步将输入图像下采样一个因子

,随后是Transformer层和向量量化器,以生成离散的潜在码。类似地,作者的解码器由多个Transformer层组成,随后是CNN解码器,这些解码器将特征上采样以获得重建图像。作者的分词器架构可以通过利用接下来两段中介绍的不同Transformer设计,适应1D和2D分词器。

基于ViT的2D tokenizer。对于2D tokenizer,其tokenizer编码器和解码器中的Transformer均采用ViT[13]架构实现。在tokenizer中,潜在特征的2D结构和token的2D结构得以保持。

1D tokenizers与Q-Former。对于1D tokenizers,作者在编码器和解码器中都实现了Transformer模块作为Q-Former [6, 34]。编码器中的Q-Former采用1D Query ,将2D输入特征转换为1D潜在token。解码器中的Q-Former使用2D Query ,将1D潜在token转换回2D特征,然后这些特征被传递给CNN解码器以重建图像。1D tokenizers消除了2D归纳偏差,并且在作者的实验中(第5.5节)表现出比2D tokenizers更好的可扩展性。

非对称编码器-解码器缩放。由于解码器面临从有损潜在码重建图像这一更具挑战性的任务,作者采用非对称设计以实现更高效的参数分配。具体而言,作者对编码器和解码器进行缩放,同时确保解码器始终大于编码器。在实践中,作者保持CNN编码器/解码器的尺寸相同且固定,仅增加Transformer模块的深度和宽度以实现缩放。

4.2. 语义正则化

在作者的初步研究(第3.2节)中,随着分词器的扩展,潜在空间的复杂度显著增加,这可能导致较大分词器在下游AR生成方面表现更差。作者假设较大分词器倾向于捕捉过多的细粒度Low-Level细节以实现更好的重建效果,从而造成潜在 Token 分布过于复杂,使得AR模型更难有效学习 Token 依赖关系。

为解决此问题,作者引入语义正则化来指导分词器编码一个更语义一致的潜在空间,该空间更简单且便于下游生成式建模。具体而言,作者在分词器训练目标中引入了一个简单的语义正则化项。该正则化项使分词器解码器的中间特征与从预训练的冻结DINOv2 [43]中提取的特征表示相一致。

从数学角度而言,设

为Transformer解码器第

层的输出特征,

为预训练图像编码器(此处为DINOv2-B [43])的语义特征,则语义正则化可表示为:

其中

是批次大小,

是图像索引,

是余弦相似度函数,

是一个将解码器特征

投影到匹配通道维度的多层感知机。

。在训练VQ tokenizers时,作者将语义正则化添加到原始VQGAN [15, 50]目标中。

在本工作中,作者经验性地将

。这里

是多种损失的组合,包括

,图像像素的

重建损失,

,感知损失 [27, 73],

,PatchGAN [26] 对抗损失,以及

,VQ码本损失。

4.3. 十亿级分词器的熵损失

在训练一个2.9B的tokenizer时,作者发现使用与622M tokenizer相同的训练方案会导致感知损失和重建损失均无法收敛,并且codebook使用率始终偏低。作者假设codebook使用率低是导致收敛困难的原因。为此,作者引入熵惩罚[66, 68]来鼓励更高的codebook利用率:

其中

表示香农熵,

是量化器的输入,被量化为

,而

是第

个码本向量。

通过惩罚量化中的不确定性来减少量化误差,而

则鼓励在整个码本中更均匀地选择码本向量。详细的推导过程可以在作者的补充材料中找到。作者发现熵惩罚解决了大型分词器的收敛困难问题。

picture.image

如图5所示,将熵损失引入2.9B分词器后,码本的使用率能够迅速达到较高水平,并且损失收敛良好。

  1. 实验

5.1. 设置

为扩展视觉分词器,作者遵循GigaTok分词器中Transformer的架构配置,如表1所示。作者从三个角度评估分词器:重建、下游AR生成和下游AR表示质量。作者使用rFID和LPIPS[73]评估重建保真度,使用gFID评估生成性能,并使用线性检测评估下游AR模型的表示质量。作者的下游AR模型是带有1D绝对位置嵌入的LlamaGen[50]。作者的扩展实验(第5.2节)和消融研究(第5.3节)使用AR检测(第3.1节中描述的111M AR模型)验证损失、gFID和线性检测来分别反映分词的学习能力、生成性能和表示质量。而在系统级比较(第5.4节)中,作者训练更大的1.4B AR模型以与前人工作进行比较。更多细节请参见补充材料。

picture.image

AR检测的有效性。如图6所示,AR检测性能包括gFID和线性检测精度与更大的LlamaGen-XL [50]模型结果一致。因此,在以下实验中,除了系统级比较外,作者始终使用AR检测。

5.2. 基于语义正则化的扩展

作者证明了所提出的语义正则化方法解决了在扩展型分词器中重建与生成之间的困境。

基于语义正则化的模型缩放。结果如图7所示。(1) 语义正则化提高了重建保真度,这体现在更低的rFID。(2) 更重要的是,在未使用语义正则化的情况下,更大规模的tokenizer会导致AR Probing验证损失和gFID下降,显示出重建与生成之间的矛盾。语义正则化解决了这一矛盾,验证损失相对受限且gFID持续下降证明了这一点。(3) 线性 Prob 结果表明,随着tokenizer模型的规模扩大,语义正则化有助于AR模型学习到更好的表征。

picture.image

可视化分词器特征空间。作者可视化了一组图像中来自第一个Transformer解码器层的分词器特征的前3个PCA分量。如图8所示,作者发现普通的分词器编码了一个语义一致性有限的潜在空间,这可能会损害其对于下游AR模型的可学习性。相比之下,GigaTok呈现了语义一致的模式(图8),表明这是一个有意义且一致的潜在空间。

picture.image

5.3. 非对称1D Tokenizer更具可扩展性

分词器解码器需要更多参数。为了确定在扩展规模时应该优先考虑解码器还是编码器,作者在表2中比较了

和 B-S 分词器,两者均在相同设置下训练了100个epoch。作者的结果表明,扩展解码器而非编码器能够带来在重建和下游生成方面的更大改进,这表明解码器扩展应该优先考虑。

picture.image

扩展分词器编码器同样重要。虽然优先扩展分词器解码器能带来显著效益,但作者也发现扩展分词器编码器能够进一步提升下游模型的表现。在表2中,作者展示了与S-L分词器相比,B-L分词器获得了显著改进。因此,作者建议同时扩展编码器和解码器,并保持解码器比编码器更大,以实现最佳性能。

5.4. 系统级比较

一维分词器比二维分词器更具可扩展性。作者使用相同的语义正则化设置训练了S-S、S-B和B-L的一维/二维分词器。如图9所示,一维分词器始终在rFID和AR Probing线性检测准确率上优于二维分词器。在AR Probing gFID方面,一维分词器表现出更陡峭的扩展趋势,随着模型规模的增大,最终超越了二维分词器。作者将一维分词器的可扩展性优势归因于其减少的归纳偏差。

picture.image

实验设置。使用GigaTok进行分词,作者将LlamaGen [50] AR模型在

ImageNet训练集上进行300个epoch的训练,以与其他方法进行比较。作者不使用AdaLN [44, 53],因为它专门用于类条件生成。作者提供了使用DINO判别器 [36, 53] 训练的B-L分词器的结果,以公平地比较rFID。但在实践中作者发现,DINO判别器对LPIPS的提升有限,并可能影响十亿规模分词器的训练稳定性。因此,作者从主要设计中排除了它。

结果。如表3所示,作者的2.9B GigaTok在所有离散分词器中实现了最先进的重建性能(rIFD)。此外,使用作者的2.9B分词器,下游的1.4B AR模型在LLM风格的自动回归下一词预测模型中实现了最先进的图像生成性能(gFID)。VAR [53] 采用下一尺度预测而非下一词预测来预测图像,这与语言模型不太兼容。GigaTok在不包含VAR等视觉特定设计的简单LLM风格下游AR生成器的情况下,实现了与VAR [53] 相当的gFID。此外,在2.9B分词器上训练的1.4B AR模型通过视觉生成预训练实现了最先进的线性检测精度5。这表明作者的GigaTok有助于下游生成模型学习更好的表示。从生成预训练中学习的高质量表示也可能有助于未来原生多模态模型统一生成和理解。

picture.image

5.5. 讨论与消融实验

关于生成成本讨论。在生成图像时,AR模型需要多次预测token,而tokenizer只需要一次前向传递。因此,与AR模型相比,将token解码为图像的时间消耗相对较小。作者在表4中记录了不同tokenizer/AR模型在tokenizer解码上花费的时间比例。对于1.4B的AR模型,作者最大的2.9B tokenizer仅占总推理时间的

picture.image

搜索最佳语义正则化层。作者搜索

,即Transformer解码器在提取中间特征以计算公式1中的语义正则化之前的层索引。如表5所示,改变

在AR检测的gFID和线性精度之间存在权衡。较小的

意味着对潜在空间的正则化更严格,因此下游生成模型能学习到更好的表示。然而,较小的

也牺牲了生成质量。作者选择

,以在所有分词器中获得更平衡的rFID、gFID和线性检测精度。

picture.image

探索预训练语义编码器的选择。作者将CLIP-B (DFN) [16, 46]、SigLIP-400M [71]和DINOv2-B [43]作为S-B分词器的语义正则化来源进行比较。如表6所示,使用DINOv2-B作为语义编码器进行正则化,能够产生最佳的重建分词器、下游条件生成和表示质量。

picture.image

探索语义正则化的权重。作者研究了不同正则化权重

(公式2)的影响,范围从0.25到1.00。如表7所示,较大的

(0.75, 1.00)会损害重建质量,但有利于线性检测精度,而较小的

(0.25)会导致次优的rFID和线性检测精度。作者选择更均衡的

作为所有分词器的默认值。

picture.image

  1. 结论

在本工作中,作者研究和解决了视觉 Token 器扩展中的重建与生成困境。作者发现这一困境源于较大 Token 器中潜在空间的复杂性增加。作者提出了语义正则化方法,通过注入预训练表示来有效正则化 Token 器潜在空间,使其与训练中的 Token 器特征对齐。语义正则化方法,结合作者探索的若干关键实践,成功研发出首个30亿参数 Token 器GigaTok,该 Token 器在重建、下游AR生成以及下游AR表示质量方面均达到了当前最佳水平。有关局限性和未来工作的讨论,请参阅补充材料。

参考

[1]. GigaTok: Scaling Visual Tokenizers to 3 Billion Parameters for Autoregressive Image Generation

picture.image

扫码加入👉「集智书童」交流群

(备注:方向+学校/公司+昵称)

picture.image

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
CV 技术在视频创作中的应用
本次演讲将介绍在拍摄、编辑等场景,我们如何利用 AI 技术赋能创作者;以及基于这些场景,字节跳动积累的领先技术能力。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论