超越传统模型：Mask 生成 Transformer 在图像 Token 建模中的创新应用！

文本到图像的扩散模型（DMs）发展得飞快，得到了深入的理论探索和实证分析的支持。然而，DMs 与自回归模型（ARMs）之间的差异使得实现视觉和语言生成的统一目标变得复杂。
最近， Mask 生成 Transformer（MGT）作为 DMs 和 ARMs 之间的一个有前景的中间体，通过预测随机 Mask 图像 Token （即 Mask 图像建模），将 DMs 的效率与 ARMs 的离散 Token 性质相结合。
然而，作者发现关于 MGT 推理的全面分析几乎没有，因此作者旨在提出一些积极的设计选择来填补这一空白。作者修改并重新设计了一组基于 DMs 的推理技术用于 MGT，并进一步阐明了它们在 MGT 上的性能。作者还讨论了纠正 Token 分布以提高推理的方法。
通过大量实验和实证分析，作者得到了具体而有效的设计选择，这些设计选择可以合并以实现进一步的性能提升。
例如，在提高推理方面，作者与纯采样相比，在 HPS v2 上的获胜率达到了约 70%，与最近的 SOTA MGT Meissonic 相当。作者的贡献有可能进一步提高 MGT 的能力和未来发展。

1 Introduction

生成模型的快速发展在计算机视觉和自然语言处理领域成功地激发了深度学习的创新。自然语言处理中大型语言模型（LLMs）的出现，以及它们在跨域和任务上的强泛化能力[27, 43, 46]，得益于自回归模型（ARM）与 Transformer 解码块[52]。相比之下，文本到图像（T2I）合成的主导范式是扩散模型（DM），它采用多步去噪过程从高斯噪声中合成高质量图像。ARMs 和 DMs 在训练和推理之间的显著变化阻碍了计算机视觉和自然语言处理中的生成范式的统一。近年来，一些 ARMs 在视觉生成方面的成就，如 LlamaGen [42]，Lumina-mGPT [21]和 Fluid [8]，表明 DM 并不是实现图像生成的唯一选择。这种范式可以合成极其高质量的图像，但合成一张图像需要数百或数千次函数评估（NFEs）[44]。相反， Mask 生成 Transformer（MGTs）[5]在每个正向传播中预测多个 Mask Token ，从而在 DM 和 ARM 之间取得权衡。这种方法保留了 DM 的效率，同时稳定了图像向离散 Token 的转换，与 LLM[34]的部分特征描述相吻合。

最近，Meissonic [3] 进行的最新 MGT，首次在 1024 × 1024 分辨率下实现高质量图像合成，在 HPS v2 [49] 上达到最先进性能，比 SD XL [29] 优 0.69。这一现象证实了 MGT 合成高分辨率图像的能力，并提出了开发商业级生成模型（如 FLUX [16]）的潜力。然而，与 DM 领域广泛的理论研究和实证分析相比，学术界对高分辨率 MGT 的探索和理解仍属空白，这阻碍了 MGT 在训练和推理方面的进一步发展 [3, 5, 6]。

为了填补这一差距，本文专注于 MGT 推理阶段，旨在通过在高分辨率图像合成场景中的实证分析，识别出能增强视觉质量和实现高效采样的有效设计选择。首先，作者阐明了在 DM 中广泛应用的无训练方法，并阐述需要重新设计以确保其有效性。如图 2 和表 1 所示，MGT 的采样过程与 DM 的采样过程具有显著相似性，因此将 DM 中的算法直接应用于 MGT 是合理的。作者在此背景下探索了 DPM-Solver [23, 24]，TomeSD [4]，以及 Z-Sampling [2]算法，但发现这三个算法都需要针对 MGT 的特征进行特定修改，以降低 NFE、加速推理或实现性能改进。

