LoopDiT: Loop Transformers for Diffusion Models

Three-segment LoopDiT三段式 LoopDiT

Four numbers describe a model: p, s, and e are block counts, l is the loop count. p2 s2 l4 e2 means 2 prefix blocks, 2 shared blocks looped 4 times, and 2 suffix blocks: effective depth 2 + 2×4 + 2 = 12 blocks, six of them carrying their own weights. Because the letters alone are easy to misread, each segment carries an underline color: prefixsharedloops (repeats of the shared segment, dashed)suffix. Chart labels, legends, and the text use the same rendering.

全文用四个数字描述一个模型：p、s、e 是 block 数，l 是循环次数。p2 s2 l4 e2 即 2 个 prefix block，2 个 shared block 循环 4 轮，2 个 suffix block；有效深度 2 + 2×4 + 2 = 12 个 block，其中 6 个拥有独立权重。字母本身易混淆，因此每段以下划线颜色标识：prefixsharedloops（shared 的重复轮数，虚线）suffix。图表里的行标签、legend 和正文都按同一套渲染。

Summary

All numbers are FID on ImageNet 256 with 50k samples and CFG 2.0, after 240 epochs unless stated. Our DiT-B/4 baseline is 11.91 (11.53 on a repeat); the MeanFlow reference implementation reports 11.35 for the same setting, so read gaps under 0.6 as noise. The mainline configuration is p5 s2 e5; other configurations are ablations, and each figure states its training budget on the source row.

所有结果为 ImageNet 256、50k 样本、CFG 2.0 的 FID；未标注时为 240 epoch。我们的 DiT-B/4 基线为 11.91（重复一次 11.53），MeanFlow 官方实现在同样设置下报告 11.35，因此 0.6 以内的差距视为噪声。主配置为 p5 s2 e5，其余为消融实验；每张图的训练预算标注在来源行。

1. The penalty for sharing shrinks as the shared segment grows.At matched compute the dense twin always wins, but its margin collapses when the shared segment goes from six blocks to twelve. Compute is never saved. 共享段越大，参数共享的代价越小。等计算量下稠密对照始终领先，但共享段从 6 个块增至 12 个块时，领先幅度大幅收窄。计算量则始终没有节省。
2. Independent blocks after the shared segment contribute more than blocks before it, and a fixed budget of them is best split evenly.Suffix blocks refine the representation the loop has built; prefix blocks only prepare its input. 共享段之后的独立块贡献大于之前的独立块；独立块总数固定时均分最优。suffix 块在精炼循环已构建的表示，prefix 块只负责准备循环的输入。
3. Conditioning the shared blocks on the loop index through adaLN lowers FID by 1 to 2 at negligible parameter cost.Without it the shared blocks apply the same function on every loop; with it one set of weights acts as L slightly different functions. 通过 adaLN 将循环索引注入共享块，以可忽略的参数开销降低 1～2 FID。没有它，共享块每轮计算同一个函数；有了它，同一组权重在 L 轮中表现为 L 个不同的函数。
4. Returns from additional loops diminish and reverse beyond about ten.Two shared blocks support only so many distinct functions; beyond about ten iterations the extra passes add no information and begin to interfere. The curve is from an 80-epoch budget and a full-budget recheck is in progress. 循环轮数的边际收益递减，超过约十轮后转为负。两个共享块能承担的功能有限；超过十轮后额外的迭代不再增加信息，反而相互干扰。该曲线来自 80 epoch 预算，满预算复核正在进行。
5. Under the same epoch budget, training with four loops first and eight later beat every fixed loop count; the comparison is not compute-matched.Few loops keep early training cheap; depth is needed only late, and the reverse order confirms it. 相同 epoch 预算下，先 4 轮后 8 轮的两阶段训练优于所有固定轮数；该比较未对齐计算量。少轮数使早期训练开销低，深度只在后期需要；相反顺序的结果印证了这一点。
6. With mixture-of-experts in the shared blocks, FID tracks the number of active experts per token and is flat in the size of the expert pool.Four and sixteen experts perform identically; top-1 versus top-2 routing accounts for the whole difference. A dense MLP of the same active width does not reproduce it. 共享块采用混合专家后，FID 随每 token 的激活专家数变化，对专家池大小不敏感。4 个与 16 个专家结果相同；top-1 与 top-2 之间是全部差异。等激活宽度的稠密 MLP 无法复现这一收益。
7. Input injection, noise initialization, and stochastic loop counts, all effective in looped language models, degrade FID here.Each was designed to keep a language model's recurrence from collapsing; a diffusion loop of a few iterations may not need that protection, and the extra signals only interfere. 在循环语言模型中有效的 input injection、噪声初始化和随机循环轮数，在此均使 FID 变差。三者都用于防止语言模型的循环状态塌缩到不动点；几轮的扩散循环可能不需要这类保护，额外信号只构成干扰。
8. Under a fixed inference budget, additional sampling steps are more effective than additional loops.Multi-step sampling wins outright at small budgets; a looped model catches up only at 20 forward passes. 固定推理预算下，增加去噪步比增加循环轮数更有效。小预算下多步采样直接胜出；循环模型要到 20 次前向才追平。

