一句话总结

微软 Mage 团队的 Mage-Flow 是一个高效的 4B 模型系列，它联合设计了一个低成本基于扩散的 tokenizer（Mage-VAE）并结合锚点潜在正则化，以及一个通过整流流匹配训练的原生分辨率多模态 Transformer，通过分辨率打包和内核融合将吞吐量提升 $2.5\times$2.5×，并生成 Turbo 变体，在单个 A100 GPU 上实现 $1024^2$10242 分辨率下 0.59 秒生成和 1.02 秒编辑。

核心贡献

• Mage-VAE 是一个轻量级潜在 tokenizer，结合了一步扩散式的编码/解码与锚点潜在正则化，重建保真度媲美强大的公开 VAE，同时将 tokenization 成本降低超过一个数量级。
• 该堆栈集成了原生分辨率多模态扩散 Transformer（NR-MMDiT）与原生分辨率打包和融合 CUDA 内核，实现灵活分辨率训练，并将端到端训练吞吐量提升约 2.5 倍。
• 构建了一个完整的生成与编辑模型系列，包含基础版、RL 对齐版和 Turbo 变体。Diffusion-NFT 作为后训练阶段应用于增强提示遵循、文本渲染、美学质量和编辑保真度，并结合基于视觉基础模型的对抗感知引导进行解耦 DMD 蒸馏，得到 4 步 Turbo 模型；这些模型在标准基准上具有竞争力，并在单个 NVIDIA A00 GPU 上以 1024² 分辨率实现 0.59 秒生成和 1.02 秒编辑。

引言

当前视觉生成模型需要合成高保真图像、遵循组合文本提示、渲染多语言文本，并在设计、原型制作和交互式编辑等应用中保持空间布局。闭源系统性能强大但缺乏透明度，而开源替代方案依赖从 6B 到 80B 参数规模的主干网络，这使得在有限算力预算下进行微调、受控消融研究和部署变得不切实际。作者引入了 Mage-Flow 堆栈，一个紧凑的 4B 规模基础模型，将轻量级一步扩散 tokenizer 与原生分辨率多模态扩散 Transformer 和融合 CUDA 内核相结合，从而以比大得多的模型显著更低的延迟和内存实现具有竞争力的文本到图像生成和基于指令的编辑。

数据集

作者从两个互补的流程构建了一个大规模训练语料库：一个文本到图像生成数据集和一个基于指令的图像编辑数据集。两个流程均旨在提升视觉质量、安全性、语义对齐、多样性以及对代表性不足的能力的覆盖。

文本到图像生成数据集

• 原始来源：从开源数据集中收集约 100 亿个图文对。
• 样本级过滤：移除损坏文件、低分辨率图像、极端宽高比和方向错误的样本；然后对图像的亮度、饱和度、模糊度、熵、水印、美学质量、OCR 内容和 NSFW 安全性进行评分。阈值在训练阶段（256² → 512² → 1024² → SFT）逐步收紧，使得早期阶段保持广泛覆盖，而后期阶段使用更干净的数据。
• 跨样本去重：每张图像使用副本检测描述符（SSCD）进行编码，并用 FAISS 建立索引。在数据集内部，余弦相似度高于 0.9 的对被分组，仅保留质量最佳的代表样本。跨数据集时，使用持久索引拒绝近似重复项。还匹配了一个留出的基准索引以减少评估污染。
• 多粒度描述：保留的图像由大型视觉语言模型在四个层次上生成描述：短语、实体、组合和摄影描述。对于文本丰富的图像，描述器还将可见文本转换为渲染指令。在训练期间，模型从这些通道中采样，以学习遵循不同长度和具体性的提示。
• 概念感知合成与平衡：为长尾概念生成补充数据，如长文本渲染、稀有物体、罕见属性、结构化布局和代表性不足的风格（例如，具有多样化字体、布局、语言、颜色和背景的合成文本渲染样本）。所有补充图像经过相同的质量过滤，并合并到主语料库中。短语级描述用于估计概念分布，训练期间的概念感知采样降低了频繁物体、自然场景和常见逼真风格的支配地位。
• 最终规模：保留了约 13 亿个高质量图文对；从该池中按阶段抽取预训练子集，并采用逐步收紧的过滤和更强的重新加权。

