论文

环形扩散Transformer

Looped Diffusion Transformer

模型架构递归架构

摘要

传统上,改进文本到图像模型依赖于增加模型大小或去噪步骤的数量。在这项工作中,我们探索了一种扩展计算的替代方法,通过在每个去噪步骤中重复运行共享 Transformer 块,有效增加计算深度,同时保持参数计数固定。这种循环计算可以在没有显式推理标记的情况下迭代细化内部表示。然而,简单的循环无法持续提高图像质量。我们将这个问题追溯到中间循环的弱监督和不受监管的注意力更新,这些更新逐渐侵蚀了本地信息。为了克服这些挑战,我们提出了循环扩散Transformer(Looped-DiT),它将跨中间循环的深度监督与自调节注意力相结合,以稳定循环特征更新。在匹配参数和匹配计算设置下,Looped-DiT 始终优于非循环基线。值得注意的是,2.6 亿参数的循环模型在多个文本到图像基准测试中可以超过 6.5 倍大的模型,同时需要的推理计算量降低 4.9 倍。除了这种性能增益之外,我们发现循环计算可以为扩散模型提供更有效的迭代计算形式,与在固定推理预算下添加更多去噪步骤相比,增加循环深度会产生更大的增益。此外,更深的循环可以逐步纠正早期循环中所犯的错误,表现出暗示潜在推理的行为。总之,这些结果表明循环计算为扩展视觉生成模型提供了一种有前途的方法。

环形扩散Transformer:图1:循环计算以较少参数扩展文生图模型。(a)增加循环深度显著改善文生图推理表现。(b)这些收益使用更少参数和更低推理成本获得。(c)循环之间出现自我纠错。
图1:循环计算以较少参数扩展文生图模型。(a)增加循环深度显著改善文生图推理表现。(b)这些收益使用更少参数和更低推理成本获得。(c)循环之间出现自我纠错。