The penalty for sharing shrinks as the shared segment grows共享段越大，参数共享的代价越小

All three figures are 240-epoch models. The two comparisons answer different questions. Matched on depth, a looped model has fewer distinct functions than the dense one: twelve dense blocks are twelve functions, while six blocks looped twice are six, and one block looped twelve times is a single map iterated to a fixed point. Every split loses to DiT-B/4, and the loss grows as the shared segment shrinks. Matched on parameters, the same twelve blocks run twice and reach a lower FID, but they also spend twice the compute.

The third comparison, matched on compute, was run afterwards: a dense model of independent blocks against the same shape looped twice at half the parameters. Run at two scales it gives the figure below. The dense twin wins both times, so sharing never buys compute, but the margin it wins by is not a constant: with six blocks shared it is large, and with twelve it is close to the noise between two seeds. The more distinct blocks a loop has to work with, the less the missing parameters seem to cost.

Read within a row. The lower row was measured after the original storage was retired and the pipeline rebuilt, which moved this looped configuration by more than the gap the row is being used to measure, and its two runs used different seeds.

三张图均为 240 epoch 的模型。两组比较回答不同的问题。对齐有效深度时，循环模型的独立函数少于稠密模型：12 个稠密块是 12 个函数，6 个块循环 2 轮是 6 个，1 个块循环 12 轮则是同一映射的迭代。每种配置都落后于 DiT-B/4，共享段越小落后越多。对齐参数量时，同样 12 个块执行两遍得到更低的 FID，但 FLOPs 翻倍。

第三组比较即对齐计算量，已于其后完成：独立块的稠密模型对同样形状、参数量减半的循环模型。在两个规模上各做一次，即下图。稠密对照两次都赢，因此循环从未换来计算量；但它领先的幅度并非常数：共享 6 个块时差距明显，共享 12 个块时已接近两个种子之间的噪声。循环能支配的独立块越多，缺失的参数量看起来代价越小。

应按行内比较，不要跨行比较。下面一行是在原存储退役、管线重建之后测得的，重建对该循环配置造成的偏移大于这一行要测量的差距，且两次运行使用了不同的种子。

Independent blocks are worth more after the loop than before it独立块放在共享段之后优于之前

Conditioning on the loop index helps at negligible cost以可忽略的开销将循环索引注入共享块

Every configuration from here on uses loop embedding, so it comes first. Without it the shared block sees the same condition vector on every loop and cannot tell which loop it is in. Loop embedding writes the loop index into the adaLN condition. Each block reads c = t_emb + h_emb + y_emb, where t is the MeanFlow time, h = t − r is the interval of the average velocity, and y is the class. Loop k (k = 0 … L−1) adds a scalar t_loop = 1 − k/(L−1), mapped by a TimestepEmbedder of the same form as the time embedder (256-dimensional sinusoidal features followed by a two-layer MLP) to a hidden-size vector and added to c. The shared weights do not change; only the modulation differs per loop. Prefix and suffix blocks still read the original c. The embedder adds about 0.8M parameters, under 1% of the model. The direction matches denoising: 1.0 on the first loop, 0.0 on the last, and a fixed 1.0 when L = 1.

