在本系列的第一篇文章中，我们介绍了目标：完全从头开始、以开放方式、大规模地训练一个具有竞争力的文生图基础模型。我们主要聚焦于架构选择，并阐述了模型 PRX 背后的核心设计决策。我们还发布了一个早期的小型（12 亿参数）版本模型，作为我们正在构建内容的预览（如果你还没试过，快去试试吧😉）。
在这篇文章中，我们将重点从架构转向训练。目标是记录在尝试让模型训练更快、收敛更可靠、学习到更好的表征时，哪些方法真正起到了作用。该领域发展迅速，“训练技巧”的清单也在不断增长，因此我们并未试图进行全面调查，而是将其构建为一本实验日志：我们复现（或改编）了一系列近期思路，在一致的设置中实现它们，并报告它们在实际中如何影响优化和收敛。最后，我们不仅孤立地报告这些技术，还探索了哪些技术在组合使用时仍然有效。
在下一篇文章中，我们将把完整的训练方案（包括本文中的实验）以代码形式发布。我们还将进行并报告一次公开的“速通”，将最佳方案整合到单一配置中，并对其进行端到端的压力测试。这项工作既可作为对当前训练流程的压力测试，也可作为具体示例，展示在严格限制下精心设计的训练能达到何种效果。如果你还没加入，我们邀请你加入我们的 Discord 以继续讨论。这个项目的很大一部分内容是通过与社区成员的交流形成的，我们非常重视外部反馈、消融实验以及对结果的不同解读。
基线
在介绍任何训练效率技术之前，我们首先建立一个干净的参考运行。这个基线故意设计得很简单。它使用标准组件，避免辅助目标，并且不依赖架构捷径或技巧来节省计算资源。其作用是作为所有后续实验的稳定比较基准。
具体来说，这是一个纯 Flow Matching（Lipman 等人，2022）训练设置（如第一部分所述），没有额外的目标函数，也没有架构层面的速度技巧。我们将使用本系列第一篇文章中介绍的小型 PRX-1.2B 模型（对图像 token 和文本 token 采用全局注意力的单流架构）作为基线，并在 Flux VAE 潜空间中进行训练，除非另有说明，否则在所有对比中保持配置不变。
基线训练设置如下：
| 设置项 | 数值 |
|---|---|
| 训练步数 | 10 万 |
| 数据集 | 使用 MidJourneyV6 生成的公开 100 万张合成图像 |
| 分辨率 | 256×256 |
| 全局批次大小 | 256 |
| 优化器 | AdamW |
| 学习率 | 1e-4 |
| 权重衰减 | 0.0 |
| epsilon | 1e-15 |
| betas | (0.9, 0.95) |
| 文本编码器 | GemmaT5 |
| 位置编码 | 旋转位置编码（RoPE） |
| 注意力掩码 | 填充掩码 |
| 指数移动平均（EMA） | 禁用 |
这个基线配置提供了一个透明且可复现的锚点。它使我们能够将观察到的改进和退化归因于特定的训练干预措施，而不是变化的超参数或隐藏的设置改动。在本文的剩余部分中，每一项技术都将对照这个参考基准进行评估，并始终围绕一个核心问题：
与基线相比，这项修改是否改善了收敛性或训练效率？
基线模型在 10 万训练步数后的生成示例。
基准测试指标
为了使本文内容扎实可靠，我们依赖一小套指标来随时间监控检查点。这些指标中没有一个是感知图像质量的完美代理，但它们共同提供了一个实用的评分板，供我们在迭代过程中参考。
Fréchet Inception Distance（FID）：（Heusel 等人，2017）使用 Inception-v3 特征统计量（均值和协方差）衡量生成图像分布与真实图像分布之间的接近程度。数值越低通常表示样本保真度越高。
CLIP Maximum Mean Discrepancy（CMMD）：（Jayasumana 等人，2024）使用 CLIP 图像嵌入和最大均值差异（MMD）来衡量真实图像分布与生成图像分布之间的距离。与 FID 不同，CMMD 不假设高斯特征分布，并且可以更高效地利用样本；在实践中，它通常比 FID 更能反映感知质量，尽管它仍然是一个不完美的代理指标。
DINOv2 最大均值差异（DINO-MMD）：与 CMMD 相同的基于 MMD 的距离度量，但计算对象是 DINOv2（Oquab 等人，2023）的图像嵌入向量而非 CLIP。这提供了在自监督视觉骨干网络下对分布偏移的补充视角。
网络吞吐量：每秒处理的平均样本数（样本/秒），作为端到端训练效率的衡量指标。
在定义了评分板之后，我们现在可以深入探讨所探索的方法，这些方法分为四大类：表征对齐、训练目标、Token 路由与稀疏化，以及数据。
表征对齐
扩散模型和流模型通常使用单一目标进行训练：从被破坏的输入中预测一个类似噪声的目标（或向量场）。在训练早期，这一个目标同时承担两项任务：它必须构建有用的内部表征，并在此基础上学习去噪。表征对齐通过保留去噪目标并添加一个辅助损失函数来明确这一点，该损失函数使用一个强大且冻结的视觉编码器直接监督中间特征。这往往会加速早期学习，并使模型的特征更接近现代自监督编码器的特征。其结果是，达到相同质量通常需要更少的算力。
一个有用的理解方式是，将去噪器分解为一个生成中间隐藏状态的隐式编码器，以及一个将这些状态映射到去噪目标的解码器。其核心观点是，表征学习是瓶颈：扩散和流 Transformer 确实能学习判别性特征，但在算力受限的训练条件下，它们落后于基础视觉编码器。因此，借用强大的表征空间可以使去噪问题变得更容易。
REPA（Yu 等人，2024）
使用预训练视觉编码器的表征对齐。图来自 arXiv:2410.06940。
REPA 在基础的流匹配目标之上添加了一个表征匹配项。
设 \(x_{0} \sim p_{\text{data}}\) 为干净样本，\(x_{1} \sim p_{\text{prior}}\) 为噪声样本。模型在插值状态 \(x_{t}\)（其中 \(t \in [0,1]\)）上训练，并预测向量场 \(v_{\theta}(x_{t}, t)\)。在 REPA 中，预训练的视觉编码器 \(f\) 处理干净样本 \(x_{0}\)，生成图像块嵌入向量 \(y_{0} = f(x_{0}) \in \mathbb{R}^{N \times D}\)，其中 \(N\) 是图像块 token 数量，\(D\) 是教师模型嵌入维度。与此同时，去噪器处理 \(x_{t}\) 并生成中间隐藏 token \(h_{t}\)（每个图像块对应一个 token）。一个小型投影头 \(h_{\phi}\) 将这些学生模型的隐藏 token 映射到教师模型的嵌入空间，并通过一个辅助损失函数最大化教师与学生对应 token 之间的图像块级相似度：
\[ \mathcal{L}_{\text{REPA}}(\theta, \phi) = - \mathbb{E}_{x_{0}, x_{1}, t} \left[ \frac{1}{N} \sum_{n=1}^{N} \text{sim} \left( y_{0,[n]}, \; h_{\phi}(h_{t,[n]}) \right) \right] \]
这里 \(n \in \{1, \ldots, N\}\) 是图像块 token 的索引，\(y_{0,[n]}\) 是第 \(n\) 个图像块的教师模型嵌入向量，\(h_{t,[n]}\) 是时间 \(t\) 对应的学生模型隐藏 token，而 \(\text{sim}(\cdot, \cdot)\) 通常为余弦相似度。
该项与主流程匹配损失函数结合：
\[ \mathcal{L} = \mathcal{L}_{\text{FM}} + \lambda \; \mathcal{L}_{\text{REPA}} \]
其中 \(\lambda\) 控制两者之间的权衡。
在实践中，学生模型被训练为从 \(x_{t}\) 生成对噪声鲁棒且与数据一致的图像块表示，这样后续层可以专注于预测向量场和生成细节，而无需从头重新发现语义骨架。
In practice, the student is trained to produce noise-robust, data-consistent patch representations from $x_{t}$x t, so later layers can focus on predicting the vector field and generating details rather than rediscovering a semantic scaffold from scratch.
我们的观察结果
我们在基线 PRX 训练的基础上运行了 REPA，使用了两个冻结的教师模型：DINOv2 和 DINOv3（Siméoni 等人，2025）。结果模式非常一致：加入对齐能提升质量指标，更强的教师模型效果更好，但代价是速度略有下降。
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批/秒 ↑ |
|---|---|---|---|---|
| 基线 | 18.2 | 0.41 | 0.39 | 3.95 |
| REPA-Dinov3 | 14.64 | 0.35 | 0.3 | 3.46 |
| REPA-Dinov2 | 16.6 | 0.39 | 0.31 | 3.66 |
在质量指标上，两个教师模型都优于基线。DINOv3 的效果最强，在此次运行中取得了整体最佳数值。
REPA 并非没有代价：我们需要额外进行一次冻结教师模型的前向传播以及补丁级相似度损失计算，这体现在吞吐量从 3.95 批/秒下降到了 3.66（DINOv2）或 3.46（DINOv3）。换句话说，DINOv3 以更慢的训练速度为代价，优先追求最高的表征质量；而 DINOv2 则提供了更高效的权衡方案，在速度下降更小的同时，仍能带来显著的性能提升。
我们的实际结论是，REPA 是文本到图像训练中的一个有力手段。在我们的设置中，吞吐量的权衡是真实存在的，而净加速效果（达到给定图像质量所需的时间）感觉不如论文作者在 ImageNet 风格、类别条件生成任务上报告的那样显著。尽管如此，质量提升仍然非常明显。从定性角度看，我们也很早就看到了差异：大约在 10 万步之后，使用对齐训练的样本倾向于锁定更清晰的全局结构和更连贯的布局，这很容易理解为什么 REPA（以及更广泛的对齐变体）已成为现代 T2I 训练方案中的首选成分。
| 基线 | Repa-DinoV2 | Repa-DinoV3 |
|---|---|---|
 |  |  |
