本文目录

1 JiT：让去噪生成模型真正回归 "去噪"
(来自 MIT：Tianhong Li, Kaiming He)
1 JiT 论文解读
1.1 去噪模型，正一步步偏离 "去噪" 本质
1.2 JiT 预测空间，Loss 空间和推理空间
1.3 Toy Experiment：验证流形假设
1.4 Just Image Transformers
1.5 JiT 做法：预测 \bm{x} ，以 \bm{v} 为 Loss
1.6 实验结果

太长不看版

本文是 MAR 一作 Tianhong Li 在 MIT Kaiming He 老师团队里，2025.11 份的大作。

本文属于是生成式模型探索的文章。

本文的核心观点是：现在的扩散模型可能走错了路，与其预测噪声啊，速度场啊这类的 "noised quantity"，不如说直接预测原始图片这样的 "clean data" (也就是本文的题目：让去噪模型还是做去噪的任务吧)。

这里先列举一下本文依赖的假设和作者的观点 (注意，这里是假设和作者观点，不等于客观事实，也不等于公理。靠文章中的大量实验来佐证。)：

• 前提1：预测噪声，速度场啊这类的 "noised quantity"，预测原始图片这样的 "clean data" 是两码事。
• 前提2：流形假设：自然界的数据 "clean data" 是分布在高维空间当中的低维流形上面。但是带了噪声之后的 "noised quantity" 就不是了，分布在高维。

基于以上两个大前提，本文作者认为扩散模型应该预测 "clean data"，而不是 "noised quantity"。

作者训练了一个：

1. no tokenizer
2. no pre-training
3. no extra loss

的 "三无" Transformer。值得一提的是 patch 比较大 (256px 是 16，512px 是 32)。这个 "三无" Transformer：

• 预测 "clean data"：性能依旧能打。
• 但预测 "noised quantity"：直接失败。

这个点很有意思：说明一个明显 under-capacity 的模型，该预测的，应该是 "clean data"。

下面是对本文的详细介绍。

JiT：让去噪生成模型真正回归 "去噪"

• 论文名称：Back to Basics: Let Denoising Generative Models Denoise

• 论文地址：https://arxiv.org/pdf/2511.13720

• 代码链接：https://github.com/LTH14/JiT

01 JiT 论文解读：

1.1 去噪模型，正一步步偏离 "去噪" 本质

生成式模型，或者说去噪扩散模型的发展历程中的两座里程碑，分别是：

1. 预测 noise (\bm{\epsilon} - prediction)：在 diffusion model 里很常见。
2. 预测速度场 (\bm{v} - prediction)：在 flow-based model 里很常见，速度场可以视为一种 noised quantity。

但是很少有工作研究到底模型该去预测什么，似乎已经默认了模型能够完美执行要预测的任务。

本文作者认为：预测 clean data 和预测 noised quantity 是不一样的，原因是根据流形假设：高维的数据分布在低维的流形上面。那也就是说：

• clean data：分布在低维流形上。
• noised quantity：分布在完整的高维空间。

假设现在有一个场景，高维的数据分布在低维的流形上面：

• 那我现在如果说用神经网络去预测 noise 的话，那这网络就得 high-capacity。因为这 noise 维度很高，你想把信息预测全面的话，网络自身的能力就得比较强。
• 那我现在如果说用神经网络去预测 clean data 的话，那这网络就可以 limited-capacity。因为这 clean data维度很低，你想把信息预测全面的话，网络自身的能力就不必那么高。

那为什么在之前那么多预测 noise 或者速度场的工作？作者认为通过预训练的 VAE 把 image 映射到 latent space，然后在 latent space 上完成任务的方式降低了预测的难度。这个问题其实是被隐藏了，而没有被真正地解决。

因此，在本文中做的事情是：

• 去掉了 VAE
• 让模型 (single Transformer) 直接预测 clean data

本文的核心发现是：

当 Transformer 用了比较大的 patch 的时候，比如 512px 用了 32，也就意味着每个 token 的维度超高，预测   是可以的；预测   或者   不 work。

