Back to Basics: Let Denoising Generative Models Denoise

Status: completed

Authors: Tianhong Li, Kaiming He

Venue / Year: CVPR 2026，pp. 36115–36125；arXiv v1 于 2025-11-17 提交，v2 于 2026-01-07 修订

Affiliations: Massachusetts Institute of Technology (MIT)

Links: CVF Open Access | arXiv | PDF | Official GitHub

Tags: [[diffusion-models]] [[flow-matching]] [[pixel-space-generation]] [[vision-transformers]] [[x-prediction]] [[manifold-learning]] [[image-generation]] [[ImageNet]]

One-Sentence Summary

JiT 的核心不是一套复杂新架构，而是把高维像素扩散模型的网络直接输出从噪声 $\bm{\epsilon}$ 或速度 $\bm{v}$ 政策性地改成干净图像 $\bm{x}$：在带有强 patch bottleneck 的纯 ViT 中，这让网络只需表达近似低内在维的数据流形，而不必把高维噪声逐维穿过瓶颈，从而使大 patch 的像素空间 Transformer 在无 tokenizer、无预训练、无辅助损失的条件下达到有竞争力的 ImageNet 生成质量。

阅读结论先行

这是一篇概念极简、实验现象很强、理论解释尚未闭合的论文。

• 最值得记住的贡献不是“$\bm{x}$-prediction 是新方法”——它并不新——而是作者把网络直接输出空间和训练损失所在空间严格拆开，并证明在高维、强瓶颈 ViT 中，直接输出什么会造成灾难性差异。
• 最强证据是同架构下的 $3\times3$ 输出空间/损失空间控制实验：JiT-B/16 在 ImageNet $256\times256$ 上，只有直接 $\bm{x}$-prediction 的三列能够工作。
• 最需要克制的地方是“根本不同”这一表述。对无限容量、以 $t$ 为条件且达到逐点 MSE Bayes 最优的函数，九种参数化可以导向同一个条件均值去噪器；实际差异来自有限容量、瓶颈、跨时间参数共享、数值条件和优化几何。
• 论文证明了“一个特定且有代表性的架构失效机制”，但没有证明所有 U-Net、带长跳连的 Transformer、latent diffusion 或所有自然数据域都应改用 $\bm{x}$-prediction。
• 公开 PyTorch 代码能验证最终 JiT-B/L/H 的主路径，但不包含 JiT-G、JiT/64、原始 JAX/TPU 实现，也不包含关键 $3\times3$ 消融的另外八种配置；截至 2026-07-28，公开 checkpoint 下载还出现了不可用报告。

论文故事线与章节任务

整体故事： 现代 diffusion/flow 模型通常让网络直接预测 $\bm{\epsilon}$ 或 $\bm{v}$，再通过解析变换得到干净图像。作者认为“损失可以互相重加权”并不意味着“有限容量网络直接输出什么都一样”。如果自然图像近似落在低内在维流形上，而噪声占满环境空间，那么强瓶颈网络可以丢掉法向噪声后直接预测 $\bm{x}$，却无法为 $\bm{\epsilon}$ 或 $\bm{v}$ 保留全部高维细节。基于这个判断，作者构造 JiT：把 raw pixel 切成极大的 patch，用低秩 patch embedding 压缩后交给普通 Transformer，直接输出干净 patch，并在 $\bm{v}$-loss 下训练。实验从二维流形玩具问题、九种参数化、噪声调度、瓶颈维度、架构组件一直推进到 ImageNet 256/512/1024，说明这种简单设计可以扩展。

章节	任务	完成情况
Introduction / Related Work	指出“损失重加权等价”掩盖了有限网络的输出空间问题	叙事清晰；“根本不同”的适用范围写得偏宽
Sec. 3 Prediction Outputs	建立 $\bm{x},\bm{\epsilon},\bm{v}$ 三空间和九种组合的统一代数	代数正确；未充分讨论 Bayes 最优意义下的等价性
Sec. 3.3 Toy Experiment	用可控内在维与环境维分离现象	方向性强；缺少重复试验、误差条和容量边界曲线
Sec. 4 JiT	给出极简 pixel ViT、$\bm{x}$-prediction 与 $\bm{v}$-loss	方法完整，工程细节在附录中较充分
Sec. 4.2 Ablations	证明输出空间而非损失加权、噪声偏移或宽度是关键变量	是全文最有说服力的部分
Sec. 5 Comparisons	展示 256/512/1024、四个模型规模和外部基线	扩展性好；外部比较并非严格同预算，G 模型还使用早停
Appendix	披露优化器、EMA/CFG 搜索、预条件器、附加指标和代码细节	信息丰富；超参数选择与公开代码默认值仍有复现落差

Problem

Core Challenge

论文讨论的是条件图像生成。给定干净图像

$$\bm{x}\sim p_{\mathrm{data}}(\bm{x}),$$

以及独立高斯噪声

$$\bm{\epsilon}\sim p_{\mathrm{noise}}(\bm{\epsilon}) =\mathcal{N}(\bm{0},\bm{I}),$$

作者采用与通常 diffusion 记号方向相反的时间约定：$t=0$ 是纯噪声，$t=1$ 是数据。线性插值为

$$\bm{z}_t=t\bm{x}+(1-t)\bm{\epsilon}, \qquad t\in[0,1].$$

对应的常速度目标为

$$\bm{v} =\frac{\mathrm{d}\bm{z}_t}{\mathrm{d}t} =\bm{x}-\bm{\epsilon}.$$

采样时从

$$\bm{z}_0\sim p_{\mathrm{noise}}$$

出发，求解概率流 ODE

$$\frac{\mathrm{d}\bm{z}_t}{\mathrm{d}t} =\bm{v}_{\theta}(\bm{z}_t,t),$$

一直积分到 $t=1$。

经典做法通常让网络直接输出 $\bm{\epsilon}_{\theta}$ 或 $\bm{v}_{\theta}$。本文追问的是：

即使三种预测量可以解析互换，有限容量网络直接输出 $\bm{x}$、$\bm{\epsilon}$ 或 $\bm{v}$，是否真的一样容易？

Motivation

作者依赖流形假设：自然图像在高维像素环境空间中集中于低内在维结构，而独立高斯噪声沿所有环境维扩散。若单个 patch 的环境维度为

$$q=3p^2,$$

而网络先把它压到 $d'\ll q$，那么：

• 直接预测 $\bm{x}$ 时，瓶颈可以保留与数据流形相关的自由度，滤去大量法向噪声；
• 直接预测 $\bm{\epsilon}$ 或 $\bm{v}$ 时，输出仍包含高维、样本特定的噪声信息，强瓶颈可能无法近似恒等地传递这些细节。

这个动机对 JiT 架构非常贴切：JiT-B/32 的单 patch 像素维数为 $3072$，但 bottleneck 只有 $128$ 维，隐藏宽度只有 $768$。问题因此不再是“Transformer 能否在像素上生成”，而是“让它把什么量穿过信息瓶颈”。

Previous Work Gaps

