2025-03-20 14:07 北京
FlowMo，一种基于Transformer的扩散自编码器。
机器之心报道
机器之心编辑部
当我们看到一张猫咪照片时，大脑自然就能识别「这是一只猫」。但对计算机来说，它看到的是一个巨大的数字矩阵 —— 假设是一张 1000×1000 像素的彩色图片，实际上是一个包含 300 万个数字的数据集（1000×1000×3 个颜色通道）。每个数字代表一个像素点的颜色深浅，从 0 到 255。
为了更加高效地从成千上万张图像中学习，AI 模型需要对图片进行压缩。比如当前最先进的图像生成模型，第一步就是一个名叫 tokenization 的操作，用于执行此操作的组件叫 tokenizer。tokenizer 的主要目标是将原始图像压缩到一个更小、更易处理的潜在空间，使得生成模型能够更高效地学习和生成。因此，如何得到更好的 tokenizer 是该领域的研究者非常关心的问题。
在一篇新论文中，来自斯坦福大学李飞飞、吴佳俊团队的研究者提出了一种名叫「FlowMo」的改进方案（论文一作是斯坦福大学计算机科学博士生 Kyle Sargent）。
FlowMo 的训练分为两个阶段：第一阶段先学习如何全面捕捉图像的多种可能重建结果，第二阶段则学习如何从这些可能中选择最接近原图的重建方案。这种方法既保证了图像重建的多样性，又确保了重建质量，使得 FlowMo 在 ImageNet-1K 数据集上展现出了领先的重建性能。
论文标题：Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization
研究背景
自从 VQGAN 和潜在扩散模型等视觉生成框架问世以来，最先进的图像生成系统通常采用两阶段设计：先将视觉数据压缩到低维潜在空间进行 tokenization，再学习生成模型。
Tokenizer 训练一般遵循标准流程，即在均方误差（MSE）、感知损失和对抗损失的组合约束下压缩并重建图像。扩散自编码器曾被提出作为学习端到端感知导向图像压缩的方法，但在 ImageNet-1K 重建这一竞争性任务上尚未达到最先进水平。
李飞飞团队提出了 FlowMo，一种基于 Transformer 的扩散自编码器，它在多种压缩率下实现了图像 tokenization 的新性能标准，且无需使用卷积、对抗损失、空间对齐的二维潜在编码，或从其他 tokenizer 中蒸馏知识（这与传统的基于 GAN 的 tokenizer，如 VQGAN，非常不同）。
研究的关键发现是 FlowMo 训练应分为模式匹配预训练阶段和模式寻求后训练阶段。此外，研究者进行了广泛分析，并探索了基于 FlowMo tokenizer 的生成模型训练。
图 1：无论是在低比特率训练（FlowMo-Lo）还是高比特率训练（FlowMo-Hi）下，FlowMo 模型都实现了最先进的 image tokenization 性能。
作者强调，尽管基于 GAN 的 tokenizer 在图像 tokenization 任务上已经取得了很好的性能，但 FlowMo 提供了一种简单且不同的方法。
FlowMo 方法
众所周知，基于 Transformer 的扩散自编码器包含编解码结构，因此 FlowMo 也是由编码器 e_θ 和解码器 d_θ 组成，其核心架构遵循了 MMDiT，在 Transformer 的架构中学习一维潜在空间。
图 2：FlowMo 架构概览
首先，编码器将输入图像转换为潜空间向量 c，然后解码器则根据潜空间 c 学习重建图像的条件分布，与旨在产生单一确定性输出的传统方法不同，FlowMo 的解码器会生成可能的重建分布，从而更好地捕捉图像重建中固有的模糊性。
FlowMo 架构的主要包括以下四点：
基于 Transformer 的设计：编码器和解码器都使用 Transformer 架构，从而能够更有效地处理图像数据。
一维潜空间表示：FlowMo 产生紧凑的潜在表示，使其适用于下游生成建模任务。
量化层：编码器的输出被量化以创建离散 token，从而实现更高效的压缩。
扩散过程：解码器使用扩散过程逐渐将随机输入去噪为高质量重建。
在 FlowMo 架构中的一个核心创新点是其两阶段训练策略，这一策略使得重建分布偏向于原始图像具有高度感知相似性的模式。
阶段 1A：模式匹配预训练
图 3：FlowMo 的训练过程结合了基于流的损失和感知损失，以引导模型实现高质量的重建。
在阶段 1A 中，FlowMo 通过联合训练编码器与解码器，以实现两个核心目标：最大化潜在编码的信息量，并使其重建分布与真实分布相匹配。这一训练过程巧妙地结合了多种损失函数，展现出其独特的技术优势：
修正流损失（Rectified flow loss）：引导扩散过程向目标图像分布靠拢，确保生成结果的准确性；
感知损失（Perceptual loss）：保证了重建图像在视觉上与原始图像高度相似；
熵损失（Entropy loss）：鼓励生成多样化的潜在编码，避免模式单一化；
承诺损失（Commitment loss）：使得编码器输出与量化表示尽可能接近，进一步优化了模型的稳定性与效率。
具体而言，FlowMo 作为扩散自动编码器进行端到端训练，以优化解码器输出上的修正流损失 L_flow，在过程中使用了 L_perc 来监督图像生成中的去噪预测，同时在潜空间 c 上，作者还结合了 LFQ 的熵损失和承诺损失来进行训练。其中损失函数的数学表达式如下所示：
结合这些损失函数，并最终得到了第一阶段的损失表达式：
阶段 1B：模式探索后训练
在第二阶段中，FlowMo 的核心目标是优化解码器分布 pθ(x∣c)，以寻找那些在感知上与原始图像高度相似的模式。为实现这一目标，FlowMo 采用了创新的训练策略：首先冻结编码器，随后在 Lflow 的基础上，联合训练解码器，并引入受扩散模型训练后的 x_0 来生成目标 Lsample。这一过程通过以下步骤实现：
概率流 ODE：通过少量步骤的概率流常微分方程（ODE）集成；
感知损失计算：在生成样本后，模型会计算其与原始图像之间的感知损失，确保重建结果在视觉上与原始图像保持一致；
解码器参数更新：基于感知损失，FlowMo 对解码器参数进行优化。
图 4：模式搜索训练过程，编码器处于冻结状态（雪花表示），而解码器则进行微调以提高感知质量。
如上图所示，其中 FlowMo 通过冻结编码器，集中精力优化解码器，使其在重建图像时更加注重感知相似性，从而进一步提升生成图像的质量与真实感。对概率流 ODE 进行积分的 n 步样本感知损失 Lsample 如下所示：
第二阶段模式探索损失如下所示：
采样过程
为了生成重构图像，FlowMo 通过求解概率流 ODE，对给定一维潜空间 c 的重建图像的多模态分布进行采样
FlowMo 采样方法的一项关键创新是使用「移位」采样器。FlowMo 不使用统一的时间步长间隔，而是采用可调的移位超参数，将采样步骤集中在扩散过程的某些区域，从而提高感知质量。
采样过程需要多次前向通过解码器模型，这在计算上很昂贵，但可以产生高质量的结果。
实验结果分析
主要结果
FlowMo 在多个比特率设置下（0.07 BPP 和 0.22 BPP）与当前最先进的 tokenizer 进行了比较，在重建 FID（rFID）、PSNR 和 SSIM 指标上均取得了最佳结果。在 0.07 BPP 设置下，FlowMo-Lo 的 rFID 为 0.95，相比 OpenMagViT-V2 的 1.17 有显著提升；在 0.22 BPP 设置下，FlowMo-Hi 的 rFID 为 0.56，略优于 LlamaGen-32 的 0.59。
表 1. tokenization 结果。
消融实验分析
研究团队进行了大量消融实验，分析了 FlowMo 设计中的关键决策：噪声调度、量化策略、模型架构和后训练策略等。结果表明，thick-tailed logit-normal 噪声分布、shifted sampler 和后训练阶段对模型性能至关重要。
图 5：噪声调度导致失真的可视化案例。
特别是，没有模式寻求后训练阶段，FlowMo-Lo 的 rFID 会从 0.95 下降到 1.10，FlowMo-Hi 的 rFID 会从 0.56 下降到 0.73。
表 2：后训练消融实验结果。
生成任务验证
在生成任务中，基于 FlowMo 训练的 MaskGiT 在某些指标上表现优于基于 OpenMagViT-V2 训练的模型，但在 FID 上略逊一筹（4.30 vs 3.73）。这表明 tokenizer 质量与下游生成模型质量之间存在复杂关系，需要进一步研究。
表 3：生成模型指标对比。
图 6：生成图像对比。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com
