又一爆款！港大&字节开源GigaTok首次破解图像重建与生成矛盾，30亿参数登顶SOTA！

文章链接：https://arxiv.org/pdf/2504.08736 
项目链接：https://silentview.github.io/GigaTok/ 
代码链接：https://github.com/SilentView/GigaTok 

亮点直击

• 分词器规模扩展中存在的重建与生成困境源于更大规模分词器隐空间复杂度的提升。本文提出语义正则化方法，有效缓解该困境从而实现分词器的规模扩展。
• 探索了分词器规模扩展的最佳实践方案，包括：采用混合CNN-Transformer架构的一维分词器、非对称编码器-解码器扩展策略，以及针对十亿级分词器的熵损失优化。
• GigaTok成为首个规模扩展至30亿参数的分词器，在ImageNet数据集上同步实现了重建质量、下游自回归生成和下游自回归表征三项指标的突破性性能。

2.9 B GigaTok 在 ImageNet 256 × 256 分辨率上以 1.4 B AR 模型实现了 SOTA 自回归图像生成

总结速览

解决的问题

• 视觉分词器（visual tokenizer）缩放困境：现有方法在扩大视觉分词器规模时，虽然提升了图像重建质量，但导致下游自回归生成质量下降（即“重建 vs. 生成”矛盾）。
• 隐空间复杂性增长：分词器规模扩大时，隐空间复杂度增加，使得下游自回归模型难以有效学习离散token的分布。
• 大规模分词器训练稳定性：亿级参数分词器的训练面临收敛困难问题。

提出的方案

• 语义正则化（Semantic Regularization）：

a.利用预训练视觉编码器（如DINOv2）的特征对齐分词器特征，约束隐空间复杂度，避免生成过于复杂的token依赖关系。

• 分词器缩放关键策略：

a.1D分词器：相比2D分词器，1D结构更具可扩展性。

b.非对称模型缩放：优先扩展解码器而非编码器。

c.熵损失（Entropy Loss）：用于稳定亿级参数分词器的训练。

应用的技术

• 混合架构：基于CNN-Transformer的向量量化（VQ）分词器，支持1D/2D token生成。
• 自回归探测（AR Probing）：通过轻量级下游AR模型监控分词器训练过程，分析token分布学习难度。
• 预训练模型特征对齐：使用DINOv2等模型的语义特征作为正则化目标。

达到的效果

• 性能提升：

a.重建质量：达到SOTA图像重建效果。

b.生成质量：下游自回归生成任务性能显著提升。

c.表征学习：下游AR模型的表征能力增强。

• 规模化能力：成功将GigaTok扩展至30亿参数，突破“重建 vs. 生成”矛盾。
• 训练稳定性：通过熵损失实现亿级参数分词器的稳定训练。

先导研究

本文首先引入自回归探测(AR Probing)作为有效监控分词器对下游生成效果的代理方法，随后通过先导实验研究当简单扩展视觉分词器时遇到的重建与生成挑战。

用于分词器评估的自回归探测

在自回归视觉生成中，分词器和下游AR模型的训练分两个独立阶段进行。第一阶段训练视觉分词器将图像压缩为离散token，以重建目标进行优化。第二阶段基于预训练分词器产生的离散token训练下游生成模型。然而，在第一阶段重建效果良好的分词器，并不一定能带来下游生成模型的更好表现。因此，在评估训练好的分词器时，除了重建质量外，评估其对下游生成的有效性同样至关重要。

尽管很重要，但评估分词器如何影响下游生成模型的计算成本可能很高。例如，充分训练一个3.43亿参数的下游AR生成器需要在64块V100 GPU上耗时170小时。为解决这一挑战，受表征学习中线性探测(Linear Probing)的启发，引入了自回归探测方法。其核心思想是使用小型AR模型的性能作为反映大规模AR模型性能趋势的代理指标。

本文使用分词器训练一个小型Llama风格模型(1.11亿参数)50个周期，并评估其gFID、验证损失和线性探测准确率，以便在不同分词器之间进行公平比较。训练所提出的自回归探测模型来评估分词器的效率比训练原始3.43亿下游AR模型高10倍。下图6实验表明，自回归探测观察到的趋势与充分训练后大规模AR模型的性能表现一致。

gFID。自回归探测的生成FID表明两阶段框架的整体图像生成性能。它既反映了分词器的重建保真度，也反映了下游自回归探测模型学习视觉token依赖关系的能力(即token分布的可学习性)。

