目录
一、摘要
二、引言
三、相关工作
四、方法
基于扩散先前的离散标记化器利用广告
架构
阶段 1A：模式匹配预训练
阶段 1B：模式搜索后训练
采样
第二阶段：潜在生成建模
五、Coovally AI模型训练与应用平台
六、实验
主要结果
分析
七、结论
论文题目：Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization
论文链接：
https://arxiv.org/pdf/2503.11056v1
一、摘要
自VQGAN和潜在扩散模型等流行的视觉生成框架出现以来，最先进的图像生成系统一般都是两阶段系统，首先将视觉数据标记化或压缩到低维潜在空间，然后再学习生成模型。标记化训练通常采用标准方法，即根据MSE、实际损失和对抗损失的组合对图像进行压缩和重建。扩散自动编码器在之前的工作中已被提出，作为一种学习端到端感知导向图像压缩的方法，但在ImageNet-1K重构这一竞争性任务中尚未显示出最先进的性能。我们提出的FlowMo是一种基于变换器的扩散自动编码器，它能以多种压缩率实现最新的图像标记化，而无需使用卷积、对抗损失、空间对齐的二维拉编码或从其他标记化器中提炼。我们的主要见解是，FlowMo训练应分为模式匹配前训练阶段和模式搜索后训练阶段。此外，我们还进行了广泛的分析，并探索了在FlowMo标记器上训练生成模型的方法。
二、引言
扩散模型和离散自回归模型等生成模型在图像和视频内容创建方面的应用引人注目。虽然级联系统和单级系统等不太常见的系统也被探索过，但最先进的视觉生成系统一般包括两个阶段：首先，学习一个“标记器”，将像素数据压缩到一个更小、更容易处理的离散或连续潜空间；其次，在这个压缩的潜空间上训练一个生成模型。由于这种两阶段模式的优势，标记化本身已成为一个活跃的研究领域，许多重要的研究工作都非常关注标记化的训练和设计。在这项工作中，我们重点研究具有离散潜空间的标记化器。
自VQGAN以来，离散图像标记符的主流架构和训练方案已经出现。最先进的图像标记器通常是经过训练的卷积自动编码器，先将视觉数据下采样为二维、空间对齐的潜码，然后再将潜码上采样为经过精算和感知损失正则化的重构。有人提出了与这一设置不同的方法：TiTok使用了基于变压器的架构和一维潜码，但在初始蒸馏阶段依赖于传统的基于CNN的标记化器；ViT-VQGAN使用了基于变压器的编码器和解码器。但总的来说，图像标记化的基本设置仍然占主导地位。
在这项工作中，我们提出了FlowMo（Flow to the Mode，流向模式），它在技术上大大突破了当前图像标记化的技术水平。首先，在处理图像标记化任务时，我们建议使用解码器的整流目标来模拟多模态再构造分布。其次，我们使用完全基于变压器的架构，对一维潜码进行编码和解码。第三，我们对FlowMo进行了端到端优化，而不是从已有的二维标记符中进行提炼，或在二维标记符的潜在空间之上进行编码。
最重要的是，FlowMo通过一个关键的洞察力实现了最先进的性能。FlowMo是一种扩散编码器，这一系统已在之前的工作中进行过探索。然而，在最直观的感知重建基准ImageNet-1K方面，最先进的技术仍然被更传统的基于CNN或GAN的标记化器所主导。FlowMo通过以下关键见解实现了最先进的标记化：对于感知重建任务来说，对重建分布中实际上接近原始图像的模式进行采样比试图匹配所有模式更好。因此，我们将FlowMo的训练分为模式匹配前训练阶段和模式搜索后训练阶段，前者是对系统进行端到端训练，解码器上有扩散损失，就像典型的扩散自动编码器一样；后者是对系统进行训练，有选择性地放弃重建分布中不匹配的模式。
在这一阶段，系统会接受训练，选择性地放弃重建分布中不接近原始图像的模式。我们将在第3节解释这两个阶段。
尽管FlowMo在方法上与之前的研究有很大的不同，但与现有最强的多BPP标记化器相比，FlowMo仍是最先进的。我们的主要贡献如下：
-
我们基于扩散自动编码器和多模态扩散图像转换器 (MMDiT)，提出了一种简单但新颖的图像标记化架构。
-
我们为扩散自编码器提出了一种新颖的两阶段训练方案，包括模式匹配前训练和模式搜索后训练。
-
在每像素0.07比特和每像素0.22比特两种情况下，我们在感知图像标记化的rFID、PSNR和其他指标方面都达到了最新水平。我们还表明，使用FlowMo标记化器训练的生成模型可以与使用传统标记化器训练的生成模型相媲美（尽管不能超越）。
-
我们对FlowMo中的系统设计选择进行了广泛的分析，概述了噪声调度、采样器设计、模型设计、量化策略和后期训练中几个微妙而关键的决定。
三、相关工作
图像标记化。最先进的视觉生成系统一般由两个阶段组成（少数例外情况是尝试直接在像素空间中学习）。第一阶段是“标记化”，通过投射到连续或离散的潜在空间，重新生成像素数据的空间或时空维度。这一范例的主导地位促使许多作品将标记化器本身作为重要组件进行研究。在这项工作中，我们研究的是具有离散潜变量的标记化器。与之前的工作不同，FlowMo是第一个基于扩散自动编码器的架构，在ImageNet-1K重建中达到了最先进的性能。扩散自动编码器 扩散模型在视觉生成中非常流行，而简化框架（如整流模型）已被进一步采用。利用扩散解码器端对端学习自动编码器的想法最早是在中提出的。随后又有许多研究跟进了这一想法，研究了用于表征学习的扩散自动编码器，特别是用于感知导向的图像压缩。扩散后训练。各种研究都试图通过专门的后训练策略来提高扩散模型样本的质量，其目的通常是灌输所需的属性，如美学质量。DDPO和DPOK探索的是强化学习目标，而DRAFT和AlignProp探索的是通过采样链进行反向传播。FlowMo采用基于反向传播的后训练策略，以适应整流的连续噪声时间表和图像标记化的独特设置。
并行工作DiTo在与我们同时进行的工作中研究了用于连续图像标记化任务的扩散自动编码器。与DiTo不同，我们专注于离散图像标记化，并与现有最强大的离散标记化器进行比较。同时，Flex-Tok提出了一种系统，在该系统中，扩散标记化器是在传统连续变异自动编码器（VAE）的潜空间之上学习的，该自动编码器经过感知损失、对抗损失和重建损失的训练。FlowMo不依赖辅助VAE。
四、方法
现有的先进标记化器在视觉生成方面取得了相当大的进步。不过，它们也有很多缺点。首先，它们需要对抗损失，这可能不稳定且难以调整。它们几乎都需要在至少一个训练阶段使用CNN，因此很难利用变压器的硬件效率和广为人知的扩展行为。最后，它们通常利用以前训练过的标记化器进行提炼，以达到最先进的性能。
FlowMo有几个关键目标。我们提出的标记符是：
-
基于扩散先前的离散标记化器利用广告
-
基于扩散。先前的离散标记化器利用广告
这就需要自适应梯度尺度计算、LeCam正则化，或对损失权重进行仔细调整以确保稳定性。相反，我们将使用扩散解码器，因为扩散模型已被证明简单、非常适合多模态建模，而且在大规模应用中非常可靠。
-
纯变换器，一维潜码。
几乎所有最先进的图像标记化框架都采用了基于CNN的架构，或预测了一个局部对齐的二维潜码，但TiTok除外，它仍然依赖于从一个基于CNN的预训练标记化器中提炼出二维潜码。虽然这些选择可能会提供有用的归纳偏差，但我们认为基于变换器的一维潜码标记器最终可能会在大型数据和模型规模上提供更高的效率和灵活性。FlowMo在基于转换器的架构中学习一维潜码，这与MMDiT类似。
-
最新技术。
在扩散自动编码器方面已经有相当多的研究，但FlowMo是第一个在ImageNet-1K上实现最先进的感知重建的技术。
FlowMo实现了这些目标，我们现在将通过对系统的阐述以及随后（在第 4 节中）通过大量的实验、消融和分析加以说明。作为一种扩散自动编码器，鉴于潜在代码c中的信息有限，重建图像x的分布（表示为p(x|c)）必然是多模态的。我们的主要见解是，为了利用扩散自动编码器实现最先进的标记化，应采取多种措施使 p(x|c) 偏向于与原始图像具有高感知相似性的模式。我们通过两个关键要素来实现这一目标： (1) 3.3节中解释的模式搜索后训练，以及 (2) 3.4节中解释的移位采样器。
现在我们将详细介绍FlowMo。这是一种基于MMDiT的模拟变压器架构。然后，我们将解释两个训练阶段：模式匹配预训练（阶段1A）和模式搜索微调（阶段1B）。最后，我们将讨论在FlowMo潜在空间上的生成建模（阶段 2）和采样。
-
架构
FlowMo的结构图如图3所示。FlowMo是一个扩散自动编码器。解码器是一个条件扩散模型，用于学习重建图像的条件分布p(x|c)。
FlowMo由编码器eθ和解码器dθ组成。两者都是基于多模态扩散图像变换器（MMDiT）的变换器。给定一个经过修补的图像x∈Rn和一个初始潜码c0∈Rd（一个全为零的向量），编码器产生一个潜标记序列
cˆ然后通过无查找量化（LFQ）的量化操作q进行元素二值化，得到
按照整流，解码器被训练成从噪声到数据的速度场v建模，定义为
解码器处理xt和c的方式与编码器处理x和c0的方式相同，但通过AdaLN调制]，解码器还接受时间（或噪声水平）参数t来调节每个MMDiT块。编码器eθ和解码器dθ在结构上是对称的，但大小不同，解码器更大更深。我们对所有模型都使用了μP参数化，以简化探索性配置和放大配置之间的超参数转移。
-
阶段 1A：模式匹配预训练
在第一个训练阶段，我们的目标是联合训练编码器eθ和解码器dθ，使量化后的c对x具有最大的信息量，从而使pθ(x|c)与真实分布p(x|c)相匹配，由于量化后的c只包含有限的信息，因此真实分布必然是多模态的。图 4 给出了阶段1A的完整示意图。关于理论依据，我们参考了之前在扩散自动编码器方面的研究，其中指出端到端扩散目标对应于非高斯解码器 p(x|c) 的修正变分下限。
FlowMo被训练成端到端扩散自动编码器。具体来说，FlowMo编码器和解码器是端对端训练，以优化解码器输出上的整流损失Lflow。给定噪声z∼N (0, I )、数据 x∼px 和时间（或噪声水平）t∼pt,t∈[0,1]，我们定义
并优化流量匹配目标
我们还使用学习到的感知距离dperc，通过以下方式对网络的1步去噪预测进行优化
最后，在潜码c上，我们按照的方法使用LFQ的熵损失和承诺损失，即
为简单起见，这两个损失可以用有限标量量化（FSQ）代替，但我们发现LFQ 的性能稍好一些。我们在1A阶段的训练损耗为
损失权重和更多细节见补充材料。图像被重新调整为位于[-1, 1]范围内。噪声水平t是按照 Stable Diffusion 3 ，从厚尾logit-normal分布中采样的。
-
阶段 1B：模式搜索后训练
在这一阶段，我们的目标是优化解码器分布pθ(x|c)，以寻求在感知上与原始图像高度相似的模式。为了实现这一目标，我们冻结了编码器，并在Lflow和后训练目标Lsample的基础上对解码器进行联合训练，后训练目标Lsample的灵感来源于之前在扩散模型后训练方面的工作。
Lsample只是对概率流ODE的n步采样的感知损失，我们通过它进行区分。我们随机采样积分的时间步t1、......、tn，以便在测试时使用不同的采样计划进行实验，并使用梯度检查点来反向传播采样链。图 5 给出了阶段1B的完整示意图。
让dti (xt)表示流量采样更新函数，即
然后，我们定义
第 1B 阶段的全部损失为
我们发现本阶段的λsample值尤为重要，因此使用λsample=0.01。Lsample的权重太小会导致rFID效果不佳，因为网络会遗忘在第1A阶段获得的实际特征；而权重太大则会导致所谓的“奖励劫持”（即解码器dθ与dperc过度拟合）或训练发散。我们在第1A阶段使用LPIPS-VGG网络作为感知网络，但在第 1B 阶段，我们使用ResNet作为感知网络。我们在补充材料中解释了这一设计选择和其他设计选择，并通过消融法证明，仅仅在Lperc中的1步去噪预测中使用ResNet网络是无效的；因此有必要引入后训练阶段。这一阶段的计算成本很高，因为它需要通过整个采样链进行反向传播。我们使用n=8，并发现它在采样时能很好地适用于其他 n。
-
采样
给定一个量化的潜码c，通过求解概率流ODE来采样给定c的重建图像的多模态分布（表示为 p(x|c)）。
给定初始 x1∼N (0, I)。另外，FlowMo 也可用于计算样本对数似然，方法是按照流匹配反向求解流 ODE。
在推理时，我们用可调整的移动超参数给出的时间步距对整流ODE进行积分。图5.第1B阶段。冻结编码器eθ将输入图像编码为c，作为解码器 dθ 的条件，解码器dθ通过整个采样链的反向传播进行训练。我们还使用 Lflow 进行协同训练，这与图4中的方法相同。
ρ 表示采样时间步向低噪声水平集中。对于时间步距：
设置ρ=1相当于通常的线性间隔整流ODE采样器。在极端情况下，让ρ→∞意味着在t=1时迈出一大步，相当于给定c对x进行回归。我们使用ρ=4，这相当于在接近 t=1时迈出一大步，并将采样集中到 p(x|c) 的平均值上，同时在低噪声水平上仍然花费了大量的采样FLOP，而之前的研究表明，这一选择对rFID至关重要。我们的采样器大大提高了rFID和PSNR。
-
第二阶段：潜在生成建模
与其他使用离散空间训练自动编码器的工作一样，我们验证了我们的标记器可以用来训练高质量的第二阶段生成模型。我们使用MaskGiT，这一阶段的设置主要来自MaskGiT和TiTok。
五、Coovally AI模型训练与应用平台
如果你也想要进行模型训练或模型改进，Coovally平台满足你的要求！
Coovally平台整合了国内外开源社区1000+模型算法和各类公开识别数据集，无论是YOLO系列模型还是Transformer系列视觉模型算法，平台全部包含，均可一键下载助力实验研究与产业应用。
而且在该平台上，无需配置环境、修改配置文件等繁琐操作，一键上传数据集，使用模型进行训练与结果预测，全程高速零代码！
具体操作步骤可参考：YOLO11全解析：从原理到实战，全流程体验下一代目标检测
平台链接：https://www.coovally.com
如果你想要另外的模型算法和数据集，欢迎后台或评论区留言，我们找到后会第一时间与您分享！
六、实验
-
主要结果
标记化。对于标记化这一主要任务，所有标记化器都采用图像输入，将其编码为量化潜像，然后重建图像。所有标记化器都是在ImageNet-1K上训练的。重建质量用rFID 、PSNR、SSIM和LPIPS等指标来衡量。我们在分辨率为256×256的ImageNet-1K验证集上进行评估。
标记化器量化潜码所包含的信息量以每像素比特（BPP）为单位。
其中 S是潜码序列长度，V是标记词词汇量，标记化器的BPP计算公式为：
在不同BPP下训练的标记化器不能进行苹果对苹果的比较，因为获取更多比特会提高性能。因此，我们训练了两种模式，以匹配现有最先进的标记符。
FlowMo-Lo以0.0703 BPP进行训练，以匹配OpenMagViT-V2的BPP。FlowMo-Hi以0.219 BPP进行训练，以匹配LlamaGen-32的BPP。除了BPP之外，FlowMo-Lo和FlowMo-Hi都是完全相同的FlowMo架构。为了与我们进行比较的tokens生成器的BPP相匹配，我们修改了FlowMo tokens词汇的大小或tokens数量。就=rFID、PSNR和SSIM而言，我们在两个BPP上都取得了最先进的结果。FlowMo性能优异的唯一例外是LPIPS指标，FlowMo在该指标上仍然表现不佳。可在补充材料和我们的网站上查看具体的可视化比较。
生成。在生成任务中，我们对不同标记化器生成的编码标记序列训练生成模型MaskGiT。然后，我们通过各自的标记化器对生成的标记序列进行解码，并用图像生成指标来衡量图像质量，即FID、sFID、Inception Score以及Precision和Recall。生成指标在ImageNet-1K基准上进行评估。尽管实现了更强的标记化，但使用FlowMo作为下游生成模型的标记化器，只能改善部分生成指标，而不能改善所有生成指标。标记化质量和生成质量之间存在有趣而复杂的相互作用，我们希望在未来的工作中改进结果。
表2中的模型大小、批量大小和训练长度等训练超参数与托克生成器相同。具体来说，我们在批量大小为128、学习率为1×10-4的两个token- enizers上都训练了300个epoch。其他超参数设置见补充材料。虽然我们没有足够的资源来训练最先进的生成模型（例如，需要1,000 - 1,080个历元的训练，批量大小为1,024 - 2,048），但我们在这里的目标并不是为生成模型设定最先进的水平，而是从生成模型所提供的潜在空间的角度对标记生成器进行公平的比较。
其他比较。在表3中，我们还对FlowMo和DiTo进行了标记化比较。为了进行比较，我们使用连续潜空间训练标记化器，将Lay- erNorm作为最终编码层，并使DiTo中提出的 “噪声同步 ”增强技术。我们还均衡了总体潜空间大小，使用256个标记，标记维度为16。实验细节。我们所有的模型都是μP参数化的，因此超参数可以在探索性配置和放大配置之间“μ-转移”。我们的大部分实验是在混合硬件（A6000、L40S、A100、H100 或 H200 GPU）上进行的，ViT补丁大小为8，隐藏维度为768；这也是表4中用于消融研究的配置。
我们的最终实验（FlowMo-Lo、FlowMo-Hi）是在8×H100节点上进行的，使用的是放大模型，补丁大小为4，仅解码器的隐藏维度增加到 1152。所有其他超参数都直接进行了μ转移。我们对FlowMo-Lo进行了100个epoch的预训练，对FlowMo-Hi进行了80个epoch的预训练，然后对两者进行了约 1个epoch的后训练。在预训练中，rFID的性能没有达到饱和。更长时间的训练会耗费更多资源，但可能会进一步提高性能。
我们对FlowMo-Lo和FlowMo-Hi进行了批量大小为128、学习率为10-4的训练。我们按照EDM2，每一步对编码器和解码器的MLP块中的权重矩阵进行强制归一化，以防止激活和权重矩阵爆炸。所有模型都在PyTorch中以bfloat16的精度进行训练。我们使用Adam优化器，其值为 (β1，β2 )=(0.9,0.95)，因为我们注意到，可能由于bfloat16精度或变压器序列长度较长，β2越高越不稳定。我们将编码器的学习率设置为0，训练步数为200,000 步。我们使用的是指数移动平均值，速率为0.9999。
-
分析
消融研究。我们对标记化性能进行了消融研究，分析了不同决策对rFID和其他指标的影响。实验设置与我们主要的ImageNet-1K标记化实验相同。
我们首先对表4中第 1A 阶段的设计决策进行消融研究。我们对这些决策进行分析，并在适用时参考先前的扩散自动编码器作品。在这些实验中，我们使用了减小μP宽度的配置，以提高实验效率。
双倍贴片尺寸。在视觉转换器和去肋架构中，“补丁尺寸”决定了图像转换序列的长度，补丁尺寸越小，模型的计算量越大。之前关于像素空间中基于变换器的扩散模型的研究已经指出了小补丁尺寸的重要性。我们证实了这一点：增大补丁尺寸会影响所有指标。
使用MSE训练的编码器。之前的扩散自动编码器研究已经探索了首先用MSE或基于LPIPS的回归目标来训练自动编码器，也有可能用对抗损失来训练，然后用训练得到的冻结特征作为扩散解码器的条件。虽然这种训练方式可能会提高PSNR，但为了实现最佳的rFID，必须在所有噪声水平下对整个系统进行端到端训练，以确保潜码包含有助于在所有噪声水平下进行速度估计的特征。
无感知损失。之前在扩散自动编码器方面的研究已经注意到感知损失对扩散解码器1步去噪预测的重要性。如果没有感知损失，性能就会下降。
FSQ量化。与我们的默认量化策略LFQ相比，我们发现FSQ的最终训练损失和成对度量性能略胜一筹，但却以最关键的rFID为代价。我们还尝试了其他量化技术，但LFQ的性能最佳。
对数正态噪声。用于整流训练的logit-normal噪声计划将 t∈{0, 1} 的概率质量分配为0。虽然这种方法适用于t=0（此时整流目标变得奇异），但由于我们的调节信号非常强，因此在t=1（纯噪声）时对v(xt,t)的估计至关重要。使用通常的logit-normal时间表会导致PSNR和rFID下降。取而代之，我们使用厚尾 logit-normal噪声计划：我们从 [0, 1] 上的均匀分布中抽取10%的噪声水平样本。如果不做这种修改，t=1时的速度估计值就会不准确，从而导致低PSNR和变色。示例见图 7。
非移位采样器。如第3.4节所述，我们使用时间步距t对整流ODE进行积分，时间步距按比例向t=0集中。使用线性间隔的时间步会降低性能。
无引导。与GAN标记化器相比，FlowMo的一个优势是能够利用无分类器引导和引导间隔等技术来提高感知质量。如果不在有限的时间间隔内应用分类器引导，rFID就会恶化。
表5是纳入阶段1B的消融研究。如果没有这一关键阶段，所有指标的性能都会下降。有趣的是，尽管只在Lsample上进行了训练，但这一阶段也提高了 PSNR。我们在图8中说明了FlowMo解码器在后期训练后如何保持多模态。
局限性。FlowMo的主要局限在于推理时间。FlowMo需要多次模型前向传递（我们在工作中使用了n=25步）才能从给定量化潜码的解码器中生成一个样本。相比之下，几乎所有其他标记化器都是基于GAN的，因此只需一次前向传递即可重新构建给定图像。尽管如此，将扩散模型提炼为单步或几步系统的工作仍然相当多，这些技术也适用于FlowMo。
七、结论
我们介绍了基于变压器的扩散自动编码器FlowMo，它利用扩散自动编码器的训练和采样应着眼于寻求感知高质量模式这一关键见解，实现了最先进的图像标记化。我们在256×256分辨率的竞争性ImageNet-1K基准上展示了最先进的性能，与之前的许多工作不同，我们没有使用二维空间对齐潜码、对抗损失、来自辅助标记化器的代理目标，也没有进行压缩。