以 Z-Sampling 为例，作者发现，如果将 DDIM Inversion [26]（在 DM 中使用），并采用随机 Mask ，在某些情况下将不起作用，除非将 Mask 限制在低置信预测 Token 。特别是，作者的实验结果表明，在这三个算法中，DPM-Solver 和 TomeSD 对 MGT 的影响相对较小，而像 Z-Sampling 这样的重排算法可以获得显著的性能提升。此外，作者还调查了 MGT 中的噪声计划，类似于 EDM [14]，并发现，在某些条件下，余弦噪声计划是次优的。这些发现表明，DM 和 MGT 的训练和推理机制不一致，可能导致在 DM 上增强的推理算法对于 MGT 并不有效。

第二，作者关注由 Backbone （即 Transformer）生成的概率分布，进而发展出几种“便宜”（即没有显著计算开销）但有效分布校正算法，包括噪声正则化和（低熵）差分采样。具体来说，噪声正则化根据在应用软 max 之前的时间步数，动态地（高斯）扰动 Backbone 输出，以增强合成图像的多样性。另一方面，差分采样计算两个相邻时间步输出之间的 Kullback-Leibler（KL）离差，并重新采样与 Transformer 输出过于相似的 Token ，从而避免信息冗余并提高视觉质量。

第三，作者还研究了 Meissonic 上的模型量化以实现高效的内存使用。作者的结果表明，Weight4Activation16（W4A16）量化在实际中无法减少内存使用，而 W4A8 量化会导致推理崩溃。为解决这个问题，作者只对幅度较小的激活值层进行量化，将内存占用从 11.98 GB 减少到 4.57 GB（无显著性能损失）。

第四，作者使用各种指标（包括 ImageReward [50]，HPS v2 [48]，PickScore [10, 19]，和 AES [17]）来验证作者提出的布局选择在 HPD v2 基准测试中的有效性。如图 1 所示，这些策略显著提高了合成的图像质量。此外，作者引入了 Challengebench，包括 220 个 Prompt 来评估 Meissonic [3]在具有挑战性的案例中的鲁棒性。结果表明，与 SD XL [29]相比，Meissonic 在 Challengebench 上的性能提高了，而不是在 HPD v2 基准测试上。

2 Preliminaries

作者首先回顾了在视觉合成领域不断增长的三个生成模型：扩散模型（DM）[12]，自回归模型（ARM）[42]，以及 Mask 生成 Transformer（MGT）[3]。然后，作者概述了 MGT 的简单采样过程，该过程由 MaskGIT[5]引入。

视觉生成模型的推理机制

3 Enhanced Inference

第 3 节和第 4 节将分别讨论作者关于增强和高效的推理的研究。在本节中，增强推理涉及对已知 DM 方法的适应，以及根据 MGT 属性设计的新算法的应用。请注意，为了更合理地介绍作者的研究，作者将以渐进探索的形式介绍作者的方法和实验。对于后续内容，作者采用表 1 中的定义。除非另有说明，所有实验都是在 HPD v2 子集（见附录 7）上的 Meissonic-1024×1024 进行的。

Convexity Exploration of Noise Schedule

Masked Z-Sampling for MGT

Zigzag 扩散采样（Z-Sampling）[2, 25]的核心思想是通过在采样过程中提前引入“未来”语义信息，使用“zigzag”路径来提高 MCMC 的采样质量。作者的目标是将已在 MCMC 中证明有效的此算法扩展到 MGT，以提高合成的图像保真度。Z-Sampling 的逻辑可以通过图 5 顶部给出的方程进行说明。在获得潜在的Zi后，它使用“特定”的 Mask 算法（对应于 MCMC 中的 DDIM 反向）回溯到t=i-1，并从 t=i-1 到 t=i 再次进行采样。