验证损失。本文使用自回归探测模型的验证损失作为潜在token可学习性的解耦衡量指标。验证损失计算为ImageNet 50k验证集上基于下一token预测范式中各token交叉熵损失的平均值。在相同词汇量、相同视觉token数量和结构、相同自回归探测模型的情况下，更大的验证损失表明隐空间对AR模型来说更难学习。因此，我们用验证损失来反映AR模型的隐空间复杂度和可学习性。

线性探测准确率。除了视觉生成质量外，本文还研究扩展分词器是否会促使AR模型学习到更好的视觉表征，这可能为未来使用AR模型进行统一多模态理解和生成的研究提供启发。为评估AR模型的表征质量，本文采用标准做法，使用自回归探测模型中间Transformer层的特征评估线性探测准确率。

简单扩展分词器不可行

为研究简单扩展视觉分词器时面临的挑战，在ImageNet 256×256分辨率下训练了三个规模递增的向量量化分词器。如下图3所示，随着分词器规模增大，虽然重建质量(rFID)持续提升，但AR生成性能(gFID)显著下降。这凸显了分词器扩展中的重建与生成困境。此外，我们观察到自回归探测的验证损失随分词器规模扩大持续增加，表明更大分词器会导致更复杂的token依赖关系，使AR模型更难学习。这一发现促使我们设计语义正则化方法来约束分词器的隐空间复杂度，从而解决重建与生成困境。

GigaTok

本节介绍可扩展视觉分词器GigaTok的模型结构和训练策略。首先提出支持1D和2D token结构的分词器主干，并讨论编码器-解码器的非对称扩展策略。然后引入语义正则化，通过用预训练视觉表征约束隐空间复杂度来解决重建与生成困境。最后展示熵损失如何促进十亿级规模分词器的收敛。

架构

当前研究中，CNN架构因其捕捉细粒度局部细节的有效性成为图像分词器的主流选择。而Transformer则是更具扩展性且归纳偏置更少的架构。因此，本文设计了结合CNN和Transformer的混合架构向量量化(VQ)分词器主干(下图4)。编码器包含系列CNN块(将输入图像逐步下采样p倍)，后接Transformer层和向量量化器生成离散潜在编码。解码器包含多个Transformer层，后接CNN解码器(将特征上采样获得重建图像)。通过采用下文介绍的两种Transformer设计，我们的分词器架构可适配1D和2D分词器。

2D分词器(ViT实现) ：对2D分词器，编解码器的Transformer均采用ViT架构实现，全程保持潜在特征和token的2D结构。

1D分词器(Q-Former实现) ：对1D分词器，编解码器的Transformer模块均实现为Q-Former。编码器的Q-Former使用1D查询，将2D输入特征转为1D潜在token；解码器的Q-Former使用2D查询，将1D潜在token转回2D特征后输入CNN解码器重建图像。1D分词器消除了2D归纳偏置，实验表明其扩展性优于2D分词器。

语义正则化

在先导研究中，发现随着分词器规模的扩大，隐空间复杂度显著增加，这可能导致更大规模的分词器在下游AR生成任务中表现更差。假设，更大的分词器倾向于捕捉过多细粒度的低级细节以获得更好的重建效果，从而导致潜在token分布过于复杂，使得AR模型更难有效学习token依赖关系。

为解决这一问题，本文引入了语义正则化方法，引导分词器编码出语义更一致的隐空间——这种空间复杂度更低，更有利于下游生成建模。在分词器训练目标中加入了一个简单的语义正则化项。该正则化将分词器解码器的中间特征与预训练冻结DINOv2模型提取的特征表示对齐。

十亿级分词器的熵损失

在训练 29 亿参数的分词器时，发现使用与 6.22 亿参数分词器相同的训练方案会导致感知损失和重建损失无法收敛，并持续出现码本使用率低下的问题。我们假设低码本使用率是造成收敛困难的原因。为解决这个问题，本文引入熵惩罚来提升码本利用率：

实验

实验设置

为扩展视觉分词器，本文采用下表1总结的GigaTok分词器Transformer架构配置。从三个维度评估分词器：重建质量、下游AR生成质量和下游AR表征质量。使用rFID和LPIPS评估重建保真度，gFID评估生成性能，线性探测评估下游AR模型的表征质量。我们的下游AR模型采用具有1D绝对位置编码的LlamaGen。扩展实验和消融研究使用自回归探测的验证损失、gFID和线性探测分别反映token可学习性、生成性能和表征质量。系统级对比则训练14亿参数AR模型与之前工作进行对比。

自回归探测的有效性。如前面图6所示，自回归探测的gFID和线性探测准确率等指标与更大规模的LlamaGen-XL模型结果一致。因此除系统级对比外，后续实验均采用自回归探测。