图像编辑数据集

• 原始三元组：约 9000 万个（源图像，编辑指令，目标图像）三元组。大约 5000 万来自开源基于指令的编辑数据集；其余约 4000 万为内部合成——1000 万低层级图像处理三元组和 3000 万通用语义编辑三元组。
• 基于 VLM 的过滤：每个三元组由三个具有不同评判标准的独立视觉语言模型专家评估。一位专家检查编辑是否被正确应用，无关区域是否被保留，以及结果是否在视觉上合理。只有当至少两位专家投票通过时，三元组才被保留。过滤后，约 2000 万开源三元组和 2500 万合成三元组存活，得到 4500 万三元组的保留池。
• 编辑类型标记与平衡：定义了一个包含 19 个编辑类别的手工分类法。数据单元（整个数据集、子数据集或共享明确编辑类型字段的样本）被映射到这些类别。然后调整采样率，使训练混合在整个分类法中提供广泛且比例适当的覆盖，防止频繁操作占主导地位，同时保持稀有编辑类型得到充分代表。
• 使用：最终平衡后的编辑集直接用于编辑模型训练。

方法

作者提出 Mage-Flow 堆栈，围绕两个核心组件构建：Mage-VAE，一个轻量级高保真潜在 tokenizer，和一个采用整流流匹配训练的 4B 原生分辨率多模态扩散 Transformer（NR-MMDiT）。Mage-VAE 提供一个紧凑的、可用于生成的潜在空间，而 NR-MMDiT 则以灵活的分辨率和纵横比对打包的潜在序列建模。

为解决传统 VAE 在高分辨率下的延迟瓶颈，作者将 Mage-VAE 设计为基于像素扩散的轻量级 tokenizer。它用轻量级一步扩散式的编码和解码取代了繁重的高分辨率自编码器组件。该架构采用对称设计：解码器使用堆叠的卷积扩散块和一个解耦的像素扩散头，直接从潜在表示重建 RGB 像素，而编码器则作为以像素为条件的潜在生成器。为确保潜在空间可用于生成，作者将后验分布正则化到由 FLUX.2-VAE 诱导的锚点潜在分布。

训练流程分为三个阶段。第一阶段，编码器和解码器分别作为多步扩散模型预训练。第二阶段，使用重建损失、DMD 损失和 DINOv2 投影的 GAN 损失将解码器蒸馏为一步模型。第三阶段，对一步编码器和解码器进行联合微调，并加入锚点潜在 KL 正则项。

生成主干是一个 4B 参数的 NR-MMDiT。与标准的分桶训练不同，NR-MMDiT 直接在原生分辨率序列上训练。不同分辨率和纵横比的图像由 Mage-VAE 编码，展平为可变长度的潜在 token 序列，并在固定的 token 预算下打包到单个连续批次中。相同的打包原则也适用于文本条件。这种原生分辨率公式消除了单桶限制，让模型接触到更丰富的分辨率分布，并使单个检查点能够泛化到灵活的输出尺寸，高度和宽度均可从 512 到 2048。

文本条件由冻结的 Qwen3-VL 文本编码器提供。模型在 Mage-VAE 潜在空间中采用整流流匹配目标进行训练：

$$\mathcal{L}(\theta) = \mathbb{E}_{(x, \tau), t, \epsilon} \left[ \| v_{\theta}(z_t, t, \tau) - (z - \epsilon) \|_2^2 \right]$$L(θ)=E(x,τ),t,ϵ​[∥vθ​(zt​,t,τ)−(z−ϵ)∥22​]

其中 $z_t = (1 - t)z + t\epsilon$zt​=(1−t)z+tϵ 且 $\epsilon \sim \mathcal{N}(0, I)$ϵ∼N(0,I)。对于编辑任务，模型在 2D 旋转位置嵌入的基础上增加了一个额外的帧维度，以区分源和目标视觉 token。