iREPA（Singh 等人，2025）
REPA 的一个自然延伸问题是：我们究竟应该对齐什么？iREPA 认为答案是空间结构，而非全局语义。在横跨 27 个视觉编码器的大规模实验中，作者发现，ImageNet 风格的“全局”质量（例如，对 patch token 的线性探测准确率）与 REPA 下的下游生成质量仅有微弱关联，而 patch token 空间自相似性的简单度量则与 FID 的相关性要强得多。基于这一诊断，iREPA 对 REPA 方案做了两个微小但有针对性的改动，以更好地保留和传递空间信息：
- 将通常使用的 MLP 投影头替换为一个轻量级的 3×3 卷积投影，作用于 patch 网格上。
- 对教师模型的 patch token 应用空间归一化，移除全局叠加层（跨空间位置的均值），以增强局部对比度。
尽管这些调整“只涉及不到 4 行代码”，但它们能持续加速收敛，并在不同编码器、模型规模乃至与 REPA 相近的训练方案中提升质量。
我们的观察
在我们的实验设置中，当在 DINOv2 之上应用 iREPA 的空间调整时，我们观察到了类似的提升效果：收敛过程更加平滑，指标在前 10 万步内提升得也更稳定。有趣的是，当在 DINOv3 教师模型之上应用同样的改动时，效果并未能干净地迁移，反而倾向于降低性能而非提升。我们不想过度解读这一结果：这很可能与我们特定的架构、分辨率/分块方式、损失权重，甚至是一些微小的实现细节有关。尽管如此，鉴于不同教师模型之间存在这种不一致性，我们很可能不会将这些调整纳入默认方案，尽管在针对特定设置进行调优时，它们仍是一个值得重新审视的有趣选项。
关于在完整训练过程中使用 REPA
论文《REPA 有效，直到失效：早期停止的全局对齐加速扩散训练》（Wang 等人，2025）强调了一个关键注意事项：REPA 是一种强大的早期加速器，但在训练后期可能会陷入平台期，甚至成为阻碍。作者描述了容量不匹配的问题。一旦生成模型开始拟合完整的数据分布（尤其是高频细节），强制其保持接近冻结的识别编码器的低维嵌入流形就会变得具有约束性。他们的实践结论很简单：在对齐的“预热”阶段保持对齐，然后通过分阶段调度将其关闭。
我们在自己的训练中也观察到了相同的定性模式。在训练预览模型时，在大约 20 万步后移除 REPA，显著改善了图像质量的整体观感，纹理、微对比度和精细细节持续锐化，而不是显得略微暗淡。因此，我们也建议将表示对齐视为一种临时支架。用它来获得快速的早期进展，然后在模型自身的生成特征跟上之后，过一段时间再将其移除。
Token 潜在空间中的对齐
到目前为止，“对齐”意味着在将分词器/潜在空间视为固定的情况下，对生成器的内部特征进行正则化，使其与冻结的教师模型保持一致。一个更直接的杠杆是塑造潜在空间本身，使得呈现给流式主干的表示在本质上更易于建模，同时不牺牲编辑和下游工作流所需的重建保真度。
REPA-E（Leng 等人，2025）将这一点具体化。它的出发点是一个失败模式：如果你简单地将扩散/流损失反向传播到 VAE 中，分词器会很快为去噪器学习到一个病态简单的潜在表示，这甚至可能降低最终的生成质量。REPA-E 的解决方案是一种双信号训练方案：
- 保留扩散损失，但应用停止梯度，使其仅更新潜在扩散模型（而不更新 VAE）；
- 使用端到端的 REPA 对齐损失同时更新 VAE 和扩散模型。
借助这两个技巧，分词器被明确优化，以生成能带来更高对齐度和经验上更优生成结果的潜变量。
与此同时，Black Forest Labs 的 FLUX.2 AE 工作将潜变量设计视为可学习性、质量和压缩之间的权衡。他们的核心论点是，提高可学习性需要将语义结构注入到表征中，而不是将分词器视为一个纯粹的压缩模块。这促使他们重新训练潜空间，以明确瞄准“同时实现更好的可学习性和更高的图像质量”。他们没有分享完整的配方，但清楚地陈述了关键思想：通过添加语义或表征对齐来使自编码器的潜空间更具可学习性，并明确指出他们基于并整合到 FLUX.2 AE 中的机制是 REPA 风格的对齐（与一个冻结的视觉编码器）。
我们的观察
为了探究潜空间中的对齐，我们将两个预训练的自编码器作为即插即用的分词器，用于同一个流模型主干进行比较：一个是 REPA-E-VAE（我们按照论文添加了 REPA 对齐目标），另一个是 Flux2-AE（我们遵循其建议，没有添加 REPA）。老实说，结果在定量和定性上都极其令人印象深刻。在样本中，差距一目了然：生成结果展现出更连贯的整体结构和更清晰的布局，并且“早期训练”伪影显著减少。
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批/秒 ↑ |
|---|---|---|---|---|
| 基线 | 18.20 | 0.41 | 0.39 | 3.95 |
| Flux2-AE | 12.07 | 0.09 | 0.08 | 1.79 |
| REPA-E-VAE | 12.08 | 0.26 | 0.18 | 3.39 |
首先一个引人注目的点是，两种潜空间干预都将 FID 降低了约 6 个点（从 18.20 降至约 12.08），这比我们通常通过“仅仅”对齐中间特征所获得的提升要大得多。这有力地支持了核心观点：如果分词器产生的表征在本质上更具可学习性，那么流模型将在各个方面受益。
这两个自编码器在细节上的表现差异很大。Flux2-AE 在大多数指标上占据优势（CMMD 和 DINO_MMD 非常低），但代价是吞吐量大幅下降：每秒批处理数从 3.95 降至 1.79。在我们的案例中，这种速度下降可以用他们同样强调的实际因素来解释：模型本身更重，并且产生的潜变量也更大（32 个通道），这增加了扩散主干网络每一步所需的工作量。
REPA-E-VAE 是“均衡”选项：它在达到与 Flux2-AE 基本相同的 FID 值的同时，吞吐量更接近基线水平（3.39 批次/秒）。
| 基线 | Flux2-AE | REPA-E-VAE |
|---|---|---|
 |  |  |