语义正则化的扩展效果

本文证明所提出的语义正则化解决了分词器扩展中的重建与生成困境。

带语义正则化的模型扩展。结果如下图7所示：(1)语义正则化通过降低rFID改善了重建保真度；(2)更重要的是，无语义正则化时，自回归探测验证损失和gFID随分词器增大而恶化，显示重建与生成困境。语义正则化通过相对受限的验证损失和持续下降的gFID解决了该困境；(3)线性探测结果表明，语义正则化帮助AR模型在分词器扩展时学习到更好的表征。

分词器特征空间可视化。本文对一组图像的首个Transformer解码层特征进行前三主成分可视化。如下图8所示，普通分词器编码的隐空间语义一致性有限，可能损害下游AR模型的可学习性；而GigaTok展现出语义一致的模式(图8)，表明其隐空间具有意义和一致性。

非对称1D分词器更具可扩展性

解码器需要更多参数。为确定扩展时应优先扩展编码器还是解码器，在下表2中比较了S-B4和B-S分词器（两者在相同设置下训练100个epoch）。结果表明，相比编码器，扩展解码器能带来更大的重建和下游生成性能提升，说明应优先扩展解码器。

编码器扩展同样重要。虽然优先扩展解码器收益显著，但发现扩展编码器也能进一步提升下游模型性能。上表2显示，B-L分词器相比S-L分词器获得显著提升。因此建议同时扩展编码器和解码器，但保持解码器规模大于编码器以获得最优性能。

1D分词器比2D分词器更具扩展性。在相同设置下训练了S-S、S-B和B-L规模的1D/2D分词器（均使用语义正则化）。如下图9所示，1D分词器始终获得更好的rFID和自回归探测线性分类准确率。在自回归探测gFID指标上，1D分词器展现出更陡峭的扩展趋势，最终在模型规模扩大时超越2D分词器。我们将1D分词器的优越扩展性归因于其更低的归纳偏置。

系统级对比

实验设置：使用GigaTok进行token化，在256×256 ImageNet训练集上对LlamaGen AR模型进行300个epoch的训练，与其他方法进行对比。未采用AdaLN方法，因其专为类别条件生成设计。为公平比较rFID，提供了使用DINO判别器训练的B-L分词器结果。但实际发现DINO判别器对LPIPS提升有限，且可能影响十亿级分词器的训练稳定性，故未将其纳入主要设计。

实验结果：如下表3所示，29亿参数GigaTok在所有离散分词器中实现了最优的重建性能(rIFD)。此外，配合29亿参数分词器，下游14亿参数AR模型在LLM风格的自回归下一token预测模型中取得了最佳的图像生成性能(gFID)。VAR采用下一尺度预测而非下一token预测，与语言模型兼容性较差。本文的模型仅使用简单的LLM风格下游AR生成器（不含VAR等视觉专用设计）就达到了与VAR相当的gFID。更重要的是，基于29亿分词器训练的14亿AR模型通过视觉生成式预训练获得了最优的线性探测准确率，表明GigaTok能帮助下游生成模型学习更好的表征。这种通过生成式预训练获得的高质量表征，可能有助于未来原生多模态模型统一生成与理解能力。

讨论与消融研究

生成成本讨论

在图像生成过程中，AR模型需多次前向传播预测token，而分词器仅需单次前向传播。因此，相较于AR模型，分词器将token解码为图像的时间占比极小。表4记录了不同分词器/AR模型中分词器解码的时间占比：对于14亿参数AR模型，我们最大的29亿参数分词器仅占总推理时间的约10%。

语义正则化最佳层搜索

探索了Transformer解码器中用于计算语义正则化（公式1）的中间

预训练语义编码器选择

本文比较了CLIP-B (DFN)、SigLIP-400和DINOv2-B作为S-B分词器的语义正则化源。下表6显示：

• 采用DINOv2-B作为语义编码器时，分词器在重建、下游类别条件生成和表征质量上均表现最佳

语义正则化权重探索

结论

本文研究并解决了视觉分词器扩展过程中的重建与生成困境。发现该困境源于更大规模分词器中不断增加的隐空间复杂度。通过注入预训练表征并与分词器特征对齐，提出语义正则化方法来有效约束分词器隐空间。这种语义正则化结合我们探索的若干关键实践，最终实现了首个30亿参数的分词器GigaTok，在重建质量、下游AR生成性能和下游AR表征质量上均达到SOTA水平。

本文转自AI生成未来 ，作者：AI生成未来

原文链接:​​https://mp.weixin.qq.com/s/qRevutdSGTOBrdqYOmeAxQ​​