Looped Diffusion Transformer
Yong Xien Chng, Tianyi Chen, Wenwen Tong, Haiwen Diao, Zhongang Cai, Lei Yang, Ziwei Liu, Lewei Lu, Dahua Lin, Gao Huang
cs.CV, cs.LG
2026-10-01
在每个去噪步内把共享 Transformer 块循环跑 N 次,配合深监督和自调制注意力,260M 参数文生图模型平均分超过 1.7B 模型,推理算力还省 4.9 倍。
文生图模型的常规扩法有两条:堆参数(Qwen-Image、FLUX.2 已经到了数十亿量级),或者加去噪步数。两条都直接推高部署成本。还有第三条路来自语言模型那边:looped Transformer,把同一组中间块反复套在隐藏状态上,算力深度上去了,参数和序列长度都不动。
这件事对图像生成有额外吸引力。生成一张连贯的图不只是照抄 prompt,模型要推断隐含的视觉内容、解一组相互牵连的约束、并在生成过程中发现并修正不一致。这些过程天然是视觉的,在隐藏表示里迭代做,不需要显式的文字推理链。
但 naive 的循环不work。在 MiniT2I 这个极简 MMDiT 上直接加循环,浅层循环时性能可能低于不循环的基线,加循环先饱和再下降,超过训练循环数后继续崩。作者用一个 ridge 回归探针去测每个循环深度下隐藏状态还能不能线性解码出图像 token 的二维位置:R² 从第 1 轮后的 0.865 掉到第 8 轮后的 0.562。反复施加同一个变换,注意力更新会逐渐覆写、削弱图像 token 里的局部信息。
诊断出两个病根:中间循环没有直接监督信号(梯度要从最后一轮绕回来,路径太长),以及注意力更新强度不随循环调节。Looped-DiT 对症下两味药。
把 17 个 MMDiT 块切成三段:前段 A 跑一次,中间 5 块 B 循环跑 N 次(训练时 N=4),后段 C 跑一次。B 的参数在循环间共享,所以加循环只加计算深度,不加参数。循环发生在每个去噪步内部,和采样步数解耦,可以独立调。
两个关键组件:
消融显示 XSA 比 Gated 更强,并且有直接证据支撑机理:加 XSA 后注意力更新范数相对残差流最小、位置可解码性保持最高,没有 XSA 时后几轮会继续改一张已经正确的图并多画一个物体。XSA 对循环模型的增益(+1.6)也远大于对不循环模型的(参数匹配 −0.1、算力匹配 +0.7),说明调制专门服务于重复更新这个场景。
深监督还带来一个副产品:模型对推理循环数不敏感。训练在 4 轮,推理时 1 到 8 轮都稳定,单轮设置都超过不循环基线。这让自适应循环成为可能:一个轻量门控网络(4 个可学习 query 的 cross-attention 加 MLP 头)预测继续循环的期望收益,阈值 λ 一调就能在平均循环数 1.25 到 2.70 之间滑动,不用重训,而且在相同平均循环数下优于固定循环。
六个 benchmark,平均分对比:
| 模型 | 参数量 | GenEval | DPG | PRISM | CoRe | Spatial | TIIF-Short | 平均 |
| MiniT2I-L/16 | 1.6B | 90.3 | 83.4 | 58.5 | 45.0 | 51.6 | 78.2 | 67.8 |
| InternVL-U | 1.7B | 85.0 | 85.2 | 63.5 | 48.2 | 54.5 | 77.7 | 69.0 |
| Uni-CoT(CoT 类最佳) | 6.6B | 81.2 | 84.1 | 58.8 | 45.8 | 53.8 | 76.8 | 66.8 |
| Looped-DiT B/16 | 0.26B | 87.4 | 87.0 | 67.0 | 53.5 | 54.6 | 79.7 | 71.5 |
对次优的 InternVL-U,平均分高 2.5,参数少 6.5 倍,推理算力低 4.9 倍,延迟低 6.5 倍。六个基准里五个第一,GenEval 输给 MiniT2I-L/16。
收益不能被三种替代方案解释。参数和算力都匹配的对照(B/32):比参数匹配基线高 4.0 分;训练算力也匹配的 deeper 基线(加了深监督)58.1 vs 59.1。同样推理预算下,把算力分给循环深度比分给更多去噪步涨得更多(25 到 50 步设置下一致)。和文字 CoT 重写 prompt 对比:循环在约束满足类子任务上增益大(PRISM/Long Text +9.1),CoT 在推断隐含内容类子任务上增益大(CoRe/Generalization +22.2),两者结合总分最高。
循环数从 1 加到 4,平均分单调上升。定性例子显示逐轮补齐缺失内容、重排物体满足空间约束、修正渲染错误,而更大的非循环模型和显式 CoT 基线在同样 prompt 下留着这些错误。
这是第一个对文生图做循环计算的系统性研究,给出的证据链比较完整:收益在参数、推理算力、训练算力三种匹配设置下都成立,也排除了加深、加宽、加去噪步、加 CoT 这些解释。对从业者,循环多了一个和模型规模、去噪步数正交的推理时扩展旋钮,而且训练一次就能通过循环数在延迟和质量之间滑动。
XSA 是个零参数改动,单独搬走用在别的重复计算场景也值得试。自适应循环那个「不重训调算力」的性质,对需要按请求分级的服务有直接用处。
也要诚实:这是一个基于 MiniT2I 的受控研究,规模在 260M,结论能不能搬到工业级数十亿参数模型上,论文没有验证。
论文自己承认的:naive 循环不可靠,有效循环需要中间监督和自适应调节,这不是免费的午餐;证据是「suggestive of latent reasoning」,作者自己也用了「行为提示」这样的措辞,没有给出循环内部到底在算什么的机制性解释。
读下来的几点疑问。第一,主对比里 Looped-DiT 用的是自己训的 MiniT2I 系基线,和 InternVL-U、CogView4 这些模型的训练数据、数据量都不同,6.5 倍参数差距的对比不是纯粹的架构对照,更像「受控设置内的结论 + 跨模型参考点」。第二,主结果固定 4 循环,对更长循环(8 轮以上)只有 B/32 的探针分析,主模型在大循环数下是否还有收益没有展示。第三,「循环胜过去噪步」的结论在 25 到 50 步区间成立,更少步数(蒸馏到一步的 regime)下未测。第四,结论全部基于 MiniT2I 这一个 backbone,循环和别的架构(比如更深的工业模型)的交互没有碰。