训练目标：超越基础流匹配
架构决定了容量，但训练目标决定了如何使用这些容量。在实践中，损失函数的微小改动往往会对收敛速度、条件保真度以及模型“锁定”全局结构的速度产生巨大影响。在下面的章节中，我们将介绍在基线修正流设置基础上测试的各种目标，首先从一个简单但效果出奇好的修改开始：对比流匹配。
对比流匹配（Stoica 等人，2025）
流匹配在无条件生成情况下有一个很好的特性：隐式地鼓励轨迹保持唯一性（流不应相交）。但一旦我们转向条件生成（类别或文本条件），不同的条件仍然可能导致流重叠，这在经验上表现为“平均化”行为：条件特异性减弱，全局结构变得模糊。对比流匹配通过添加一个对比项来直接解决这个问题，该对比项将条件流与批次中的其他流推开。
对比流匹配使类别条件流更加独特，减少了标准流匹配中出现的重叠现象，并生成更高质量、更能代表每个类别的图像。图片来自 arXiv:2506.05350。
对于给定的训练三元组 $(x, y, \epsilon)$，标准条件流匹配训练模型速度 $v_{\theta}(x_t, t, y)$ 以匹配目标传输方向。对比流匹配保留了这一正项，但额外从批次中采样一个负样本对 $(\tilde{x}, \tilde{y}, \tilde{\epsilon})$，如果模型预测的流也与该其他轨迹兼容，则对其进行惩罚。用论文的符号表示，这变为：
$$ \mathcal{L}_{\Delta \text{FM}}(\theta) = \mathbb{E} \left[ \| v_{\theta}(x_t, t, y) - (\dot{\alpha}_t x + \dot{\sigma}_t \epsilon) \|^2 - \lambda \| v_{\theta}(x_t, t, y) - (\dot{\alpha}_t \tilde{x} + \dot{\sigma}_t \tilde{\epsilon}) \|^2 \right] $$
其中 $\lambda \in [0, 1)$ 控制“推开”项的强度。直观理解：匹配自己的轨迹，同时与别人的轨迹不兼容。
作者证明，对比流匹配能产生更具判别性的轨迹，这进而转化为质量和效率的提升：在 ImageNet（Deng 等人，2009）和 CC3M（Sharma 等人，2018）的实验上，收敛速度更快（据报告，达到相似 FID 所需的训练迭代次数最多减少 9 倍），采样步数更少（据报告，去噪步数最多减少 5 倍）。
一个关键优势在于，该目标函数几乎可以即插即用：你保留通常的流匹配损失，然后利用同一批次中的其他样本作为负样本，添加一个单一的对比“推开”项，这在不引入额外模型前向传播的情况下提供了额外的监督信号。
我们的观察结果
方法
| FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批处理量/秒 ↑ | 基线 |
|---|---|---|---|---|
| 18.20 | 18.20 | 0.41 | 0.39 | 3.95 |
| 对比流匹配 | 20.03 | 0.40 | 0.36 | 3.75 |
在此次运行中，对比流匹配在表征驱动指标上带来了微小但可测量的改进：CMMD 从 0.41 降至 0.40，DINO-MMD 从 0.39 降至 0.36。改进幅度小于论文在 ImageNet 上报告的结果，这并不太令人意外：文本条件比离散类别复杂得多，且训练数据分布的可分离性可能不如 ImageNet，这使得对比信号更难被利用。
在此特定实验中，我们并未看到 FID 的改善（反而略有恶化），但实际运行中吞吐量成本可以忽略不计（3.95 → 3.75 批次/秒）。考虑到这一改动非常简单，且条件/表征指标持续向正确方向移动，我们很可能仍会将对比流匹配作为低成本正则化手段保留在训练流程中。
JiT（Li 和 He，2025）
《回归本源：让去噪生成模型去噪》可能是我们在扩散模型领域最喜爱的近期论文之一，因为它并非新技巧，而是一次重置：停止要求网络预测流形外的量（噪声或速度），只需让它去噪。大多数现代扩散模型和流模型都训练网络预测噪声 $\epsilon$ 或混合量（如速度 $v$）。根据流形假设，自然图像位于低维流形上，而 $\epsilon$ 和 $v$ 本质上处于流形之外，因此预测它们可能是一个比表面看起来更困难的学习问题。
根据流形假设，干净图像位于数据流形上，而噪声和速度则不在流形上。因此，训练模型预测干净图像从根本上比训练它预测类似噪声的目标更容易。图来自 arXiv:2511.13720。
作者用干净图像 $x$ 与噪声 $\epsilon$ 之间的标准线性插值来表述该问题：
$$ z_{t} = t \textrm{ } x + \left(\right. 1 - t \left.\right) \textrm{ } \epsilon , $$
z t=t x+(1−t)ε，以及对应的流速度：
$$ v = \frac{d z_{t}}{d t} = x - \epsilon . $$
v=d t d z t=x−ε。
模型不直接输出 $v_{\theta}$，而是预测一个干净图像的估计值：
$$ x_{\theta} \left(\right. z_{t} , t \left.\right) : = \left(n e t\right){\theta} \left(\right. z{t} , t \left.\right) , $$
x θ(z t,t):=net θ(z t,t)，然后我们通过以下方式将其转换为速度预测：
$$ v_{\theta} \left(\right. z_{t} , t \left.\right) = \frac{x_{\theta} \left(\right. z_{t} , t \left.\right) - z_{t}}{1 - t} . $$
v θ(z t,t)=1−t x θ(z t,t)−z t。
这样我们就可以在 v 空间中保持完全相同的流式目标：
$$ \mathcal{L}{v} = \mathbb{E}{t , x , \epsilon} \left[\right. \left(\parallel v_{\theta} \left(\right. z_{t} , t \left.\right) - v \parallel\right)_{2}^{2} \left]\right. \text{with} v = x - \epsilon . $$
L v=E t,x,ε[∥v θ(z t,t)−v∥2 2] 其中 v = x − ε。
这种公式化使得学习问题在高维空间中变得容易得多：网络不是预测噪声或速度（它们在像素空间中基本上不受约束），而是预测干净图像 $x$，即位于数据流形上的东西。在实践中，这使得直接在像素上训练大块 Transformer 成为可能，无需 VAE 或 tokenizer，同时保持优化稳定且 token 总数可控。
我们的观察结果
我们首先在与其余目标实验相同的设置下评估了 x-prediction，即在 FLUX 潜空间中训练，分辨率为 256×256。
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批/秒 ↑ |
|---|---|---|---|---|
| 基线 | 18.20 | 0.41 | 0.39 | 3.95 |
| X-Pred | 16.80 | 0.54 | 0.49 | 3.95 |
在这种设定下，x-prediction 的优势并不明确。虽然 FID 相比基线略有改善，但 CMMD 和 DINO-MMD 都明显下降，且吞吐量没有变化。这表明，当在已经结构良好的潜空间中工作时，预测干净图像而非速度并不能持续优于基线目标，甚至可能损害表征层面的对齐。
话虽如此，这个实验并非 x-prediction 真正大放异彩的领域。
令人兴奋之处在于，x 预测稳定了高维训练，使得使用更大的图像块并直接在像素空间（无需 VAE）以更高分辨率进行去噪成为可能。利用 JiT，我们直接在 1024×1024 的图像上使用 32×32 的图像块训练了一个模型，而不是在压缩的潜在空间中操作。尽管分辨率高得多且没有分词器，优化过程仍然稳定且快速。我们达到了 FID 17.42、DINO_MMD 0.56 和 CMMD 0.71，吞吐量为 1.33 批次/秒。
这些结果非常显著：直接在 1024×1024 图像上训练，与在 256×256 潜在空间中训练相比，速度仅慢约 3 倍，且是在原始像素上操作。这有力地支持了“回归基础”的核心主张：让模型预测干净图像使得学习问题显著更容易，并为无需分词器的高分辨率文本到图像训练打开了大门，且无需高昂的计算成本。
因此，我们计划将此公式作为即将进行的极速实验的骨干，以探索当与上述其他效率和稀疏化技术结合时，我们能将其推进到何种程度。目前的主要缺点是，这种方法无法让我们受益于 FLUX.2 VAE 的优良特性；探索某种形式的对齐或混合训练是否能使这两者兼容，是我们计划进一步研究的开放方向。
通过 Token 路由与稀疏化降低计算成本
到目前为止，我们讨论的大多数技术都侧重于让每个训练步骤更有效：改进目标函数、塑造表征或加速收敛。下一个杠杆是正交的：让每一步更廉价。
对于扩散模型和流式 Transformer，主要成本是在大量图像/潜在 token 上运行深层 Transformer 堆栈，其中注意力机制随序列长度扩展性不佳。Token 稀疏化方法直接针对这一点，确保只有一部分 token 在网络中计算密集的部分承担全部计算代价，同时仍保留足够的信息流以保持高质量。
大多数掩码方法通过在前向传播中移除 token，然后让模型从学习到的占位符中“幻觉”出缺失内容，从而加速训练。这种方法效果出奇地好，但它违背了迭代去噪的本质。我们并非在每一步中细化所有内容，而是从头开始重建部分信息。
近期两篇论文展示了一种更简洁的替代方案：不是删除信息，而是重新组织计算资源的分配。TREAD 和 SPRINT 有着相同的宏观目标——避免在每个层面对每个 token 都进行全深度计算，但它们通过互补的策略来实现这一目标。
TREAD（Krause 等人，2025）的核心思想是：将原本通过信息丢失（如丢弃或掩码 token）来减少计算的方式，替换为通过 token 路由进行信息传输来减少计算。它引入了一条路由：对于每个训练样本，随机选择一部分 token，暂时跳过连续的一段层，然后在后续再将这些 token 重新注入。这些 token 并未被丢弃，而是避免了支付全深度计算的代价。具体来说，对于一个由模块堆叠 $L_{0} , \ldots , L_{B - 1}$ 组成的去噪器，TREAD 定义了一条路由 $r_{i t o j}$（起始层 $i$，结束层 $j$）。一部分 token 沿着廉价路径（恒等映射）穿过 $L_{i} , \ldots , L_{j}$，而其余 token 则进行正常的全量计算。然后，这两条流在 $L_{j}$ 处再次合并。在实践中，该论文表明，路由高达 50% 的 token 仍然有效，而更高的比例则会开始降低质量。
TREAD 通过将 token 路由绕过某些层来提升训练效率。图源：arXiv:/2501.04765。
SPRINT（Park 等人，2025）通过在网络计算开销最大的部分引入稀疏性，同时保留一条密集的信息通路，扩展了这一方法。其方案结构明确：先对所有 token 运行密集的早期层，以构建可靠的底层特征；然后在计算量最大的稀疏中间层仅保留部分 token；最后，在生成输出前，通过将稀疏深层特征与来自早期层的密集残差流进行重新扩展与融合，再次恢复为密集状态。与 TREAD 的关键区别在于鲁棒性的来源：TREAD 让 token 保持“存在”但更浅（路由），而 SPRINT 则允许大量 token 在中间块中缺失，依靠密集残差路径来保留全分辨率信息。这正是在实践中能够实现更激进稀疏化的原因。该论文探索了约 75% 的丢弃比例，而 TREAD 约为 50%。
SPRINT 超越了 TREAD，它在丢弃中间层大部分 token 的同时，保留了一条密集残差路径以维持全分辨率信息。图片来自 arXiv:/2510.21986。
我们的观察结果
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批次数/秒 ↑ |
|---|---|---|---|---|
| 基线 | 18.20 | 0.41 | 0.39 | 3.95 |
| TREAD | 21.61 | 0.55 | 0.41 | 4.11 |
| SPRINT | 22.56 | 0.72 | 0.42 | 4.20 |
在我们标准的 256×256 潜在空间设置下，两种方法都实现了我们主要追求的性能提升。TREAD 从 3.95 批次数/秒提升至 4.11 批次数/秒，而 SPRINT 则进一步将其推高至 4.20 批次数/秒。代价是，根据我们的评估协议，这种额外的吞吐量提升伴随着明显的质量下降：FID 从 18.20 上升至 21.61（TREAD）和 22.56（SPRINT），CMMD 和 DINO-MMD 也呈现出相同的趋势。
从表面数据来看，路由带来了约 7–9% 的适度吞吐量提升，但在此基准测试中，其代价是指标恶化，其中 SPRINT（更激进的方案）对质量的损害略高于 TREAD。
一个重要的注意事项是，token 稀疏 / 路由模型在标准无分类器引导（CFG）下往往得分较低，在我们的设置中，由于这些运行仍处于相对欠训练状态，这种效应可能会被放大。《引导 Token 稀疏扩散模型》（Krause 等人，2025）的作者认为，这在一定程度上是评估不匹配：路由改变了模型的有效容量，而简单的“有条件 vs. 无条件”CFG 通常会效果不佳，从而人为地降低了质量。我们刻意没有使用专门的引导方案，以保持我们的基准测试在各方法间的一致性，并且在现阶段，将稀疏模型视为引导方面的“劣化版本”也没有太大意义。因此，我们认为这些数字具有方向性参考价值，但仍偏悲观，值得谨慎解读。
在 256×256 分辨率下，路由带来的提升较为有限，因为模型处理的 token 相对较少。在 1024×1024 分辨率下，情况则完全不同。处理 1024 个 token 时，路由终于瞄准了主要成本，结果令人瞩目。
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批次/秒 ↑ |
|---|---|---|---|---|
| 基线 | 17.42 | 0.71 | 0.56 | 1.33 |
| TREAD | 14.10 | 0.46 | 0.37 | 1.64 |
| SPRINT | 16.90 | 0.51 | 0.41 | 1.89 |
TREAD 和 SPRINT 相比密集基线都带来了显著的吞吐量提升，其中 SPRINT 在速度上最为突出。更重要的是，这一次的提升并非以牺牲质量为代价，恰恰相反。尤其是 TREAD，其 FID 大幅下降（17.42 → 14.10），同时 CMMD 和 DINO-MMD 也有显著改善。SPRINT 则更为激进，质量上略有不稳定，但相比基线仍有明显提升，同时是最快的选项。
简而言之，这正是 token 路由大放异彩的场景：高分辨率、大量 token，以及 JiT 风格的像素空间训练。在此场景下，路由不再是一种边际优化——而是同时提升模型训练速度和效果的关键杠杆。
数据
在介绍了表征对齐、核心训练目标和 token 路由之后，我们转向了在实践中始终至关重要的第四个维度：数据。我们发现，训练数据的选择，包括如何通过描述性文字（caption）来表述数据，对训练过程的影响程度不亚于优化技术。以下是我们实验中持续带来显著效果的三个具体数据实验。
长描述 vs 短描述
描述性文字是训练集的关键组成部分：对于文生图模型而言，它们不仅仅是元数据，更是监督信号。DALL·E 3（Betker 等人，2023）的研究论文表明，更丰富的描述性文字可以成为提升训练信号质量和提示词遵循能力的最有力手段之一。为了在我们的设置中隔离其影响，我们保持其他所有因素不变，仅改变描述性文字的样式进行比较：
- 长描述性文字（我们的基线）：包含多个从句的描述，提及构图、属性、光照、材质和物体间关系。
示例
一张照片展示了一只毛茸茸的垂耳兔，它坐在户外一块风化的木制表面上。这只兔子以白色为主，头部和耳朵上带有浅棕色和棕褐色的斑块。它的耳朵明显下垂，毛发看起来柔软而浓密。兔子的眼睛乌黑且富有表现力。它略微偏离中心位置，面朝画面的左侧。在兔子身后，略微失焦的地方，有一个深红色的微型金属手推车。兔子左侧放着一个部分可见的橙色苹果。以红褐色为主的秋日落叶散落在兔子和苹果周围的木制表面上。背景是模糊但仍可辨认的广阔绿草地，暗示着户外环境。光线柔和而自然，很可能是漫射的日光，没有投下任何强烈的阴影。整体氛围宁静、平和，充满秋意。其美学风格质朴而迷人，以兔子作为主要焦点。色彩搭配柔和而自然，主要由白色、棕色、橙色和绿色构成。风格写实而直接，没有任何明显的艺术加工。整体感觉温柔而暖心。
- 简短的单行标题：结构极简的描述。
示例
“一只兔子坐在一张木桌上。”
我们的发现
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批次数/秒 ↑ |
|---|---|---|---|---|
| 基线 | 18.20 | 0.41 | 0.39 | 3.95 |
| 简短标题 | 36.84 | 0.98 | 1.14 | 3.95 |
结果一目了然：改用简短标题严重损害了所有指标上的收敛效果。长标题提供了更丰富的监督信号：除了提示词遵循度之外，还有一个非常实际的优化原因。更多的模型 token 通常意味着更多的信息，因此也为去噪器提供了更多的学习信号。当条件文本指定了构图、属性、光照、材质和关系时，模型就能获得一个更清晰的“目标”，明确去噪轨迹应该保留和优化什么，尤其是在训练早期阶段。
一个有趣的悖论是，这些额外的细节往往让学习问题变得更简单，而非更困难：直觉上，人们可能会认为更长的提示词，包含更多属性、约束和关系，会增加复杂度并加重模型负担。但实际上，情况恰恰相反。简短的描述留下了大量未指定的自由度，迫使模型在模糊性下学习，并隐式地在多种可能的解释之间取平均。而长描述则通过将隐式选择转化为显式约束来消除这种不确定性，使去噪器能够将能力集中在优化一个定义明确的问题上，而不是去猜测哪些信息重要。
长描述是一种强大的训练加速手段，但我们仍然希望模型在短提示词上表现良好，因为这才是人们实际使用这些系统的方式。一个简单的解决方法是，在训练结束时，用长短描述混合的数据集进行一个短期的微调阶段。这样既能保留早期丰富监督信号带来的好处，又能教会模型在条件信息稀疏时保持鲁棒性。
利用合成图像进行自举训练
我们探索的另一个与数据相关的研究问题是，与同等规模的真实语料库相比，低成本的合成语料库能否加速早期训练。为了进行这项基准测试，我们在一个从 Pexels 收集的真实图像数据集上训练了一个模型，并将其与我们的基线模型进行了比较，后者是在使用 MidjourneyV6 生成的合成数据上训练的，两个数据集均包含约 100 万张图像。我们针对同一个 Unsplash 参考集（完全由真实图像组成）对两个训练结果进行了评估。
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ | 批/秒 ↑ |
|---|---|---|---|---|
| 合成图像 | 18.20 | 0.41 | 0.39 | 3.95 |
| 真实图像 | 16.6 | 0.5 | 0.46 | 3.95 |
在合成数据上训练的模型在 CMMD 和 DINO-MMD 指标上得分更高，而在真实图像上训练的模型则获得了更低的 FID 值。这种差异并非矛盾，而主要反映了这些指标各自侧重的评估方向。
FID 对低层级图像统计特征尤为敏感：精细纹理、高频细节、噪声模式以及真实摄影中微妙的非规则性。由于我们的评估参考集由真实图像构成，在真实照片上训练的模型自然更贴近这些统计特征，从而获得更优的 FID 分数。相比之下，合成图像通常表现出略有不同的高频特征、更清晰的边缘、更平滑的微观纹理以及更均匀的噪声——这些差异在定性层面几乎难以察觉，但仍会受到 FID 这类分布度量指标的惩罚。
从定性角度看，这种差异很容易识别。基于合成数据训练的模型倾向于生成全局结构更清晰、构图与物体一致性更强的图像，但同时也呈现出更明显的合成感，表现为纹理更平滑、摄影噪声减少。相比之下，基于真实图像训练的模型能更好地捕捉自然照片中典型的非规则精细纹理，不过通常需要更多训练才能达到相当的全局结构水平。
合成数据在训练初期依然如此有效的一个合理解释是：它让模型接触到更广泛的构图碰撞——即物体、属性、风格和视角之间那些在自然数据集中极少同时出现的非寻常组合。虽然这可能在纹理层面损害真实感，但它迫使模型去解释更广阔的组合空间，这似乎有助于早期的解耦与结构学习。
综合来看，这提示了一种简单而实用的策略：合成数据是快速启动训练并锁定全局结构的有效方式；而如果匹配摄影纹理统计特征是优先目标，那么真实图像在训练后期仍然至关重要。
| 使用真实数据训练的模型 |  |  |  |  |  |  |
|---|---|---|---|---|---|---|
| 使用合成数据训练的模型 |  |  |  |  |  |  |
使用 Alchemist 进行 SFT：小数据集，真实影响
最后，我们使用 Alchemist（Startsev 等人，2025）进行了一次有针对性的监督微调（SFT）实验。Alchemist 是一个精心策划的紧凑型数据集，专为高影响力而设计。Alchemist 在设计上规模很小（3350 个图像-文本对），但通过一个复杂的筛选流程构建而成，该流程从一个网络规模的池开始，逐步提炼出视觉上卓越的样本。
在我们的设置中，我们在 Alchemist 数据集上对预览模型进行了 20K 步的微调。尽管数据集规模很小，但它产生了超乎寻常的效果：它增加了一个独特的“风格层”，带来了更好的构图、更精美的摄影质感以及更丰富的场景，同时对泛化能力没有明显影响。
下面的样本展示了同一基础模型在 Alchemist 微调前后的生成结果对比。
| 无 SFT |  |  |  |  |  |  |  |  |
|---|---|---|---|---|---|---|---|---|
| 有 SFT |  |  |  |  |  |  |  |  |
更多有用的训练技巧
最后但同样重要的是，我们将简要介绍两个实际训练细节，它们的重要性超出了我们的预期。这些因素很容易被忽视，但在我们的案例中，它们对收敛速度和最终图像质量产生了明显影响。
Muon 优化器
在我们的基准测试中，我们通常默认使用 AdamW，因为它可预测且便于在不同运行之间进行比较。然而，最近我们看到人们对那些试图表现得像良好预处理器、但又没有二阶方法全部开销的优化器重新产生了兴趣。最近的一个例子是 Muon（Jordan 等人，2024），从高层次上讲，它试图通过应用条件更好的更新步骤来改进优化，这通常意味着更快的收敛速度和训练早期更清晰的进展。
在我们的设置中，Muon 是少数几个改变优化器就能立即对指标产生可观察效果的案例之一。
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ |
|---|---|---|---|
| 基线 | 18.20 | 0.41 | 0.39 |
| Muon | 15.55 | 0.36 | 0.35 |
在这个实验中，我们使用了 Muon 的官方 PyTorch 实现，该实现目前仅支持分布式数据并行（DDP）训练。如果你正在运行完全分片数据并行（FSDP），社区有可用的变体实现；例如，可以在这里找到。
尽管我们不会仅凭单一基准就得出广泛结论，但这些结果表明，优化器的选择不仅关乎训练稳定性，还能在达到目标质量所需的时间上带来切实收益。
精度陷阱：BF16 中的类型转换与权重存储
我们最终发现了一个设置错误：去噪器权重曾一度被错误地以 bfloat16 格式存储。
需要明确的是，使用 BF16 自动混合精度是很好的做法。以前向和后向传播采用 BF16 或混合精度进行运算是标准操作，通常也是兼顾速度与内存的理想选择。问题出在将参数本身以 BF16 精度进行存储，这会对数值敏感的操作产生负面影响。
在实践中，某些层和操作对参数精度降低的容忍度要低得多：
- 归一化层（例如 LayerNorm / RMSNorm 的统计量），
- 注意力 softmax / logits 路径，
- RoPE，
- 优化器的内部状态 / 更新动态。
| 方法 | FID ↓ | CMMD ↓ | DINO-MMD ↓ |
|---|---|---|---|
| 基线 | 18.20 | 0.41 | 0.39 |
| BF16 权重（错误） | 21.87 | 0.61 | 0.57 |
因此，我们现在严格遵守的规则是：计算时使用 BF16 自动混合精度，但将权重（以及优化器状态）保持在 FP32 精度，或者至少确保数值敏感的模块保持 FP32。
这并非什么炫酷的技巧，但正是那种“无声”的细节——如果不尽早发现，可能会让你白白浪费数天的工作。
总结
我们对 PRX 训练进行了一组系统的消融实验，结合质量指标和吞吐量，将一系列优化、表征、效率和数据选择方案与一个干净的流匹配基线进行了比较。
最大的提升来自对齐：REPA 能加速早期收敛（最佳用法是作为预热阶段，然后关闭），而更好的潜在表示/分词器（REPA-E/FLUX2-AE）则在质量上带来巨大飞跃，但伴随明显的速度权衡。目标函数的调整效果不一——对比流匹配略有帮助，而 x 预测的重要性主要体现在它能实现稳定的 1024² 像素训练。Token 路由（TREAD/SPRINT）在 256² 分辨率下效果甚微，但在高分辨率下则成为主要优势。数据和实践细节也很关键：长描述至关重要，合成数据与真实数据会改变纹理与结构的侧重，小规模 SFT 能起到润色作用，Muon 优化器有帮助，而以 BF16 存储权重则会悄然损害训练效果。
下一步是什么？
第二部分就到这里！如果你想体验本系列早期公开的检查点，PRX-1024 T2I 测试版仍可在此处获取。
我们对下一步感到非常兴奋：未来几周内，我们将发布 PRX 训练框架的完整源代码，并开展一次公开的 24 小时“极速挑战”，将本文中的最佳思路整合到一次运行中，看看这套完整方案在一天内能推进到何种程度。
如果您读到了这里，首先非常感谢您的关注。此外，我们诚挚邀请您加入我们的 Discord 社区，在这里我们讨论 PRX 的进展与成果，以及所有与扩散模型和文生图模型相关的话题。
In the first post of this series, we introduced our goal: training a competitive text-to-image foundation model entirely from scratch, in the open, and at scale. We focused primarily on architectural choices and motivated the core design decisions behind our model PRX. We also released an early, small (1.2B parameters) version of the model as a preview of what we are building (go try it if you haven't already 😉).
In this post, we shift our focus from architecture to training. The goal is to document what actually moved the needle for us when trying to make models train faster, converge more reliably, and learn better representations. The field is moving quickly and the list of “training tricks” keeps growing, so rather than attempting an exhaustive survey, we structured this as an experimental logbook: we reproduce (or adapt) a set of recent ideas, implement them in a consistent setup, and report how they affect optimization and convergence in practice. Finally, we do not only report these techniques in isolation; we also explore which ones remain useful when combined.
In the next post, we will publish the full training recipe as code, including the experiments in this post. We will also run and report on a public "speedrun" where we put the best pieces together into a single configuration and stress-test it end-to-end. This exercise will serve both as a stress test of our current training pipeline and as a concrete demonstration of how far careful training design can go under tight constraints. If you haven’t already, we invite you to join our Discord to continue the discussion. A significant part of this project has been shaped by exchanges with community members, and we place a high value on external feedback, ablations, and alternative interpretations of the results.
The Baseline
Before introducing any training-efficiency techniques, we first establish a clean reference run. This baseline is intentionally simple. It uses standard components, avoids auxiliary objectives, and does not rely on architectural shortcuts or tricks to save compute resources. Its role is to serve as a stable point of comparison for all subsequent experiments.
Concretely, this is a pure Flow Matching (Lipman et al., 2022) training setup (as introduced in Part 1) with no extra objectives and no architectural speed hacks. We will use the small PRX-1.2B model we presented in the first post of this series (single stream architecture with global attention for the image tokens and text tokens) as our baseline and train it in Flux VAE latent space, keeping the configuration fixed across all comparisons unless stated otherwise.
The baseline training setup is as follows:
| Setting | Value |
|---|---|
| Steps | 100k |
| Dataset | Public 1M synthetic image generated with MidJourneyV6 |
| Resolution | 256×256 |
| Global batch size | 256 |
| Optimizer | AdamW |
| lr | 1e-4 |
| weight_decay | 0.0 |
| eps | 1e-15 |
| betas | (0.9, 0.95) |
| Text encoder | GemmaT5 |
| Positional encoding | Rotary (RoPE) |
| Attention mask | Padding mask |
| EMA | Disabled |
This baseline configuration provides a transparent and reproducible anchor. It allows us to attribute observed improvements and regressions to specific training interventions, rather than to shifting hyperparameters or hidden setup changes. Throughout the remainder of this post, every technique is evaluated against this reference with a single guiding question in mind:
Does this modification improve convergence or training efficiency relative to the baseline?
Examples of baseline model generations after 100K training steps.
Benchmarking Metrics
To keep this post grounded, we rely on a small set of metrics to monitor checkpoints over time. None of them is a perfect proxy for perceived image quality, but together they provide a practical scoreboard while we iterate.
Fréchet Inception Distance (FID): (Heusel et al., 2017) Measures how close the distributions of generated and real images are, using Inception-v3 feature statistics (mean and covariance). Lower values typically correlate with higher sample fidelity.
CLIP Maximum Mean Discrepancy (CMMD): (Jayasumana et al., 2024) Measures the distance between real and generated image distributions using CLIP image embeddings and Maximum Mean Discrepancy (MMD). Unlike FID, CMMD does not assume Gaussian feature distributions and can be more sample-efficient; in practice it often tracks perceptual quality better than FID, though it is still an imperfect proxy.
DINOv2 Maximum Mean Discrepancy (DINO-MMD): Same MMD-based distance as CMMD, but computed on DINOv2 (Oquab et al. 2023) image embeddings instead of CLIP. This provides a complementary view of distribution shift under a self-supervised vision backbone.
Network throughput: Average number of samples processed per second (samples/s), as a measure of end-to-end training efficiency.
With the scoreboard defined, we can now dive into the methods we explored, grouped into four buckets: Representation Alignment, Training Objectives, Token Routing and Sparsification, and Data.
Representation Alignment
Diffusion and flow models are typically trained with a single objective: predict a noise-like target (or vector field) from a corrupted input. Early in training, that one objective is doing two jobs at once: it must build a useful internal representation and learn to denoise on top of it. Representation alignment makes this explicit by keeping the denoising objective and adding an auxiliary loss that directly supervises intermediate features using a strong, frozen vision encoder. This tends to speed up early learning and bring the model’s features closer to those of modern self-supervised encoders. As a result, you often need less compute to hit the same quality.
A useful way to view it is to decompose the denoiser into an implicit encoder that produces intermediate hidden states, and a decoder that maps those states to the denoising target. The claim is that representation learning is the bottleneck: diffusion and flow transformers do learn discriminative features, but they lag behind foundation vision encoders when training is compute-limited. Therefore, borrowing a powerful representation space can make the denoising problem easier.
REPA (Yu et al., 2024)
Representation alignment with a pre-trained visual encoder. Figure from arXiv:2410.06940.
REPA adds a representation matching term on top of the base flow-matching objective.
Let $x_{0} sim p_{\text{data}}$x 0∼p data be a clean sample and $x_{1} sim p_{\text{prior}}$x 1∼p prior be the noise sample. The model is trained on an interpolated state $x_{t}$x t (for $t \in \left[\right. 0 , 1 \left]\right.$t∈[0,1]) and predicts a vector field $v_{\theta} \left(\right. x_{t} , t \left.\right)$v θ(x t,t). In REPA, a pretrained vision encoder $f$f processes the clean sample $x_{0}$x 0 to produce patch embeddings $y_{0} = f \left(\right. x_{0} \left.\right) \in \mathbb{R}^{N \times D}$y 0=f(x 0)∈R N×D, where $N$N is the number of patch tokens and $D$D is the teacher embedding dimension. In parallel, the denoiser processes $x_{t}$x t and produces intermediate hidden tokens $h_{t}$h t (one token per patch). A small projection head $h_{\phi}$h ϕ maps these student hidden tokens into the teacher embedding space, and an auxiliary loss maximizes patch-wise similarity between corresponding teacher and student tokens:
$$ \mathcal{L}{\text{REPA}} \left(\right. \theta , \phi \left.\right) = - \mathbb{E}{x_{0} , x_{1} , t} \left[\right. \frac{1}{N} \sum_{n = 1}^{N} \text{sim} \left(\right. y_{0 , \left[\right. n \left]\right.} , \textrm{ } h_{\phi} \left(\right. h_{t , \left[\right. n \left]\right.} \left.\right) \left.\right) \left]\right. $$
L REPA(θ,ϕ)=−E x 0,x 1,t[N 1n=1∑Nsim(y 0,[n],h ϕ(h t,[n]))] Here $n \in \left{\right. 1 , \ldots , N \left.\right}$n∈{1,…,N} indexes patch tokens, $y_{0 , \left[\right. n \left]\right.}$y 0,[n] is the teacher embedding for patch $n$n, $h_{t , \left[\right. n \left]\right.}$h t,[n] is the corresponding student hidden token at time $t$t, and $\text{sim} \left(\right. \cdot , \cdot \left.\right)$sim(⋅,⋅) is typically cosine similarity.
This term is combined with the main flow-matching loss:
$$ \mathcal{L} = \mathcal{L}{\text{FM}} + \lambda \textrm{ } \mathcal{L}{\text{REPA}} $$
L=L FM+λ L REPA
with $\lambda$λ controlling the trade-off.
In practice, the student is trained to produce noise-robust, data-consistent patch representations from $x_{t}$x t, so later layers can focus on predicting the vector field and generating details rather than rediscovering a semantic scaffold from scratch.
What we observed
We ran REPA on top of our baseline PRX training, using two frozen teachers: DINOv2 and DINOv3 (Siméoni et al., 2025). The pattern was very consistent: adding alignment improves quality metrics, and the stronger teacher helps more, at the cost of a bit of speed.
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Baseline | 18.2 | 0.41 | 0.39 | 3.95 |
| REPA-Dinov3 | 14.64 | 0.35 | 0.3 | 3.46 |
| REPA-Dinov2 | 16.6 | 0.39 | 0.31 | 3.66 |
On the quality metrics, both teachers improve over the baseline. The effect is strongest with DINOv3, which achieves the best overall numbers in this run.
REPA is not free: we pay for an extra frozen teacher forward and the patch-level similarity loss, which shows up as a throughput drop from 3.95 batches/s to 3.66 (DINOv2) or 3.46 (DINOv3). In other words, DINOv3 prioritizes maximum representation quality at the cost of slower training, while DINOv2 offers a more efficient tradeoff, still delivering substantial gains with a smaller slowdown.
Our practical takeaway is that REPA is a strong lever for text-to-image training. In our setup, the throughput trade-off is real and the net speedup (time required to reach a given level of image quality) felt a bit less dramatic than what the authors of the paper report on ImageNet-style, class-conditioned generation. That said, the quality gains are still clearly significant. Qualitatively, we also saw the difference early: after ~100K steps, samples trained with alignment tended to lock in cleaner global structure and more coherent layouts, which makes it easy to see why REPA (and alignment variants more broadly) have become a go-to ingredient in modern T2I training recipes.
| Baseline | Repa-DinoV2 | Repa-DinoV3 |
|---|---|---|
| | |
iREPA (Singh et al., 2025)
A natural follow-up to REPA is: what exactly should we be aligning?iREPA argues that the answer is spatial structure, not global semantics. Across a large sweep of 27 vision encoders, the authors find that ImageNet-style “global” quality (e.g., linear-probe accuracy on patch tokens) is only weakly predictive of downstream generation quality under REPA, while simple measures of patch-token spatial self-similarity correlate much more strongly with FID. Based on that diagnosis, iREPA makes two tiny but targeted changes to the REPA recipe to better preserve and transfer spatial information:
- Replace the usual MLP projection head with a lightweight 3×3 convolutional projection operating on the patch grid.
- Apply a spatial normalization to teacher patch tokens that removes a global overlay (mean across spatial locations) to increase local contrast.
Despite representing “less than 4 lines of code”, these tweaks consistently speed up convergence and improve quality across encoders, model sizes, and even REPA-adjacent training recipes.
What we observed
In our setup, we observed a similar kind of boost when applying the iREPA spatial tweaks on top of DINOv2: convergence was a bit smoother and the metrics improved more steadily over the first 100K steps. Interestingly, the same changes did not transfer as cleanly when applied on top of a DINOv3 teacher and they tended to degrade performance rather than help. We do not want to over-interpret that result: this could easily be an interaction with our specific architecture, resolution/patching, loss weighting, or even small implementation details. Still, given this inconsistency across teachers, we will likely not include these tweaks in our default recipe, even if they remain an interesting option to revisit when tuning for a specific setup.
About Using REPA During the Full Training:
The paper REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training (Wang et al., 2025) highlights a key caveat: REPA is a powerful early accelerator, but it can plateau or even become a brake later in training. The authors describe a capacity mismatch. Once the generative model starts fitting the full data distribution (especially high-frequency details), forcing it to stay close to a frozen recognition encoder’s lower-dimensional embedding manifold becomes constraining. Their practical takeaway is simple: keep alignment for the “burn-in” phase, then turn it off with a stage-wise schedule.
We observed the same qualitative pattern in our own runs. When training our preview model, removing REPA after ~200K steps noticeably improved the overall feel of image quality, textures, micro-contrast, and fine detail continued to sharpen instead of looking slightly muted. For that reason, we also recommend treating representation alignment as a transient scaffold. Use it to get fast early progress, then drop it after a while once the model’s own generative features have caught up.
Alignment in the Token Latent Space
So far, “alignment” meant regularizing the generator’s internal features against a frozen teacher while treating the tokenizer / latent space as fixed. A more direct lever is to shape the latent space itself so the representation presented to the flow backbone is intrinsically easier to model, without sacrificing the reconstruction fidelity needed for editing and downstream workflows.
REPA-E (Leng et al., 2025) makes this concrete. Its starting point is a failure mode: if you simply backprop the diffusion / flow loss into the VAE, the tokenizer quickly learns a pathologically easy latent for the denoiser, which can even degrade final generation quality. REPA-E’s fix is a two-signal training recipe:
- keep the diffusion loss, but apply a stop-gradient so it only updates the latent diffusion model (not the VAE);
- update both the VAE and the diffusion model using an end-to-end REPA alignment loss.
Thanks to these two tricks, the tokenizer is explicitly optimized to produce latents that yield higher alignment and empirically better generations.
In parallel, Black Forest Labs’ FLUX.2 AE work frames latent design as a trade-off between learnability, quality, and compression.Their core argument is that improving learnability requires injecting semantic structure into the representation, rather than treating the tokenizer as a pure compression module. This motivates retraining the latent space to explicitly target “better learnability and higher image quality at the same time". They do not share the full recipe, but they do clearly state the key idea: make the AE’s latent space more learnable by adding semantic or representation alignment, and explicitly point to REPA-style alignment with a frozen vision encoder as the mechanism they build on and integrate into the FLUX.2 AE.
What we observed
To probe alignment in the latent space, we compared two pretrained autoencoders as drop-in tokenizers for the same flow backbone: a REPA-E-VAE (where we do add the REPA alignment objective, as in the paper) and the Flux2-AE (where we do not add REPA, following their recommendation). The results were, honestly, extremely impressive, both quantitatively and qualitatively. In samples, the gap is immediately visible: generations show more coherent global structure and cleaner layouts, with far fewer “early training” artifacts.
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Baseline | 18.20 | 0.41 | 0.39 | 3.95 |
| Flux2-AE | 12.07 | 0.09 | 0.08 | 1.79 |
| REPA-E-VAE | 12.08 | 0.26 | 0.18 | 3.39 |
A first striking point is that both latent-space interventions lower the FID by ~6 points (18.20 to ~12.08), which is a much larger jump than what we typically get from “just” aligning intermediate features. This strongly supports the core idea: if the tokenizer produces a representation that is intrinsically more learnable, the flow model benefits everywhere.
The two AEs then behave quite differently in the details. Flux2-AE dominates most metrics (very low CMMD and DINO_MMD, but it comes with a huge throughput penalty: batches/sec drops from 3.95 to 1.79. In our case this slowdown is explained by practical factors they also emphasize: the model is simply heavier, and it also produces a larger latent (32 channels), which increases the amount of work the diffusion backbone has to do per step.
REPA-E-VAE is the “balanced” option: it reaches essentially the same FID as Flux2-AE while keeping throughput much closer to the baseline (3.39 batches/sec).
| Baseline | Flux2-AE | REPA-E-VAE |
|---|---|---|
| | |
Training Objectives: Beyond Vanilla Flow Matching
Architecture gets you capacity, but the training objective is what decides how that capacity is used. In practice, small changes to the loss often have outsized effects on convergence speed, conditional fidelity, and how quickly a model “locks in” global structure. In the sections below, we will go through the objectives we tested on top of our baseline rectified flow setup, starting with a simple but surprisingly effective modification: Contrastive Flow Matching.
Contrastive Flow Matching (Stoica et al., 2025)
Flow matching has a nice property in the unconditional case: trajectories are implicitly encouraged to be unique (flows should not intersect). But once we move to conditional generation (class- or text-conditioned), different conditions can still induce overlapping flows, which empirically shows up as “averaging” behavior: weaker conditional specificity, and muddier global structure. Contrastive flow matching addresses this directly by adding a contrastive term that pushes conditional flows away from other flows in the batch.
Contrastive flow matching makes class-conditional flows more distinct, reducing overlap seen in standard flow matching, and produces higher-quality images that better represent each class. Figure from arXiv:2506.05350.
For a given training triplet $\left(\right. x , y , \epsilon \left.\right)$(x,y,ε), standard conditional flow matching trains the model velocity $v_{\theta} \left(\right. x_{t} , t , y \left.\right)$v θ(x t,t,y) to match the target transport direction. Contrastive flow matching keeps that positive term, but additionally samples a negative pair $\left(\right. \overset{\sim}{x} , \overset{\sim}{y} , \overset{\sim}{\epsilon} \left.\right)$(x,y,ε~) from the batch and penalizes the model if its predicted flow is also compatible with that other trajectory. In the paper’s notation, this becomes:
$$ \mathcal{L}{\Delta \text{FM}} \left(\right. \theta \left.\right) = \mathbb{E} \left[\right. \parallel v{\theta} \left(\right. x_{t} , t , y \left.\right) - \left(\right. \left(\overset{\cdot}{\alpha}\right){t} x + \left(\overset{\cdot}{\sigma}\right){t} \epsilon \left.\right) \parallel^{2} \textrm{ }\textrm{ } - \textrm{ }\textrm{ } \lambda \parallel v_{\theta} \left(\right. x_{t} , t , y \left.\right) - \left(\right. \left(\overset{\cdot}{\alpha}\right){t} \overset{\sim}{x} + \left(\overset{\cdot}{\sigma}\right){t} \overset{\sim}{\epsilon} \left.\right) \parallel^{2} \left]\right. $$
L Δ FM(θ)=E[∥v θ(x t,t,y)−(α˙tx+σ˙tε)∥2−λ∥v θ(x t,t,y)−(α˙tx+σ˙tε)∥2]
where $\lambda \in \left[\right. 0 , 1 \left.\right)$λ∈[0,1) controls the strength of the “push-away” term. Intuitively: match your own trajectory, and be incompatible with someone else’s.
The authors show that contrastive flow matching produces more discriminative trajectories and that this translates into both quality and efficiency gains: faster convergence (reported up to 9× fewer training iterations to reach similar FID) and fewer sampling steps (reported up to 5× fewer denoising steps) on ImageNet (Deng et al. 2009) and CC3M(Sharma et al., 2018) experiments.
A key advantage is that the objective is almost a drop-in replacement: you keep the usual flow-matching loss, then add a single contrastive “push-away” term using other samples in the same batch as negatives which provides the extra supervision without introducing additional model passes.
What we observed
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Baseline | 18.20 | 0.41 | 0.39 | 3.95 |
| Contrastive-FM | 20.03 | 0.40 | 0.36 | 3.75 |
On this run, contrastive flow matching yields a small but measurable improvement on the representation-driven metrics: CMMD goes from 0.41 → 0.40 and DINO-MMD from 0.39 → 0.36. The magnitude of the gain is smaller than what the paper reports on ImageNet, which is not too surprising: text conditioning is much more complex than discrete classes, and the training data distribution is likely less “separable” than ImageNet, making the contrastive signal harder to exploit.
We do not see an improvement in FID in this specific experiment (it slightly worsens), but the throughput cost is negligible in practice (3.95 → 3.75 batches/sec). Given the simplicity of the change and the consistent movement in the right direction for the conditioning/representation metrics, we will likely still keep contrastive flow matching in our training pipeline as a low-cost regularizer.
JiT (Li and He, 2025)
Back to Basics: Let Denoising Generative Models Denoise is probably one of our favorite recent papers in the diffusion space because it is not a new trick but a reset: stop asking the network to predict off-manifold quantities (noise or velocity) and just let it denoise. Most modern diffusion and flow models train the network to predict noise $\epsilon$ε or a mixed quantity like velocity $v$v. Under the manifold assumption, natural images live on a low-dimensional manifold, while $\epsilon$ε and $v$v are inherently off-manifold, so predicting them can be a harder learning problem than it looks.
Under the manifold assumption, clean images lie on the data manifold while noise and velocity do not. Thus training the model to predict clean images is fundamentally easier than training it to predict noise-like targets. Figure from arXiv:2511.13720.
The authors frame the problem with the standard linear interpolation between the clean image $x$x and the noise $\epsilon$ε:
$$ z_{t} = t \textrm{ } x + \left(\right. 1 - t \left.\right) \textrm{ } \epsilon , $$
z t=t x+(1−t)ε, and the corresponding flow velocity:
$$ v = \frac{d z_{t}}{d t} = x - \epsilon . $$
v=d t d z t=x−ε.
Instead of outputting $v_{\theta}$v θ directly, the model predicts a clean image estimate:
$$ x_{\theta} \left(\right. z_{t} , t \left.\right) : = \left(n e t\right){\theta} \left(\right. z{t} , t \left.\right) , $$
x θ(z t,t):=net θ(z t,t), and we convert it to a velocity prediction via:
$$ v_{\theta} \left(\right. z_{t} , t \left.\right) = \frac{x_{\theta} \left(\right. z_{t} , t \left.\right) - z_{t}}{1 - t} . $$
v θ(z t,t)=1−t x θ(z t,t)−z t.
Then we can keep the exact same flow-style objective in v-space:
$$ \mathcal{L}{v} = \mathbb{E}{t , x , \epsilon} \left[\right. \left(\parallel v_{\theta} \left(\right. z_{t} , t \left.\right) - v \parallel\right)_{2}^{2} \left]\right. \text{with} v = x - \epsilon . $$
L v=E t,x,ε[∥v θ(z t,t)−v∥2 2]with v=x−ε.
This formulation makes the learning problem substantially easier in high dimensions: instead of predicting noise or velocity (which are essentially unconstrained in pixel space), the network predicts the clean image $x$x, i.e., something that lies on the data manifold. In practice, this makes it feasible to train large-patch Transformers directly on pixels without a VAE or tokenizer while keeping optimization stable and the total number of tokens manageable.
What we observed
We first evaluated x-prediction in the same setting as the rest of our objective experiments, namely training in the FLUX latent space at 256×256 resolution.
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Baseline | 18.20 | 0.41 | 0.39 | 3.95 |
| X-Pred | 16.80 | 0.54 | 0.49 | 3.95 |
In this regime, the benefit of x-prediction is unclear. While FID improves slightly compared to the baseline, both CMMD and DINO-MMD degrade noticeably, and throughput is unchanged. This suggests that, when working in an already well-structured latent space, predicting clean images instead of velocity does not consistently dominate the baseline objective, and can even hurt representation-level alignment.
That said, this experiment is not where x-prediction really shines.
The exciting part is that x-prediction stabilizes high-dimensional training, making it feasible to use larger patches and denoise directly in pixel space, without a VAE, at much higher resolutions. Using JiT, we trained a model directly on 1024×1024 images with 32×32 patches, instead of operating in a compressed latent space. Despite the much higher resolution and the absence of a tokenizer, optimization remained stable and fast. We reached FID 17.42, DINO_MMD 0.56, and CMMD 0.71 with a throughput of 1.33 batches/sec.
These results are remarkable: training directly on 1024×1024 images is only about 3× slower than training in a 256×256 latent space, while operating on raw pixels. This strongly supports the core claim of Back to Basics: letting the model predict clean images makes the learning problem significantly easier, and opens the door to high-resolution, tokenizer-free text-to-image training without prohibitive compute costs.
As a result, we plan to use this formulation as the backbone of our upcoming speedrun experiments, to see how far we can push it when combined with the other efficiency and sparsification techniques discussed above. The main downside for now is that this approach does not let us benefit from the very nice properties of the FLUX.2 VAE; exploring whether some form of alignment or hybrid training could make these two worlds compatible is an open direction we plan to investigate further.
Token Routing and Sparsification to Reduce Compute Costs
So far, most of the techniques we discussed focus on making each training step more effective: improving the objective, shaping the representations, or accelerating convergence. The next lever is orthogonal: make each step cheaper.
For diffusion and flow transformers, the dominant cost is running deep transformer stacks over a large set of image/latent tokens where attention scales poorly with sequence length. Token sparsification methods target this directly by ensuring that only a subset of tokens pays the full compute price in the expensive parts of the network, while still preserving enough information flow to keep quality high.
Most masking approaches accelerate training by removing tokens from the forward pass, then asking the model to hallucinate the missing content from learned placeholders. That works surprisingly well, but it violates the spirit of iterative denoising. Instead of refining all the content in each step, we are reconstructing parts from scratch.
Two recent papers illustrate a cleaner alternative: instead of deleting information, they reorganize where compute is spent. TREAD and SPRINT share the same high-level objective of avoiding full-depth computation for every token at every layer, but they pursue it through complementary strategies.
TREAD's (Krause et al., 2025) core idea is to replace compute reduction through information loss, such as dropping or masking tokens, with compute reduction through information transport using token routing. It introduces a route: for each training sample, it randomly selects a fraction of tokens and temporarily bypasses a contiguous chunk of layers, then re-injects those tokens later. Tokens are not discarded. Instead, they avoid paying the cost of full depth. Concretely, for a denoiser with a stack of blocks $L_{0} , \ldots , L_{B - 1}$L 0,…,L B−1, TREAD defines a route $r_{i t o j}$r i t o j (start layer $i$i, end layer $j$j). A subset of tokens follows the cheap path (identity) across $L_{i} , \ldots , L_{j}$L i,…,L j, while the rest follows the normal full computation. Then both streams merge again at $L_{j}$L j. In practice, the paper shows that routing up to 50% of tokens remains effective, while higher rates begin to degrade quality.
. TREAD enhances training efficiency by routing tokens around certain layers. Figure from arXiv:/2501.04765.
SPRINT (Park et al., 2025) extends this approach by introducing sparsity in the most computationally expensive parts of the network, while preserving a dense information pathway. Its recipe is intentionally structured: run dense early layers over all tokens to build reliable low-level features, then keep only a subset of tokens through the sparse middle layers where compute is heaviest, and finally go dense again by re-expanding and fusing sparse deep features with a dense residual stream from the early layers, before producing the output. The key distinction from TREAD is where robustness comes from: TREAD keeps tokens “present” but shallower (routing), whereas SPRINT allows many tokens to be absent in the middle blocks, relying on the dense residual path to preserve full-resolution information. This is what enables more aggressive sparsification in practice. The paper explores drop ratios around 75%, versus ~50% for TREAD.
SPRINT goes beyond TREAD by dropping most tokens in the middle layers while keeping a dense residual path to preserve full-resolution information. Figure from arXiv:/2510.21986.
What we observed
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Baseline | 18.20 | 0.41 | 0.39 | 3.95 |
| TREAD | 21.61 | 0.55 | 0.41 | 4.11 |
| SPRINT | 22.56 | 0.72 | 0.42 | 4.20 |
Under our standard 256×256 latent setup, both methods deliver the primary benefit we were targeting. TREAD goes from 3.95 → 4.11 batches/sec, and SPRINT pushes it a bit further to 4.20 batches/sec. The cost is that under our evaluation protocol, this extra throughput comes with a clear loss in quality: FID rises from 18.20 to 21.61 (TREAD) and 22.56 (SPRINT), with the same pattern observed in CMMD and DINO-MMD.
Taken at face value, routing yields a modest ~7–9% throughput gain, but it comes with worse metrics in this benchmark, with SPRINT (the more aggressive scheme) degrading quality slightly more than TREAD.
One important caveat is that token-sparse / routed models tend to score worse under vanilla Classifier-Free Guidance (CFG), and this effect is likely amplified here because these runs are still relatively undertrained in our setting. The authors of Guiding Token-Sparse Diffusion Models (Krause et al., 2025) argue this is partly an evaluation mismatch: routing changes the model’s effective capacity, and plain “conditional vs. unconditional” CFG often becomes less effective, which can artificially reduce quality. We deliberately did not use specialized guidance schemes to keep our benchmark consistent across methods, and at this stage it would also not be very meaningful to treat the sparse model as a “bad version of itself” for guidance. As a result, we consider these numbers directionally useful, but still pessimistic and worth interpreting with caution.
At 256×256, routing only gave modest gains because the model processes relatively few tokens. At 1024×1024, the picture changes completely. With 1024 tokens, routing finally targets the dominant cost, and the results are striking.
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Baseline | 17.42 | 0.71 | 0.56 | 1.33 |
| TREAD | 14.10 | 0.46 | 0.37 | 1.64 |
| SPRINT | 16.90 | 0.51 | 0.41 | 1.89 |
Both TREAD and SPRINT deliver large throughput gains over the dense baseline, with SPRINT pushing speed the furthest. More importantly, this time the gains do not come at the expense of quality but quite the opposite. TREAD in particular stands out, with a dramatic drop in FID (17.42 → 14.10) alongside strong improvements in CMMD and DINO-MMD. SPRINT is slightly more aggressive and a bit noisier in quality, but still clearly improves over the baseline while being the fastest option.
In short, this is the regime where token routing really shines: high resolution, many tokens, and JiT-style pixel-space training. Here, routing is no longer a marginal optimization—it’s a major lever that improves both how fast and how well the model trains.
Data
After covering representation alignment, the core training objective, and token routing, we turned to the fourth axis that kept constantly mattered in practice: data. We found that the choice of training data, including how it is described through captions, can influence the trajectory of a training run as much as optimization techniques. Below are three concrete data experiments that consistently moved the needle in our setup.
Long vs. Short Captions
Captions are an essential part of the training set: for a text-to-image model, they are not just metadata, they are the supervision. The DALL·E 3 (Betker et al., 2023) research paper showed that richer captions can be one of the strongest levers for improving training signal and prompt-following. To isolate the effect in our setup, we kept everything else fixed and changed only the caption style to compare:
- Long, descriptive captions (our baseline): multi-clause captions that mention composition, attributes, lighting, materials, and relationships.
Example
"A photograph depicts a fluffy lop-eared rabbit sitting on a weathered wooden surface outdoors. The rabbit is predominantly white with patches of light brown and tan fur, particularly on its head and ears. Its ears droop noticeably, and its fur appears soft and thick. The rabbit's eyes are dark and expressive. It is positioned slightly off-center, facing towards the left of the frame. Behind the rabbit, slightly out of focus, is a miniature dark red metal wheelbarrow. A partially visible orange apple sits to the left of the rabbit. Fallen autumn leaves, predominantly reddish-brown, are scattered around the rabbit and apple on the wooden surface. The background is a blurred but visible expanse of green grass, suggesting an outdoor setting. The lighting is soft and natural, likely diffused daylight, casting no harsh shadows. The overall atmosphere is calm, peaceful, and autumnal. The aesthetic is rustic and charming, with a focus on the rabbit as the main subject. The color palette is muted and natural, consisting mainly of whites, browns, oranges, and greens. The style is naturalistic and straightforward, without any overt artistic manipulation. The vibe is gentle and heartwarming."
- Short, one-line captions: minimal descriptions with much less structure.
Example
"A rabbit sitting on top of a wooden table."
What we observed
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Baseline | 18.20 | 0.41 | 0.39 | 3.95 |
| Short-Captions | 36.84 | 0.98 | 1.14 | 3.95 |
The outcome was unambiguous: switching to short captions severely hurt convergence across all metrics. Long captions provide a richer supervision signal: beyond prompt adherence, there is a very practical optimization reason. More tokens usually means more information, and therefore more learning signal for the denoiser. When the conditioning text specifies composition, attributes, lighting, materials, and relationships, the model gets a sharper “target” for what the denoising trajectory should preserve and refine, especially early in training.
The fun paradox is that this extra detail often makes the learning problem easier, not harder: intuitively, one might expect longer prompts, with more attributes, constraints, and relationships, to increase complexity and burden the model. In practice, the opposite happens. Short captions leave many degrees of freedom unspecified, forcing the model to learn under ambiguity and implicitly average across multiple plausible interpretations. Long captions collapse that uncertainty by turning implicit choices into explicit constraints, allowing the denoiser to focus its capacity on refining a well-posed solution instead of guessing what matters.
Long captions are a strong training-time accelerator, but we still want the model to behave well on short prompts because that is how people actually use these systems. A simple workaround is to end training with a short fine-tuning stage on a mixture of long and short captions. That keeps the benefits of rich supervision early, while teaching the model to stay robust when conditioning is sparse.
Bootstrapping With Synthetic Images
Another data-related research question we explore is whether a low-cost synthetic corpus can accelerate early training compared to a real corpus of similar size. For this benchmark, we trained a model on a dataset of real images collected from Pexels and compared it with our Baseline which was trained on synthetic data generated with MidjourneyV6, both of which have around 1M images. We evaluated both runs against the same Unsplash reference set, composed exclusively of real images. 
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ | batches/sec ↑ |
|---|---|---|---|---|
| Synthetic images | 18.20 | 0.41 | 0.39 | 3.95 |
| Real images | 16.6 | 0.5 | 0.46 | 3.95 |
The synthetic-trained model scores better on CMMD and DINO-MMD, while the model trained on real images achieves a lower FID. Rather than a contradiction, this split mostly reflects what these metrics emphasize.
FID is particularly sensitive to low-level image statistics: fine textures, high-frequency detail, noise patterns, and the subtle irregularities of real photography. Since our evaluation reference is composed of real images, a model trained on real photos naturally matches those statistics more closely, which translates into a better FID. Synthetic images, by contrast, often exhibit slightly different high-frequency signatures, cleaner edges, smoother micro-textures, more uniform noise, which are barely noticeable qualitatively but still get penalized by distributional metrics like FID.
Qualitatively, this difference is easy to spot. Models trained on synthetic data tend to produce images with cleaner global structure and stronger compositional and object coherence, but also exhibit a more synthetic appearance, characterized by smoother textures and reduced photographic noise. In contrast, models trained on real images better capture the irregular, fine-grained textures typical of natural photographs, though they often require more training to achieve comparable global structure.
One plausible explanation synthetic data remains so effective early on is that it exposes the model to a wider range of compositional collisions: unusual pairings of objects, attributes, styles, and viewpoints that rarely co-occur in natural datasets. While this can hurt realism at the texture level, it forces the model to explain a broader space of combinations, which appears to help with early disentanglement and structure learning.
Considered jointly, this suggests a simple but practical strategy: synthetic data is an efficient way to bootstrap training and lock in global structure quickly, while real images remain important later on if matching photographic texture statistics is the priority.
| Model trained with real data | | | | | | |
|---|---|---|---|---|---|---|
| Model trained with synthetic data | | | | | | |
SFT With Alchemist: Small Dataset, Real Impact
Finally, we experimented with a targeted Supervised Fine-Tuning (SFT) pass using Alchemist (Startsev et al., 2025), a compact dataset explicitly curated for high-impact. Alchemist is small by design (3,350 image–text pairs), but is constructed through a sophisticated curation pipeline that starts from a web-scale pool and progressively distills it down to visually exceptional samples.
In our setup, we fine-tuned our preview models for 20K steps on Alchemist. Despite the dataset’s small size, it had an outsized effect: it adds a distinct “style layer” with better composition, more photographic polish, and richer scenes without a clear impact on generalization.
The samples below show a side-by-side comparison of generations from the same base model, before and after the Alchemist fine-tuning pass.
| Without SFT | | | | | | | | |
|---|---|---|---|---|---|---|---|---|
| With SFT | | | | | | | | |
More Useful Tips for Training
Last but not least, we will briefly cover two practical training details that turned out to matter more than we expected. These factors are easily overlooked and in our case they had a clear impact on convergence and final image quality.
Muon Optimizer
We generally default to AdamW for our benchmarks because it’s predictable and easy to compare across runs. However, lately, we have seen a renewed interest in optimizers that try to behave more like a good preconditioner without the full overhead of second-order methods. One recent example is Muon (Jordan et al., 2024), which, at a high level, tries to improve optimization by applying a better-conditioned update step, often translating into faster convergence and cleaner progress early in training.
In our setup, Muon was one of the rare cases in which a change of optimizer produced an immediately observable effect on the metrics.
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ |
|---|---|---|---|
| Baseline | 18.20 | 0.41 | 0.39 |
| Muon | 15.55 | 0.36 | 0.35 |
For this experiment, we used the official PyTorch implementation of Muon, which at the moment supports Distributed Data Parallel (DDP) training only. If you’re running Fully Sharded Data Parallel (FSDP), there are community variants available; for example here.
While we refrain from broad conclusions based on a single benchmark, these results indicate that optimizer choice extends beyond stability considerations and can yield tangible gains in time-to-quality.
Precision Gotcha: Casting vs. Storing weights in BF16
We eventually identified an error in our setup, where the denoiser weights were mistakenly stored in bfloat16 for a period of time.
To be clear, using the BF16 autocast is great. Running the forward and backward passes in BF16 or mixed precision is standard and usually what you want for speed and memory. The problem arises from keeping the parameters in BF16 precision, which negatively impacts numerically sensitive operations.
In practice, some layers and operations are much less tolerant to reduced parameter precision:
- normalization layers (e.g. LayerNorm / RMSNorm statistics),
- attention softmax/logits paths,
- RoPE,
- optimizers’ internal state / update dynamics.
| Method | FID ↓ | CMMD ↓ | DINO-MMD ↓ |
|---|---|---|---|
| Baseline | 18.20 | 0.41 | 0.39 |
| BF16 weights (bug) | 21.87 | 0.61 | 0.57 |
So the rule we now follow very strictly is: use BF16 autocast for compute, but keep weights (and optimizer state) in FP32 or at least ensure numerically sensitive modules stay FP32.
It is not a glamorous trick but it is exactly the kind of “silent” detail that can cost you multiple days of work if you do not notice it early.
Summary
We ran a systematic set of ablations on PRX training, comparing a range of optimization, representation, efficiency, and data choices against a clean flow-matching baseline using both quality metrics and throughput.
The biggest gains came from alignment: REPA boosts early convergence (best used as a burn-in, then turned off), and better latents/tokenizers (REPA-E/FLUX2-AE) give a large jump in quality with clear speed trade-offs. Objective tweaks were mixed—contrastive FM helped slightly, while x-prediction mattered most by enabling stable 1024² pixel training. Token routing (TREAD/SPRINT) is minor at 256² but becomes a major win at high resolution. Data and practical details also mattered: long captions are critical, synthetic vs. real data shifts texture vs. structure, small SFT adds polish, Muon helped, and BF16-stored weights quietly hurt training.
What's next?
That’s it for Part 2! If you want to play with an earlier public checkpoint from this series, the PRX-1024 T2I beta is still available here.
We are really excited about what’s next: in the coming weeks we will release the full source code of the PRX training framework, and we will do a public 24-hour “speedrun” where we combine the best ideas from this post into a single run and see how far the full recipe can go in one day.
If you made it this far, first of all thank you very much for your interest. Furthermore, we would love to have you join our Discord community where we discuss PRX progress and results, along with everything related to diffusion and text-to-image models.