斯坦福大学李飞飞团队新突破！FlowMo 革新图像 Tokenizer

当我们悠然刷着手机，看到一张可爱猫咪的照片时，大脑会瞬间识别出「这是一只猫」，这一切不过是电光火石间的事儿。但在计算机的 “眼中”，情况却复杂得超乎想象。假设这是一张1000×1000像素的彩色照片，在计算机的世界里，它摇身一变，成了一个包含300万个数字的庞大数据集（1000×1000×3个颜色通道）。每一个数字都承担着重要使命，它们代表着对应像素点颜色的深浅程度，取值范围从0到255，0是深沉的黑色，255则是明亮的白色，而中间的数值就是各种过渡的灰色调。这些数字密密麻麻地排列在一起，就像一幅神秘的数字拼图，等待计算机去解读。在人工智能领域，随着技术的飞速发展，图像数据如同潮水般涌来。AI模型要想高效地从成千上万张这样复杂的图像中学习，就必须找到一种巧妙的方法来处理这些庞大的数据。就好比一个人面对堆积如山的书籍，要想快速获取知识，就得有一套高效的阅读和整理方法。对AI模型而言，这个方法就是对图片进行压缩。在当前最先进的图像生成模型里，第一步操作至关重要，它有个专业的名字叫tokenization，执行这一操作的组件就叫做tokenizer，中文可译为 “分词器”，不过在图像领域，它的作用远超普通的分词概念。 tokenizer堪称图像领域的 “神奇压缩大师”，它的主要目标是将原始图像，也就是那一大串复杂的数字矩阵，压缩到一个更小、更易于处理的潜在空间中。想象一下，把一个庞大的图书馆压缩成一个小巧的移动硬盘，却又能保留其中的关键信息。经过tokenizer处理后，图像被转化成了一种更紧凑的表示形式，这些紧凑的 “token” 就像是经过提炼的知识精华，让后续的生成模型能够更高效地学习图像的特征和模式，进而更加快速、准确地生成新的图像。比如在生成逼真的风景图像时，生成模型可以依据tokenizer处理后的信息，快速勾勒出大致的山川、河流形状，再逐步细化色彩和纹理。因此，如何研发出性能更优、效果更好的tokenizer，成了图像生成、处理等相关领域研究者们心心念念、全力攻克的重要课题。最近，来自斯坦福大学的李飞飞、吴佳俊团队在这个难题上取得了新突破。他们在一篇新论文中，提出了一种名为「FlowMo」的创新改进方案，论文的第一作者是斯坦福大学计算机科学专业的博士生Kyle Sargent。这一方案就像一把新钥匙，为解决图像tokenizer的优化问题，打开了一扇充满可能的新大门。

FlowMo：重塑图像 tokenization 的新范式

在图像生成与处理领域，模型的训练与优化一直是研究的核心。斯坦福大学李飞飞团队提出的 FlowMo，为这一领域带来了极具创新性的解决方案。FlowMo 的训练匠心独运，分为两个关键阶段，每个阶段都有着明确的目标与作用。

第一阶段：博采众长，捕捉多元重建可能

在第一阶段，FlowMo 致力于全面捕捉图像的多种可能重建结果。想象一下，面对一幅风景图像，不同的人对画面的理解和重构方式可能千差万别。FlowMo 就像一位求知若渴的学习者，通过大量的数据训练，学习从各个角度、各种风格对图像进行重建。在此过程中，它广泛探索图像潜在的各种表达方式，不放过任何一种可能的组合。这一阶段的训练，赋予了 FlowMo 丰富的 “想象力”，为后续生成多样化的图像重建结果奠定了基础。

第二阶段：精益求精，筛选最优重建方案

有了第一阶段积累的大量重建可能性，第二阶段的 FlowMo 开始学习如何从众多可能中挑选出最接近原图的重建方案。这就好比在众多候选答案中，精准找出最正确的那一个。通过细致的学习与对比，FlowMo 不仅能够保证图像重建的多样性，还能显著提升重建质量。正是这种独特的两阶段训练方式，让 FlowMo 在 ImageNet-1K 数据集上大显身手，展现出了领先的重建性能，超越了许多同类模型。

