何凯明团队提出JiT方法，用最朴素的Vision Transformer在像素级空间实现高质量图像生成
在人工智能生成内容迅猛发展的今天，去噪扩散模型已经成为图像生成领域的主流技术。从最初的DDPM到如今的Stable Diffusion，这些模型通常不是直接预测干净图像，而是预测噪声或噪声相关量。但麻省理工学院的Tianhong Li和Kaiming He的最新研究提出了一个根本性问题：「我们是否走错了方向？」
流形假设：干净数据与噪声的本质区别
根据机器学习中的流形假设，自然数据实际上存在于一个高维空间中的低维流形上。简单来说，虽然一张图片可能有数百万个像素点，但其有效信息可能只集中在某个低维子空间中。
这一假设揭示了关键问题：干净图像（x）位于低维流形上，而噪声（ε）和流速（v）则分布在整个高维空间中。当神经网络需要预测这些高维噪声量时，就需要足够大的容量来保留所有噪声信息。相比之下，预测干净图像只需要网络保留低维有用信息，同时过滤掉噪声。
"这一发现解释了为什么在低维潜在空间中，传统的噪声预测方法能够工作，但在高维像素空间中会遭遇灾难性失败。"研究团队在论文中表示。
传统扩散模型的演变与局限
回顾扩散模型的发展历程，2015年Jascha Sohl-Dickstein等人首次提出了扩散模型的基本思想，但直到2020年DDPM的出现才真正引起广泛关注。DDPM的一个关键发现是让网络预测噪声（ε-prediction）比预测干净图像能产生更好的生成质量。
随后，流动匹配方法引入了流速预测（v-prediction），将扩散模型与基于流的方法联系起来。如今，实践中的扩散模型通常预测噪声或噪声相关量。
然而，这些方法在高维像素空间中面临严重挑战。研究团队通过一个简单的玩具实验证明了这一点：当观测维度增加时，ε-prediction和v-prediction方法会迅速失效，而x-prediction仍能保持稳定性能。
JiT方法：简单而强大的解决方案
受到这些发现的启发，研究团队提出了"Just image Transformers"（JiT）方法——本质上就是直接在像素块上应用标准的Vision Transformer，并采用x-prediction目标。
JiT的设计极其简洁：
-
将图像划分为非重叠的块（如16×16或32×32像素）
-
通过线性嵌入层将每个块投影为token
-
使用标准的Transformer块处理token序列
-
输出层将每个token投影回图像块空间
这种方法不依赖任何预训练、标记器或额外损失函数，实现了完全自包含的生成范式。
预测目标与损失空间的组合分析
研究团队系统地分析了九种不同的预测目标和损失空间组合（如下表所示），发现在高维设置下，只有x-prediction能够稳定工作。
| 预测目标 | x-loss | ε-loss | v-loss |
|---|---|---|---|
| 「x-pred」 | 10.14 | 10.45 | 8.62 |
| 「ε-pred」 | 379.21 | 394.58 | 372.38 |
| 「v-pred」 | 107.55 | 126.88 | 96.53 |
表1：在ImageNet 256×256上使用JiT-B/16的不同组合的FID结果（越低越好），红色标注表示灾难性失败
值得注意的是，即使使用v-loss（流速损失），只要网络直接预测干净图像（x-prediction），模型仍能取得良好性能。这表明「损失加权本身并不是解决问题的关键」，网络预测的目标空间才是决定性因素。
瓶颈设计的反直觉优势
更令人惊讶的是，研究团队发现「在网络中引入瓶颈结构甚至是有益的」。通过将线性块嵌入层替换为低秩线性层对，即使将维度降低到极小的16维，模型也不会出现灾难性失败，反而在一定范围内能提升生成质量。
这一发现与经典流形学习观点一致：瓶颈结构鼓励网络学习数据的内在低维表示，过滤掉无关噪声。
高分辨率生成的突破性表现
JiT方法在高分辨率图像生成上展现了显著优势。研究团队在ImageNet 512×512甚至1024×1024分辨率上进行了实验，使用相应的大块大小（32×32和64×64），结果令人印象深刻。
| 分辨率 | 模型 | 块维度 | 参数量 | FID |
|---|---|---|---|---|
| 256×256 | JiT-B/16 | 768 | 131M | 4.37 |
| 512×512 | JiT-B/32 | 3072 | 133M | 4.64 |
| 1024×1024 | JiT-B/64 | 12288 | 141M | 4.82 |
表2：JiT在不同分辨率下的表现，尽管块维度显著增加，模型仍能保持稳定性能
尤为重要的是，这些模型在参数和计算量上几乎没有增加，仅通过调整块大小就适应了不同分辨率，避免了传统方法中计算成本随分辨率平方增长的问题。
与现有方法的对比
研究团队将JiT与当前主流方法进行了全面对比，结果显示JiT在保持简洁性的同时，实现了有竞争力的性能。
| 方法 | 预训练 | 参数量 | FID↓ |
|---|---|---|---|
| DiT-XL/2 | SD-VAE+VGG | 724M | 2.27 |
| SiT-XL/2 | SD-VAE+VGG | 724M | 2.06 |
| ADM-G | 无 | 559M | 7.72 |
| 「JiT-H/16」 | 「无」 | 「953M」 | 「1.86」 |
| 「JiT-G/16」 | 「无」 | 「2B」 | 「1.82」 |
表3：ImageNet 256×256上的结果对比，JiT在不依赖任何预训练组件的情况下达到领先水平
JiT的优势在于其「完全自包含」的特性——不需要预训练的VAE标记器、不需要对抗性损失、不需要感知损失，也不需要表示对齐的预训练模型。
通用Transformer架构的优势
JiT方法的另一个关键优势是能够直接受益于通用Transformer架构的进步。研究团队集成了多种最初为语言模型开发的技术：
-
SwiGLU激活函数
-
RMSNorm归一化
-
RoPE位置编码
-
qk-norm注意力规范化
这些通用组件的引入进一步提升了模型性能，证明了任务无关架构设计的价值。
实际训练与推理算法
JiT的训练和推理过程异常简洁：
「训练过程」：
-
采样时间步t、干净图像x和噪声ε
-
计算噪声样本z_t = t·x + (1-t)·ε
-
计算目标流速v = x - ε
-
网络直接预测干净图像x_θ = net_θ(z_t, t)
-
推导预测流速v_θ = (x_θ - z_t)/(1-t)
-
最小化‖v_θ - v‖²
「推理过程」：
-
从噪声分布初始化z_0
-
通过ODE求解器逐步去噪：dz_t/dt = v_θ(z_t, t)
-
使用50步Heun方法获得最终图像
这种简洁性使得JiT易于实现和调整，为后续研究提供了坚实基础。
跨领域应用的潜力
JiT方法的「领域无关性」为其在科学计算等领域的应用开辟了道路。在蛋白质结构预测、分子设计、气象模拟等任务中，设计有效的标记器往往很困难，而JiT的直接像素级（或对应领域的原始数据）处理方法显示出独特优势。
研究团队强调："通过最小化领域特定的设计，我们希望起源于计算机视觉的通用'Diffusion+Transformer'范式能够找到更广泛的应用领域。"
技术细节与优化策略
在实际实现中，研究团队采用了几项重要技术：
「时间步分布」：使用logit-normal分布采样时间步t，通过调整均值μ控制噪声水平。
「高分辨率适配」：对于512×512和1024×1024分辨率，分别将噪声尺度放大2倍和4倍，保持信噪比稳定。
「上下文类别条件」：在序列前添加多个类别标记（默认32个），并在特定Transformer块开始引入上下文条件。
「训练优化」：对大模型使用dropout防止过拟合，在监控到FID退化时实施早停。
生成样本展示
JiT在ImageNet数据集上生成的样本显示出了「出色的多样性和真实性」。从鸟类到动物，从自然场景到人造物体，模型都能捕捉到关键特征并生成连贯图像。
这些样本是在达到报告FID值的CFG参数下生成的，没有经过精心挑选，体现了模型的稳定可靠性。
未来展望与研究意义
JiT方法的意义不仅在于其当前实现的性能，更在于它为生成式AI研究提供了新的方向：
「理论层面」，这项工作强调了神经网络容量分配的重要性，提醒我们在追求更复杂概率公式的同时，不应忽视模型的基本能力限制。
「实践层面」，JiT的简洁性和有效性使得扩散模型能够更易于应用到缺乏大量预训练资源的领域。
「生态层面」，减少对特定领域设计的依赖有助于构建更加统一和可迁移的AI基础设施。
结语
麻省理工学院的这项研究通过回归基本原理，揭示了去噪扩散模型的一个根本性洞见：「让网络直接预测干净数据，而非噪声」。基于这一理念提出的JiT方法，用极其简洁的设计在像素级空间实现了强大的生成性能。
这项工作的核心价值在于其「简约性与有效性」的完美结合。在这个越来越复杂的AI时代，有时候最重要的突破恰恰来自于回归基础、重新思考根本假设。
正如论文作者所总结的："噪声与自然数据有着本质区别。多年来，扩散模型的发展主要关注概率公式，而对神经网络能力（和限制）的关注较少。然而，神经网络的能力不是无限的，它们应该更好地利用自身容量来建模数据而非噪声。"
JiT的出现标志着生成式AI研究的一个重要转折点——在追求更高性能的同时，不应忽视模型的简洁性、通用性和可解释性。这一理念将为实现更可持续、更可访问的人工智能技术奠定基础。
以上，既然看到这里了，如果觉得不错，随手点个赞、在看、转发三连吧，如果想第一时间收到推送，也可以给我个星标⭐～谢谢你看我的文章，我们，下次再见。