后续配置均使用 loop embedding，故先说明。没有它时，共享块在每轮接收完全相同的条件向量，无法区分当前处于第几轮。Loop embedding 将循环索引编码进 adaLN 的条件向量：每个 block 原本读取 c = t_emb + h_emb + y_emb，其中 t 是 MeanFlow 的时间，h = t − r 是平均速度的区间长度，y 是类别。第 k 轮（k = 0…L−1）另给一个标量 t_loop = 1 − k/(L−1)，经与时间步同款的 TimestepEmbedder（256 维正弦编码接两层 MLP）映射为 hidden 维向量，加到 c 上得到 c_loop；共享块的权重不变，仅每轮的调制不同。prefix 和 suffix 仍读取原始 c。l_embedder 约 0.8M 参数，不到模型的 1%。方向与去噪一致：第一轮 1.0，最后一轮 0.0，L = 1 时固定为 1.0。

for k in range(n_loops):
    t_loop = 1.0 - k / max(n_loops - 1, 1)          # 1 → 0
    c_loop = c + l_embedder(jnp.full((B,), t_loop))  # same form as t_embedder
    for block in shared_blocks:
        x = block(x, c_loop)                        # adaLN reads c_loop

The same design fixes how the loop count changes at inference. Training with L loops and evaluating with L′ < L keeps the training step 1/(L−1) and takes the first L′ values instead of respacing [1, 0]: trained with 6 and evaluated with 4, t_loop runs 1.0, 0.8, 0.6, 0.4. Fewer loops then means an early exit rather than a spacing the block never saw.

这一设计同时决定了推理时如何改变循环轮数。训练 L 轮、评估 L′ < L 轮时，不把 [1, 0] 重新均分，而是保持训练时的步长 1/(L−1)，只取前 L′ 个值：训练 6 轮、评估 4 轮，t_loop 依次为 1.0、0.8、0.6、0.4。这样减少轮数等价于提前退出，而不是让共享块面对训练时未见过的条件间隔。

Returns from more loops diminish, then reverse循环轮数的收益递减并转负

Hold p5 s2 e5 fixed, keep loop embedding on, and vary only how many times the two shared blocks run. Quality improves up to about ten loops and then slides back. With the loop index as the only thing that changes between passes, two blocks can be asked to play only so many distinct roles; beyond that, extra passes revisit the same computation and the gradient through a long chain of shared weights gets harder to use within 80 epochs. Whether the ceiling is capacity or optimization is not separated here.

A later run at the full 240-epoch budget indicates that much of this curve is budget-limited rather than a property of the loop: the same configuration at ten loops improves by more than the entire spread of the curve above. The full-budget run at sixteen loops, which decides whether the reversal survives a longer schedule, has not finished.

固定 p5 s2 e5 并保留 loop embedding，只改变两个共享块的循环轮数。FID 在约十轮前持续改善，之后回落。每轮之间唯一变化的是循环索引，两个块能承担的功能有限；超过之后，额外迭代只是重复同样的计算，穿过长链共享权重的梯度在 80 epoch 内也更难优化。上限来自容量还是优化，此处未加区分。

后续的满预算运行表明这条曲线很大程度上受预算限制，而非循环本身的性质：同一个十轮配置的提升幅度超过了上图整条曲线的跨度。决定这一反转能否在更长训练下成立的十六轮满预算运行尚未结束。

The same rise appears without any independent blocks, at both two and one shared block, so it is a property of the loop rather than of this particular split.

无独立块时同样存在这一趋势，一个和两个共享块均如此，说明这是循环本身的性质而非该配置的特例。

Four loops first and eight later beat every fixed loop count at equal epochs先 4 轮后 8 轮：等 epoch 下更优

If few loops train fast and many loops train well, use each where it helps: p5 s2 e5 spends the first third of training with four loops and the rest with eight. Under the same epoch budget this beats either count held fixed, and the reverse order loses, so the gain comes from having depth late in training rather than from the switch itself. The four-loop phase also runs fewer blocks per step, so the schedule used less compute than a fixed eight; a compute-matched run is pending.