为使原生分辨率的 4B 规模学习切实可行，作者实现了堆栈级内核融合。他们将 Mage-VAE、Qwen3-VL 文本编码器和 NR-MMDiT 中重复模块内的主要算子链融合，将中间结果保存在片上内存中，以减少激活内存移动和内核启动。

Mage-Flow 生成语料库由约 100 亿个原始图文对精心整理而成。该流程包括样本级过滤、跨样本去重、多粒度描述和概念感知合成。样本级过滤器去除损坏、低质量或不安全的图像。跨样本去重使用 SSCD 描述符抑制近重复的视觉模式。多粒度描述分配短语级、实体级、组合级和摄影描述，以标准化文本监督。

概念感知合成为长尾图像提供补充。最终的数据构成确保了在物体、场景、设计和合成领域的广泛覆盖，同时概念感知采样在训练期间降低了频繁物体的支配地位。

作者使用统一的方案训练 Mage-Flow 模型系列。预训练遵循渐进式课程：从 256x256 分辨率开始，用于广泛的视觉-语言对齐；然后转移到 512 像素原生宽高比模式，最后到 1024 像素模式以捕捉精细细节。随后在一个高质量子集上进行监督微调。对于基于指令的编辑，模型从文本到图像基础检查点初始化，并使用编辑三元组与生成对的平衡混合进行适配。

后训练阶段应用 Difusion-NFT，这是一种在流匹配生成器的前向过程上操作的负向感知微调方法。它利用在线展开组和能力路由的奖励评估器，以隐式正负策略优化奖励加权的流匹配目标。最后，为降低推理成本，作者利用解耦 DMD 结合对抗感知引导，将对齐后的检查点蒸馏为 4 步 Turbo 模型。

实验

Mage-VAE 被评估为一个从 FLUX.2-VAE 蒸馏而来的轻量级 tokenizer，结果表明锚点潜在监督在保持重建质量和潜在空间兼容性的同时，将编码成本降低了超过一个数量级。Mage-Flow 和 Mage-Flow-Edit 模型还通过解耦 DMD 与对抗感知引导进一步蒸馏为 4 步 Turbo 变体，该变体保持了生成和编辑质量，对抗性的收益在生成和文本编辑方面尤为明显。总体而言，4B 的 Mage-Flow 在与更大模型的比较中取得了有竞争力的文本到图像和基于指令的编辑性能，而 4 步版本在很大程度上保持了这种能力。

Mage-VAE 在 CLIC 2020 上的重建保真度与最强的 FLUX.2-VAE 基线相当，同时在 FFHQ 上达到最佳或接近最佳的质量。它将编码的计算复杂度降低了约一个数量级，解码复杂度降低超过二十倍，从而能够高效处理其他 VAE 变得异常缓慢或超出内存限制的高分辨率图像。轻量级设计保留了从昂贵教师蒸馏而来的生成兼容性潜在结构，提供了卓越的质量-效率权衡。Mage-VAE 在 CLIC 2020 上的 PSNR、SSIM 和 LPIPS 与 FLUX.2-VAE 相当，但参数量和每像素 kMAC 值更低。在 1024×1024 的 FFHQ 上，Mage-VAE 在所有比较的 VAE 中达到了最高的 PSNR 和 SSIM，超越了 SD 变体和教师 FLUX.2-VAE。与 FLUX.2-VAE 相比，编码复杂度降低约 12 倍，解码复杂度降低约 22 倍，同时仍保持高重建质量。在 4096×4096 分辨率下，基线 VAE 要么耗尽内存，要么变得极慢，而 Mage-VAE 在编码和解码方面均保持高效。尽管计算量大为节省，但蒸馏后的 tokenizer 保留了 FLUX.2 风格扩散生成器所需的潜在几何结构，下游任务中的跨 tokenizer 兼容性也证明了这一点。

