引言
欢迎回来 👋
在之前的两篇文章（第一部分和第二部分）中，我们探讨了扩散模型在架构和训练技巧方面的广泛内容。我们尝试孤立地评估每个想法，测量吞吐量、收敛速度和最终图像质量，并试图理解哪些因素真正起到了关键作用。
在这篇文章中，我们想回答一个更实际的问题：
当我们把所有有效的技巧结合起来时，会发生什么？
我们不再一次只优化一个维度，而是将最有前景的要素叠加在一起，看看在严格的计算预算下，能将性能推到何种程度。
为了具体说明，我们进行了一次 24 小时极速挑战：
- 32 块 H200
- 总计算预算约 1500 美元（每 GPU 每小时 2 美元）
这与早期扩散模型时代相去甚远，当时训练有竞争力的模型可能需要花费数百万美元。这里的目标是展示该领域已经取得了多大的进步，以及通过精心的工程设计，在仅仅一天的训练中能取得怎样的成果。
这次极速挑战不仅仅是一个有趣的实验。它很可能将成为我们未来大规模训练方案的基础。
除了结果之外，我们还开源了我们的代码（GitHub 链接），其中包含：
- 本次极速挑战使用的训练代码
- 上一篇博文中的实验框架
这样你就可以自行复现、修改和扩展所有内容。
训练方案
现在让我们来详细看看这次 24 小时运行中包含了哪些内容。
X 预测与像素空间训练
我们使用了来自《回归基础：让去噪生成模型去噪》[Li and He, 2025] 的 x 预测公式。如第二部分所述，这可以直接在像素空间中进行训练，并完全消除了对 VAE 的需求。我们使用 32 的补丁大小，并在初始 token 投影层中使用 256 维的瓶颈。这种设计将序列长度控制在合理范围内，使得即使在更高分辨率下，像素空间训练在计算上也是可行的。
在 512 像素下，序列长度为：
$$ \left(\right. 512 / 32 \left.\right)^{2} = 256 $$
(512/32)²=256
在 1024 像素下，序列长度变为：
$$ \left(\right. 1024 / 32 \left.\right)^{2} = 1024 $$
(1024/32)²=1024
我们没有遵循常规的 256px → 512px → 1024px 训练策略，而是直接从 512px 开始训练，然后在 1024px 分辨率下进行微调。
在可控的 token 数量和现代硬件条件下，像素空间训练已不再难以实现。它只是一种更简洁、更直接的建模方式。
感知损失
在像素空间中直接预测 $x_{0}$ 的一个非常好的附带效果是，我们可以复用经典计算机视觉领域的整套工具箱。
当模型输出的是隐空间向量时，感知监督会变得很棘手。你要么需要解码回像素空间，要么在一个可能与人类感知对齐也可能不对齐的、经过学习的隐空间中定义损失函数。一旦你直接预测像素，一切又变得简单直接。你可以完全按照感知损失最初的设计方式来使用它们。
我们从论文《PixelGen: Pixel Diffusion Beats Latent Diffusion with Perceptual Loss》[Ma et al.] 中汲取了灵感，该论文的作者在扩散损失的基础上引入了额外的感知目标。他们表明，加入感知信号可以显著提升收敛速度和最终的视觉质量。
在这次 24 小时的运行中，我们增加了两个辅助损失：
- LPIPS（[Zhang et al.]）
- 基于 DINO 的感知损失（我们使用了 DINOv2 [Oquab et al.]）
思路很简单：除了标准的流匹配目标函数外，我们还鼓励预测出的干净图像在感知特征空间上与目标图像相匹配。LPIPS 捕捉低层次的感知相似性，而 DINO 特征则提供了更强的语义信号。
我们保留了论文中的整体思路，但对一些细节进行了调整。在我们的实验中，我们根据经验发现以下做法效果更好：
- 对池化后的完整图像应用感知损失，而不是基于图像块的特征
- 在所有噪声水平下都应用这些损失
这些都是实现上的小细节，但在我们的设定中，它们始终能带来更好的结果。
我们将 LPIPS 损失的权重设为 0.1，DINO 感知损失的权重设为 0.01，这与原论文中推荐的值一致。
与主 Transformer 前向传播相比，这些损失的计算量较轻，在我们的设置中，它们只增加少量开销，同时却能持续提升质量。
使用 TREAD 进行 Token 路由
为了降低每一步的计算成本，我们采用了基于 TREAD [Krause et al., 2025] 的 token 路由机制，该机制会随机选取一部分 token，让它们跳过连续的一段 Transformer 模块，之后再重新注入，从而确保没有任何 token 被丢弃。
我们选择 TREAD 而非 SPRINT（Park et al., 2025），主要是出于简洁性的考虑，并且在我们设定的场景（序列长度 64 vs. TREAD 在 512px 下的 128）中，SPRINT 额外增加的复杂度所带来的计算节省并不显著。
按照 TREAD 的方案，我们将 50% 的 token 从第 2 个模块路由到 Transformer 的倒数第二个模块。
经过路由的模型在标准 CFG 下表现可能较差，尤其是在训练不足的情况下。因此，我们实现了一种简单的自引导方案，其灵感来源于 Guiding Token-Sparse Diffusion Models（Krause et al., 2025），该方案利用密集预测与路由条件预测之间的差异进行引导，而非依赖无条件分支。
基于 REPA 和 DINOv3 的表征对齐
我们使用 REPA [Yu et al., 2024] 进行表征对齐。
在教师模型的选择上，我们采用了 DINOv3 [Siméoni et al. 2025]，因为在我们之前的实验中，它带来了最佳的质量提升。
具体来说，我们在第 8 个 Transformer 模块处应用一次对齐损失，损失权重设为 0.5。
由于我们将 REPA 与 TREAD 路由结合使用，因此我们仅对非路由 token（即实际经过我们施加损失的那些模块的 token）计算对齐损失。这保证了 REPA 信号的一致性，并避免对跳过了计算路径的 token 进行特征比较。
优化器：Muon
我们使用了 Muon 优化器，具体采用了 muon_fsdp_2 中的 FSDP 实现，因为在我们之前的运行中，它相比 Adam 有显著的改进。
Muon 仅应用于 2D 参数（主要是矩阵）。其他所有参数（偏置、归一化层、嵌入向量等）均使用 Adam 进行优化，这就是配置中包含两个参数组的原因。
| 分组 | 适用范围 | 我们使用的关键参数 |
|---|---|---|
| Muon | 2D 参数 | lr=1e-4, momentum=0.95, nesterov=true, ns_steps=5 |
| Adam | 所有非 2D 参数 | lr=1e-4, betas=(0.9, 0.95), eps=1e-8 |
训练设置
我们在三个公开可用的合成数据集上进行了训练：
- Flux 生成的数据（170 万条），lehduong/flux_generated
- FLUX-Reason-6M (6M)，LucasFang/FLUX-Reason-6M
- midjourney-v6-llava (1M)，brivangl/midjourney-v6-llava，我们使用 Gemini 1.5 对其重新进行了标注，以使提示词更加一致并减少标注噪声。
时间安排基本上是：在 512 分辨率下快速推进，然后在 1024 分辨率下进行锐化：
- 在 512px 分辨率下，以 1024 的批次大小训练 100k 步
- 在 1024px 分辨率下，以 512 的批次大小训练 20k 步，不使用 REPA。
我们还保留了权重的指数移动平均（EMA）用于采样和评估：
- 平滑系数 = 0.999
- 更新间隔 = 10ba
- EMA 起始步 = 0ba
结果与总结
以下是我们在整个训练过程中跟踪的评估曲线，以及来自最终检查点的几张示例网格图：
对于仅一天的训练来说，这已经是一个相当不错的状态了。该模型并非完美无缺（你仍然可以发现一些纹理瑕疵、偶尔出现的怪异解剖结构，并且在处理非常困难的提示词时可能会有些不稳定），但它显然是可用的。提示词遵循能力很强，整体美学风格一致，并且 1024 阶段基本达到了我们的预期：在不破坏构图的情况下锐化细节。
关键结论是，我们已经非常接近目标了。剩余的问题看起来更像是训练不足的痕迹和数据多样性有限所致，而非配方本身存在结构性缺陷。这些失败模式与人们对一个尚未见过足够多样化数据的模型所预期的表现是一致的。有了更多的算力和更广泛的数据覆盖，这个确切的设置应该会以相当可预测的方式继续改进。
从更宏观的角度来看，这次极速训练也凸显了扩散模型训练已经取得了多大的进步。通过结合像素空间训练、高效路由、表征对齐和轻量级感知引导，你现在可以在大约一天的时间内，以不久前听起来还不切实际的预算，得到一个有意义的模型。
下一步是什么？
这 24 小时的训练只是一个起点，而非终点线。接下来，我们将继续以稍大的规模推进相同的配方，并在数据集组合和标注方面进行迭代。
这次极速训练背后的所有代码和配置，以及贯穿第一部分和第二部分的完整实验框架，均可在 PRX 仓库中获取：https://github.com/Photoroom/PRX。
虽然我们不会重新分发本次运行中使用的精确训练数据集，但该流程是完全可配置的，旨在方便你适配自己的数据。你可以接入不同的数据集，调整各个组件（TREAD、REPA、感知损失、Muon 等），并以极低的摩擦成本进行受控实验。我们的目标是使其成为快速扩散研究的实用试验场，并希望社区能利用它，在自己的环境中探索、基准测试并迭代这些技术。
如果你读到了这里，感谢你的阅读。我们也非常欢迎你加入我们的 Discord 社区，在那里我们会分享 PRX 的进展和结果，并讨论任何与扩散模型和文生图相关的话题。
暂时告别，敬请期待下一轮实验！🚀
致谢。
本次速跑受到了近期几项探索快速且低成本训练扩散模型工作的启发。如果你对速跑文生图模型感兴趣，我们鼓励你查阅以下工作：
- Haridas, A., Shen, T., Yu, J. Nitro-T: 在 1 天内从零训练一个文生图扩散模型。https://rocm.blogs.amd.com/artificial-intelligence/nitro-t-diffusion/README.html
- Bhanded, S. 速跑 ImageNet 扩散模型。https://arxiv.org/abs/2512.12386
- Sehwag, V., Kong, X., Li, J., Spranger, M., Lyu, L. 把每一分钱都用在刀刃上：在微预算下从零开始训练扩散模型。https://arxiv.org/abs/2407.15811
- Yeh, S.-Y. 从零到孵化：自制扩散模型。https://arxiv.org/abs/2509.06068
Introduction
Welcome back 👋
In the last two posts (Part 1 and Part 2), we explored a wide range of architectural and training tricks for diffusion models. We tried to evaluate each idea in isolation, measuring throughput, convergence speed, and final image quality, and tried to understand what actually moves the needle.
In this post, we want to answer a much more practical question:
What happens when we combine all the tricks that worked?
Instead of optimizing one dimension at a time, we’ll stack the most promising ingredients together and see how far we can push performance under a strict compute budget.
To make things concrete, we’re doing a 24-hour speedrun:
- 32 H200
- ~$1500 total compute budget (2$/hour/GPU)
This is very far from the early diffusion days, where training competitive models could cost millions of dollars. The goal here is to demonstrate how much the field has evolved and how far careful engineering can take you in just a single day of training.
This speedrun is not just a fun experiment. It will likely serve as the foundation for our large-scale training recipe going forward.
Alongside the results, we’re also open-sourcing our code (Github link), which contains:
- The training code used for this speedrun
- The experimental framework from the previous blog post
So you can reproduce, modify, and extend everything yourself.
The Training Recipe
Now let’s walk through what went into this 24h run.
X-prediction and Training in the Pixel Space
We use the x-prediction formulation from Back to Basics: Let Denoising Generative Models Denoise [Li and He, 2025]. As seen in Part 2, this enables training directly in pixel space and eliminates the need for a VAE altogether. We use a patch size of 32 and use a 256-dimensional bottleneck in the initial token projection layer. This design keeps the sequence length under control, making pixel-space training computationally manageable even at higher resolutions.
At 512px, the sequence length is:
$$ \left(\right. 512 / 32 \left.\right)^{2} = 256 $$
(512/32)2=256
At 1024px, the sequence length becomes:
$$ \left(\right. 1024 / 32 \left.\right)^{2} = 1024 $$
(1024/32)2=1024
Instead of following the usual 256px → 512px → 1024px schedule, we start directly at 512px and then fine-tune at 1024px.
With controlled token counts and modern hardware, pixel-space training is no longer prohibitive. It is simply a cleaner and more direct formulation.
Perceptual Losses
One very nice side effect of predicting $x_{0}$x 0 directly in pixel space is that we can reuse a whole toolbox from classical computer vision.
When your model outputs latents, perceptual supervision becomes awkward. You either have to decode back to pixels or define losses in a learned latent space that may or may not align with human perception. Once you predict pixels directly, everything becomes straightforward again. You can plug in perceptual losses exactly as they were originally designed.
We take inspiration from the paper PixelGen: Pixel Diffusion Beats Latent Diffusion with Perceptual Loss [Ma et al.], where the authors introduce additional perceptual objectives on top of the diffusion loss. They show that adding perceptual signals can noticeably improve convergence speed and final visual quality.
For this 24h run, we add two auxiliary losses:
- LPIPS ([Zhang et al.])
- A DINO-based perceptual loss (we use DINOv2 [Oquab et al.])
The idea is simple: In addition to the standard flow matching objective, we encourage the predicted clean image to match the target image in a perceptual feature space. LPIPS captures low-level perceptual similarity, while DINO features provide a stronger semantic signal.
We keep the same overall idea as the paper, but we tweaked a few details. In our experiments, we empirically found that it worked better to:
- apply the perceptual losses on pooled full images instead of patch-wise features
- apply them at all noise levels
These are small implementation details, but in our setting they consistently gave better results.
We used a weight of 0.1 for the LPIPS loss and 0.01 for the DINO perceptual loss, matching the values recommended in the original paper.
These losses are lightweight compared to the main transformer forward pass, and in our setup they add only a small overhead while providing a consistent quality boost.
Token Routing with TREAD
To make each step cheaper, we use token routing with TREAD [Krause et al., 2025]), which randomly selects a fraction of tokens and lets them bypass a contiguous chunk of transformer blocks, then re-injects them later so nothing is dropped.
We picked TREAD over SPRINT (Park et al., 2025) mostly for simplicity, and because the extra complexity of SPRINT did not feel worth the fairly small additional compute savings in our setting (sequence length 64 vs. 128 with TREAD at 512px).
Following the TREAD recipe, we route 50% of the tokens from the 2nd block to the penultimate block of the transformer.
Routed models can look worse under vanilla CFG, especially when undertrained, so we implemented a simple self-guidance scheme inspired by Guiding Token-Sparse Diffusion Models (Krause et al., 2025), which guides using a dense vs. routed conditional prediction instead of relying on an unconditional branch.
Representation Alignment with REPA and DINOv3
We used REPA [Yu et al., 2024] for representation alignment.
For the teacher, we went with DINOv3 [Siméoni et al. 2025] since it gave the best quality improvements in our previous experiments.
Concretely, we apply the alignment loss once,at the 8th transformer block with a loss weight of 0.5.
Since we combine REPA with TREAD routing, we only compute the alignment loss on the non-routed tokens, meaning the tokens that actually go through the blocks where we apply the loss. This keeps the REPA signal consistent and avoids comparing features for tokens that skipped the computation path.
Optimizer: Muon
We used the Muon optimizer, using the FSDP implementation from muon_fsdp_2, since it showed a clear improvement over Adam in our previous runs.
Muon is only applied to 2D parameters (basically matrices). Everything else (biases, norms, embeddings, etc.) is optimized with Adam, which is why the config has two parameter groups.
| Group | What it applies to | Key params we used |
|---|---|---|
| Muon | 2D parameters | lr=1e-4, momentum=0.95, nesterov=true, ns_steps=5 |
| Adam | all non-2D parameters | lr=1e-4, betas=(0.9, 0.95), eps=1e-8 |
Training Settings
We trained on three publicly available synthetic datasets:
- Flux generated (1.7M), lehduong/flux_generated
- FLUX-Reason-6M (6M), LucasFang/FLUX-Reason-6M
- midjourney-v6-llava (1M), brivangl/midjourney-v6-llava which we re-captioned with Gemini 1.5 to make prompts more consistent and cut down caption noise.
The schedule is basically: go fast at 512, then sharpen at 1024:
- 512px for 100k steps with batch size 1024
- 1024px for 20k steps with batch size 512 without REPA.
We also keep an EMA of the weights for sampling and eval:
smoothing = 0.999update_interval = 10baema_start = 0ba
Results and Closing Thoughts
Below are the evaluation curves we tracked throughout the run and a few sample grids from the final checkpoint:
For a one day training run, this is already a pretty solid place to be. The model is not flawless yet (you can still spot some texture glitches, occasional weird anatomy, and it can get a bit shaky on very hard prompts), but it is clearly usable. Prompt following is strong, the overall aesthetic is consistent, and the 1024 stage mostly does what we want: sharpen details without breaking composition.
The key takeaway is that we're very close. The remaining issues look more like undertraining artifacts and limited data diversity than signs of a structural flaw in the recipe. The failure modes are consistent with what you’d expect from a model that simply hasn’t seen enough varied data yet. With more compute and broader coverage, this exact setup should continue improving in a fairly predictable way.
Zooming out, this speed run also highlights how far diffusion training has come. By combining pixel-space training, efficient routing, representation alignment, and lightweight perceptual guidance, you can now get a meaningful model in about a day on a budget that would have sounded unrealistic not that long ago.
What’s next?
This 24h run is just a starting point, not the finish line. Next, we will keep pushing the same recipe with a bit more scale and iterate on the dataset mix and captioning.
All the code and configs behind this speedrun, as well as the full experimental framework used throughout Part 1 and Part 2, are available in the PRX repository: https://github.com/Photoroom/PRX.
While we don’t redistribute the exact training datasets used in this run, the pipeline is fully configurable and designed to be easily adapted to your own data. You can plug in different datasets, tweak individual components (TREAD, REPA, perceptual losses, Muon, etc.), and run controlled experiments with minimal friction. Our goal is to make this a practical playground for fast diffusion research, and we hope the community will use it to explore, benchmark, and iterate on these techniques in their own setups.
If you made it this far, thank you for reading. We would also love to have you join our Discord community, where we share PRX progress and results, and discuss anything diffusion and text-to-image related.
Goodbye for now, and stay tuned for the next round of experiments! 🚀
Acknowledgements.
This speedrun was inspired by several recent efforts exploring fast and low-cost training of diffusion models. If you're interested in speedrunning text-to-image models, we encourage you to check out the following works:
- Haridas, A., Shen, T., Yu, J. Nitro-T: Training a Text-to-Image Diffusion Model from Scratch in 1 Day.https://rocm.blogs.amd.com/artificial-intelligence/nitro-t-diffusion/README.html
- Bhanded, S. Speedrunning ImageNet Diffusion.https://arxiv.org/abs/2512.12386
- Sehwag, V., Kong, X., Li, J., Spranger, M., Lyu, L. Stretching Each Dollar: Diffusion Training from Scratch on a Micro-Budget.https://arxiv.org/abs/2407.15811
- Yeh, S.-Y. Home-made Diffusion Model from Scratch to Hatch.https://arxiv.org/abs/2509.06068