点击下方卡片，关注「集智书童」公众号
导读
本文提出了一种基于集合 Token 化和分布建模的全新图像生成范式。与将图像序列化为具有统一压缩比固定位置潜在码的传统方法不同，作者引入了一种无序 Token 集表示，根据区域语义复杂性动态分配编码容量。这种TokenSet增强了全局上下文聚合，并提高了对局部扰动的鲁棒性。
为了解决建模离散集合的关键挑战，作者设计了一种双变换机制，该机制将集合双射转换为具有求和约束的固定长度整数序列。进一步地，作者提出了固定和离散扩散——第一个同时处理离散值、固定序列长度和求和不变性的框架，从而实现了有效的集合分布建模。
实验表明，TokenSet在语义感知表示和生成质量方面具有优越性。作者的创新，涵盖了新颖的表示和建模策略，将视觉生成推进到超越传统顺序 Token 范式的领域。
代码和模型: https://github.com/Gengzigang/TokenSet
- 引言
当代视觉生成框架[10, 24, 36, 41]主要采用两阶段范式：首先将视觉信号压缩为潜在表示，然后对低维分布进行建模。传统的 Token 化方法[10, 18, 38, 40]通常采用均匀的空间压缩比率，生成具有固定位置对应关系的序列代码。以一张海滩照片为例，上半部分是包含最少细节的天空区域；下半部分包含语义密集的前景——当前的方法对这两个区域分配相同数量的代码。这引发了一个基本问题：视觉简单区域是否应该与语义丰富区域获得相同的表现力容量？
本文介绍了一种全新的视觉压缩和分布建模方法，TokenSet（图1）。在压缩阶段，作者提出将图像分词为无序集合，而不是位置依赖的序列。与保持固定空间对应关系的序列化分词不同，作者的TokenSet公式允许根据区域语义复杂性进行动态注意力分配。这种方法增强了全局信息聚合，促进了语义感知表示，并表现出对局部扰动的优越鲁棒性。
尽管如此，与序列数据相比，对集合结构数据进行建模存在显著挑战。现有的集合建模方法可分为两类：第一类[9, 22, 42]采用基于池化操作的策略（例如，均值/总和/最大值/相似操作）来提取低维特征，这类方法缺乏对集合中每个元素的直接监督，通常导致次优结果；其他基于对应关系的方法（例如，DETR的匈牙利匹配[2]）试图通过二分图匹配来构建元素级的监督。然而，动态匹配机制的固有不稳定性导致监督信号在训练迭代中变化，从而导致次优收敛。
为了解决这个问题，作者设计了一种双重转换机制，将集合建模转换为序列建模问题。具体来说，作者计算集合中每个唯一 Token 索引的出现次数，将无序数据转换为具有以下结构的序列：
-
- 序列长度等于代码簿大小
-
- 每个元素代表非负整数计数
-
- 所有元素的总和等于集合中元素的数量
现有离散建模方法（例如，VQDiffusion [12]）虽然可以处理固定长度的整数序列，但忽略了求和约束，而连续扩散模型[26]在保持元素求和的同时，难以处理离散值表示，目前还没有任何方法同时满足这三个约束。因此，作者提出了固定求和离散扩散模型。通过引入常数和先验，作者同时满足了这三个关键属性，并实现了对这种结构化数据的有效建模。
作者的贡献可以总结如下：
-
- 作者提出了一种基于集合表示的新型图像生成范式，从根本上区别于传统的序列化表示。
-
- TokenSet展现出全局上下文 Aware ，能够在保持对局部扰动的鲁棒性的同时，实现动态的token分配到语义复杂度。
-
- 作者提出了一种通过双变换对离散集合数据进行建模的有效方法，建立了无序集合与序列化数据之间的双射关系。
作者提出了固定和离散扩散（Fixed-Sum Discrete Diffusion），这是一种针对离散数据建模的专用生成框架，它明确地强制执行求和约束，从而实现了对集合分布的优越建模。
- 相关工作
2.1 图像分词
图像分词将高维像素空间中的图像压缩成紧凑的表示，从而便于后续的理解和生成任务。早期方法如变分自编码器（VAEs）[18]将输入图像映射到低维连续的潜在分布。在此基础上，VQVAE [38] 和 VQGAN [10] 将图像投影到离散的 Token 序列，将每个图像块与一个显式的离散 Token 关联。后续工作 VQVAE-2 [32]、RQVAE [21] 和 MoVQ [44] 利用残差量化策略将图像编码成层次化的潜在表示。同时，FSQ [29]、SimVQ [46] 和 VQGAN-LC [45] 解决了在扩展代码本大小时出现的表示崩溃问题。其他创新包括 DQVAE [16] 中的动态量化策略、ImageFolder [25] 中的语义信息集成以及架构优化 [1, 39]。最近，TiTok [41] 探索了用于图像表示的 1D 潜在序列，在令人印象深刻的压缩比下实现了良好的重建。
尽管取得了这些进展，先前的方法主要将图像编码为 Token 序列，其中每个元素严格对应于固定的图像位置。本文提出将图像表示为无序的 Token 集，从而消除了位置偏差，同时有效地捕捉全局视觉语义。
2.2. 集合建模
早期基于集合的表示包括词袋模型（BoW）[17, 30, 34]及其视觉对应模型[5, 20, 35]。最近，CoC[28]提出通过聚类将图像视为点集。然而，这些基于集合的表示会丢失原始数据中的信息。相反，某些数据模态（如点云和边界框）与集合表示天生相符。这促使大量研究致力于建模排列不变的数据，但同时也带来了三个基本挑战。
首先，流行的生成范式，如自回归（AR）模型[10, 36, 37]和扩散模型[7, 11, 12, 15, 31]，旨在对序列数据进行建模，这使得它们与无序的集合结构化数据不兼容。
其次，处理排列不变数据需要严格对称的操作（例如，求和、最大值或类似操作）以避免位置依赖[9, 22, 42]。然而，这一约束限制了诸如卷积和注意力等强大工具的使用，从而限制了模型的容量。
第三，对复杂数据分布进行有效建模通常需要每个元素的监督信号，然而无序集本身缺乏这样的机制。现有的方法如DSPN [43] 使用 Chamfer 损失进行监督，而TSPN [19] 和DETR [2] 使用匈牙利匹配。然而，这些匹配过程本质上是不稳定的，往往导致训练信号不一致。替代方法如PointCloudGAN [23] 尝试直接建模全局分布，这损害了训练效果和整体性能。
本文通过双重转换操作绕过这些限制，有效地将集合转换为序列。这种转换使作者能够利用各种基于序列的建模方法来应对集合建模这一挑战性任务。
- 方法
3.1 图像集分词器
将图像分解成集合的关键在于消除视觉 Token 与图像固定位置之间的位置依赖关系。虽然先前的工作TiTok [41]通过去除二维位置关系将图像转换为1D Token 序列，但它保留了固定的1D位置对应关系。作者从这种方法出发，开发了一个完全位置无关的 Token 化框架。作者的架构在编码器和解码器组件中都使用了视觉Transformer（ViT）[8]。编码器在处理图像块的同时，处理可学习的潜在 Token ，产生连续的潜在表示，这些表示通过VQVAE [38]代码簿进行离散化。这个过程生成一个1D Token 序列
，其中
，
表示代码簿大小，
表示 Token 计数。
为了消除一维位置偏差，作者在解码过程中引入了排列不变性。具体来说，作者将
定义为
的集合表示。
在所有
的排列都被视为等价的情况下。在训练过程中，作者在保持重建目标的同时，随机排列 Token 并在解码输入前进行，如图2所示。尽管排列空间随着
的增长呈阶乘式增长，但经验结果表明，通过部分排列采样可以有效学习排列不变性。
基于集合的 Token 化方法相较于序列 Token 化具有三个主要优势：首先，通过将 Token 与固定的空间位置解耦，模型学会了根据全局图像内容动态分配 Token ，而不是依赖于局部块统计信息。其次，全局感受野通过防止过度依赖局部特征，显著提高了噪声鲁棒性。第三，通过训练， Token 自发地发展出专注于语义上不同区域（例如，物体与纹理）的专业化注意力模式。作者通过第4.2节的定量评估和可视化研究实证验证了这些特性。
3.2 双重转换
在分词后，图像被表示为一个无序的 Token 集
。使用神经网络对这类复杂集合进行建模存在显著挑战，主要由于集合固有的无序性质以及单个元素缺乏有效的监督。
现有的序列建模方法，尤其是自回归[10, 36]和扩散模型[7, 12, 31, 33]，在处理集合结构数据时存在固有的局限性。这些方法本质上依赖于元素的位置顺序，因此不适用于存在元素顺序模糊性和指数级排列可能性的排列不变集合。像PointGAN[23]这样的替代方法存在训练不稳定和缺乏对排列不变数据高效表示的问题。其他方法如DETR[2]利用匈牙利匹配来实现集合对应。然而，它受到匹配不稳定性的影响，阻碍了稳健建模。
为了解决这些挑战，作者提出了一种双向转换机制（图1），该机制可以在无序集合和结构化序列之间进行双向转换。给定一个 Token 集合
，作者通过以下方式构建一个计数向量
：
其
中
其中
表示克罗内克δ函数。通过这种方式，作者有效地将无序的 Token 集转换为序列数据，且不丢失任何信息。此外，转换后的序列数据
具有三个关键的结构先验：
- • 固定长度序列：计数向量
包含
个元素，对应于代码簿的大小，确保序列长度固定。
- • 离散计数值：每个元素
，记录码本项频率，是
区间内的整数，其中
是从编码器中提取的 Token 数。
- • 固定和约束：所有值的总和等于编码的 Token 数量：
总之，双重转换在集合和序列表示之间建立了一种双向映射，提供了两个基本优势：(1) 它将建模排列不变集合的难题简化为已研究的序列建模领域，关键性地启用了自回归和扩散框架进行建模。(2) 识别出的结构先验——固定的序列长度、离散值约束和求和守恒——提供了基于数学的规范化，引导有效的模型学习。
3.3 固定和离散扩散
鉴于作者的双变换序列数据具有三个有利的前置属性，作者系统地研究了多种建模方法。虽然自回归模型[36]和标准离散扩散方法[11, 12]对于离散值数据都有效，其中后者特别适合于固定长度的序列，但它们本身并不保证固定求和属性。
相反，连续扩散模型[31]可以通过其均值保持的均方误差损失[26]自然地保留固定长度和固定求和约束，但它们在离散分布建模[4]方面存在困难。
为了协同结合这些方法的优势，同时满足所有三个先验条件，作者提出了一种新的建模方法，称为固定和离散扩散（FSDD），如图3所示。受连续扩散方法在中间去噪步骤中强制执行求和约束的启发，此方法将一个受约束的扩散路径集成在一个离散流匹配架构中。关键创新在于确保每个中间步骤的样本严格遵循固定和约束。
3.3.1. 训练流程
扩散过程
初始噪声样本
是从长度为
的整数向量上，以固定和
的多项式分布中抽取的
给定从噪声分布中采样的
和从数据分布中采样的
，两者均满足
，作者定义约束扩散过程为：
参数
和
满足：
本设计保证了对于求和期望的约束：
然而，作者仍然不能保证每个个体样本
都满足此约束。因此，作者进行动态调整以确保其满足。
贪婪调整
作者的调整协议的核心目标是保持
的可能性。具体来说，如果
的总和超过
，作者将通过贪婪选择标准减少
的某些元素：对于每个元素，作者量化每次调整引起的可能性下降，并优先调整那些可以提高可能性或最小化其下降的元素。因此，调整后的样本
符合固定总和约束和方程（5）中的概率分布。作者在图3中提供了说明，并在算法1中提供了伪代码。在训练过程中，作者在每个扩散步骤实施这种贪婪调整策略，以整合固定总和约束。
此外，固定总和离散扩散采用标准的离散扩散损失，其中去噪网络
通过交叉熵损失训练，以从噪声输入
预测
，在保持离散状态转换的同时，保持求和不变性。
3.3.2. 推理策略
固定和扩散的推理过程遵循一个带有强制求和约束的迭代去噪方案。从满足
的噪声样本
开始，作者通过离散转换逐步细化该样本：
此处，
表示从噪声数据预测出的离散数据分布。作者采用前
个采样策略来生成
候选样本，然后通过后验项
进行处理。该项实现了截断高斯离散化：
参数定义如下：
控制采样过程中的截断比率。为确保严格遵循求和约束，作者对
应用贪婪调整，以确保
，从而有效弥合训练和推理阶段之间的潜在差距。
- 实验
4.1. 设置
作者在ImageNet数据集[6]上进行了实验，图像分辨率为
。作者报告了在50,000图像的ImageNet验证集上的结果，使用了Fréchet Inception Distance（FID）指标[14]。作者的评估协议由[7]提供。
实现细节。在分词器训练中，作者遵循了TiTok [41]中的策略，并应用了数据增强，包括随机裁剪和水平翻转。作者使用了AdamW优化器 [27]，基础学习率为1e-4，权重衰减为1e-4。模型在ImageNet上训练了
步，批大小为256，相当于200个epoch。作者实现了学习率预热阶段，随后是余弦衰减计划，梯度裁剪阈值为1.0。采用了衰减率为0.999的指数移动平均（EMA），并报告了EMA模型的结果。为了提高质量和稳定训练，作者引入了判别器损失 [10]，并在最后的
步中仅训练解码器。此外，作者遵循[41]中的MaskGIT Agent 代码 [3]以促进训练。
生成器配置与DiT [31]对齐。作者使用随机水平翻转作为数据增强。所有模型均使用AdamW [27]进行优化，学习率恒定为1e-4，批大小为256，训练200个epoch。在整个训练过程中，作者实现了带有衰减率为0.9999的EMA。对于推理，作者使用了25个采样步长，并结合无分类器引导来进一步提高图像质量。
4.2. 设置分词器
与顺序图像 Token 化方法相比，将图像表示为 Token 集引入了独特的属性，包括排列不变性、全局上下文感知以及增强对局部扰动的鲁棒性。此外，作者证明了基于集合的 Token 化可以在本质上将 Token 组织成语义上连贯的簇的同时实现精确重建。
4.2.1. 排列不变性
作者通过从不同顺序排列的编码 Token 中重建图像来测试作者分词器的排列不变性。具体来说，作者以五种不同的顺序解码 Token ：(1)原始顺序，(2)反转顺序，(3)随机打乱顺序，(4)按升序排序的 Token ，和(5)按降序排序的 Token 。如图4所示，所有重建的图像在视觉上完全相同，表明作者分词器的排列不变性。这种不变性还得到了表1中定量结果的进一步支持，这些结果在不同 Token 顺序下是一致的。这些发现表明，网络可以通过仅对排列子集进行训练来成功学习排列不变性。
4.2.2 全局上下文感知
通过强制执行排列不变性，TokenSet解耦了词元间的位置关系，从而消除了传统图像词元化中固有的由序列引起的空间偏差。这种架构设计鼓励每个词元全面整合全局上下文信息，有效地将其理论上的感受野扩展到整个特征空间。为了实证验证这一现象，作者在图5中可视化了有效的感受野，该图展示了全局注意力机制。值得注意的是，传统的基于序列的词元化器，如VQGAN[10]，在词元和固定局部区域之间存在紧密的空间耦合。相反，TokenSet从根本上消除了位置偏差，并通过具有全局感受野的词元组合来表示图像。
4.2.3. 鲁棒性
作者的集合 Token 符，虽然未绑定到特定的空间位置，但能够捕捉全局图像语义，显示出对噪声的增强鲁棒性。表2比较了不同 Token 化器对输入图像中注入的高斯噪声的鲁棒性。具体来说，作者向图像添加了不同标准差的高斯噪声，并测量了扰动图像与原始图像之间的 Token 重叠率。结果表明，与TiTok [41]和VQGAN [10]等其他 Token 化器相比，作者的 Token 化器在所有噪声水平上均能持续实现更高的重叠率。此外，尽管所有方法在噪声强度增加时都会出现性能下降，但TokenSet下降的速度较慢，这突显了其在噪声鲁棒性方面的优越性。
4.2.4. 语义聚类
给定从
个类别中抽取的
个 Token ，
的表示为
。紧凑的表示空间产生了对图像空间更高效的描述。图6 展示了编码 Token 包含特定类别的图像。有趣的是，作者观察到这些 Token 分布天生具有语义上连贯的聚类模式。例如，包含六个属于第 65 类 Token 的图像始终描绘鸟类，而十二个第 162 类 Token 则代表狗。这一现象可能表明其对基于内容的图像检索的推进有影响，其中语义上连贯的 Token 集可能实现更稳健的特征索引。此外，作者通过线性检测验证了其语义聚类能力。表3 中所示的结果表明，即使没有专门的设计，作者的 Token 器也已经实现了有希望的性能。
4.2.5 重构质量
在表4中，作者使用ImageNet验证集比较了不同分词器在重建性能上的表现。尽管TokenSet中的随机打乱操作阻止了网络利用图像的归纳位置偏差，并且从集合表示空间中获得的益处大幅减少，但作者发现这种看似不可行的方案仍然能够实现与先前主流方法[10, 21]相当的良好性能。
4.3 固定和离散扩散
4.3.1. 通过先验信息建模
在表5中，作者比较了利用不同先验的不同建模方法。所有实验均使用ViT-Small [8] 尺度的模型进行。
首先，如第4.2.1节所述，图像 Token 集的任何排列都可以等价地重建图像。因此，作者考虑随机采样两种不同的排列用于自回归建模[10]，分别称为ARorder1和AR-order2。作者发现它们都实现了几乎相同的表现。这一观察表明，排列不变性特性尚未得到充分利用。为了解决这个问题，作者提出训练一个单一的自动回归模型，以同时学习集合的所有可能排列，称为AR-random。然而，由于难以建模难以处理的大排列空间，这种方法表现出较差的生成性能。
为了克服这一限制，作者应用了双重变换，并随后对生成的序列分布进行建模。作者尝试了自回归[10]和离散扩散[11]，利用集合和离散属性，它们的性能优于AR随机。然而，它们本质上无法保证符合作者的固定和先验，导致结果平庸。同样，通过应用连续扩散[31]和后续量化来建模离散分布也被证明是无效的。
相比之下，作者提出的固定和离散扩散（FSDD）方法独特地能够同时满足所有这些期望特性，从而实现最高性能。这一发现证明了将所有必需的先验知识纳入建模方法以简化建模过程的重要性。它验证了作者的设计选择，并突出了分词器要求与作者的建模方法之间的协同作用。
4.3.2. 消融研究
先前的研究[10, 25]在图像生成的两阶段“压缩后建模”框架中识别出一个关键困境：逐步增加潜在空间容量可以持续提升重建质量，但生成质量先提升后下降。在这里，作者严格调查TokenSet是否能够解决上述困境。
作者系统地改变 Token 数和代码簿大小，以研究它们对重建和生成性能的影响。在生成评估中，作者采用一个小规模模型（36M参数）来拟合分布。如表6所示，作者发现适度增加潜在维度可以提升重建和生成质量，但超过这个范围则会降低这两个指标。这与序列化 Token 化框架中的观察结果不同，在序列化 Token 化框架中，增加潜在维度可以提升重建质量，但会损害生成性能。作者将这种差异归因于基于集合的潜在空间与图像网格之间的空间对应关系缺失。与通过空间先验实现坐标特定映射的序列化 Token 化不同，基于集合的解码器无法建立直接的空间对应机制。这与扩散模型在建模过于复杂的潜在分布时面临的挑战相似，因此，这两种方法在维度增加过多时都会出现性能下降。关键的是，这些观察结果提出了一种解决重建-生成困境的潜在方法：通过消除解码器对低效捷径映射的依赖，作者可以使其行为与第二阶段建模的分布学习过程更加一致。
作者的分析进一步揭示，模型规模的扩展能够带来一致的性能提升，这在表7中的实证结果中得到体现。虽然这些结果暗示了进一步扩大模型规模的潜在益处，但实际限制因素限制了作者在当前实验范围之外的探索。作者将此留待未来工作。
- 结论
这项工作挑战了传统的序列化视觉表示范式，通过引入TokenSet，这是一种基于集合的框架，能够动态地为语义多样的图像区域分配表示能力。通过双重转换，作者建立了无序 Token 集和结构化整数序列之间的双射映射，使作者能够通过提出的固定和离散扩散有效地建模集合分布。实验表明，这种方法不仅实现了与区域复杂度一致的动态 Token 分配，而且增强了对抗局部扰动的鲁棒性。通过在训练和推理过程中强制执行求和约束，TokenSet解决了现有离散扩散模型中的关键局限性，并优于固定长度序列 Baseline 。
将图像作为一组进行分词相较于传统的顺序分词具有独特的优势，为图像表示和生成引入了新的可能性。这种范式转变激发了对开发下一代生成模型的全新视角。在未来工作中，作者计划进行严格的分析，以充分发挥这种表示和建模方法的全潜能。
参考
[1]. Tokenize Image as a Set
扫码加入👉「集智书童」交流群
（备注：方向+学校/公司+昵称）