字节提出FlowTok！高效文生图新框架！比PixArt 快3倍！

🌐 社群导航

🔗 点击加入➡️【AIGC/LLM/MLLM/3D/自动驾驶】 技术交流群

🔗 点击加入➡️【2-3月CCF投稿】交流群

数源AI 最新论文解读系列

论文名：FlowTok: Flowing Seamlessly Across Text and Image Tokens

论文链接：https://arxiv.org/pdf/2503.10772

开源代码：https://tacju.github.io/projects/flowtok.html

导读

弥合不同模态之间的差距对于理解代表我们世界的各种形式的数据至关重要，这既包括理解也包括生成。在多模态理解方面，大量研究集中在设计将不同模态投影到共享潜在空间的架构。这些方法通过利用模态之间的共同潜在空间，显著推动了跨模态表示学习和现实世界理解的发展。

简介

弥合不同模态之间的差距是跨模态生成的核心。传统方法将文本模态视为一种条件信号，逐步引导从高斯噪声到目标图像模态的去噪过程，而我们探索了一种更简单的范式——通过流匹配在文本和图像模态之间直接转换。这需要将两种模态投影到一个共享的潜在空间中，但由于它们本质上的不同表示方式，这带来了重大挑战：文本具有高度语义性，被编码为一维标记，而图像在空间上存在冗余，以二维潜在嵌入的形式表示。为了解决这个问题，我们引入了Flow - Tok，这是一个最小化的框架，通过将图像编码为紧凑的一维标记表示，实现了在文本和图像之间的无缝流动。与先前的方法相比，在256的图像分辨率下，这种设计将潜在空间大小减少了，无需复杂的条件机制或噪声调度。此外，FlowTok在相同的公式下自然地扩展到图像到文本的生成。凭借以紧凑的一维标记为中心的简化架构，FlowTok具有很高的内存效率，需要的训练资源显著减少，采样速度也快得多，同时其性能可与最先进的模型相媲美。

方法与模型

在本节中，我们以文本到图像生成作为主要任务来阐述 FlowTok。我们首先详细介绍如何在保留语义信息的同时，将图像和文本作为一维标记投影到统一的紧凑潜在空间中（4.1 节）。接下来，我们将 FlowTok 作为一个通用框架引入，用于实现文本和图像标记之间的无缝流，并讨论在相同公式下将其扩展到图像到文本生成的情况（4.2 节）。

图4. FlowTok概述。FlowTok是一个轻量级框架，它促进了一维文本标记和图像标记之间的无缝流动，可用于文本到图像和图像到文本的生成。顶部：对于文本到图像生成，输入文本由CLIP文本编码器编码为，投影到低维潜在空间作为文本标记，然后通过流匹配转换为相同形状的图像标记，并由一维图像变分自编码器（Variational Auto - Encoder，VAE）解码器解码以生成最终图像。底部：对于图像到文本生成，输入图像由一维图像VAE编码器编码为，通过流匹配映射到，并通过文本解码器解码为文本。与传统方法依赖噪声和图像潜在表示（例如，对于256分辨率图像为）并以文本作为条件不同，我们的直接一维转换（即）实现了的压缩率，显著降低了内存成本，加速了训练，并实现了更快的推理。

1. 统一图像和文本的潜在空间

文本和图像之间的结构差异给在流匹配的同一潜在空间中统一它们带来了重大挑战。文本本质上是语义性的，被编码为具有高维通道的一维潜在序列以保留语义，而图像包含空间冗余信息，通常表示为通道维度较低的特征图以保留空间先验信息。为了弥合这一差距，我们提出利用图像标记化方面的最新进展将图像编码为紧凑的一维标记。这种方式有助于保留文本嵌入的一维结构，只需要将它们投影到一组更紧凑的标记中，同时确保保留语义信息。下面，我们详细介绍图像和文本是如何被编码的。

将图像编码为紧凑标记。我们基于TA - TiTok [41]的核心思想进行了多项改进，以提升我们的图像标记器性能。具体来说，我们用旋转位置编码（Rotary Position Embedding，RoPE）[76]取代了原来的可学习一维位置嵌入，以提高TA - TiTok的性能。此外，我们用SwiGLU前馈网络（Feed - Forward Network，FFN）[71]取代了视觉变换器（Vision Transformer，ViT）[24]中的多层感知机（Multi - Layer Perceptron，MLP）块，这有助于学习更有效的潜在空间。为了与文本编码器的上下文标记数量相匹配，我们相应地设置了TA - TiTok中的潜在标记数量（对于CLIP [65]为）。因此，TA - TiTok的编码器将每个图像编码为一个紧凑的一维标记序列。