• DDPM / score models： 已知 $\bm{x}$、$\bm{\epsilon}$、score 可互相变换，但实践默认直接 $\bm{\epsilon}$-prediction。
• Progressive Distillation / $\bm{v}$-prediction： 重点是不同参数化诱导的时间权重，而不是有限网络直接输出空间的表达难度。
• EDM： 以 denoiser 形式组织模型，但网络本体输出被 skip/out 系数预条件，通常并非严格的直接 $\bm{x}$。
• SiD / SiD2 / PixelFlow： 在像素空间工作，但依赖 U-Net、层级结构、小 patch 或多尺度设计，计算较重。
• PixNerd： 使用 NeRF-style 输出头并借助表征对齐，方法不再是纯粹的通用 Transformer。
• RAE / 高维 latent diffusion： 用预训练表征编码器产生高维 latent，再设计更宽的 diffusion head；解决高维的方式与本文“直接输出流形数据”不同。

Mathematical Framework

三种预测空间的完整转换

给定 $\bm{z}_t$ 和任意一个网络直接输出，可以由

$$\bm{z}_t=t\bm{x}+(1-t)\bm{\epsilon}, \qquad \bm{v}=\bm{x}-\bm{\epsilon}$$

解出另外两个量。

网络直接输出 $\bm{x}_{\theta}$

$$\bm{x}_{\theta} =\operatorname{net}_{\theta}(\bm{z}_t,t),$$

$$\bm{\epsilon}_{\theta} =\frac{\bm{z}_t-t\bm{x}_{\theta}}{1-t},$$

$$\bm{v}_{\theta} =\frac{\bm{x}_{\theta}-\bm{z}_t}{1-t}.$$

网络直接输出 $\bm{\epsilon}_{\theta}$

$$\bm{\epsilon}_{\theta} =\operatorname{net}_{\theta}(\bm{z}_t,t),$$

$$\bm{x}_{\theta} =\frac{\bm{z}_t-(1-t)\bm{\epsilon}_{\theta}}{t},$$

$$\bm{v}_{\theta} =\frac{\bm{z}_t-\bm{\epsilon}_{\theta}}{t}.$$

网络直接输出 $\bm{v}_{\theta}$

$$\bm{v}_{\theta} =\operatorname{net}_{\theta}(\bm{z}_t,t),$$

$$\bm{x}_{\theta} =\bm{z}_t+(1-t)\bm{v}_{\theta},$$

$$\bm{\epsilon}_{\theta} =\bm{z}_t-t\bm{v}_{\theta}.$$

九种“损失空间 $\times$ 输出空间”

三种未加权损失为

$$\mathcal{L}_{x} =\mathbb{E}\left\| \bm{x}_{\theta}-\bm{x} \right\|_2^2,$$

$$\mathcal{L}_{\epsilon} =\mathbb{E}\left\| \bm{\epsilon}_{\theta}-\bm{\epsilon} \right\|_2^2,$$

$$\mathcal{L}_{v} =\mathbb{E}\left\| \bm{v}_{\theta}-\bm{v} \right\|_2^2.$$

设三种直接输出误差分别为

$$\Delta_x=\bm{x}_{\theta}-\bm{x}, \qquad \Delta_{\epsilon}=\bm{\epsilon}_{\theta}-\bm{\epsilon}, \qquad \Delta_v=\bm{v}_{\theta}-\bm{v}.$$

把每个损失都写回网络直接输出空间，可得到完整的时间权重：

直接输出	$\bm{x}$-loss	$\bm{\epsilon}$-loss	$\bm{v}$-loss
$\bm{x}_{\theta}$	$|\Delta_x|^2$	$\dfrac{t^2}{(1-t)^2}|\Delta_x|^2$	$\dfrac{1}{(1-t)^2}|\Delta_x|^2$
$\bm{\epsilon}_{\theta}$	$\dfrac{(1-t)^2}{t^2}|\Delta_{\epsilon}|^2$	$|\Delta_{\epsilon}|^2$	$\dfrac{1}{t^2}|\Delta_{\epsilon}|^2$
$\bm{v}_{\theta}$	$(1-t)^2|\Delta_v|^2$	$t^2|\Delta_v|^2$	$|\Delta_v|^2$

例如，本文最终选择“直接 $\bm{x}$-prediction + $\bm{v}$-loss”：

$$\bm{v}_{\theta} =\frac{\bm{x}_{\theta}-\bm{z}_t}{1-t},$$

而真实速度也可写成

$$\bm{v} =\frac{\bm{x}-\bm{z}_t}{1-t}.$$

所以

$$\mathcal{L}_{v} = \mathbb{E} \frac{1}{(1-t)^2} \left\| \bm{x}_{\theta}(\bm{z}_t,t)-\bm{x} \right\|_2^2.$$

这说明两件事必须分开：

1. Prediction space： 网络最后一层实际要表达什么函数；
2. Loss space： 相同输出误差在不同 $t$ 上被怎样加权。

论文的 $3\times3$ 实验正是通过固定其中一个、改变另一个来识别二者作用。

一个关键理论限定：Bayes 最优并未改变

对固定 $t$，只要权重 $w(t)>0$ 且不依赖具体的 $\bm{x}$，加权平方误差

$$\mathbb{E} \left[ w(t) \left\| f(\bm{z}_t,t)-\bm{x} \right\|_2^2 \right]$$

的逐点最优解仍然是

$$f^{\star}(\bm{z}_t,t) = \mathbb{E} \left[ \bm{x}\mid\bm{z}_t,t \right].$$

类似地，

$$\mathbb{E} \left[ \bm{\epsilon}\mid\bm{z}_t,t \right] = \frac{ \bm{z}_t-t\, \mathbb{E}[\bm{x}\mid\bm{z}_t,t] }{1-t},$$

$$\mathbb{E} \left[ \bm{v}\mid\bm{z}_t,t \right] = \frac{ \mathbb{E}[\bm{x}\mid\bm{z}_t,t]-\bm{z}_t }{1-t}.$$

因此，在无限容量、完美优化、连续时间条件充分表达的理想条件下，三种输出空间可以描述同一个 Bayes 去噪器。本文观察到的巨大差异应更精确地解释为：

$$\boxed{ \mathrm{parameterization} \;+\; \mathrm{information\ bottleneck} \;+\; \mathrm{finite\ capacity} \;+\; \mathrm{optimization} }$$

之间的交互，而不是三个统计问题拥有不同的理想解。

这个限定并不削弱论文的工程价值；恰恰相反，它说明论文发现了一种非常实用的有限模型归纳偏置。

Method

Architecture Overview

对输入

$$\bm{x}\in\mathbb{R}^{B\times3\times H\times W},$$

JiT 的数据流为：

$$\mathrm{raw\ pixels} \rightarrow \mathrm{nonoverlapping\ patches} \rightarrow \mathrm{low\ rank\ patch\ embedding} \rightarrow \mathrm{Transformer\ blocks} \rightarrow \mathrm{linear\ patch\ predictor} \rightarrow \bm{x}_{\theta}.$$

不存在 VAE、VQ tokenizer、U-Net 多尺度塔、感知损失或对抗损失。

Patchification 与维度

