单GPU从零训练2.1亿参数文生图DiT：三个关键实测发现

一位开发者用一块 RTX PRO 6000、耗时 3.5 天，从零训练了一个 2.1 亿参数的文本到图像扩散 Transformer（DiT），训练数据为 420 万张 256² 分辨率图像。他的目标不是刷 SOTA，而是完整走通训练流程，并把过程中三个「别处很少被明确说出来」的实测结论分享出来。相比于展示生成样本，这些工程细节对想理解扩散模型训练机制的人更有价值。

学习到的 null 注意力槽成了「注意力汇聚点」

第一个发现关乎模型内部的注意力分布。作者借鉴了 register token 的思路，在图像流中加入 16 个 register token，并在每个交叉注意力层追加 2 个可学习的 key/value 槽（learned null slots）。

实测结果颇为反直觉：在中等噪声水平、模型中间层的位置，这 2 个可学习槽吸收了约 90% 的交叉注意力质量。而在传统交叉注意力模型中通常充当「汇聚点（sink）」的 EOS token，占比骤降到约 4%；真正的内容词（content words）各自只保留几个百分点，但注意力锐利地聚焦在对应物体上。

更值得关注的是范数变化：到模型中间层时，register 向量的范数增长到图像 token 的 4–13 倍。这说明这些人为引入的「垃圾桶」token 在训练中自发承担了吸收冗余注意力的功能，让内容词的注意力更加干净、聚焦。这为 register token 机制在扩散 Transformer 中的作用提供了一个可量化的证据。

Register token 的概念最初由 Darcet 等人在 Vision Transformer 研究中提出，用于解决 ViT 中部分 patch token 出现异常高激活值（artifact）的问题。其核心思路是在输入序列中插入若干不对应任何实际图像区域的「寄存器」token，让模型在计算全局注意力时有专属的「草稿纸」来承接那些与局部内容无关的全局信息。在语言模型中有类似作用的是注意力汇聚（attention sink）现象——序列开头的 token（如 BOS）往往会吸收大量注意力权重，即便它本身并不携带重要语义，这一机制被认为有助于维持 softmax 注意力的数值稳定性。本文发现的现象则是两者的结合：在交叉注意力中，人为加入的可学习 null 槽完全取代了原本的 EOS sink 角色，且吸收比例高达 90%，说明模型主动将「无处安放的注意力」集中分配给这些专用槽，从而使内容词的注意力信号更加纯净。

Flow-matching 损失是「健康信号」而非「质量信号」

第二个结论直击训练监控中的常见误区：不要用训练损失去判断生成质量。

整个训练过程中，flow-matching 损失仅从 0.805 缓慢下降到 0.754。如果只盯着这个数字，很容易误以为模型几乎没有进步。但真正的质量指标却大幅改善：

• 保留集 FID 从 33.7 降到 27.0
• FD-DINOv2 从 570 降到 218
• 基于检测器的物体准确率从 65% 提升到 90%

作者解释，高噪声阶段的大部分损失来自速度目标（velocity target）本身的不可约方差（irreducible variance），这部分噪声无论如何训练都无法消除。一个佐证是：训练损失与保留集损失在长达 24 个 epoch 里保持到小数点后第三位都相等——说明模型没有过拟合，损失曲线的平坦更多反映的是任务的内在噪声，而非训练停滞。这提醒实践者：评估扩散模型必须依赖 FID、FD-DINOv2、物体准确率等下游指标，而非单纯看 loss。

Flow matching 是近年扩散生成模型的主流训练目标之一，其核心思想是学习一个将噪声分布「流动」到数据分布的速度场（velocity field），而非像 DDPM 那样预测噪声本身。在 rectified flow 框架下，训练目标是让模型预测从纯噪声到真实图像的直线插值速度。这一目标的内在挑战在于：对于任意给定的中间时间步，速度目标本身具有不可约方差——同一噪声输入在条件概率分布下对应多种合理的真实图像，模型只能学到期望值，无法消除这部分固有的不确定性。这与分类任务的交叉熵损失有本质区别；后者在模型完美拟合时可以趋近于零，而 flow matching 损失在模型完全收敛后仍会保留一个由数据多样性决定的下界。FID（Fréchet Inception Distance）和 FD-DINOv2 则通过比较生成图像与真实图像的特征分布差异来评估感知质量，能捕捉到 loss 无法体现的结构与语义改善。