将文本转换为紧凑的标记。我们使用预训练的文本编码器[65]来提取初始文本嵌入，其中表示通道数。值得注意的是，通常比图像潜在空间大小大得多，因为它携带了更丰富的语义信息。由于我们的目标是将文本嵌入直接流入图像潜在空间，我们需要确保两种嵌入具有相同的形状。虽然我们已经通过仔细调整图像标记器使与对齐，确保图像编码后与文本标记的长度相匹配。但剩下的挑战在于对齐通道数（即和），我们使用文本投影器来解决这个问题。由于只需要调整的通道数，同时保留其一维形状，我们采用了几个简单的Transformer块作为投影器。为了在从相同文本生成图像时引入可变性，我们通过应用KL散度正则化将投影后的文本潜在空间建模为高斯分布。

文本到图像生成的一个关键方面是确保生成的图像准确反映输入文本描述。由于通过可学习的投影器减少文本嵌入的通道维度可能会导致语义信息丢失，我们引入了辅助文本对齐损失来保持语义一致性。具体来说，我们使用一个轻量级的多层感知器（MLP）将投影到一个新的空间进行对齐。然后，我们沿着通道维度将和展平并归一化，并计算它们之间的对比损失，这受到了CLIP[65]的启发。具体而言，我们使用可学习的温度参数计算缩放后的成对余弦相似度，然后计算对称交叉熵损失：

其中表示转置操作，CE表示交叉熵损失，标签根据其批次索引分配，确保每个文本标记都被明确训练为与同一批次中对应的CLIP文本嵌入对齐。我们还探索了其他保留语义信息的方法，例如与平均池化后的文本嵌入对齐或使用带边界的余弦相似度损失。然而，我们发现CLIP风格的损失取得了最佳性能。更多细节在第5.3节提供。

通过上述设计，FlowTok在保留语义信息的同时，有效地将文本标记化到相同的低维潜在空间。这种与标记化后的图像潜在空间的对齐为压缩后的文本标记和图像标记之间的直接流动奠定了基础。值得注意的是，当使用CLIP作为文本编码器时，与传统的二维流匹配方法相比，FlowTok有效地减小了潜在空间大小。在图像分辨率为256时，潜在空间大小从减小到，实现了的压缩。这种减小显著降低了内存需求并加速了训练，提高了框架的效率和可扩展性。

2. FlowTok：跨文本和图像标记的无缝流通用框架

文本到图像生成。如图4（上）所示，在图像和文本都映射到同一潜在空间的情况下，Flow - Tok通过堆叠DiT模块[61]利用了普通的流匹配方法[51]。值得注意的是，源模态（即文本）被直接视为流匹配的源分布，从而消除了DiT模块内的拼接或交叉注意力机制的需求。这种设计选择进一步简化了整体框架，并优化了文本到图像的生成过程。结合第4.1节中引入的紧凑一维标记，FlowTok实现了高内存效率，在上支持的批量大小。此外，它还能实现快速采样，比现代文本到图像扩散模型[69]快以上，显著降低了训练大规模文本到图像生成模型的计算门槛。

图像到文本生成。如图4（下）所示，FlowTok还可以在相同的公式下使用紧凑的一维图像和文本标记无缝扩展到图像到文本生成。具体来说，图像标记流向文本标记，经过训练的文本解码器将作为输入并输出标记器索引，然后可以将其解码回相应的文本描述。

实验与结果

在本节中，我们首先提供FlowTok的实现细节（第5.1节），然后给出文本到图像和图像到文本生成的主要结果（第5.2节）。最后，我们进行消融实验，以更好地理解FlowTok在文本到图像生成中的设计选择（第5.3节）。

1. 实现细节

图像标记器。我们在官方TA - TiTok [41]代码库的基础上进行了最小的修改来构建我们的图像标记器。编码器使用ViT - B [24]，解码器使用ViT - L，两者的补丁大小均为。为了与CLIP文本编码器的输出序列长度对齐，我们将一维潜在标记的数量设置为77，标记维度设置为16。此外，我们使用RoPE [76]和SwiGLU FFN [71]对标记器进行了增强。值得注意的是，我们增强后的标记器在ImageNet验证集的零样本评估中实现了1.02的FID，与原始的具有128个标记的TA - TiTok性能相当。