研究背景：图像生成技术的演进与挑战

随着 VQGAN 和潜在扩散模型等视觉生成框架的相继问世，图像生成系统迎来了全新的发展阶段。当下，最先进的图像生成系统普遍采用两阶段设计：首先，将视觉数据压缩到低维潜在空间，完成 tokenization 操作；随后，基于这些压缩后的数据，学习并构建生成模型。这一设计思路大大提高了模型处理图像数据的效率，使得生成高质量图像成为可能。

在 Tokenizer 的训练过程中，通常遵循一套标准流程。模型在均方误差（MSE）、感知损失和对抗损失的共同约束下，对图像进行压缩与重建。然而，扩散自编码器虽然曾被提出用于学习端到端感知导向图像压缩，但在 ImageNet-1K 重建这一极具竞争性的任务中，始终未能达到最先进水平，在重建精度、多样性等方面存在一定的不足。

FlowMo：创新架构引领新突破

李飞飞团队提出的 FlowMo，全称基于 Transformer 的扩散自编码器，为图像 tokenization 带来了全新的技术路径。与传统方法不同，FlowMo 另辟蹊径，无需使用卷积、对抗损失、空间对齐的二维潜在编码，也无需从其他 tokenizer 中蒸馏知识。而传统的基于 GAN 的 tokenizer，如 VQGAN，在训练过程中高度依赖这些技术。FlowMo 凭借独特的架构和训练方法，在多种压缩率下实现了图像 tokenization 性能的显著提升，树立了新的性能标准。

论文与项目信息

论文标题：Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization
论文地址：https://arxiv.org/pdf/2503.11056v1
项目主页：Flow to the Mode: Mode-Seeking Diffusion Autoencoders for State-of-the-Art Image Tokenization

因此，FlowMo 的出现，为图像生成与处理领域注入了新的活力，有望推动该领域朝着更高精度、更具多样性的方向发展，为后续研究和应用提供了极具价值的参考与借鉴。

研究背景

在人工智能图像生成领域，早期的图像生成模型受限于算法和算力，生成的图像质量与真实感都不尽如人意。直到VQGAN和潜在扩散模型等视觉生成框架横空出世，彻底改变了这一局面，引领图像生成技术迈向新的阶段。如今，最前沿的图像生成系统普遍采用两阶段设计。在第一阶段，系统将高维度、数据量庞大的视觉数据，压缩到低维潜在空间，完成tokenization操作，这就好比把杂乱无章的原始信息进行初步梳理，整理成有序的、更易处理的形式。在第二阶段，基于经过tokenization处理的数据，模型学习如何生成逼真的图像。在训练tokenizer时，研究人员一般会遵循一套标准流程。模型在均方误差（MSE）、感知损失和对抗损失的共同约束下，对图像进行压缩，随后再重建图像。均方误差能让模型关注像素层面的差异，感知损失确保重建图像在人类感知层面的合理性，对抗损失则推动模型生成更具真实感的图像。扩散自编码器作为一种新兴的技术，旨在学习端到端感知导向的图像压缩方法。它试图通过一种更直接的方式，让模型从整体上理解图像，实现高效的图像压缩与重建。然而，在ImageNet-1K重建这项竞争异常激烈的任务中，扩散自编码器始终难以跻身最先进技术行列，重建的图像在细节还原、视觉效果等方面与顶尖模型存在差距。面对这一困境，斯坦福大学李飞飞团队提出了FlowMo，一种基于Transformer的扩散自编码器。与传统方法截然不同，FlowMo另辟蹊径，无需依赖卷积、对抗损失、空间对齐的二维潜在编码，也无需从其他tokenizer中获取知识。而传统的基于GAN的tokenizer，例如VQGAN，在训练时对这些技术高度依赖。FlowMo凭借独特的架构和训练方式，在多种压缩率下，实现了图像tokenization性能的大幅提升，为这一领域树立了新的标杆。