这说明：

当预测的东西的维度很高的时候 (比如大 patch 场景，或者不用 tokenizer 时)，让模型预测低维流形的数据 (更简单，维度降下来了)，模型能 work。但是让模型预测高维的数据 (更困难，维度降不下来)，模型就不能 work。

所以，本文号召大家在训练扩散模型的时候，去预测 clean data 这种低维流形的东西，让扩散模型回归去噪任务吧。

1.2 JiT 预测空间，Loss 空间和推理空间

需要说明的是：

1. 预测空间：模型在哪个空间做预测任务，即预测出来的结果代表什么含义。
2. Loss 空间：在哪个空间计算 Loss。
3. 推理空间：采样的时候在哪个空间进行。

预测空间

预测空间是指模型在哪个空间做预测任务，即预测出来的结果代表什么含义。

 之间满足关系：

有如下的排列组合：

根据图 2，当我们预测出了   中的任意一个量时，可以推断出另外两个。

Loss 空间

Loss 空间是指在哪个空间计算 Loss。Loss 虽然经常在   空间中定义，但是实际计算的时候可以在   空间中进行。比如上图2第3行（a），  －prediction，  －loss 的情况下   ，那 Loss 变为：

也就是  －loss 的重新加权形式。

推理空间

推理空间是指采样的时候在哪个空间进行。

在推理的时候，无论网络预测   中的任意一个，都一律转化到   空间中进行推理（图 2 第 3 行）。然后解 ODE 方程   来进行采样。

1.3 Toy Experiment：验证流形假设

不要忘记，本文有两个大前提：

• 前提1：预测噪声，速度场啊这类的 "noised quantity"，预测原始图片这样的 "clean data" 是两码事。
• 前提2：流形假设：自然界的数据 "clean data" 是分布在高维空间当中的低维流形上面。但是带了噪声之后的 "noised quantity" 就不是了，分布在高维。

Toy Experiment 的目的就是验证这两个前提。

假设现在是有   维数据，以流形的形式分布在   维空间中。这样的话，用一个列正交的投影矩阵   ，就可以把   维数据   映射为   。

这里取   。使用  －loss，分别做   projection。

训练一个 5－layer ReLU MLP 作为生成器，再使用   矩阵映射回   维。结果如图3所示。

注意这里模型的 hidden dimension 是 256 ，当   的时候，模型宽度不够了，对于   预测而言，都失败了。但是   预测依旧有效。

这说明了什么问题？ 数据   是分布在   维流形的，而   很低，  比较高。

• 对于   预测而言：要预测的东西到了 512 维，太高了，hidden dimension 是 256 的模型宽度不够。
• 对于   预测而言：要预测的东西仅 2 维，不高了，hidden dimension 是 256 的模型宽度够。

因此预测   仍然可以很好地生成数据。

1.4 Just Image Transformers

Just Image Transformers 完全 follow DiT 的架构，如图4所示，但是：

1. 去掉了 tokenizer。
2. 使用了比较大的 Patch：256px 为 16 (每个 Patch 的维度是 16×16×3=768)，512px 为 32 (每个 Patch 的维度是 32×32×3=3072)。

1．4．1 验证  －prediction

为了验证  －prediction，作者进行了这样的实验：

图 5(a) 是 ImageNet 256 的实验，使用大小为 16×16 的 Patch，图 5(b) 是 ImageNet 64 的实验，使用大小为 4×4 的 Patch。这两组实验的序列长度都是 256。然后用的模型都是 JiT-B (hidden size 为 768)。

结果非常耐人寻味：

图 5（a）中，大小为   的 Patch（Patch 维度是 768），仅有   预测成功了，   预测都失败了。在此情况下，JiT－B 的 hidden size 虽然是 768，但是作者认为需要一定余量用以处理位置编码信息。这样一来，模型就没有足够的能力进行 noised quantity 的处理了。