patch 大小为 $p\times p$，token 数量为

$$N = \frac{H}{p}\frac{W}{p},$$

每个 raw patch 的环境维度为

$$q=3p^2.$$

设置	$H\times W$	$p$	$N$	raw patch 维度 $q$	JiT-B 隐藏维度
JiT-B/16	$256\times256$	16	256	768	768
JiT-B/32	$512\times512$	32	256	3072	768
JiT-B/64	$1024\times1024$	64	256	12288	768

作者让 patch size 与分辨率同比增长，因此三个分辨率始终只有 $16\times16=256$ 个图像 token。Transformer 主干 FLOPs 基本固定，分辨率差异主要落在输入/输出投影。

Low-Rank Patch Embedding

最终模型不是直接用一个 $q\times d$ 矩阵，而是使用两个连续线性映射：

$$\mathbb{R}^{q} \xrightarrow{\bm{W}_1} \mathbb{R}^{d'} \xrightarrow{\bm{W}_2} \mathbb{R}^{d}.$$

其中：

• B/L 模型默认 $d'=128$；
• H/G 模型默认 $d'=256$；
• $\bm{W}_1$ 在代码中实现为 kernel size 与 stride 都等于 $p$ 的卷积；
• $\bm{W}_2$ 是 $1\times1$ 卷积；
• 两层之间没有非线性，因此整体等价于秩至多为 $d'$ 的线性 patch embedding。

这一步是论文机制成立的关键，而不只是普通实现细节。以 JiT-B/32 为例：

$$3072\rightarrow128\rightarrow768.$$

若任务要求逐维传递噪声，第一步已经造成不可逆信息压缩；若任务只需恢复低内在维图像结构，这个压缩反而可以正则化模型。

Transformer Block

最终 JiT 使用：

• 固定二维 sine-cosine absolute positional embedding；
• RoPE 作用于 query/key；
• 每个 head 上的 qk-RMSNorm；
• RMSNorm；
• SwiGLU MLP；
• adaLN-Zero 注入时间和类别条件；
• attention 与 MLP 两个残差分支各自带可学习门；
• 最终 RMSNorm、adaLN modulation 和线性 unpatchify 头。

对第 $l$ 个 block，抽象更新可写为

$$\bm{h}'_l = \bm{h}_l + g^{\mathrm{attn}}_l(\bm{c}) \odot \operatorname{Attn} \left( \operatorname{Mod} \left( \operatorname{RMSNorm}(\bm{h}_l), \bm{c} \right) \right),$$

$$\bm{h}_{l+1} = \bm{h}'_l + g^{\mathrm{mlp}}_l(\bm{c}) \odot \operatorname{SwiGLU} \left( \operatorname{Mod} \left( \operatorname{RMSNorm}(\bm{h}'_l), \bm{c} \right) \right),$$

其中条件向量为

$$\bm{c} = \operatorname{Embed}_t(t) + \operatorname{Embed}_y(y).$$

所有 adaLN modulation 的末层和最终输出线性层均零初始化，因此网络从近恒等残差主干和零输出开始。

In-Context Class Conditioning

除了 adaLN-Zero，作者还在网络中部插入 32 个类别 token：

$$\bm{s}_j = \operatorname{Embed}_y(y)+\bm{p}^{\mathrm{class}}_j, \qquad j=1,\ldots,32.$$

这些 token：

• 不是从第一个 block 就加入；
• B/L/H/G 分别从第 4/8/10/10 个 block 附近开始；
• 与图像 token 一起参加后续 self-attention；
• 最终输出前被移除。

因此“Just image Transformers”在最终强模型上并非完全等同于原始 ViT：它仍包含 adaLN-Zero、RoPE、qk-norm、SwiGLU、低秩 patch embedding 和中途插入的多类别 token。

Model Scale

模型	深度	隐藏维度	heads	bottleneck	dropout
JiT-B	12	768	12	128	0
JiT-L	24	1024	16	128	0
JiT-H	32	1280	16	256	0.2
JiT-G	40	1664	16	256	0.2

H/G 的 dropout 只放在中间一半 blocks，同时作用于 attention probability/projection 与 MLP。G 模型仍会过拟合，因此作者在约 320 epochs 处按 FID 早停。

Algorithm

Training Forward Pass

训练时：

1. 图像被缩放到 $[-1,1]$。
2. 采样

$$s\sim\mathcal{N}(\mu,\sigma^2), \qquad t=\operatorname{sigmoid}(s),$$

默认

$$\mu=-0.8, \qquad \sigma=0.8.$$

3. 按分辨率采样噪声

$$\bm{\epsilon} \sim \mathcal{N} \left( \bm{0}, \alpha^2\bm{I} \right), \qquad \alpha=\frac{H}{256}.$$

所以 256、512、1024 分辨率分别使用 $\alpha=1,2,4$。 4. 构造

$$\bm{z}_t = t\bm{x} + (1-t)\bm{\epsilon}.$$

5. 以概率 $0.1$ 丢弃类别标签，训练 classifier-free guidance 的无条件分支。
6. JiT 直接输出

$$\bm{x}_{\theta} = \operatorname{JiT}_{\theta} \left( \bm{z}_t,t,y \right).$$

7. 转到速度空间并计算 MSE。

端点裁剪后的真实训练目标

论文正文写作

$$\bm{v}_{\theta} = \frac{\bm{x}_{\theta}-\bm{z}_t}{1-t}, \qquad \bm{v} = \frac{\bm{x}-\bm{z}_t}{1-t}.$$

实现中为避免 $t\rightarrow1$ 时除零，定义

$$D_{\tau}(t) = \max(1-t,\tau), \qquad \tau=0.05.$$

代码实际计算

$$\widetilde{\bm{v}}_{\theta} = \frac{\bm{x}_{\theta}-\bm{z}_t}{D_{\tau}(t)},$$

$$\widetilde{\bm{v}} = \frac{\bm{x}-\bm{z}_t}{D_{\tau}(t)}.$$

两者相减后，真实优化目标精确等于

$$\mathcal{L}_{\mathrm{impl}} = \mathbb{E} \left[ \frac{ \left\| \bm{x}_{\theta}-\bm{x} \right\|_2^2 }{ D_{\tau}(t)^2 } \right].$$

因此：

• 当 $t\leq0.95$ 时，它与理论 $\bm{v}$-loss 诱导的 $\bm{x}$ 权重相同；
• 当 $t>0.95$ 时，权重上限为

$$\frac{1}{\tau^2}=400,$$

不再随 $t\rightarrow1$ 发散；

• 同时，裁剪后的训练 target 不再严格等于恒定的 $\bm{v}=\bm{x}-\bm{\epsilon}$，而是乘上 $(1-t)/\tau$。

这是数值稳定性所需的合理修改，但分析公式和真实实现之间应显式区分。

Optimization

• AdamW，$\beta_1=0.9$，$\beta_2=0.95$；
• 全局 batch size 1024；
• base learning rate $5\times10^{-5}$，按 batch size 线性缩放后实际为

$$2\times10^{-4};$$