文本投影器。我们为文本到图像生成训练了一个文本投影器，它将CLIP文本嵌入转换为形状为的潜在表示，与我们的图像标记器编码的图像潜在空间对齐。文本投影器由六个Transformer [81]模块组成，每个模块包括一个多头自注意力机制和一个多层感知机（MLP），两者都通过跳跃连接[35]进行了增强，以确保训练的稳定性。

文本解码器。我们为图像到文本生成训练了一个文本解码器，与文本投影器类似，它由六个Transformer [81]模块组成。解码器将文本潜在表示作为输入，并输出相应的CLIP文本标记器索引，这些索引可以使用CLIP文本标记器进一步转换为文本。

FlowTok。我们采用DiT [61]模块作为FlowTok的基本构建单元来建模标记之间的交互。具体来说，我们遵循DiT架构实现了用于高效消融实验的FlowTok - B和用于增强性能的FlowTok - XL。为了进一步提升性能，我们增加了深度、宽度和注意力头的数量，构建了具有11亿参数的FlowTok - H。详细的模型配置见表1。

数据集。我们采用开源数据集 [41] 以促进 FlowTok 简单框架的可复现性。具体而言，我们的图像分词器在 DataComp - 1B [30] 上进行训练，文本分词器在 COCO [50] 上进行训练。对于文本到图像生成，受近期研究  的启发，我们采用两阶段训练策略：预训练和微调。预训练阶段结合使用 DataComp - 1B [30]、CC12M [15] 和 LAION - aesthetic [1]，而微调阶段纳入了额外的高质量数据集，包括 LAION - art [3]、LAION - pop [4]、JourneyDB [77] 和 DALLE3 - 1M [26]。对于图像到文本生成，我们遵循 COCO [20] 的 Karpathy 划分 [40] 来划分训练集和验证集。详细的数据集信息见附录。

训练。FlowTok 的训练目标主要集中在预测流匹配中的速度，记为 。对于文本到图像生成，我们引入了两个额外的损失：KL 散度损失  用于强制文本标记服从高斯分布，以及文本对齐损失  用于保留第 4.1 节中讨论的语义信息。形式上，总体训练目标为：

其中  和  控制损失的权重。默认情况下，对于文本到图像生成，我们将  设置为 ，将  设置为 1，而对于图像到文本生成，两者都设置为 0。

评估。我们遵循标准评估实践，报告文本到图像和图像到文本生成的相关指标。具体而言，对于文本到图像生成，我们报告在 COCO [50] 上的 FID - 30K 和在 MJHQ - 30K [46] 上的 FID。对于图像到文本生成，我们报告在 COCO Karpathy 划分 [40] 上的 BLEU - 4 [60]、METEOR [11]、ROUGE [49]、CIDEr [82] 和 SPICE [8]。为了在 FlowTok 中纳入无分类器引导（CFG）[36]，我们遵循 CrossFlow [53] 并使用 CFG 指标。除非另有说明，由于 FlowTok 的一维潜在空间较小，我们发现仅使用 20 步采样就足够了。这显著加快了推理过程，在不影响性能的情况下实现更快的生成。

2. 主要结果

文本到图像生成。我们在表2中报告了在COCO [50]和MJHQ -  [46]数据集上的零样本文本到图像生成结果。比较的方法分为两组：文本作为条件，即文本作为图像生成的引导信号；文本作为源分布，即在生成过程中文本直接被建模为一种分布。可以观察到，在COCO FID - 30K指标上，FlowTok在这两类方法中都取得了与先前方法相当的性能。具体而言，与同样将文本作为源分布的CrossFlow [53]相比，FlowTok - H的FID - 30K达到了9.67，与CrossFlow大致相当。当进一步在MJHQ - 30K上评估FlowTok以评估生成图像的美学质量时，我们发现，尽管FlowTok仅在公开可用的数据集上进行训练，没有使用高质量的专有数据，但FlowTok - XL已经超越了其他最先进的模型，证明了其生成多样化、高质量图像的能力。此外，FlowTok - H进一步将FID分数提高到7.15，凸显了其卓越的图像生成能力。

除了性能之外，与现有的最先进模型相比，FlowTok所需的训练资源显著减少。具体来说，FlowTok - XL仅需20.4个A100 GPU天即可完成训练，而FlowTok - H将预算略微增加到26.1个A100 GPU天。相比之下，效率最高的文本作为条件的模型PixArt -  [17]仍需要94.1个A100 GPU天。与同样将文本作为源分布且需要78.8个A100 GPU天的CrossFlow [53]相比，FlowTok的效率要高得多。

