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，其余配置是 ablation；每张图的训练预算写在来源行。

1. Looping does not save compute.At the depth of DiT-B/4 no looped split reaches it. The same twelve blocks looped twice do, but at twice the compute, and the dense model of that depth has not been run, so nothing about parameter efficiency follows yet. 循环省不下算力。在 DiT-B/4 的深度上没有一种循环切法追得上它。同样十二个块循环两轮追上了，但算力翻倍，而同深度的稠密模型还没有跑，所以关于参数效率还说不出任何结论。
2. Independent blocks after the shared segment contribute more than blocks before it, and a fixed budget of them is best split evenly.Blocks after the loop refine a representation the loop has already built; blocks before it only prepare its input. 放在共享段之后的独立层贡献大于之前的，固定数量的独立层对半分最好。循环之后的层是在精修循环已经建好的表示，之前的层只是在准备输入。
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 个略有不同的函数。
4. Returns from additional loops diminish and reverse beyond about ten.Two shared blocks can play only so many distinct roles; past about ten the loops stop adding and start blurring. 循环轮数的收益递减，超过十轮左右转为负。两个共享块能扮演的角色有限；十轮之后再循环不再增加信息，反而互相干扰。
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 make early training cheap; the deep model is only needed late, and the reverse order confirms it. 同样的 epoch 预算下，先四轮后八轮胜过所有固定轮数；这一比较未对齐算力。少轮让早期训练便宜，深模型只在后期才需要；反过来的顺序验证了这一点。
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 tie; one versus two active experts is the whole difference, which is a difference in per-token compute. 共享块换成混合专家后，FID 随每 token 激活的专家数变化，对专家池大小不敏感。四个和十六个专家打平；一个和两个激活专家之间才是全部差异，而那是每 token 算力的差异。
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 passes may not need that protection, and the extra signals only disturb it. 在循环语言模型里有效的 input injection、噪声初始化和随机轮数，在这里都使 FID 变差。三者都是为了防止语言模型的循环塌缩而设计的；几轮的扩散循环可能不需要这种保护，额外的信号只是干扰。
8. Under a fixed inference budget, additional sampling steps are more effective than additional loops.Small budgets go to multi-step sampling outright; a looped model catches up only at twenty passes. 固定推理预算下，增加采样步比增加循环轮数更有效。预算小时多步采样直接胜出；循环模型要到二十次前向才追平。

Looping does not save compute循环省不下算力

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 dense 24-block model at the same width, which would separate the effect of sharing from the effect of depth, has not been run. The one clean statement is the first: looping does not save compute.

三张图都是 240 epoch 的模型。两组比较回答的是不同的问题。对齐深度时，循环模型的不同函数比稠密模型少：十二个稠密块是十二个函数，六个块循环两轮是六个，一个块循环十二轮则是同一个映射迭代到不动点。每种切法都输给 DiT-B/4，共享段越小输得越多。对齐参数时，同样十二个块跑两遍得到更低的 FID，但算力也翻了倍；能把共享的作用和深度的作用分开的同宽度稠密 24 层模型还没有跑。能干净地说出口的只有第一条：循环省不下算力。

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，所以先讲它。没有它时，shared block 在每一轮看到的条件向量完全相同，无法区分自己处在第几轮。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；shared block 的权重不变，只是每轮的调制不同。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。这样少跑几轮等价于提前退出，而不是让 block 面对训练时没见过的条件间隔。

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.

固定 p5 s2 e5，保留 loop embedding，只改两个共享块跑多少轮。质量改善到十轮左右，然后回落。每轮之间唯一变化的是循环序号，两个块能扮演的角色有限；再往后，额外的轮次只是重复同样的计算，穿过一长串共享权重的梯度在 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, eight later: better at equal epochs先四轮后八轮：同 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 前三分之一用四轮，其余用八轮。同样的 epoch 预算下这胜过任一固定轮数，反过来则输，说明收益来自训练后期的深度而不是切换本身。四轮阶段每步跑的块更少，所以这个方案的总算力也低于固定八轮；算力对齐的运行待补。

MoE: FID follows active experts, not pool size混合专家：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, not how many MLPs there are to choose from. The natural control is a dense MLP with doubled width; it has not been run, so the gain is not yet separated from plain widening.

激活两个专家而不是一个，在两种训练长度下都有帮助；专家池从四个扩到十六个则毫无变化。不管 router 学到了什么，收益跟着每个 token 得到多少 MLP 走，而不是有多少 MLP 可选。自然的对照是 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 的结果直接比较。每一个去噪步看到的是新的、噪声更少的输入；每一轮循环只是用同样的权重再处理同一个输入。预算小时前者值钱得多，把预算全花在循环上的模型过了五轮反而变差。只有预算大到去噪步的收益已经饱和，循环模型才在十步时略微超过稠密模型。

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 或多步扩散训练的关系没有考察。

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/blogs/loopdit.html}
}