• 5 epochs warmup，之后 constant learning rate；
• weight decay 为 0；
• 训练 200 epochs 做消融，最终模型标称 600 epochs；
• 同时维护多个 EMA 候选；
• H/G 使用 dropout，G 使用基于 FID 的早停。

Sampling

从

$$\bm{z}_0 \sim \mathcal{N} \left( \bm{0}, \alpha^2\bm{I} \right)$$

开始，使用 50-step Heun ODE solver。Euler predictor 为

$$\widehat{\bm{z}}_{t_{k+1}} = \bm{z}_{t_k} + \Delta t_k \bm{v}_{\theta} \left( \bm{z}_{t_k},t_k \right),$$

Heun corrector 为

$$\bm{z}_{t_{k+1}} = \bm{z}_{t_k} + \frac{\Delta t_k}{2} \left[ \bm{v}_{\theta} \left( \bm{z}_{t_k},t_k \right) + \bm{v}_{\theta} \left( \widehat{\bm{z}}_{t_{k+1}},t_{k+1} \right) \right].$$

最后一步使用 Euler，避免在 $t=1$ 再评估速度场。

CFG 在速度空间进行：

$$\bm{v}_{\mathrm{cfg}} = \bm{v}_{\mathrm{uncond}} + \omega(t) \left( \bm{v}_{\mathrm{cond}} - \bm{v}_{\mathrm{uncond}} \right).$$

若使用 CFG interval，则仅在

$$t\in(0.1,1.0)$$

内令 $\omega(t)>1$，区间外取 1。作者在每个模型上扫描 CFG scale 与 EMA 后，再用选中配置报告 FID。

Evidence

Toy Experiment：控制内在维与环境维

作者从二维 spiral 数据

$$\widehat{\bm{x}}\in\mathbb{R}^{d}, \qquad d=2,$$

出发，用固定随机列正交矩阵

$$\bm{P}\in\mathbb{R}^{D\times d}, \qquad \bm{P}^{\top}\bm{P}=\bm{I}_d$$

嵌入环境空间：

$$\bm{x} = \bm{P}\widehat{\bm{x}} \in \mathbb{R}^{D}.$$

生成器是 5-layer ReLU MLP，隐藏宽度 256；比较

$$D\in\{2,8,16,512\}.$$

所有模型都使用同一个 $\bm{v}$-loss，只改变直接输出空间。结果是：

• $D=2$ 时三种输出都可工作；
• $D$ 增长后 $\bm{\epsilon}$/$\bm{v}$-prediction 快速退化；
• $D=512>256$ 时，只有 $\bm{x}$-prediction 仍能恢复二维 spiral。

证据强度：中到强。 它把内在维、环境维和网络宽度显式拆开，方向上支持瓶颈解释。但论文没有报告多随机种子、定量距离、不同 $d$、不同宽度或非线性嵌入，所以它更像机制演示而不是完整定理验证。

九种组合：全文最关键实验

ImageNet 256 上，JiT-B/16 的 raw patch 与 hidden 都是 768 维；ImageNet 64 上，JiT-B/4 的 raw patch 只有 48 维。两者 token 数都为 256。

ImageNet $256\times256$，JiT-B/16，200 epochs，FID-50K

损失空间	$\bm{x}$-pred	$\bm{\epsilon}$-pred	$\bm{v}$-pred
$\bm{x}$-loss	10.14	379.21	107.55
$\bm{\epsilon}$-loss	10.45	394.58	126.88
$\bm{v}$-loss	8.62	372.38	96.53

ImageNet $64\times64$，JiT-B/4，200 epochs，FID-50K

损失空间	$\bm{x}$-pred	$\bm{\epsilon}$-pred	$\bm{v}$-pred
$\bm{x}$-loss	5.76	6.20	6.12
$\bm{\epsilon}$-loss	3.56	4.02	3.76
$\bm{v}$-loss	3.55	3.63	3.46

解释：

1. 高维 patch 下，决定性变量是网络直接输出空间，不是 loss space；
2. 低维 patch 下，九种组合都可工作，主要差异回到时间权重；
3. $\bm{x}$-prediction + $\bm{v}$-loss 是高维设置中的最好组合；
4. 这组实验比只比较最终 JiT 和外部 baseline 更能支撑核心 claim，因为它保持架构、数据、训练时长与采样器一致。

主要保留意见： 没有 seed 方差，也没有公开另外八种配置和 checkpoint；因此“FID 300 级灾难性失败”尚缺独立复现。

Noise-Level Shift 不能救回错误输出空间

JiT-B/16、ImageNet 256、$\bm{v}$-loss、200 epochs：

$\mu$	$\bm{x}$-pred	$\bm{\epsilon}$-pred	$\bm{v}$-pred
0.0	14.44	464.25	120.03
-0.4	9.79	372.91	109.93
-0.8	8.62	372.36	96.53
-1.2	8.99	355.25	106.85

更负的 $\mu$ 让训练采样偏向更小 $t$，也就是更高噪声。适当 shift 能改善已经可工作的 $\bm{x}$-prediction，但无法把 $\bm{\epsilon}$/$\bm{v}$-prediction 从灾难状态救回。

这排除了“作者只是选错了高分辨率噪声分布”这一简单解释，但只扫描了四个 $\mu$，不能排除更复杂的 weighting/schedule。

Bottleneck Ablation

JiT-B/16 的 raw patch 为 768 维。作者把 bottleneck $d'$ 从 768 降到 16：

• $d'=32$ 到 $512$ 的广泛范围内，FID 相比无 bottleneck 最多改善约 1.3；
• $d'=16$ 仍未灾难性失败；
• 最终 B/L 使用 $d'=128$。

这个结果与论文叙事高度一致：直接预测 $\bm{x}$ 时，低秩输入投影不仅可容忍，甚至是有利正则化。

但需要注意，它并未直接估计 ImageNet patch 的内在维度。$d'=16$ 只是一个中间线性秩；后续 Transformer 利用 256 个 token 的全局上下文，单 patch 局部细节可以由邻域和类别信息补全。

“Just Advanced” Components

ImageNet 256、200 epochs、bottleneck 128：

配置	JiT-B/16 FID	使用 CFG interval	JiT-L/16 FID	使用 CFG interval
SwiGLU + RMSNorm	7.48	6.32	-	-
+ RoPE + qk-norm	6.69	5.44	-	-
+ 32 in-context class tokens	5.49	4.37	3.39	2.79

最终性能并非只来自 $\bm{x}$-prediction。RoPE/qk-norm、CFG interval 和额外 class tokens 的总增益很大；因此“只是把输出改为 $\bm{x}$ 就得到 FID 1.8”是不准确的。

Resolution Scaling

JiT-B 保持 256 token：

分辨率	模型	patch dim	hidden	参数量	单次前向 GFLOPs	FID
$256\times256$	B/16	768	768	131M	25	4.37
$512\times512$	B/32	3072	768	133M	26	4.64
$1024\times1024$	B/64	12288	768	141M	30	4.82

这张表真正证明的是：当直接预测 $\bm{x}$ 且同步放大噪声时，环境维度远大于 hidden width 不会导致灾难性崩溃。它没有证明 1024 生成达到当时 SOTA，也没有与 1024 baseline 做同预算比较。