少轮数训练快、多轮数效果好，因此分阶段使用：p5 s2 e5 前三分之一训练用 4 轮，其余用 8 轮。相同 epoch 预算下优于任一固定轮数，相反顺序则更差，说明收益来自训练后期的深度而非切换本身。4 轮阶段每步执行的块更少，因此该方案的总计算量低于固定 8 轮；计算量对齐的对照待补。

MoE: FID follows the number of active experts混合专家：FID 取决于激活专家数

Activating two experts instead of one helps at both training lengths, while widening the pool from four to sixteen changes nothing. Whatever the router learns, the benefit tracks how much MLP each token gets. The natural control, a dense MLP at double width and therefore the active parameter count of two ratio-4 experts, was trained afterwards and is the second row of the figure. Widening alone does not reproduce the routing gain. It also lands behind the ratio-4 dense MLP, which is not how width usually behaves, so treat the control as evidence that the two are not interchangeable rather than as a measurement of what routing is worth.

top-2 路由在两种训练预算下均优于 top-1；专家池从 4 个扩至 16 个则无变化。无论 router 学到什么，收益取决于每 token 分配到的 MLP 计算量。自然的对照是 MLP 宽度翻倍的稠密模型，其激活参数量与两个 ratio-4 专家相当；该对照后续训练完成，即图中第二行。单纯加宽并不能复现路由带来的收益。但它同时落后于 ratio-4 的稠密 MLP，而加宽通常不会使结果变差，因此这一对照只说明两者不可互换，不足以用来衡量路由本身的价值。

Three recipes from looped language models degrade FID here循环语言模型的三种技术在此均使 FID 变差

At inference, sampling steps beat loops for the same budget等推理预算下去噪步优于循环

This is an inference-time question and does not contradict the first finding: there parameters are shared at training time, here a fixed inference budget is split between loops and denoising steps, and every model in this figure is an 80-epoch checkpoint, so the numbers are not comparable to the 240-epoch results above. A denoising step sees a new, less noisy input each time; a loop re-processes the same input with the same weights. At small budgets the first is worth far more, and a model that spends its whole budget on loops gets worse past five. Only when the budget is large enough that extra denoising steps have saturated does a looped model, at ten steps, edge past the dense one.

这是推理阶段的问题，与第一条结论不矛盾：前者是训练时的参数共享，此处是将固定推理预算在循环轮数与去噪步数之间分配，且本图中的模型均为 80 epoch 的 checkpoint，数值不能与前文 240 epoch 的结果直接比较。每个去噪步接收一个新的、噪声更低的输入；每轮循环只是以相同权重重新处理同一输入。小预算下前者收益远高，将全部预算用于循环的模型超过 5 轮后反而变差。只有当预算大到去噪步的收益饱和时，循环模型才在 10 步时略微超过稠密模型。

Limitations局限

Two limits. The model is small: everything here is DiT-B/4 on ImageNet 256, and nothing says how looping behaves at the sizes where it would matter. The objective is fixed: MeanFlow was the only training target, and whether loops interact differently with flow matching or with multi-step diffusion training was not studied.

两点局限。模型规模小：所有实验均为 ImageNet 256 上的 DiT-B/4，循环在更大规模上的表现无法从中推断。训练目标单一：仅使用 MeanFlow，循环与 flow matching 或多步扩散训练的相互作用未加考察。

Takeaways核心结论

Looping DiT-B/4 under the MeanFlow one-step objective: at equal compute the looped model holds no advantage, and extra denoising steps pay off more than extra loops.

在 MeanFlow 一步生成目标下循环 DiT-B/4：同等计算量下循环模型不比稠密模型占优，增加去噪步数比增加循环次数更有收益。

Acknowledgements致谢

The author is pleased to acknowledge that the work reported on in this post was substantially performed using the Princeton Research Computing resources at Princeton University. Princeton Research Computing is a consortium of groups including the Princeton Institute for Computational Science and Engineering (PICSciE) and Research Computing at Princeton University.

Citation

To cite this post:

引用本文请使用：

@misc{chai2026loopdit,
  title        = {LoopDiT: Loop Transformers for Diffusion Models},
  author       = {Chai, Wenhao},
  year         = {2026},
  month        = sep,
  howpublished = {Blog post},
  url          = {https://wenhaochai.com/exps/loopdit.html}
}