在保持下游主干不变的情况下，将 tokenizer 在 Mage-VAE 和 FLUX.2-VAE 之间互换，在生成和编辑基准上性能几乎相同，显示出强大的跨 tokenizer 兼容性。结果证实，轻量级 Mage-VAE 保留了 FLUX.2 风格扩散生成器所需的潜在几何结构，使其成为高效的即插即用替代方案，且不会牺牲输出质量。在 Mage-Flow-Turbo 主干下，两种 tokenizer 的 GenEval 得分相同（0.88），DPG-Bench 仅相差 0.01，其他指标如 OneIG-EN 和 LongText-CN 也匹配或几乎匹配。对于 FLUX.2-Klein-4B 主干，GenEval 保持不变为 0.83，DPG-Bench 最多变化 0.01，LongText-EN 和 LongText-CN 的得分无论使用哪种 tokenizer 都保持相同。跨主干的稳定基准表明，Mage-VAE 可以在不显著影响生成或基于指令的编辑质量的情况下，取代计算量更大的 FLUX.2-VAE。

将无分类器引导的条件分支和无条件分支打包到单个批次中，消除了冗余计算，并在不改变去噪轨迹的情况下降低了整体推理延迟。在文本到图像和基于指令的编辑变体中，这种优化在单个 NVIDIA A100 GPU 上将吞吐量持续提升 9–15%，基础生成模型受益最大。在所有测试的模型和步数配置下，打包 CFG 的加速比从 1.09 倍到 1.15 倍不等。采用 20 步的基础文本到图像模型获得了最大的相对增益，而每步时间较长的编辑变体则获得了较温和的 1.09–1.10 倍提升。所有测量均在单块 A100 GPU 上进行，确保报告的加速反映了实际的推理效率提升。

从 FLUX.2-VAE 切换到轻量级 Mage-VAE tokenizer，即刻将每步训练时间缩短约 29%，并将模型 FLOP 利用率从 33.20% 提升至 45.63%，而 GPU 峰值内存基本不变。为 VAE 和文本编码器添加融合 CUDA 内核仅带来少量额外加速，而融合 NR-MMDiT 主干则提供了最主要的改进，将内存降至 141.44 GB，并将整体加速比提升至 2.49 倍。完整系统表明，紧凑的 tokenizer 和堆栈级内核融合对于消除重复模块中的内存受限开销都是必不可少的。仅用 Mage-VAE 替换 FLUX.2-VAE 即可实现 1.41 倍加速，并将 MFU 提升超过 12 个百分点，而峰值内存不变。融合 VAE 和 Qwen3-VL 文本编码器各自带来增量加速（最高 1.46 倍），而融合 NR-MMDiT 模块则使整体加速比大幅跃升至 2.49 倍。DiT 内核融合还将每个 GPU 的峰值内存从约 175 GB 降至 141 GB，表明内存受限开销集中在扩散主干中。完整配置将 MFU 提高了一倍多（从 33.20% 到 77.26%），证实堆栈级算子融合有效地将硬件能力转化为可用吞吐量。

从初始预训练阶段到监督微调阶段，过滤阈值逐渐变得更加严格。最低分辨率和美学质量要求提高，而允许的水印概率急剧下降，重心从广泛的视觉覆盖转向更高分辨率、更干净、更具视觉吸引力的图像。纵横比、文件大小和 NSFW 阈值在所有阶段保持不变。水印容忍度从 256² 阶段的低于 0.5 急剧收紧至 SFT 阶段的低于 0.05，实际上在后期训练阶段消除了带水印的图像。美学分数下限在阶段间从 4.5 上升至 6.5，确保只有视觉上吸引人的样本才能通过最终过滤步骤。

论文评估了 Mage-VAE 作为一个轻量级 tokenizer，证实其重建质量与 FLUX.2-VAE 相当，同时将编码和解码复杂度降低超过一个数量级，并且可以在下游生成中无缝替代教师模型，对质量毫无影响。打包的无分类器引导去除了冗余计算，将推理加速高达 15%，而融合 NR-MMDiT 主干使训练吞吐量翻倍以上，同时峰值内存减半。数据过滤流程从初始预训练到监督微调应用越来越严格的美学和水印阈值，确保后期阶段仅使用干净、高质量的图像。