PRX 第三部分：在 24 小时内训练文本到图像模型
像素空间训练与 X 预测
模型采用 x-prediction 公式，能够直接在像素空间进行训练，从而无需变分自编码器（Variational Autoencoder，VAE）。为使序列长度在计算上保持可控，团队使用了 32 的 patch 大小以及在初始 token 投影层中的 256 维瓶颈。
在 512px 分辨率下，序列长度为 256 个 token；在 1024px 时，序列长度增至 1024 个 token。训练计划直接从 512px 开始，最终在 1024px 上进行微调。
感知损失用于视觉质量
直接预测像素使得可以使用经典的计算机视觉感知损失来提升收敛速度和视觉质量。团队在标准的 flow matching 目标上额外实现了两种辅助损失：
- LPIPS: 使用权重 0.1 来捕获低层次感知相似性。
- DINO-based perceptual loss: 使用 DINOv2，权重为 0.01，以提供更强的语义信号。
这些损失是对所有噪声水平下的整幅图像进行池化后计算的，而不是对 patch 级特征计算，从而获得了更好的经验效果。
使用 TREAD 的高效 Token 路由
为了降低每一步的计算成本，团队采用了 TREAD（Token Routing for Efficient Architecture-agnostic Diffusion Training）。该方法随机选择一部分 token 绕过一段连续的 transformer 块后再重新注入。
具体而言，50% 的 token 从第二个块路由到倒数第二个块。为缓解在普通 Classifier-Free Guidance（CFG）下可能出现的质量下降，团队实现了一种自我引导方案，以密集条件预测与路由条件预测进行对比引导。
表征对齐与优化器
表征对齐使用 REPA 完成，教师模型为 DINOv3。对齐损失在第 8 层 transformer 上计算一次，损失权重为 0.5。为保持与 TREAD 路由的一致性，对齐损失仅在未路由的 token 上计算。
在优化方面，团队对二维参数（矩阵）使用 Muon optimizer，对所有非二维参数（偏置、归一化、嵌入）使用 Adam。
| Optimizer Group | Target Parameters | Key Hyperparameters |
|---|---|---|
| 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 generated (1.7M images)
- FLUX-Reason-6M (6M images)
- midjourney-v6-llava (1M images)，使用 Gemini 1.5 重新生成字幕以保持一致性。
训练计划为：先在 512px 上进行 100k 步，批量大小为 1024；随后在 1024px 上进行 20k 步，批量大小为 512（不使用 REPA）。
结果与意义
虽然最终模型仍存在一些纹理瑕疵和偶发的解剖错误，但团队描述其为“明显可用”，在提示遵循和整体美感上表现强劲。作者认为剩余问题更可能源于训练不足和数据多样性受限，而非配方本身的结构缺陷。
Photoroom 已通过 PRX GitHub repository 开源了训练代码和实验框架。