Model Scaling

模型	ImageNet 256，200 ep	ImageNet 256，长训练	ImageNet 512，200 ep	ImageNet 512，长训练
JiT-B	4.37	3.66	4.64	4.02
JiT-L	2.79	2.36	3.06	2.53
JiT-H	2.29	1.86	2.51	1.94
JiT-G	2.15	1.82	2.11	1.78

表中“长训练”列标作 600 epochs，但附录说明 JiT-G/16 与 G/32 约在 320 epochs 因 FID 变差而早停。故：

• B/L/H 的数字可理解为 600-epoch；
• G 的数字更准确地说是“最多 600 epochs、按 FID 选择的早停 checkpoint”；
• 不能把 G 相对 H 的差值完全归因于参数 scaling，因为模型选择策略也改变了。

External Baselines

ImageNet $256\times256$

方法	空间	额外预训练/组件	参数量	GFLOPs	FID
DiT-XL/2	latent	SD-VAE + VGG	675M + 49M	119	2.27
REPA SiT-XL/2	latent	SD-VAE + VGG + DINOv2	675M + 49M	119	1.42
RAE DiT-DH-XL/2	latent	RAE + VGG + DINOv2	839M + 415M	146	1.13
SiD2 UViT/1	pixel	无 tokenizer	N/A	653	1.38
PixelFlow XL/4	pixel	无 tokenizer	677M	2909	1.98
JiT-H/16	pixel	无预训练、无辅助损失	953M	182	1.86
JiT-G/16	pixel	无预训练、无辅助损失	2B	383	1.82

ImageNet $512\times512$

方法	空间	参数量	GFLOPs	FID
DiT-XL/2	latent	675M + 49M	525	3.04
RAE DiT-DH-XL/2	latent	839M + 415M	642	1.13
SiD2 UViT/2	pixel	N/A	653	1.48
JiT-H/32	pixel	956M	183	1.94
JiT-G/32	pixel	2B	384	1.78

JiT 的优势不是绝对最低 FID，而是：

• 不依赖 tokenizer、VGG perceptual loss 或 self-supervised encoder；
• 分辨率翻倍时 token 数不变，主干计算几乎不变；
• 相比 PixelFlow 等像素方法，单步 FLOPs 很低；
• 架构更接近通用 Transformer。

公平性限制：

• 表中 GFLOPs 只计一次生成器前向；latent tokenizer 的成本未计入 GFLOPs，但 decoder 参数计入参数量；
• 不同方法的训练 epochs、训练总 FLOPs、数据增强、采样 NFE、CFG/EMA 搜索范围不同；
• JiT 是 50-step Heun，单步 FLOPs 不能直接等于端到端生成成本；
• 表格是参考性横向定位，不是严格 compute-matched benchmark。

Additional Metrics

ImageNet 256：

方法	FID	IS	Precision	Recall
DiT-XL/2	2.27	278.2	0.83	0.57
SiT-XL/2	2.06	277.5	0.82	0.59
RAE	1.13	262.6	0.78	0.67
JiT-B/16	3.66	275.1	0.82	0.50
JiT-L/16	2.36	298.5	0.80	0.59
JiT-H/16	1.86	303.4	0.78	0.62
JiT-G/16	1.82	292.6	0.79	0.62

JiT scaling 明显改善 recall，但 G 相对 H 的 FID 只改善 0.04，IS 反而下降。论文对过拟合的解释与此一致：大模型的边际收益已经很小，模型选择对结果敏感。

Preconditioner Experiment

EDM-style 预条件器写为

$$\bm{x}_{\theta} = c_{\mathrm{skip}}\bm{z}_t + c_{\mathrm{out}} \operatorname{net}_{\theta}(\bm{z}_t,t).$$

作者把 EDM 的 VE 记号近似映射到本文插值：

$$\sigma_t=\frac{1-t}{t},$$

$$c_{\mathrm{skip}} = \frac{1}{t} \frac{ \sigma_{\mathrm{data}}^2 }{ \sigma_{\mathrm{data}}^2+\sigma_t^2 },$$

$$c_{\mathrm{out}} = \frac{ \sigma_{\mathrm{data}}\sigma_t }{ \sqrt{ \sigma_{\mathrm{data}}^2+\sigma_t^2 } }, \qquad \sigma_{\mathrm{data}}=0.5.$$

另一个线性版本为

$$c_{\mathrm{skip}}=t, \qquad c_{\mathrm{out}}=1-t.$$

JiT-B/16、ImageNet 256：

loss	直接 $\bm{x}$-pred	EDM-style	linear preconditioner
$\bm{x}$-loss	10.14	28.94	39.50
$\bm{\epsilon}$-loss	10.45	72.05	67.56
$\bm{v}$-loss	8.62	35.49	46.25

预条件器优于纯 $\bm{\epsilon}$/$\bm{v}$-prediction，却明显差于严格直接 $\bm{x}$ 输出。这个实验支持“直接输出定义重要”，但它不是对 EDM 的完全等价复现：作者明确承认 VE schedule 到当前 schedule 无法完全转换。因此它不能被读作“EDM 预条件器普遍无效”。

Claims $\rightarrow$ Evidence Mapping

Claim	类型	支撑证据	强度	风险
高维强瓶颈下，直接 $\bm{x}$-prediction 显著优于 $\bm{\epsilon}$/$\bm{v}$	Empirical	Fig. 3 toy；Table 2 的 $3\times3$ 实验	强	只覆盖 MLP 与 JiT；没有 U-Net/长跳连对照
Loss weighting 不能解释全部差异	Algebraic + Empirical	九种组合代数；Table 2 同列跨 loss	强	有限训练下的结论，不等于 Bayes 最优不同
Noise shift 不能单独修复失败	Empirical	Table 3，四个 $\mu$	中	搜索范围有限
网络 hidden 不必大于 patch 环境维度	Empirical	B/32、B/64；Table 5	强	依赖 $\bm{x}$ 输出与固定 token 数
Bottleneck 可以有益	Empirical	Fig. 4，$d'=16$ 到 768	中到强	无 seed 方差，未估计真实内在维
Plain Transformer 可在 pixel space 达到竞争结果	Empirical	Tables 6–8，FID/IS/PR	强	最强结果包含多项现代组件、CFG/EMA sweep 和早停
分辨率翻倍不导致主干计算四倍增长	Architectural	固定 256 tokens；GFLOPs 25 到 26	强	更大 patch 可能牺牲局部细节；未报告完整吞吐
方法可推广到蛋白质、分子、天气等 raw natural data	Speculative	Conclusion	弱	没有跨域实验

Theory and Mathematical Rigor

正确且有价值的部分

1. $\bm{x}$、$\bm{\epsilon}$、$\bm{v}$ 的线性转换正确。
2. 作者明确区分 loss space 与 network output space，这比把二者统称为“prediction objective”更严谨。
3. 九种目标没有被错误说成训练函数完全相同；不同输出参数化确实改变有限网络的函数类与梯度。
4. 高维瓶颈的直觉与 toy 和 ImageNet 实验方向一致。
5. 端点除法使用 $\tau=0.05$ 裁剪，避免了 $t\rightarrow1$ 的数值爆炸。