很遗憾，在实验中，将随机 Mask （即图 5 中的基本 Z-Sampling）应用于模拟 DDIM 反转在受限制的 DM 中降低了推理性能。作者认为这是由于随机 Mask 错误地从潜在空间中移除了一些对合成图像有显著贡献的 Token 。例如，图 5 中的第一轮前采样获得的紫色 Token 可能被掩盖，尽管这些紫色 Token 通常代表最“未来”的信息。因此，作者采用了一种与采样机制一致的新 Mask 流程进行回溯，即在 i 步时用低对数概率（即图 5 中的被掩盖的 Z-Sampling） Mask 预测 Token 的一部分。作者还需要提到一个重要参数：反向分类器无关的指导（CFG）规模，它指的是在 Mask 阶段用于生成选择低置信度的 Token 的 CFG 规模。作者研究了反向 CFG 规模如何影响合成图像的质量。如[2]所述，选择一个理想的反向 CFG 规模可以最大化语义信息注入的积极影响，从而产生一个“恰到好处”的 CFG 间隙。

作者在图 6（左）中呈现了消融结果。从黑色虚线（即平均指标）的变化可以看出，反向 CFG 规模在-1 和 9 附近表现最好。为了降低计算成本，作者在实验中设置反向 CFG 规模为 0（即不使用 CFG），将标准 CFG 规模设置为 9，从而通过减少 NFE 避免额外的计算开销。

作者进一步在 HPD v2 数据集上验证了 Mask Z-Sampling，结果如图 6（右）所示。可以看出，作者的算法在几乎所有领域和指标上都显著优于普通采样，这说明 Mask Z-Sampling 可以持续提高 MGT 的性能。

Noise Regularization

根据作者的研究，MGT 与 DMs 之间一个显著的区别是，MGT 可以通过调整模型输出的概率分布来提高合成的图像的视觉质量和多样性。因此，作者提出了噪声正则化和差分采样的方法。在这里，作者首先介绍一种简单而有效的噪声正则化方法，该方法可以描述为

如图 9（顶部）所示，这种方法可以显著降低相似性的传播，并有效地引入多样性到采样过程中。此外，作者对超参数 z 进行了消融实验，结果如图 10 所示。当 z 从 0 增加到 100（z=0 表示普通采样）时，采样性能起初提高，然后下降。有趣的是，即使应用于所有 Token ，差分采样仍然优于普通采样，突显了差分采样的鲁棒性。此外，实证结果表明，当 z 设置为 75 时，性能最佳，作者在附录 9.4 中给出了差分采样在 HPD v2 上的性能。

4 Efficient Inference

另一种作者探索的路径是确定如何实现对 MGT 的高效推理。作者考虑模型量化[13]， Token 合并[4]，以及类似于 DM[23, 24, 30]中使用的调度策略。

Secondary Calibration for Quantization

实现内存效率的最有效方法是将生成模型的 Backbone 应用模型量化，这种技术在 FLUX [16]和 SD 3.5 [40]中已经成功使用。然而，这种方法在 Meissonic-1024 × 1024 上不起作用，原因有以下两点：

Meissonic-1024 × 1024 的模型参数数量有限（仅 10 亿个），其压缩层会积极减少 Token 数量至 1024。这些限制导致在应用 W4A16 后量化（PTQ）时出现问题，使得合成正常图像的能力受限。
由于 Meissonic 引入了多模型 Transformer 块，过于复杂的架构设计在实际中无法显著减少量化内存。一个简单的解决方案是量化激活值。然而，这种操作将进一步降低模型性能。

为解决这些问题，作者提出了二次量化校准（SCQ）方法。作者的核心贡献包括：

1）使用 Meissonic 合成的图像进行量化感知训练（QAT），以校正量化的值范围；

2）引入一种二次校准策略，记录每个层次在第一次校准后的幅度，然后只对幅度较小的激活值进行量化，进一步校准它们。

在作者的实验中，作者默认只量化激活值的 1/3，从而将 Backbone 网络的内存使用从 3.34 GB 降低到 2.24 GB。

