PRX Part 3 – Training a Text-to-Image Model in 24h
2 months ago
- #text-to-image
- #diffusion-models
- #machine-learning
- 在32块H200 GPU上以1500美元预算进行24小时极限训练，融合了扩散模型的架构技巧与训练策略
- 采用x-prediction像素空间训练方案（无需VAE），设置32的补丁尺寸与256维瓶颈层
- 引入感知损失（LPIPS+DINO）加速收敛并提升画质，在所有噪声级别对完整图像进行池化处理
- 通过TREAD令牌路由机制，从第2到倒数第2个transformer块跳过50%令牌计算以降低开销
- 使用REPA表示对齐技术，以DINOv3作为教师模型，在第8个transformer块施加权重0.5的对齐损失
- 二维参数采用Muon优化器，其余参数使用Adam，并为不同参数组定制学习率与动量配置
- 训练数据采用合成数据集（Flux生成/FLUX-Reason-6M/midjourney-v6-llava），512px训练10万步后接1024px训练2万步
- 效果呈现优秀的提示跟随与美学一致性，仅存在因训练不足导致的纹理瑕疵和结构错误等轻微缺陷
- 已在GitHub开源训练代码与实验框架供社区使用适配