重要理论缺口

1. “预测噪声需要保留全部噪声”不是无条件命题

平方误差训练的网络并不需要重建监督样本中的不可预测随机噪声；其 Bayes 解是条件均值

$$\mathbb{E} \left[ \bm{\epsilon} \mid \bm{z}_t,t \right].$$

在本文线性插值下，它由条件干净图像均值解析决定。真正的问题是：JiT 的 bottleneck 是否能高效表示从 $\bm{z}_t$ 到这个条件均值的参数化，以及输出端是否有便宜的 $\bm{z}_t$ 跳连。

论文的直觉在“无长跳连的强瓶颈网络”上成立，但措辞有时把架构依赖现象提升为输出随机变量本身的普遍难度。

2. 缺少直接的架构因果实验

最关键的缺失 baseline 是：

$$\bm{v}_{\theta} = \frac{ g_{\theta}(\bm{z}_t,t)-\bm{z}_t }{ 1-t }$$

或等价地，在 $\bm{v}$-prediction 网络外显式加入无瓶颈的 $\bm{z}_t$ 跳连。官方 GitHub issue #8 中作者也承认，这种直接的、随 $t$ 缩放的残差连接可能隐式学习同样变换。

如果加入 cheap skip 后 $\bm{v}$-prediction 恢复，则结论应定位为：

输出参数化必须与网络的信息通路匹配。

这比“所有高维 diffusion 都应直接 $\bm{x}$-prediction”更精确。

3. 流形假设没有被定量检验

论文没有：

• 估计 ImageNet patch 在不同噪声水平下的局部内在维；
• 测量瓶颈前后的 mutual information 或 Jacobian rank；
• 分解切向误差和法向误差；
• 验证 $\bm{x}$-prediction 是否真的优先保留 manifold tangent；
• 证明所需宽度随内在维而非环境维增长。

因此“manifold”当前是一个高度一致的解释，而不是已验证机制。

4. 自然图像并不严格落在光滑低维流形

传感器噪声、纹理、量化、裁剪和数据增强会让像素分布具有非零厚度。更稳妥的说法是图像概率质量集中在低复杂度集合附近，而非严格位于零测度流形上。

5. 端点裁剪改变了理论 ODE

采样代码也使用

$$\frac{1}{D_{\tau}(t)}$$

而非严格的 $1/(1-t)$。作者报告这种 clipping 对 ImageNet FID 有益，但它意味着最后区间使用的是被正则化的速度场。论文没有分析由此引入的积分偏差。

后续理论如何补足

2026 年的 k-Diff 工作《Revisiting Diffusion Model Predictions Through Dimensionality》进一步建立广义预测 target，并把环境维与内在维的比例联系到最优参数化。这类工作说明 JiT 的现象具有理论延展性，但也反向说明原论文自身主要提供的是经验发现与机制假说，而非闭合证明。

Experimental Design Assessment

做得好的地方

• $3\times3$ 因子实验真正拆开 prediction space 与 loss space；
• ImageNet 64 与 256 保持 token 数一致，隔离 patch 环境维；
• $\mu$ 扫描排除简单 noise shift 解释；
• bottleneck sweep 直接测试论文最反直觉的预测；
• 多模型规模、多分辨率和多个指标支持结果不是单点偶然；
• 附录披露 EMA/CFG sweep、早停、训练参数和预条件器。

重大缺失

1. 没有随机种子与方差。 FID 差 0.04 或 0.1 时无法判断统计显著性。
2. 核心八个失败消融未公开。 官方代码只实现最终 $\bm{x}$-prediction + $\bm{v}$-loss。
3. 没有架构交叉。 缺 U-Net、无 bottleneck ViT、显式 input-output skip、不同 patch mixer。
4. 没有训练总成本。 只报告单次 forward GFLOPs，不报告 TPU/GPU days、总 NFE、吞吐或能耗。
5. 模型选择预算很大。 每个消融都在 CFG 1.0 到 4.0、步长 0.1，以及多个 EMA 上用 8K samples 选最佳，再评 50K；这提高了结果，也引入隐含验证预算。
6. G 模型早停。 “600-epoch”列中的 G 实际约 320 epochs，跨规模比较不是完全同训练协议。
7. 只用 ImageNet class-conditional。 对 text-to-image、无条件生成、非视觉数据和真实高维科学场的外推未经验证。
8. FID 对 ImageNet training set。 这是领域惯例，但仍可能奖励训练分布拟合，尤其作者已观察到 G 过拟合。

Cherry-Picking 风险

没有证据表明作者隐瞒负结果；相反，附录披露了过拟合、预条件器失败和分类损失探索。但以下选择会放大最优数字：

• 每模型独立 CFG/EMA sweep；
• G 模型按监控 FID 早停；
• 外部 baseline 来自各论文最优报告，训练预算不统一；
• 主表强调 FID，precision/recall 只在附录。

因此更可信的结论是“JiT 建立了强且高效的像素 Transformer baseline”，而不是“JiT 全面优于 latent/pixel 现有方法”。

Contribution Positioning

Innovation Type

• [x] 概念创新： 把有限网络的直接输出空间从 loss weighting 中分离出来。
• [ ] 理论创新： 没有正式定理；流形解释是机制假说。
• [x] 方法创新： $\bm{x}$-prediction、低秩大 patch 与普通 ViT 的组合。
• [x] 经验创新： 首次系统展示高维 patch 下三种输出空间的灾难性分化。
• [x] 工程创新： 固定 token 数使 256 到 1024 的主干计算近似不变。

与直接竞品的增量

工作	核心路线	相对 JiT 的优势	相对 JiT 的劣势
SiD2	分层 UViT 的 pixel diffusion	FID 更低，像素建模成熟	FLOPs 更高，架构更专用
PixelFlow	多尺度 pixel-space flow	强高分辨率生成	单步计算极高，系统复杂
PixNerd	ViT + neural field head + representation alignment	局部连续表示强	需要专用 head 与预训练表征
RAE	预训练表征编码器 + 宽 diffusion head	FID 1.13，语义 latent 强	不是自包含，需要预训练 encoder/decoder
DDT	解耦宽 head 解决高维 latent	针对高维 latent 的结构更直接	仍依赖 tokenizer/representation encoder
JiT	大 patch pixel ViT + 直接 $\bm{x}$ 输出	最简、自包含、分辨率扩展便宜	训练/采样仍重，绝对 FID 非最佳
k-Diff	数据驱动学习连续 target 参数	理论定位更完整，不固定三选一	后续工作，系统更复杂

Innovation Score: 7.5/10

理由：

• $\bm{x}$-prediction 本身是旧概念；
• 但“输出空间 $\times$ 损失空间”的完整因子化及其高维失效实验很有洞察；
• JiT 把一个被认为需要专用架构的问题，降成了目标参数化选择；
• 贡献很可能改变 pixel diffusion 与科学数据 diffusion 的默认设计习惯；
• 理论因果、跨架构普适性和复现资产仍不足以给到 9 分。

