Hugging Face 分享如何在有限时间内训练可用的文生图模型,覆盖数据处理、模型选择、优化策略等核心步骤。
在前两篇文章(第1部分和第2部分)中,我们探索了扩散模型的各种架构和训练技巧。我们尝试单独评估每个想法,测量吞吐量、收敛速度和最终图像质量,并试图理解什么真正推动了进展。
在这篇文章中,我们想回答一个更实际的问题:
我们不再一次优化一个维度,而是将最有希望的组成部分堆叠在一起,看看在严格的计算预算下我们能把性能推到多远。
为了更具体,我们在做一个24小时的速跑:
~1500美元的总计算预算(2美元/小时/GPU)
这与早期扩散模型的日子相差甚远,当时训练竞争力的模型可能需要花费数百万美元。这里的目标是演示该领域进化了多少,以及仔细的工程在仅仅一天的训练中能为您带来多远。
这个速跑不仅仅是一个有趣的实验。它很可能会成为我们今后大规模训练方案的基础。
除了结果,我们还开源了我们的代码(Github链接),其中包含:
所以你可以自己复制、修改和扩展所有内容。
现在让我们看看这24小时的运行涉及了什么。
我们使用来自《回到基础:让去噪生成模型去噪》[Li and He, 2025]的x预测公式。如第2部分所示,这能够直接在像素空间中训练,完全消除了对VAE的需求。我们使用32的patch大小,并在初始token投影层中使用256维的瓶颈。这种设计将序列长度保持在控制范围内,即使在更高的分辨率下也能使像素空间训练在计算上易于管理。
在512px处,序列长度为:
(512/32)^2 = 256
在1024px处,序列长度变为:
(1024/32)^2 = 1024
我们没有遵循通常的256px → 512px → 1024px计划,而是直接从512px开始,然后在1024px处进行微调。
有了受控的token计数和现代硬件,像素空间训练不再是禁止的。它只是一个更干净、更直接的表述。
在像素空间中直接预测x0的一个很好的副作用是,我们可以重复使用来自经典计算机视觉的整个工具箱。
当您的模型输出潜在变量时,感知监督变得尴尬。您要么必须解码回像素,要么在可能与也可能不与人类感知相一致的学习潜在空间中定义损失。一旦您直接预测像素,一切就变得简单明了了。您可以完全按照其最初设计的方式插入感知损失。
我们从论文PixelGen获得灵感:像素扩散用感知损失战胜潜在扩散[Ma et al.],其中作者在扩散损失的基础上引入了额外的感知目标。他们表明添加感知信号可以明显改善收敛速度和最终的视觉质量。
对于这个24小时的运行,我们添加了两个辅助损失:
这个想法很简单:除了标准的流匹配目标,我们鼓励预测的干净图像在感知特征空间中与目标图像匹配。LPIPS捕捉低级感知相似性,而DINO特征提供更强的语义信号。
我们保持与论文相同的总体思路,但我们调整了一些细节。在我们的实验中,我们通过实证发现最好是:
这些是小的实现细节,但在我们的设置中,它们始终给出更好的结果。
我们为LPIPS损失使用0.1的权重,为DINO感知损失使用0.01,与原始论文中推荐的值匹配。
这些损失与主transformer前向传递相比是轻量级的,在我们的设置中,它们只增加少量开销,同时提供一致的质量提升。
为了使每一步更便宜,我们使用TREAD的token路由[Krause et al., 2025]),它随机选择一小部分token,让它们绕过transformer块的连续块,然后稍后重新注入它们,所以没有任何东西被丢弃。
我们选择TREAD而不是SPRINT (Park et al., 2025)主要是为了简单起见,以及因为在我们的设置中SPRINT的额外复杂性不值得相当小的额外计算节省(序列长度64 vs. 128,在512px处为TREAD)。
按照TREAD配方,我们将50%的token从第2个块路由到transformer的倒数第二个块。
路由模型在vanilla CFG下可能看起来更差,特别是当训练不足时,所以我们实现了一个简单的自引导方案,灵感来自指导token稀疏扩散模型(Krause et al., 2025),它使用密集与路由条件预测进行指导,而不是依赖无条件分支。
我们使用REPA [Yu et al., 2024]进行表示对齐。
对于教师模型,我们选择了DINOv3 [Siméoni et al. 2025],因为它在我们之前的实验中给出了最佳的质量改进。
具体来说,我们在第8个transformer块应用一次对齐损失,损失权重为0.5。
由于我们将REPA与TREAD路由结合,我们只在非路由token上计算对齐损失,这意味着实际通过我们应用损失的块的token。这保持REPA信号的一致性,并避免比较跳过计算路径的token的特征。
我们使用Muon优化器,使用muon_fsdp_2中的FSDP实现,因为它在我们之前的运行中相比Adam显示出明显的改进。
Muon只应用于2D参数(基本上是矩阵)。所有其他内容(偏差、范数、嵌入等)都用Adam优化,这就是为什么配置有两个参数组。
我们在三个公开可用的合成数据集上进行了训练:
计划基本上是:在512处快速进行,然后在1024处锐化:
我们还为采样和评估保持权重的EMA:
update_interval = 10ba
以下是我们在整个运行过程中跟踪的评估曲线和最终检查点的一些样本网格:
对于一天的训练运行,这已经是一个相当扎实的地方了。该模型还不是完美的(您仍然可以看到一些纹理故障、偶尔的奇怪解剖学结构,并且对于非常困难的提示可能会有点摇晃),但它显然是可用的。提示跟随能力强,整体美感一致,1024阶段基本上做了我们想要的:锐化细节而不破坏构成。
关键的收获是我们非常接近。剩下的问题看起来更像是欠训练的伪影和有限的数据多样性,而不是配方中结构缺陷的迹象。失败模式与您对只是还没有看到足够多样化数据的模型的期望相一致。有了更多的计算和更广泛的覆盖,这个确切的设置应该以相当可预测的方式继续改进。
从宏观角度看,这个速跑也突出了扩散训练已经进展到了多远。通过结合像素空间训练、高效路由、表示对齐和轻量级感知指导,您现在可以在不久前听起来不切实际的预算内在大约一天内获得一个有意义的模型。
这个24小时的运行只是一个起点,不是终点。接下来,我们将继续使用稍微更多的规模推动相同的配方,并迭代数据集混合和字幕。
这个速跑背后的所有代码和配置,以及整个第1部分和第2部分中使用的完整实验框架,都可以在PRX存储库中获得:https://github.com/Photoroom/PRX。
虽然我们不重新分配此运行中使用的确切训练数据集,但管道是完全可配置的,设计用于轻松适应您自己的数据。您可以插入不同的数据集、调整单个组件(TREAD、REPA、感知损失、Muon等),并以最小的摩擦运行受控实验。我们的目标是使其成为快速扩散研究的实用游乐场,我们希望社区将使用它在自己的设置中探索、基准测试和迭代这些技术。
如果您读到这里,感谢您的阅读。我们也很希望您加入我们的Discord社区,我们在那里分享PRX的进展和结果,讨论任何与扩散和文本到图像相关的内容。
暂时再见,敬请期待下一轮实验!🚀
这个速跑受到了最近几项探索快速和低成本训练扩散模型的工作的启发。如果您对文本到图像模型的速跑感兴趣,我们鼓励您查看以下作品:
提问1: 24小时非常令人印象深刻,但看到32个H200使其更有意义xD。你使用Muon太酷了!并且,你计划发布重新字幕的brivangl/midjourney-v6-llava吗?
回复: 是的,在消费者硬件上肯定不可能...还没有:D这也是我们添加$1500标签的原因。
关于发布字幕,我们还没有真正考虑过。如果您认为这对社区有价值,我们也可以努力发布它们。
提问2: 嗨,这篇文章很鼓舞人心。我想听听您对感知损失的想法。在(原始)图像恢复的研究实习中——并尝试探索感知/扭曲权衡——我做了几次VGG和Dino损失的尝试。我总是最终得到棋盘图案(无论vgg还是dino)。如果您查看Johnson的原始论文,它充满了这些伪影。您是否面临过任何此类问题,或者可以分享有关该主题的有用见解?
非常感谢。Balthazar
回复: 我认为我们并没有真正观察到这样的伪影,或者至少不是由于LPIPS。我们在使用JIT在欠训练模型上时确实看到了没有LPIPS的类似情况,但我认为我们不再观察到它们。可能只是因为模型还没有收敛。
我们开始观察棋盘图案的一种情况是在这些模型之一上使用GRPO或SRPO进行奖励对齐时。在这种情况下,伪影更加明显,但我认为它们与LPIPS无关,因为我们当时没有使用它。
您使用的是什么类型的模型?