训练期时间步 shift 比翻倍采样步数更值

第三个发现关于推理效率与采样策略。作者在 2456 个保留 prompt 上用最终权重做了对比：

配置	FID
20 步 + shift 2.8	27.0
50 步 + shift 2.8	26.6
8 步 + shift 2.8	28.4
20 步 + 无 shift	27.3（FD-DINOv2 从 218 升到 228）

关键结论是：把采样步数从 20 提到 50，FID 只改善 0.4；但去掉 shift，质量反而明显退化。也就是说，恰当的时间步 shift 带来的收益，超过把采样步数翻倍还多。

这里的 shift 值 2.8 并非拍脑袋得来，而是遵循 SD3/RAE 的规则 √(32·32·32/4096) 计算而来，对应 FLUX.2 使用的 32 通道潜空间。这一细节说明，扩散采样的调度参数应当随潜空间维度理论性地推导，而不是靠盲目堆叠推理步数换质量。

时间步 shift（timestep shift）是针对 flow matching 采样调度的一种校正手段。在标准 rectified flow 中，推理时各时间步均匀分布在 [0, 1] 区间，但研究发现这种均匀分布并不最优——高噪声阶段（接近 t=1）对最终图像结构的影响远大于低噪声阶段（接近 t=0），因此将更多采样步数分配到高噪声区间通常能显著提升质量。Shift 参数通过对均匀时间步施加一个偏移变换来实现这种重分配，shift 值越大，采样越集中在高噪声端。SD3 和 FLUX 等模型的实践表明，shift 值应与潜空间的通道数和空间分辨率挂钩，公式 √(C·H·W / token_count) 提供了一种理论推导路径，避免了纯凭经验调参的随意性。这一结论的工程含义是：在推理预算有限时，优先校准调度参数而非简单堆叠步数，是更高性价比的优化方向。

训练配置一览

为了让结果可复现，作者给出了完整的技术栈：

• 架构：交叉注意力 DiT（896 宽 × 16 层），2D RoPE、QK-norm、SwiGLU、adaLN-single
• 目标：rectified flow，logit-normal 时间步采样 + 上述 shift；配合余弦速度损失和 dispersive 辅助损失
• 多分辨率：从第一步起就用 5 个约 256 token 的宽高比 bucket
• 文本编码：冻结的 flan-t5-base；每张图长/短/空 caption 按 50/40/10 采样
• 数据构成：Pexels 2.8M（60%）、经质量筛选的 FLUX-Reason-6M 1.2M 切片（25%）、带 GPT-4V caption 的 COCO（15%）
• 训练：batch 256、40 万步、EMA 0.9999、最后四分之一线性衰减学习率、torch.compile（相比 eager 提速 2.4×）

作者已开源全部代码、权重、写作说明与在线 demo（tinydit 项目），并在文末抛出下一阶段的问题：准备用 Flow-GRPO 在此基础上做强化学习，读者会优先选择 PickScore/HPSv2、基于检测器的物体奖励，还是像「计数」这样可验证的奖励作为起点？

对独立研究者的启示

这个项目最有价值的地方，不在于 2.1 亿参数模型本身的能力（毕竟规模有限），而在于它把从零训练文生图模型的每一个决策都摊开、量化并给出出处。对预算有限的独立研究者和爱好者而言，它证明了单卡数天就能跑通完整的扩散 Transformer 训练管线，并沉淀出可迁移的经验：注意力汇聚点会被可学习槽接管、loss 不等于质量、调度 shift 优于暴力增步。这类「诚实的负结果与中间态测量」往往比华丽的样本图更能推动社区理解模型的实际工作机制。