想象一下你正在尝试拼凑一个巨大的拼图,但你没有看到包装盒上的图案,而是必须从零开始猜测每一块碎片的颜色和形状。在人工智能的世界里,这就是计算机试图创建图像或理解复杂场景时的情景。长期以来,科学家们一直在使用两种主要的工具:一种像作家打字写句子一样,通过逐块构建图像(称为自回归模型);另一种则从一个静态的、充满噪声的模糊图像开始,逐渐将其清理干净,直到图像显现出来(称为扩散模型)。最近,一种名为“掩码离散扩散”(Masked Discrete Diffusion)的新型混合方法成为了明星。它将文本和图像都视为一长串拼图碎片(即标记/tokens),计算机通过观察已知的碎片来学习预测缺失的部分。这非常强大,因为它允许人工智能一次性观察整个画面,而不仅仅是下一个碎片,这使得它在图像编辑或解决数学问题方面表现出色。然而,这里有一个问题:为了清理图像,计算机在每一步中都必须重新检查每一块拼图碎片,即使是那些已经解决或仍然隐藏的碎片。这就像一位厨师在炖汤时,每分钟都要品尝一下锅里的每一种食材,即使有些食材已经调味完美,他也得尝一下,只为了确保味道正确。这使得整个过程既缓慢又耗费计算资源。
于是,Sparse-LaViDa 登场了,它像是一个为这些 AI 模型量身定制的聪明“副厨”。这项工作的研究人员意识到,计算机并不需要每次都盯着整个拼图看;它只需要专注于那些当前正在被解决的碎片。他们构建了一个系统,能够动态地“截断”(truncate),或者说切掉那些不必要的掩码标记(即隐藏或已解决的碎片)在烹饪过程中的存在。但巧妙之处在于:你不能直接把这些碎片扔掉,否则人工智能可能会忘记整个图像的上下文。因此,Sparse-LaViDa 引入了特殊的“寄存器标记”(register tokens)。你可以把它们想象成微小的、神奇的书签,代表着缺失的碎片。它们是紧凑的摘要,告诉人工智能:“嘿,这里还有 1,000 块碎片,但你现在不需要逐一观察它们;只需信任这个书签即可。”这使得人工智能能够在保持全局观的同时,跳过冗余的工作。
论文证明了这种方法效果惊人。通过使用这些寄存器标记和一种特殊的内存组织方式(称为 KV 缓存/KV caching),这个新模型 Sparse-LaViDa 在不损失任何质量的前提下,显著提高了处理速度。在测试中,它使文本生成图像的速度提升了近 2 倍(具体为 1.96 倍加速),使图像编辑速度提升了近 3 倍(2.84 倍加速),并使视觉数学推理速度提升了 2.80 倍。至关重要的是,作者展示了这不仅仅是适用于简单任务的小技巧;它保留了进行复杂任务的能力,如“局部重绘”(inpainting,填充图像缺失部分)和“外延绘制”(outpainting,扩展图像),而其他快速方法往往会破坏这些能力。研究人员指出,虽然我们通过微调现有模型实现了这些结果,但该方法在理论上未来也能够用于从头开始训练大规模模型。最终,Sparse-LaViDa 表明我们可以鱼与熊掌兼得:既拥有超快的生成速度,又不会牺牲那些让模型在理解和创造视觉世界方面如此出色的双向思维能力。
技术摘要:Sparse-LaViDa
问题陈述
掩码离散扩散模型(Masked Discrete Diffusion Models, MDMs)已成为处理多模态任务(包括图像理解、生成和编辑)的一种强大的统一框架。与自回归(AR)模型不同,MDMs 利用双向上下文和并行解码,在图像修复(inpainting)和文本填空(text infilling)等任务中具有优势。然而,现有的 MDMs 面临显著的推理效率问题:
- 冗余处理: 在每个采样步骤中,标准的 MDMs 必须处理完整的 token 序列,其中包括大量不携带实质性信息的冗余掩码 token。
- 缺乏 KV-Cache 支持: 由于依赖全注意力机制而非因果注意力机制,标准的 MDMs 在推理过程中无法利用 KV-cache 进行加速。
- 现有加速器的局限性: 先前的 MDM 加速尝试(如 Block Diffusion)采用了从左到右的块因果顺序(block-causal order)。虽然这实现了 KV-cache,但牺牲了对图像生成和编辑至关重要的双向上下文,使得图像修复等任务变得困难甚至无法实现。
方法论
作者提出了 Sparse-LaViDa,这是一种旨在通过动态截断不必要的掩码 token 来加速 MDM 采样,同时保持生成质量和双向上下文的新型建模框架。该方法引入了三个核心创新:
1. 稀疏参数化 (Sparse Parameterization)
Sparse-LaViDa 不再将所有掩码 token 实例化到输入序列中,而是仅使用以下内容来表示部分掩码序列:
- 提示词 token (Prompt tokens) 和 先前已解码的 token(存储在 KV cache 中)。
- 为当前步骤选定的用于解码的 掩码 token 子集。
- 寄存器 Token (Register Tokens): 引入特殊的 token 作为被截断(移除)的掩码 token 的紧凑表示。这些寄存器有助于恢复因截断而损失的建模能力。在最终设计中,使用了 64 个寄存器 token,并将其置于序列末尾。
2. 步进因果注意力掩码 (Step-Causal Attention Mask)
为了确保训练与推理之间的一致性,作者设计了一种专门的注意力掩码,用以模拟推理过程中观察到的增量 token 缓存行为:
- 训练阶段: 序列被划分为若干块。提示词 token 被分配到第 0 块。清晰 token 被分配到第 1…M 块,掩码 token 被分配到第 M+1…M+N 块。
- 注意力规则: 特定块中的 token 可以关注其自身块或之前的块。至关重要的是,特定块中的掩码 token 不能关注其他掩码块中的掩码 token,但它们可以关注所有的提示词和清晰 token。
- 寄存器集成: 寄存器 token 被附加到每个掩码块中,并遵循相同的注意力规则。
这种设计允许高效的并行训练,同时能够忠实地匹配推理过程中仅观测到部分 token 的过程。
3. 采样策略 (Sampling Strategy)
该框架支持根据任务采用不同的去掩码(unmasking)策略:
- 视觉生成(文生图、编辑): 利用预生成的 2D 去掩码顺序(例如,分层随机采样),以在不依赖置信度分数的情况下保持高生成质量。
- 文本生成/理解: 采用半自回归策略,即按从左到右的顺序对 token 块进行采样,但在块内部进行动态去掩码。与 Block Diffusion 不同,这种方法在块之间保持了双向上下文。
核心贡献
- 高效的参数化: Sparse-LaViDa 是首个支持 KV-cache 和任意 token 截断的方法,且无需假设从左到右的解码顺序,也不会牺牲双向上下文。
- 寄存器 Token: 引入寄存器 token 作为截断 token 的压缩表示,有效地缓解了由于激进的 token 移除导致的性能下降。
- 训练-推理一致性: 步进因果注意力掩码弥合了训练与推理之间的差距,使模型能够直接学习稀疏参数化,而无需进行蒸馏或启发式修改。
- 统一框架: 该方法基于最先进的统一 MDM —— LaViDa-O 构建,使其能够在单一加速框架内处理多样化任务(生成、编辑、推理)。
实验结果
作者在广泛的多模态基准测试上对 Sparse-LaViDa 进行了评估,并将其与基线 LaViDa-O 及其他最先进的模型进行了对比。
- 文生图生成: 在 GenEval 基准测试中,Sparse-LaViDa 取得了与 LaViDa-O 相当的性能(总分 0.78 vs 0.77),同时将推理延迟从 21.27s 降低至 10.86s,实现了 1.96× 的加速。在 DPG-bench 和 MJHQ-30k 上,它保持或略微提升了感知指标(例如,在相同数据子集训练下,FID 从 8.11 提升至 7.63)。
- 图像编辑: 在 ImgEdit 基准测试中,Sparse-LaViDa 将总分从 3.71 提高到 3.79,同时将延迟从 63.98s 降低至 22.55s,实现了 2.83× 的加速。
- 视觉数学推理: 在 MathVista 基准测试中,模型实现了 2.80× 的加速(3.72s vs 10.41s),同时保持了与基线相当的准确率(56.7 vs 56.9)。
- 消融实验:
- 移除寄存器 token 会导致细粒度对齐(DPG)和图像质量指标(FID, HPS v3)下降,证实了它们在保留低级视觉细节方面的作用。
- 移除步进因果注意力掩码或跳过微调会导致性能显著下降,凸显了所提训练流水线的必要性。
意义与主张
本文声称 Sparse-LaViDa 为标准 MDM 提供了一种根本性的高效参数化方案,能够在无需事后蒸馏阶段或启发式、无训练修改的情况下,实现可扩展的训练与推理。
作者强调:
- Sparse-LaViDa 在实现高达 2.84× 的大幅加速的同时,能保持与密集型基线相当的生成质量。
- 与以往的加速方法(如 Block Diffusion)不同,它保留了双向上下文,从而能够处理依赖非因果依赖关系的复杂任务,如图像修复、外扩(outpainting)和受限标题生成。
- 该方法提供了一种训练-推理一致的方法,避免了启发式加速技术中常见的不可预测的性能下降。
这项工作将 Sparse-LaViDa 定位为一种可行的路径,旨在使统一的多模态扩散模型在理论优势与实际部署所需的计算效率之间取得平衡,从而适用于实时应用场景。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。