在探索如何优化图像 tokenizer 的过程中，我看出李飞飞团队针对 FlowMo 展开了深入研究，并取得了极具价值的关键发现。其团队指出，FlowMo 的训练应分两个既相互独立又紧密关联的阶段进行，即模式匹配预训练阶段和模式寻求后训练阶段。在模式匹配预训练阶段，FlowMo 就像一位积极的学习者，通过大量数据的学习，试图捕捉图像中各种可能的重建模式。这个过程中，它通过联合训练编码器和解码器，不仅要让模型最大化潜在编码的信息量，还要让重建分布与真实图像分布尽可能匹配。具体来说，FlowMo 作为扩散自动编码器进行端到端训练，通过优化解码器输出上的修正流损失 L_flow，同时借助 L_perc 监督图像生成中的去噪预测。在潜空间层面，结合 LFQ 的熵损失和承诺损失，引导模型在复杂的图像信息中挖掘潜在的模式和规律，建立起图像特征与潜空间表示之间的联系，为后续的重建工作打下坚实基础。当顺利完成模式匹配预训练阶段后，FlowMo 便进入模式寻求后训练阶段。此时，模型的目标发生了转变，从广泛收集可能的模式，转为从已学习到的众多模式中，筛选出与原始图像在感知上最为相似的模式。为实现这一目标，研究者创新性地采用了特定训练策略。首先将编码器冻结，使模型能够集中精力对解码器分布 pθ(x∣c) 进行优化。在 Lflow 的基础上，联合训练解码器，并引入受扩散模型训练后的 x_0 来生成目标 Lsample。通过这种方式，FlowMo 不断调整解码器的参数，使其在重建图像时，能够更加精准地还原原始图像的细节和特征，显著提升重建图像的质量和真实感。此外，研究者并未局限于 FlowMo 训练过程的研究，还展开了广泛而深入的分析。他们探索了基于 FlowMo tokenizer 的生成模型训练，评估 FlowMo 在不同场景下的表现，分析其优势与不足，试图挖掘 FlowMo 在图像生成领域的最大潜力，期望通过这种探索，推动基于 FlowMo 的图像生成技术走向成熟，为该领域的发展提供更多可能。

上图反映出无论是在低比特率训练（FlowMo-Lo）还是高比特率训练（FlowMo-Hi）下，FlowMo 模型都实现了最先进的 image tokenization 性能。

作者强调，尽管基于 GAN 的 tokenizer 在图像 tokenization 任务上已经取得了很好的性能，但 FlowMo 提供了一种简单且不同的方法。

FlowMo 方法

众所周知，基于 Transformer 的扩散自编码器包含编解码结构，因此 FlowMo 也是由编码器 e_θ 和解码器 d_θ 组成，其核心架构遵循了 MMDiT，在 Transformer 的架构中学习一维潜在空间。

上图为FlowMo 架构概览