此外，FlowTok的推理速度明显更快。在分辨率下，FlowTok - XL每秒可生成22.7张图像，而FlowTok - H每秒可生成18.2张图像。相比之下，PixArt - 每秒运行速度为7.9张图像，Show - o每秒仅为1.0张图像。更值得注意的是，在文本作为源分布的类别中，与每秒仅运行1.1张图像的CrossFlow相比，FlowTok的采样时间加快了倍。这种效率源于FlowTok简化的框架及其对一维（1D）标记的有效利用，显著降低了计算开销。

图像到文本生成。我们使用Karpathy划分[40]在COCO [50]数据集上评估图像到文本生成，结果总结在表3中。为了确保公平比较，我们将方法分为两组：从图像到文本分布的直接流，这是一种利用流匹配进行直接图像到文本转换的新范式；以及其他方法，仅考虑未经过CIDEr优化训练的方法。在直接流类别中，FlowTok - XL在大多数指标上始终优于其对应方法CrossFlow [53]。具体而言，FlowTok - XL的BLEU - 4（B@4）分数达到37.1，比CrossFlow高0.7，CIDEr分数达到117.0，比CrossFlow高0.8。此外，与其他范式的最先进方法相比，FlowTok - XL表现出了有竞争力的性能，凸显了直接流匹配作为图像到文本生成的一种有前景的方法。值得注意的是，FlowTok使用紧凑的一维（1D）标记在相同的公式下进行图像到文本生成，理论上与CrossFlow采用的基于潜在变量的范式相比，需要更少的训练资源并能实现更快的采样。然而，由于CrossFlow尚未发布相应的检查点进行评估，因此无法进行直接的定量比较。

3. 消融研究

我们使用FlowTok - B对文本到图像生成进行消融研究，并在COCO数据集上评估其效率。我们的消融研究重点关注文本对齐损失的设计，因为它在保留语义信息方面起着关键作用。具体来说，我们研究了三个关键方面：文本对齐目标（表4a）、损失函数的选择（表4b）和损失权重（表4c）。以下是详细信息。

表4. 文本对齐损失的消融研究。我们对文本对齐损失的三个关键方面进行了全面的消融研究：对齐目标（表4a）、损失函数的选择（表4b）以及损失权重（表4c），旨在确定在文本到图像生成过程中，在FlowTok内保留语义信息的最有效策略。为了进行高效验证，我们报告了使用FlowTok - B在COCO上的FID - 30K，未应用CFG指标。

文本对齐目标。我们首先在表4a中研究投影文本标记的对齐目标选择。一个直接的基线方法是沿着通道维度对原始CLIP文本嵌入直接应用平均池化（第1行），将维度从768降低到16以匹配。然而，与受先前工作启发使用简单的多层感知机（MLP）来学习对齐目标（第2行）相比，这种方法的性能明显更差。我们将这种性能差距归因于CLIP文本嵌入中相邻通道不一定相关，并且简单的平均池化会丢弃过多的语义信息。相比之下，可学习的MLP减轻了这种信息损失，使其成为定义文本对齐目标的更有效选择。

文本对齐损失函数。接下来，我们在表4b中研究文本对齐损失函数。除了采用自文献[65]的对比损失外，我们还探索使用类似于文献[87]的余弦相似度损失。具体来说，我们计算文本标记与对齐目标之间的余弦相似度，并对相似度低于阈值的配对施加惩罚。我们的实验表明，虽然两种损失函数都有效，但对比损失取得了更好的性能。

文本对齐损失权重。最后，我们在表4c中研究文本对齐损失权重的影响。我们的结果表明，将设置为1.0，与流匹配损失的权重相等，足以在保持高质量图像生成的同时保留语义信息。进一步增加可能会导致文本对齐损失在训练早期阶段主导整体目标，从而可能阻碍最终性能。

结论

在本文中，我们介绍了FlowTok，这是一个简洁而强大的框架，它实现了一维文本和图像标记之间的无缝直接流。通过精心设计的关键模块和损失函数，FlowTok将两种模态投影到统一的一维潜在空间中，同时保留语义信息，从而能够在同一框架下实现文本到图像和图像到文本的生成。这种设计使FlowTok具有很高的内存效率，在训练期间仅需8块A100 GPU就能支持8K的批量大小。此外，其简洁性加速了收敛速度——在8块A100 GPU上大约20天内，FlowTok就能达到与需要显著更长训练时间的最先进模型相当的性能。简化的设计还使采样速度比现代文本到图像生成模型快倍以上。