Critical Assessment

Strengths

1. 问题选择有品味。 它挑战的是被代数等价掩盖的默认假设，而不是堆新模块。
2. 最关键实验非常干净。 九种组合把两个经常混淆的因素完全拆开。
3. 极简方法带来强可证伪预测。 bottleneck 越强不一定越差，结果确实反直觉。
4. 工程收益真实。 512 相比 256 只增加输入/输出投影，主干计算基本不变。
5. 自包含。 无 VAE、无 VGG、无 DINO、无额外 loss。
6. 写作诚实。 作者承认 $\bm{x}$-prediction 不新、G 会过拟合、预条件器转换不完全等。

Fundamental Flaws

没有发现会直接推翻全部实验结论的致命错误。

Significant Concerns

1. 核心理论语言过强。 实验支持“有限瓶颈网络的参数化差异”，不充分支持“预测随机变量在统计上根本不同”。
2. 架构依赖未隔离。 缺显式 input skip/U-Net 对照，不能确定优势在多大程度上只由 JiT bottleneck 造成。
3. 核心消融不可复现。 公开代码没有另外八种 prediction/loss 配置。
4. 统计报告不足。 没有多个 seed、置信区间或 FID 方差。
5. 最强结果依赖选择程序。 CFG/EMA sweep 和 G 早停应更突出地出现在主表。

Minor Issues

• 论文把所有模型的长训练列标作 600 epochs，但 G 实际约 320 epochs 早停；
• 1024 只展示 B 模型与 FID，没有 baseline、IS、PR 或细节敏感指标；
• “plain ViT”容易让读者忽略最终模型包含的低秩 embedding、RoPE、qk-norm、adaLN-Zero 和 32 class tokens；
• 没有 wall-clock 训练时间与 50-step 总采样延迟；
• 未报告 patch seam 或高频纹理质量的专门指标。

Reviewer Feedback

CVPR 2026 已正式接收并收录该论文，但 CVPR 的官方审稿文本与评分没有在本文提供的公开入口中展示，因此不能虚构 reviewer score 或 rebuttal 内容。

公开 GitHub 讨论提供了若干可核查补充：

讨论	作者/社区澄清	含义
Issue #3：代码是否其实是 $\bm{v}$-prediction	作者确认网络直接输出 $\bm{x}$，loss 在 $\bm{v}$ 空间	验证本文 prediction/loss 的区分
Issue #8：带残差/长跳连会怎样	作者认为显式 $\bm{z}_t$ 到输出的 $t$-dependent residual 可能学到类似变换	支持“架构交互”而非普遍不可等价
Issue #30：请求九种消融配置	截至检查时仍公开缺失	核心 claim 的复现资产不完整
Issue #42/#56：FID 复现差距	作者建议 sweep EMA/CFG；社区用 EMA 0.9996、CFG 3.6 复现 200-epoch 4.37	论文数字对推理超参数敏感，README 默认脚本不够
Issue #70：checkpoint 链接失效	截至 2026-07-28 仍为 open	当前开箱复现性下降

Recommendation

Accept，7/10。

理由：

• 核心现象新颖、明确且可能改变实践；
• 同架构控制实验足以证明 prediction space 在高维瓶颈下不是无关变量；
• 方法极简却建立了有竞争力的 pixel-space baseline；
• 理论措辞与复现缺口值得批评，但不足以否定贡献。

如果按 revision 标准，优先要求：

1. 加入 U-Net、显式 input-output skip 和无 bottleneck ViT；
2. 报告三到五个 seed；
3. 发布九种消融配置与 checkpoint；
4. 在主表注明 G 的早停和每模型独立 CFG/EMA sweep；
5. 把“根本不同”收紧为有限容量与架构条件下的参数化差异。

Field Landscape

领域路线可以分成三类：

1. 低维 latent： 用 VAE/VQ/RAE 先压缩，再在 latent 上 diffusion；效果最好但依赖 tokenizer。
2. 专用 pixel architecture： 用 U-Net、层级、多尺度、neural field 解决像素空间计算与高维问题。
3. 通用 Transformer + 合适 target： JiT 主张不改 backbone，只改网络真正需要输出的对象。

JiT 的位置可以概括为：

它不是以新生成过程取胜，而是通过把直接输出对齐到数据流形，解除大 patch Transformer 的信息瓶颈，使原本失败的纯像素 ViT 变成可扩展 baseline。

后续 k-Diff 把固定的 $\bm{x}/\bm{\epsilon}/\bm{v}$ 三选一扩展为可学习 target；2026 年关于 training-free guidance 的工作又表明，直接 $\bm{x}$ 输出可能在高噪声 guidance 中更容易保持样本位于数据流形附近。这些后续结果说明 JiT 的影响更可能体现在“target design 成为一等公民”，而不是 JiT 架构本身成为最终答案。

Research Direction and Outlook

Scientific Value

这个问题有长期价值。预测目标看似只是代数重参数化，却会改变有限网络的可表达函数、梯度尺度、数值条件和信息通路。类似问题存在于：

• weather / climate fields；
• proteins 与 molecular conformations；
• motion 与 robotics trajectories；
• audio waveform；
• scientific simulation state；
• 高维 representation latent。

Scaling Bottlenecks

• 固定 256 token 的计算优势来自 patch 随分辨率变大；这可能损失高频局部细节；
• $\bm{x}$-prediction 的端点权重需要裁剪，否则 $t\rightarrow1$ 梯度发散；
• 50-step Heun 仍需大量前向；
• 2B 参数的 JiT-G 只比 H 小幅改善 FID，说明数据规模和正则化先成为瓶颈；
• 没有 tokenizer 不等于端到端训练一定更便宜，训练 600 epochs 的像素模型成本仍高。

Best Follow-Up Questions

1. 直接给 $\bm{\epsilon}$/$\bm{v}$ 网络增加解析 input skip，能否消除灾难性差距？
2. 最优 target 是否应随 $t$、分辨率、patch 内在维和网络层深动态变化？
3. 能否从 Jacobian spectrum 直接测量瓶颈保留的是流形切向还是低频信息？
4. $\tau=0.05$ 的端点裁剪如何影响 ODE bias、细节和可逆性？
5. 对天气等连续物理场，低维“流形”是否应替换为守恒律约束的解流形？
6. 大 patch 造成的 seam 与高频损失能否由轻量局部 decoder 修复，同时保持自包含？

Personal Follow-Up Decision

• [x] 非常值得： 若研究对象是高维 raw field 或自包含 diffusion，应优先复现。
• [x] 先做小实验： 固定数据和 backbone，比较 direct $\bm{x}$、direct $\bm{v}$、$\bm{v}$ + analytic skip、连续可学习 target。
• [ ] 直接照搬 2B JiT-G。

Reusable Ideas