图5（b）中，大小为   的 Patch（Patch 维度是 48），那么   预测都成功了。在此情况下，JiT－B 的 hidden size 是 768 ，远高于 Patch 的维度。

1.4.2 改变 Noise-level 无济于事

作者也尝试了改变 Noise－level，结果如下图 6 所示。在   预测（模型本不能 work）的情况下，Noise－level Shift 无济于事。在   预测（模型本已经可以 work）的情况下，Noise－ level Shift 有点帮助。

1．4．3  －prediction 可以 handle 超大分辨率

图6里面，进一步拉大分辨率和 Patch 尺寸，比如分辨率   的 Patch，分辨率   的 Patch。前者每个 Patch 的维度是 3072dim，后者每个 Patch 的维度是 12288dim。即便如此，hidden size 仅为 768 的 JiT－B，在  －prediction 时仍然可以 work well。

这些观察表明，网络架构设计与维度是解耦的。加宽模型可以带来帮助，但并非是决定性的。

1.4.4 Bottleneck 的设计更好

与直觉相反，作者发现，在网络中引入先降维后升维的 Bottleneck 不但不会损失性能，反而更有帮助。

作者把 linear patch embedding 替换成一对先降维后升维的线性层。第一层把 hidden dimension 降为   ，第二层再把维度升为 hidden dimension。

如图 8 所示是使用 JiT-B/16 在 ImageNet 256 的 Bottleneck linear embedding 的结果。把 Bottleneck dimension 降到 16，性能也不会影响太大。把 Bottleneck dimension 调到 32 到 512，反而对性能有帮助。

作者认为，从表征学习的角度看，引入 Bottleneck 经常被用来促进学习低维表征。那么按照流形假设，自然图像分布在低维流形上，这也许是 Bottleneck 结构会帮助涨点的原因吧。这个实验也挺有意思的。

1．5 JiT 做法：预测   ，以   为 Loss

JiT 最终做法：

训练伪代码：

# net(z, t): JiT network
# x: training batch

t = sample_t()
e = randn_like(x)

z = t * x + (1 - t) * e
v = (x - z) / (1 - t)

x_pred = net_fn(z, t)
v_pred = (x_pred - z) / (1 - t)

loss = l2_loss(v - v_pred)

推理伪代码：

# z: current samples at t

x_pred = net_fn(z, t)
v_pred = (x_pred - z) / (1 - t)

z_next = z + (t_next - t) * v_pred

架构上，JiT 用了一些语言模型 Transformer 的常见改动：SwiGLU，RMSNorm，RoPE，qk-norm。

1.6 实验结果

Scalability

如下图 9 所示为 JiT 的 Scalability 结果。图 9 中模型的尺寸和 FLOPs 如图 10 和 11 所示。

下面是与其他方法的对比。值得一提的是，JiT 在 256px 的 FLOPs 与其他方法在 512px 的 FLOPs 接近。

有趣的是，随着模型尺寸增加，256px 和 512px 之间的 FID 差异减小了。

一般来讲，都是 256px的 FID < 512px 的 FID。但对于 JiT-G，512px 的 FID 结果反而更小。作者解释是因为对于大模型，ImageNet 数据集不够大，因此往往会过拟合。那么 512px 这个任务更难一点，因此过拟合风险更低了。

总结

个人认为这个文章的很大一个贡献是：图 3 的 Toy Experiment 和图 5 的发现。

当其他所有条件相同时，模型观测到的数据维度   越来越大的时候，超过或者等于模型的宽度的时候，预测原始数据是可以 work 的，但是预测速度场或者噪声就不行。这个现象惊人地与流行假设是契合的。甚至，当数据维度   很大时预测数据也是可以的。比如图 7 中， 1024 px 以及使用   的 Patch 时，预测 clean data 也能够 work。这些结果都很好地与流形假设相互印证，这个发现和倡议，个人认为才是本文贡献最大的地方。

因此，我们也许应该把关注点转回预测原始数据，即 "去噪" 任务上了。