想象一下你正在试图教一个机器人画画。为了做到这一点,机器人需要一本“字典”来理解图像是由什么组成的。在 AI 图像生成的领域中,这本字典被称为分词器(Tokenizer)。
长期以来,研究人员必须在两种截然不同的字典类型之间做出选择,而两者都不完美:
- “平滑型”字典(连续型/VAE): 把它想象成一幅水彩画。它极其细腻且平滑,捕捉到了每一个微小的细节。然而,所有的颜色都混合在一起,变成了一大滩色彩斑斓的液体。如果你想让机器人只改变“天空”而不影响“草地”,这会非常困难,因为信息全部纠缠在一起了。
- “块状型”字典(离散型/VQ): 把它想象成一套乐高积木。它将图像分解为不同的、独立的块(代码)。这有助于机器人学习模式并循序渐进地构建事物(就像写故事一样)。但这些积木往往不稳定。有时机器人会忘记如何使用某些积木,或者积木会堆叠在一起,导致画面看起来模糊或出现“坍塌”。
核心理念:MergeTok
这篇论文的作者们提出了 MergeTok,他们问道:“为什么我们不能在同一个字典中既拥有水彩画的平滑感,又拥有乐高的结构感呢?”
他们构建了一个新的系统,这个系统就像一个能同时说“平滑语”和“块状语”的双语翻译员。他们不仅仅是将两个系统拼凑在一起,而是通过一种被称为**分词合并(Token Merging)**的巧妙技巧,在两者之间架起了一座桥梁。
它是如何工作的:“分组”类比
想象一下你正在组织一场拥有 1,000 名宾客(即图像的像素)的大型派对。
- 问题所在: 如果你试图单独与每一位宾客交谈,过程会变得混乱且低效。如果你只是随机地将他们分组,你就会丢失重要的细节。
- MergeTok 的解决方案: 系统配备了一位聪明的保镖(分词合并算法),他观察着宾客并说道:“你们三个看起来都穿着红色的衬衫,而且都在聊体育运动。让我们把你们作为一个‘体育队’单元组合在一起。”
这种分组以两种方式进行,以帮助机器人学习:
- 针对“平滑侧”(VAE): 系统告诉水彩画家:“嘿,这三个人是一个‘体育队’。当你绘画时,确保让他们看起来像一个凝聚的整体,而不仅仅是随机的红色色块。”这迫使平滑的绘画进行逻辑性的自我组织,从而在以后更容易进行控制。
- 针对“块状侧”(VQ): 系统告诉乐高搭建者:“既然我们知道这三个人是一个团队,那就给他们三个不同的乐高积木,这样他们就不会看起来全都一样;但要确保没有其他团队会占用这些特定的积木。”这防止了乐高积木的坍塌或过度重复。
魔法之桥
其秘诀在于,这个“体育队”名单(称为源图谱 Source Map)是创建一次后便共享使用的。
- 它帮助平滑侧学会组织有序。
- 它帮助块状侧学会稳定与多样化。
最棒的部分在于,真正画画的机器人(生成器)甚至不知道发生了这种分组。它看到的只是一个整洁的 256 个分词列表,随时待命。这就像保镖在幕后完成了所有的繁重工作,让艺术家可以全身心地投入到绘画之中。
研究发现
研究人员在大型数据集(ImageNet)上对该系统进行了测试。结果如下:
- 更好的图像质量: 与以往最好的“平滑型”或“块状型”系统相比,该系统重建图像的质量更高(rFID 分数更低)。
- 更智能的组织: 与旧方法相比,它创建的分词能更好地理解图像的含义(例如区分猫和狗)。
- 多功能性: 由于它是一个统一的系统,它可以完美适配两种不同类型的 AI 艺术家:一种是按步骤构建图像的(自回归模型),另一种是通过清理噪声来构建图像的(扩散模型)。
总结
MergeTok 就像一位高级组织者,它将混乱的人群进行逻辑化的分组,并将结果交给艺术家。它解决了人们必须在“平滑但混乱”与“有结构但不稳定”之间做选择的老问题。通过在学习过程中将相似的信息合并在一起,它创造了一个既高质量又易于 AI 控制的超高效字典。
技术摘要:MergeTok
问题陈述
当前的图像生成视觉分词器(visual tokenizers)分为两个截然不同的家族,每一类都面临着互补的结构性局限:
- 连续型 VAE 分词器: 虽然提供了高保真度的重建能力,但其依赖的潜空间缺乏内在的组织性。在强压缩下,语义因子会纠缠在密集的潜表示中,导致生成的编码解耦性较弱,阻碍了语义可控性。
- 离散型 VQ 分词器: 虽然能够实现自回归生成,但面临优化瓶颈。量化过程的不可微特性导致梯度更新稀疏且不均匀,使得许多码本(codebook)条目接收到的监督不足。这增加了码本崩溃的风险,并降低了表示效率。
这些局限性在重建保真度、语义结构、效率和训练稳定性之间造成了持久的张力。
方法论:MergeTok 框架
MergeTok 提出了一种统一的双分支分词器,旨在单个编码器-解码器架构内联合优化连续(VAE)和离式(VQ)分词。其核心创新在于使用**令牌合并(Token Merging, ToMe)**作为两个分支之间的语义桥梁。
架构概述
给定输入图像 X,共享的 CNN 编码器 Ec 生成特征图 Z。该特征图由两个并行分支处理:
VAE 分支(连续型):
- 令牌合并: 一个配备了 ToMe 模块的注意力机制编码器将令牌序列从 L 个压缩为 K 个(K<L),通过合并相似的令牌来实现。这一过程生成了一个浓缩的表示 ZK(vae) 和一个记录原始令牌分配到合并组情况的源映射(source map) S。
- 语义对齐: 合并后的令牌通过与一个配备了 ToMe 的冻结 DINO 式教师模型进行余弦距离损失(Lalign)对齐。这使连续潜空间向具备语义感知且解耦的表示方向正则化。
- 重建: 混合解码器利用浓缩令牌和源映射 S 来“取消合并”这些组,从而重建像素空间的图像。
VQ 分支(离散型):
- 组感知量化: 该分支在正向传播期间绕过 ToMe 操作,以保留完整的序列长度(L 个令牌)。然而,它利用 VAE 分支生成的源映射 S 作为结构先验。
- 正则化: 源映射将 L 个令牌划分为 K 个语义组。针对码本分配应用了两个组感知损失:
- 组内多样性(Ldiv): 最大化每个组内的熵,以防止崩溃到单一代码。
- 组间排他性(Lcons): 最小化组间的重叠,以确保不同的语义组使用不同的代码子集。
- 重建: 标准 VQ 解码器从量化后的令牌中重建图像。
训练策略
- 统一目标函数: 模型通过结合重建、对抗和合并感知正则项的总体损失进行端到端训练:
Ltotal=λvaeLVAE+λvqLVQ+λaliLalign+λdivLdiv+λconsLcons
- 粒度感知采样: 在训练期间,合并比例(以及由此产生的保留令牌数 K)从以超参数 K∗ 为中心的离散分布中进行动态采样。这使编码器能够接触到变化的瓶颈宽度,从而鼓励在不同语义粒度下的鲁棒性。
- 推理: 在推理阶段,分词器始终输出一个包含 256 个令牌的固定序列。下游生成器(无论是自回归还是扩散模型)都在此固定序列上运行,且对内部合并结构是不可知的。
核心贡献
- 统一双分支分词器: MergeTok 是首个在单一目标下联合优化连续和离散潜空间的架构,利用令牌合并作为语义接口,增强了 VAE 训练的稳定性和 VQ 的结构优势。
- 合并感知训练约束: 作者引入了两个源自合并过程的新颖目标:
- 合并令牌对齐: 提升了连续潜空间的语义结构。
- 组感知量化: 通过组内多样性和组间排他性稳定了 VQ 训练并提高了码本利用率。
- 粒度感知合并比例采样: 一种训练时的策略,通过采样不同的令牌粒度,使模型能够学习到在不同压缩率下均具有鲁棒性的统一分词。
实验结果
实验在 ImageNet-256(256×256 分辨率)上进行,固定令牌预算为 256 个。
- 重建保真度:
- 离散型 (VQ): MergeTok-BL 实现了 0.78 的 rFID,优于强基线模型 GigaTok-BL (0.81) 和 MergeVQ-GR (1.12)。在使用冻结 DINO 分辨器的情况下,MergeTok-BL* 达到了 0.48 rFID。
- 连续型 (VAE): MergeTok-BL 实现了 0.47 的 rFID,略微优于 RAE (0.49),并显著击败了 SoftVQ-L (0.61) 和 DC-AE-f32 (0.69)。
- 表示质量:
- MergeTok-BL 达到了 78.3% 的线性探测准确率,是所列离散分词器中最高的,超过了 UniTok (70.8%) 和 MergeVQ-GR (77.9%)。这表明其具有卓越的语义组织能力。
- 下游生成:
- 自回归 (VQ 分支): 与 LlamaGen-XXL 配对,MergeTok-BL 实现了 1.93 (无 CFG) 和 2.13 (有 CFG) 的 gFID,优于 GigaTok-BL (2.03)。
- 扩散/流 (VAE 分支): 与 SiT-XL/1 配对,MergeTok-BL 实现了 1.79 (无 CFG) 和 1.29 (有 CFG) 的 gFID,达到了表中的最佳水平。
意义与主张
论文声称 MergeTok 证明了单一架构可以赋予视觉分词器既有的鲁棒语义组织和对生成器友好的离散性。
- 统一性: 它解决了 VAE 的高保真重建与自回归生成所需的离散结构之间的权衡。
- 语义桥梁: 通过利用令牌合并作为结构先验,模型成功地将语义分组信息从连续分支转移到离散分支以稳定 VQ,反之亦然。
- 兼容性: 该分词器生成的 256 令牌序列可以兼容自回归(如 LlamaGen)和扩散/流(如 DiT, SiT)生成器,无需对生成器本身进行修改。
作者将 MergeTok 定位为不仅是一个重建工具,更是一个统一的接口,能够同时支持离散自回归生成、连续扩散生成以及具有语义结构的表示学习。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。