1. 把 output parameterization 当成架构设计。 不要只问 loss 等价不等价。
2. 让网络学习结构化残差，而把高维恒等项解析传递。
3. 用强 bottleneck 做机制探针。 若目标真在低维结构附近，它应对 bottleneck 更稳健。
4. 分辨率 scaling 时固定 token 数。 适合先验证环境维与模型宽度的解耦。
5. 二维因子实验。 把 prediction space 与 loss weighting 交叉枚举，是很好的研究设计模板。
6. 端点权重必须显式写成实现公式。 理论分母和 clipped denominator 不能混为一谈。
7. 报告选择预算。 CFG、EMA、早停本质上都是模型选择的一部分。

Code Verification

Verified Version

• Repository: LTH14/JiT
• Checked commit: cbc743a2ada5e9762697da2c83f8c4f8379e8c17
• Commit date: 2025-12-07
• License: MIT
• 论文说明原实验使用 JAX + TPU；公开仓库是 PyTorch + GPU 重实现。

Claims $\rightarrow$ Code Mapping

Claim	论文描述	代码实现	状态
网络直接输出 $\bm{x}$	JiT 输出 clean image	model_jit.py 最终输出 3-channel unpatchified image；denoiser.py 命名为 x_pred	完全匹配
使用 $\bm{v}$-loss	从 $\bm{x}_{\theta}$ 转成 $\bm{v}_{\theta}$ 后 MSE	v_pred=(x_pred-z)/clamp(1-t)，对 target 做同样变换	完全匹配，含端点裁剪
Logit-normal time	$\mu=-0.8,\sigma=0.8$	sample_t() 先采正态再 sigmoid	完全匹配
分辨率噪声缩放	256/512 使用 1/2	README 与参数 noise_scale 一致	完全匹配；1024 脚本缺失
Low-rank patch embed	$q\rightarrow d'\rightarrow d$，无非线性	patch-stride Conv2d + $1\times1$ Conv2d	完全匹配
RoPE + qk-norm	attention 使用 RoPE 和 RMSNorm q/k	Attention.forward() 中先 q/k norm 再 RoPE	完全匹配
In-context class tokens	32 token，从中途 block 插入	B/L/H 分别从 4/8/10 插入	完全匹配
H/G dropout	中间一半 blocks，rate 0.2	代码支持按参数在中间一半启用	H 匹配；G 未公开
50-step Heun	最后一步 Euler	generate() 前 49 个区间用 Heun，最后 Euler	完全匹配
CFG interval	默认实验区间 $[0.1,1]$	_forward_sample() 中 interval mask	完全匹配
JiT-G	40 层、1664 hidden	JiT_models 只有 B/L/H	未找到
JiT-B/64 / 1024	patch 64，noise scale 4	model registry 和 evaluator 不支持 1024	未找到
九种消融	三输出 $\times$ 三 loss	公开 Denoiser 只有 $\bm{x}$-pred + $\bm{v}$-loss	未找到

Hyperparameter Consistency

超参数	论文	代码/README	判断
global batch	1024	8 GPU $\times$ 128	一致
learning rate	$2\times10^{-4}$	base LR $5\times10^{-5}$ 线性缩放到 $2\times10^{-4}$	一致
Adam betas	0.9, 0.95	AdamW 0.9, 0.95	一致；weight decay 为 0 时等价于 Adam
warmup	5 epochs	5	一致
LR schedule	constant	constant	一致
weight decay	0	0	一致
time sampler	$\mu=-0.8,\sigma=0.8$	P_mean=-0.8, P_std=0.8	一致
denominator clip	0.05	t_eps=5e-2	一致
EMA candidates	0.9996, 0.9998, 0.9999	默认只同时维护 0.9999 与 0.9996	部分匹配
CFG sweep	1.0 到 4.0，步长 0.1	README 给单个值，不含 sweep 脚本	部分匹配
200-ep B/16 FID 4.37	需模型特定最优 EMA/CFG	README 示例 CFG 2.9，默认 first EMA 0.9999；社区复现使用 EMA 0.9996、CFG 3.6	默认脚本不足
JiT-G early stop	约 320 epochs	无 G 配置	不可验证

Implementation Quality

优点：

• 核心模型短小，prediction/loss 逻辑直接；
• 固定随机种子按 distributed rank 偏移；
• BF16 autocast、DDP、EMA、在线 FID 均已接入；
• environment 锁定 Python 3.10、PyTorch 2.5.1、CUDA 12.4 等主要版本；
• 发布了 ImageNet 256/512 FID stats。

风险：

• 没有单元测试或集成测试；
• attention 实现硬编码 .cuda()，设备可移植性有限；
• cudnn.benchmark=True，没有启用完全确定性算法；
• 评测硬编码使用第一个 EMA，必须手工调换 decay 才能复现某些表格；
• evaluator 只支持 256/512，不支持论文中的 1024；
• 仓库历史很短，公开代码不是产生论文主结果的原始 JAX 版本；
• 核心消融、JiT-G 和完整模型选择脚本缺失；
• 截至 2026-07-28，公开 issue 报告 checkpoint URL 不可用。

Reproducibility Rating

6/10：可复现主算法，但不能完整复现论文。

有 8 张 H200/H100、ImageNet 和定制 torch-fidelity 时，可以训练 B/L/H 主配置；但要完整审计核心 claim，还缺：

• 九种输出/损失消融；
• JiT-G 与 JiT/64；
• 原始 JAX/TPU 代码；
• EMA/CFG sweep 自动化；
• 稳定可用的 checkpoint；
• 多 seed 结果。

Final Research Taste

创新性 (Novelty): ★★★★☆ (4.0/5)

• 旧 target 的新问题化方式很强；
• 简单得让人惊讶，但并非新生成理论。

严谨性 (Rigor): ★★★☆☆ (3.2/5)

• 代数和内部消融扎实；
• 缺跨架构因果、多 seed、正式理论和完整复现资产。

影响力 (Impact): ★★★★☆ (4.2/5)

• 很可能让 prediction target 重新成为高维 diffusion 的核心设计变量；
• 对科学数据和无 tokenizer 场景尤其有启发；
• JiT 本身是否成为长期主流仍取决于 text conditioning、训练成本和更强 benchmark。

Final Verdict

这篇论文在高维像素 diffusion 上通过“直接预测干净数据 + 强 bottleneck 大 patch ViT”把一个看似需要复杂像素架构的问题化约为输出参数化问题；相比 SiD2/PixelFlow，它更简单、更通用、单步计算更低，相比 RAE 它不依赖预训练 tokenizer，但绝对 FID、理论闭合度和完整复现性仍弱于最强系统。按 CVPR 标准，这是一个值得接收、很可能启发后续理论和跨域工作的强概念论文。

Follow-Up

• [ ] 复现 JiT-B/16，并固定同一 checkpoint 扫描 EMA/CFG sensitivity。
• [ ] 实现九种 prediction/loss 组合，至少跑三个 seed。
• [ ] 加入 direct $\bm{z}_t$ skip 的 $\bm{v}$-prediction baseline。
• [ ] 比较 bottleneck rank、patch size、环境维/内在维比例。
• [ ] 阅读《Revisiting Diffusion Model Predictions Through Dimensionality》并对照 k-Diff 理论。
• [ ] 在天气场或其他高维科学数据上测试 $\bm{x}$-prediction 是否保持同样优势。