表 3 中呈现了 SCQ 实验结果。其中，“A8W4-QAT & Calibration”表示对 QAT 衍生出的 Transformer（即 Backbone 网络）执行的一次校准。为了进行公平比较，在“A8W4-QAT & Calibration”中，随机选择了激活值的三分之一进行量化。此外，“CPU offloading”表示先使用 Transformer 获取所有 ZN，然后将分词器解码器加载到将 ZN 转换为合成图像。从表 3 可以得出，QAT 和辅助校准策略都是关键且有效的。

Introducing TomeSD into MGT

将 Token 合并应用于加速推理是自然的，因为 MGT 的基础是 Transformer 。不幸的是，Meissonic 只有 1024 个 Token ，比 SD XL 的注意力层中的 4096 个 Token 还要少。众所周知，注意力层的计算复杂度会随着 Token 数量呈指数增长， Token 数量较少会降低 Token 合并的潜在效益，导致作者实现的 TomeMGT 在实验中观察到的效果不明显。因此，作者专注于应用 TomeSD 到 MGT 的挑战，该方法已在 SD XL 上表现出有效性，以实现加速推理，并相应地探讨应用场景。

主要挑战有两个方面。首先，将 Token 合并到单个 transformer 可能导致推理失败，而它在多模态 transformer 中是有效的。其次，Meissonic 中的 RoPE（用于编码位置信息）也需要合并。对于前者，作者在多模态 transformer 上仅执行 Token 合并，而对于后者，作者在附录 9.5 中提供了关于作者处理 RoPE 的详细信息。在这里，作者只在表 5 中提供了合并比率的消融研究。比较实验提供在附录 9.6 中。

Momentum-based Solver

受到 DDIM[37]和 DPM-Solver[24]在 MG 上的成功启发，作者旨在在 MGT 中实现类似的机制。由于 DPM-Solver 的第一阶形式等于 DDIM，作者重点关注实现 DPM-Solver。作者将 MGT 中 DPM-Solver 的实现称为 Momentum-based Solver，因为这些算法基本上都使用动量进行加速采样[32, 38, 53]。分析图 10：差分采样的 z%的消融研究。

1st 和 2nd 阶方程可以表示为：

5 Challengebench and Additional Evaluations

作者进一步分析了 MGT 在具有挑战性的 Prompt 上的性能。作者使用 SD XL [29]合成了 150k 张图像，计算了 HPS v2 得分，并选择了最低得分的 1.5k 个 Prompt 。经过人工和 GPT-4o [1]选择后，作者在 Challengebench 上形成了 220 个语义上连贯的 Prompt ，命名为 Challengebench 。

作者在 Meissonic，SD XL [29]，FLUX.1-schnell [16]和 SD-3.5-Large [40]上进行了实验，并在表 4 中呈现了结果。作者观察到 Meissonic 在 Challengebench 上的改进（0.1838 ➡ 0.2116）大于其在 HPD v2 上的改进（0.2888 ➡ 0.2957）

原始论文表明，MGT 在具有挑战性的 Prompt 上更具鲁棒性。在 Challengebench 上，FLUX.1-schnell 和 SD-3.5-Large 的得分降低，表明其提供了对泛化的强健度量。作者通过组合设计选择在 Table 11 中获得了约 70%的获胜率，与纯随机取样的基本采样相比。更多的实验请参见附录 9.9。

6 Conclusion

作者的方法基于 Mask 生成 Transformer ，旨在确保增强和高效的推理，这是对非自回归模型的有意义探索。

在未来，作者将尝试统一并改进 Mask 生成 Transformer 的训练过程，以克服这一生成范式的 Bottleneck 。

参考文献

[0]. Bag of Design Choices for Inference of High-Resolution.

END

作者：小书童
来源：集智书童

推荐阅读

欢迎大家点赞留言，更多 Arm 技术文章动态请关注极术社区嵌入式 AI 专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。

超越传统模型：Mask 生成 Transformer 在图像 Token 建模中的创新应用 ！