受通信启发的结构化图像表示Tokenization方法 Communication-Inspired Tokenization for Structured Image Representations
提出COMiT，通过迭代局部观测和递归更新构建具有语义结构的离散视觉token序列
前置知识
Flow Matching（流匹配）
Flow Matching是一种生成式建模框架，通过学习概率分布之间的连续变换路径来生成数据。给定源分布（通常是高斯噪声）和目标分布（真实数据），Flow Matching学习一个速度场，描述从噪声到真实数据的流动轨迹。训练目标是最小化预测速度与真实速度之间的L2损失，其中噪声图像是真实图像和随机噪声的线性插值，插值参数t在[0,1]区间均匀采样。推理时通过积分ODE路径从噪声生成样本。这一框架相比扩散模型提供了更稳定和高效的训练。
COMiT使用Flow Matching作为解码器训练框架，理解它是掌握COMiT生成过程的基础，特别是时间步t和速度场的概念对于理解编码-解码统一的训练机制至关重要。
FSQ（Finite Scalar Quantization）
FSQ是一种简化版的VQ-VAE量化方法。与VQ-VAE在离散码本中查找最近邻向量不同，FSQ将连续向量映射到有限标量集的笛卡尔积上，每个维度独立量化到预定义的离散级别。例如，[8,8,8,5,5,5]表示6个维度分别量化到8或5个级别。通过投影操作实现离散化，其中先对向量进行裁剪再四舍五入。FSQ的优势是完全可微且无需维护大型码本，训练更稳定，避免了VQ-VAE中的码本崩溃问题。
COMiT使用FSQ对潜在消息进行离散化，理解FSQ的工作原理有助于理解COMiT如何将连续表示转换为离散token序列，以及为什么FSQ比传统VQ-VAE更适合COMiT的递归更新架构。
语义对齐（Semantic Representation Alignment, SREPA）
SREPA是一种知识蒸馏技术，将预训练自监督视觉模型（如DINOv2）的高级语义特征转移到目标模型的中间表示中。具体而言，在解码阶段，将COMiT网络第j层的消息token表示投影并平均池化得到向量，然后最大化其与DINOv2的CLS特征的余弦相似度。损失函数为指数负余弦相似度，这鼓励COMiT的潜在消息编码与预训练模型一致的语义信息。这种蒸馏只在解码阶段进行，避免干扰编码过程。
COMiT依赖SREPA为token序列提供语义grounding，理解SREPA有助于理解COMiT如何在无监督情况下学习到有意义的语义表示，以及为什么单独的语义对齐不足以产生结构化token序列。
组合泛化（Compositional Generalization）
组合泛化是指模型能够理解训练时未见过的元素组合的能力。例如，训练时见过狗和球、猫和飞盘的图像，测试时应该能识别猫和球或狗和飞盘的图像。在图像理解中，这要求模型学习到分离的对象表示，能够独立编码每个对象的语义信息，而不是将整个场景编码为单一混合特征。COMiT通过MSCOCO数据集评估这一能力：将图像对分割为训练集和验证集，确保对象组合在两者间不重叠，这样模型必须真正学习对象级的分离表示才能成功泛化。
组合泛化是评估视觉表示质量的重要维度，它测试模型是否真正学习到对象级结构。COMiT在这方面的显著优势是其核心贡献，理解这一概念有助于评估COMiT的实际应用价值。
研究动机
现有离散图像tokenization方法（如VQ-VAE、MaskGIT、TiTok等）主要针对重建和压缩优化，产生的token序列往往捕获局部纹理而非对象级语义结构。这些方法在单次编码pass中处理整个图像，没有显式机制来组织token的语义布局，导致语义信息在token间纠缠且定位不佳。在具体场景中，当图像包含多个对象时，现有方法将它们编码为混合特征，无法区分哪些token对应哪个对象。例如，在包含狗追逐球的图像中，现有tokenizer可能将两个对象的特征混合到同一组token中，限制下游任务如对象关系推理的性能。
本文的目标是本文的目标是设计一种新的图像tokenization框架，能够学习结构化的、以对象为中心的离散视觉token序列。该框架应该：(1)显式鼓励token序列的语义组织，使不同对象的信息分布在token序列的不同部分；(2)保持与Transformer架构的序列接口兼容性；(3)在语义编码、组合泛化和关系推理等下游任务上优于现有方法；(4)保持可接受的重建质量以支持生成应用。最终目标是让token序列能够像自然语言句子一样具有可解释的结构，便于推理和组合操作。
与已有工作不同的是，本文的独特切入角度是从人类沟通的本质中汲取灵感，将图像tokenization重新构想为一个迭代通信-重建博弈，而不是传统的编码器-解码器压缩任务。关键洞察是：当人类描述包含多个对象的场景时，会一次关注一个区域，逐步将显著信息融入消息中。这种增量过程允许听众将每次观测整合到场景的连贯心理模型中，同时逐渐减少不确定性。现有工作（如RAM、AIR）曾探索过注意力机制，但要么专注于计算效率而非表示学习，要么仅限于玩具数据集。本文首次将这种通信启发的递归观测范式应用于大规模真实图像，并将其与Flow Matching生成训练框架结合，实现端到端的结构化表示学习。
核心方法
COMiT的整体思路是将图像编码为一个离散潜在消息，通过迭代观测图像的局部区域并递归更新消息来实现。模型首先从输入图像中提取一个随机裁剪序列及其位置，计算相对偏移。潜在消息初始化为来自码本的固定token序列。在每一步，模型接收当前裁剪、时间戳（编码时为1）、偏移和前一步消息，输出更新后的消息。消息通过FSQ量化后反馈给模型形成循环。经过K步后，最终消息作为条件，通过Flow Matching解码器重建完整图像。整个流程使用单个Transformer实现编码和解码，通过组合Flow Matching重建损失和语义对齐损失端到端训练。
COMiT的核心创新点是将图像tokenization形式化为一个通信启发的迭代过程，有两个关键设计原则：(1)注意力和序列化tokenization：编码器将图像处理为局部观测序列，每步关注不同区域并逐步整合信息到离散潜在消息中；(2)同质化通信：与传统编码器-解码器使用分离网络不同，COMiT采用统一设计，同一网络既充当说话者（编码）又充当倾听者（解码），模仿人类沟通的对称性。这与现有方法的本质区别在于：现有方法在单次编码pass中产生token序列，没有显式组织token语义布局的机制；而COMiT通过递归消息更新和随机化crop数量，鼓励模型贪心地使用有限token预算，自然诱导出层次化的对象中心结构。
方法步骤详情
COMiT的完整工作流程包括：(1)输入准备：给定RGB图像，随机裁剪K个局部crop，记录其归一化中心位置。计算偏移序列。随机选择K以避免模型预分配容量。(2)消息初始化：用固定token初始化消息，其中L=256是消息长度，d=6是token维度。(3)递归编码：对k=1到K，计算更新后的消息，其中sg表示停止梯度。通过FSQ量化得到离散token。为了提高训练效率，只对最后一步更新进行反向传播。(4)流匹配解码：创建噪声图像，其中t在[0,1]均匀采样，epsilon是随机噪声。将噪声图像、最终消息和偏移输入网络，预测速度场，训练目标是L2损失。(5)语义对齐：计算SREPA损失，其中psi(x)是DINOv2的CLS特征，f-bar是网络中间层的平均池化表示。(6)总损失：组合三个损失组件，权重设置为0.5。使用Adam优化器训练200个epoch。
技术新颖性
COMiT的技术新颖性体现在多个方面：(1)首次将Flow Matching应用于同时训练编码和解码阶段，通过统一的微分流目标集成表示学习和生成，而非像现有方法那样仅用于解码。(2)引入随机化crop数量和停止梯度策略来促进贪心token使用，避免模型预分配容量给未来crop，这是tokenization中首次采用这种训练技巧。(3)提出SREPA将预训练视觉特征蒸馏到中间表示，为token序列提供语义grounding，这与传统仅重建的tokenization有本质区别。(4)使用统一Transformer模型实现编码-解码对称性，移除了传统编码器-解码器架构的冗余参数，这是生成式建模中首次尝试这种融合设计。(5)设计了专门评估语义结构的新测试套件（组合泛化、关系推理），超越了传统仅关注重建质量的评估框架。
实验结果
COMiT在三个语义探针基准测试中一致性地优于现有1D离散图像编码器。在ImageNet100分类任务上，COMiT-B达到82.91% top-1准确率，显著优于TiTok-B（19.43%）和ALIT（29.43%），与最强的FlexTok(d18-d18)（81.54%）相当。模型规模扩展到COMiT-L后进一步提升至85.80%，而扩展到COMiT-XL时略降至84.69%，表明模型容量在编码和解码间的分配存在权衡。在MSCOCO组合泛化任务上，COMiT-L达到45.31% top-5准确率，比FlexTok(d18-d18)的38.58%提升约6.7个百分点，证明COMiT学习到的token序列具有更好的对象分离特性。在Visual Genome关系推理任务上，COMiT-L达到56.42% top-1准确率，同样优于FlexTok(d18-d18)的54.46%。消融实验显示：SREPA对性能贡献显著（有SREPA时82.91% vs 无SREPA时72.26%）；使用局部crop训练即使测试时仅使用全局crop也能提升性能（82.91% vs 80.94%），证明注意力的序列化tokenization是诱导对象中心结构的关键。定性分析显示COMiT的token注意力图与目标对象紧密对齐，平均mIoU达到0.53，而ablated模型仅0.34。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| ImageNet100分类 | top-1准确率 | 85.80% | 81.54% (FlexTok-d18-18) | +4.26个百分点 |
| MSCOCO组合泛化 | top-5准确率 | 45.31% | 38.58% (FlexTok-d18-18) | +6.73个百分点 |
| Visual Genome关系推理 | top-1准确率 | 56.42% | 54.46% (FlexTok-d18-18) | +1.96个百分点 |
| ImageNet1k重建 | rFID | 3.67 (COMiT-L) | 1.61 (FlexTok-d18-d18) | 略差（表示-重建权衡） |
局限与改进
论文承认的主要局限性是重建fidelity相较于专门优化的方法略逊一筹，COMiT-L的rFID为3.67，而FlexTok-d18-d18为1.61。作者解释这是表示-重建权衡的自然结果，因为模型容量被部分用于学习语义结构。额外观察到的局限性包括：(1)计算成本：递归crop聚合增加了编码时间，特别是使用自适应crop策略时需要多次中间重建；(2)crop策略设计：论文只探索了有限的crop策略子集，最优策略（特别是任务自适应策略）尚未充分研究；(3)模型容量权衡：从L扩展到XL时语义性能略降，表明单网络中编码-解码的平衡需要更精细的架构设计；(4)视频扩展：论文专注于静态图像，时空数据的处理（如视频）尚未探索，这将是重要但具有挑战性的方向；(5)量化细节：FSQ的级别选择是经验性的，缺乏理论指导。
独立分析的弱点
COMiT存在的几个具体弱点值得改进：(1)重建质量不足：COMiT的rFID（3.67）明显高于专门优化的基线（1.61），这限制了其在生成任务中的应用。改进方向可以是引入多阶段训练流程，如先训练语义对齐，再finetune解码器以提高重建质量；或者设计自适应容量分配机制，根据输入图像复杂度动态调整编码器-解码器的容量比例。(2)编码效率：递归crop聚合需要多次前向传播，测试时间较长。改进方向包括学习更高效的crop策略（如强化学习）、并行化不相关的crop观测，或者设计更轻量的编码器网络。(3)crop策略依赖性：性能受crop选择策略影响较大，特别是组合泛化和关系推理任务。改进方向是通过元学习或强化学习自动学习最优crop策略，使其能够适应不同任务需求。(4)模型规模扩展：从L到XL时语义性能下降，表明当前架构的容量分配机制不够成熟。改进方向可以是设计模块化架构，为编码和解码分配独立的容量参数，或者引入动态路由机制根据任务调整资源分配。(5)语义对齐依赖性：SREPA依赖DINOv2等预训练模型，这限制了在全新领域的泛化能力。改进方向可以是探索自监督的语义对齐目标，或者设计多模态对齐机制从文本等其他模态获取语义信号。
未来方向
作者提出的未来工作方向主要包括：扩展COMiT到视频数据，其中时间冗余和长程结构为离散表示学习带来额外挑战和机会。基于论文成果可以延伸的方向有：(1)自适应任务依赖的视觉tokenization：通过强化学习或元学习自动学习最优crop策略，使模型能够根据不同任务（如分类、检测、分割）自动调整编码过程；(2)多模态扩展：将通信启发机制扩展到文本-图像对齐，使token序列能够同时编码视觉和语言信息，支持更丰富的推理任务；(3)层次化消息结构：设计具有显式层次结构的消息表示，使不同抽象级别的信息（如场景-对象-部分）自然分布在token序列的不同部分；(4)动态token预算：根据图像复杂度动态调整token数量，简单图像使用少量token，复杂图像使用更多token，以提高效率；(5)理论分析：深入理解递归crop聚合为什么能诱导对象中心表示，是否可以从信息论或注意力机制角度提供理论解释；(6)实时应用：优化编码效率使COMiT能够在实时系统中部署，如移动端视觉助手或AR/VR应用；(7)与其他生成范式结合：探索将COMiT与扩散模型、自回归模型等其他生成范式结合，leverage各自优势。
复现评估
COMiT的复现性较好。代码和预训练模型已在GitHub上公开（https://github.com/Araachie/comit），包含完整的训练和评估脚本。论文提供了详细的实现细节（附录A），包括模型架构（基于DiT）、训练超参数（Adam优化器，学习率3e-4，200 epoch）、crop策略（Poisson分布随机化K）和量化设置（FSQ级别[8,8,8,5,5,5]）。所有模型在ImageNet1k上训练，使用32乘GH200 GPU，训练资源需求较高但可及。评估基准（ImageNet100、MSCOCO、Visual Genome）都是公开数据集，评估协议在附录B中详细描述。probe网络设计简洁统一，公平比较不同tokenizer。唯一的潜在挑战是计算资源：训练COMiT-L需要32个GH200 GPU 200个epoch，这对于大多数研究团队来说成本较高。不过，作者提供了预训练模型，可以直接用于下游任务而无需重新训练。论文还提供了详细的超参数表格和消融实验设置，便于复现和扩展。总体而言，COMiT的可复现性在当前大规模生成模型中属于较好水平。
论文图表