在深入了解 FlowMo 在图像重建领域的卓越表现前，我们先来认识其独特的工作流程和架构设计。当一幅图像进入 FlowMo 系统，首先映入眼帘的是编码器。编码器就像一位专业的数据翻译，将输入的图像转化为潜空间向量 c。这些潜空间向量，宛如图像的 “数字密码”，蕴含了图像的关键特征与信息。随后，解码器接过 “接力棒”，基于潜空间 c 学习重建图像的条件分布。与传统方法不同，传统方法试图给出单一确定的图像重建结果，就好比在多个可能答案中只选择一个。而 FlowMo 的解码器则另辟蹊径，它会生成一系列可能的重建分布。这是因为在图像重建过程中，由于图像本身的复杂性以及数据在转换过程中的信息损失，往往存在一定的模糊性。FlowMo 这种生成多种可能重建分布的方式，就像为图像重建提供了多个 “候选方案”，能更好地捕捉这种模糊性，大幅提升重建结果的全面性与准确性。FlowMo 之所以能在图像重建领域脱颖而出，其精妙的架构设计功不可没。FlowMo 采用基于 Transformer 的设计，编码器和解码器均运用 Transformer 架构。Transformer 架构凭借强大的自注意力机制，能够对图像数据进行深入剖析，捕捉图像中不同区域之间的关联，相较于传统架构，能更高效地处理图像数据。FlowMo 产生一维潜空间表示，这种表示形式不仅紧凑，还保留了图像的关键信息，为下游生成建模任务提供了便利，就像为后续模型搭建了一个稳固的 “数据基石”。编码器的输出会经过量化层处理，量化层的作用是将编码器输出转化为离散 token，这就好比将连续的信息进行分块处理，从而实现更高效的压缩，减少数据存储和传输的负担。在解码器环节，FlowMo 采用扩散过程，从随机输入出发，通过逐步去噪，就像层层揭开迷雾，最终生成高质量的重建图像。在 FlowMo 架构众多创新点中，两阶段训练策略堪称核心。在第一阶段，模型通过大量数据的学习，探索并掌握图像的多种重建模式，尽可能全面地捕捉图像重建的各种可能性。进入第二阶段，模型基于第一阶段的学习成果，对重建分布进行优化，使其更偏向于与原始图像在感知上高度相似的模式。这种训练策略使得 FlowMo 能够在保证重建图像多样性的同时，显著提升重建图像与原始图像的相似度，在图像重建领域实现了新的突破。

阶段 1A：模式匹配预训练

上图为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 对解码器参数进行优化。

上图为模式搜索训练过程，编码器处于冻结状态（雪花表示），而解码器则进行微调以提高感知质量。

其中 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。

下表为tokenization的结果:

消融实验分析

为了深度剖析 FlowMo 各个设计要素对模型性能的影响，研究团队开展了一系列细致且全面的消融实验。消融实验就像是给模型做 “拆解手术”，通过依次去除或改变模型的特定组件或策略，观察模型性能的变化，以此来确定每个部分的重要性。在这次针对 FlowMo 的研究中，噪声调度、量化策略、模型架构和后训练策略等，都成为了研究团队重点关注的对象。噪声调度在扩散模型里，决定了噪声添加和去除的过程，对生成样本的质量和多样性影响显著。研究团队尝试不同的噪声分布类型，经过多次试验后发现，thick - tailed logit - normal 噪声分布能让模型在生成过程中，更好地捕捉数据的复杂特征，为高质量的图像重建奠定基础。量化策略决定了如何将连续的潜空间表示转换为离散的 token，这对模型的压缩效率和重建精度有重要影响。团队在实验中对多种量化策略进行测试，旨在找到最适合 FlowMo 的方案。模型架构作为 FlowMo 的 “骨架”，直接关系到模型的性能表现。无论是 Transformer 架构模块的参数设置，还是编码器、解码器之间的连接方式，都在实验的考察范围内。研究人员通过改变架构的不同参数和结构，分析模型在图像重建任务中的表现差异，进一步优化模型架构，提升其处理图像数据的能力。后训练策略同样不容忽视。在完成初步训练后，后训练阶段可以对模型进行微调，让模型朝着更优的方向发展。研究团队在实验中，尝试多种后训练方法，分析模型在重建精度、视觉效果等方面的提升情况。经过大量的消融实验，研究结果清晰地表明，thick - tailed logit - normal 噪声分布、shifted sampler 和后训练阶段，对 FlowMo 的模型性能起着至关重要的作用。thick - tailed logit - normal 噪声分布为模型的生成过程提供了更丰富的信息；shifted sampler 通过独特的采样方式，提高了模型生成样本的质量；后训练阶段则对模型进行了最后的 “打磨”，让模型在图像重建任务中展现出更卓越的性能。这些研究成果，为 FlowMo 的进一步优化和应用提供了坚实的理论依据。

下图为噪声调度导致失真的可视化的案例