扩散概率模型诞生之初,其核心物理意义与设计初衷是迭代降噪:模型学习从被噪声污染的退化图像中,逐步还原出无噪、真实的干净原图,对应的建模方式即为x-pred。然而在后续技术迭代中,两次关键范式革新彻底偏离了“直接降噪”的原始目标,逐步形成如今的行业主流,也埋下了高维建模失效的隐患:
ϵ-pred(纯噪声预测):由DDPM开创性提出,放弃直接预测干净图像,转而拟合扩散过程中人为添加的高斯噪声。该范式大幅提升了生成图像的细节丰富度与清晰度,解决了早期x-pred生成图像模糊的问题,迅速成为扩散模型的标准训练范式,被后续绝大多数像素扩散、隐空间扩散模型沿用;
v-pred(流速度预测):后续研究将扩散模型与流匹配(Flow Matching)方法统一,提出流速预测范式。其预测目标是融合了干净原图与高斯噪声的复合变量,兼具数据与噪声双重属性,进一步统一了生成模型的数学框架,成为近年高性能扩散模型的主流选择。
研究背景
过往大量研究仅聚焦于三种预测目标的数学等价变换与损失权重调优,普遍认为只要调整损失函数的加权系数,x-pred、ϵ-pred、v-pred可以实现完全等价的建模效果,模型性能差异仅来自训练超参。因此学界长期忽视预测目标本身的物理属性、空间分布与建模难度差异,从未系统性探究三者在高维空间下的固有缺陷与适用边界,这也是高分辨率像素扩散模型长期难以突破的核心理论盲区。
本文的核心理论支撑源自机器学习与流形学习领域的经典自然数据流形假设:真实世界的图像、文本、分子等自然数据,虽然观测形式为超高维空间向量(如像素空间),但其有效语义、结构、纹理信息仅分布在高维空间内的一个低维、光滑、连续的子流形上。与之相对,人工添加的高斯噪声、以及由噪声与原图融合得到的流速变量,不具备任何结构化语义约束,均匀、随机地分布在整个高维观测空间中,完全脱离数据流形。
基于该假设,三种预测目标的建模难度、信息拟合需求存在天壤之别,形成了不可逾越的建模鸿沟:
- x-pred 预测干净原图:模型无需拟合高维空间的全部随机信息,仅需要捕捉、还原低维数据流形上的有效结构化信息,同时自动过滤流形外的无效噪声信息。该任务对网络容量要求极低,即便参数有限、存在结构瓶颈的轻量化网络,也能精准完成降噪还原任务;
- ϵ/v-pred 预测含噪变量:噪声与流速变量遍布全高维空间,无低维流形约束。模型必须拟合高维空间中所有随机、无规律的信息,对网络隐藏层维度、参数量、表达能力提出极高要求。一旦输入补丁维度超过网络容量,模型无法承载海量高维信息,会出现梯度失效、拟合崩溃、生成完全失真的灾难性失效现象。
现有的方案存在三种缺陷:
- 隐空间扩散高度依赖外部预训练组件:当前主流的隐式扩散模型(如LDM、DiT、SiT)均依赖预训练VAE分词器实现像素空间到低维隐空间的压缩,同时需要感知损失、对抗损失、自监督对齐损失辅助优化生成细节。整套系统无法独立训练、自包含建模,高度依赖外部预训练权重与领域先验,跨领域泛化能力严重受限;
- 传统像素空间扩散存在严重设计冗余:为勉强适配ϵ/v-pred的高维拟合需求,传统像素扩散模型普遍采用密集卷积、极小图像补丁、超大通道数、超长跳跃连接等冗余设计,以此弥补网络容量不足的缺陷。这类方案计算成本极高、算力开销巨大,且只是“缓解”高维退化问题,无法从机理上彻底解决;
- 高维建模存在固有认知误区:行业长期固化“输入维度必须匹配网络隐藏层维度”的认知,认为高维像素补丁输入必须配备超大容量网络,否则必然拟合失败。这一误区导致高分辨率像素扩散模型普遍臃肿、算力成本高昂,且始终无法突破高维空间的建模瓶颈。
基于以上难题,文章提出了JiT模型。

准备工作
为厘清三种预测目标的本质区别,本文从流匹配与扩散模型的统一ODE框架出发,严格推导了预测空间(网络直接输出)与损失空间(监督信号计算)的全部组合关系,穷尽9种合法建模组合,构建了一套完备、自洽的扩散模型数学理论体系。通过严格的公式推导与逻辑证明,首次明确:不同预测目标并非简单的权重缩放差异,而是建模空间、信息拟合范围、任务难度的本质差异,是解释高维扩散模型性能分化的核心底层逻辑。
基础扩散流模型
本文摒弃传统DDPM的方差爆炸调度,采用更简洁、可解释性更强的线性流匹配调度范式构建扩散过程。对于任意时间步 ,对干净图像施加线性插值加噪,得到对应的含噪样本 ,实现从纯干净图像到纯噪声的平滑过渡:
式中 代表原始无噪干净图像,是唯一具备流形约束的有效数据; 为标准高斯白噪声,无任何结构化信息,均匀分布在全维度空间。
基于含噪样本的时序演化规律,定义时间流速度v为含噪样本对时间的一阶导数,表征扩散过程中图像的动态变化速率,通过公式推导可得到贯穿全文的核心恒等式,建立起干净图像、噪声、流速三者的强关联:
在模型推理采样阶段,生成过程等价于求解一阶常微分方程(ODE),通过持续迭代更新样本状态,实现从纯噪声到真实图像的逆向生成:
训练完成后,从初始纯噪声分布 开始积分,直至时间步 最终得到生成图像。本文统一采用50步Heun求解器完成ODE数值求解,在保证生成质量的同时兼顾推理效率,是实验的统一基准设置。
三大预测目标的转换关系
在统一的ODE与流匹配框架下,网络的最终推理目标均为流速,但网络训练时的直接输出对象可自由选择,分为三大范式:直接预测干净图像(x-pred,输出 )、直接预测高斯噪声(ϵ-pred,输出 )、直接预测流速(v-pred,输出 )。依托基础恒等式,三者可实现精准的相互推导转换,完整的范式差异与转换逻辑如下表所示,清晰体现不同预测方式的建模逻辑区别:

损失函数定义
本文通过穷尽9种预测目标+损失空间的组合消融实验,最终筛选出最优建模组合:x-pred(网络直接预测干净图像)+ v-loss(基于流速空间计算监督损失),该方案完美兼顾训练稳定性、收敛速度与最终生成画质:

论文严格证明:9种建模组合不存在任何数学等价性,低维场景下性能差异微弱可被超参掩盖,但在真实高维像素建模场景中,不同组合的性能差距会被极致放大,呈现“有效与完全失效”的两极分化,预测目标的选择成为决定模型成败的核心关键。
时间采样分布
为平衡不同噪声等级样本的训练权重,避免模型偏向单一噪声区间,本文训练阶段摒弃均匀时间采样,采用logit-normal时间采样策略,令时间步满足 。实验最优默认超参为 ,其中可灵活调整整体噪声等级(越小、整体噪声越强),控制时间步分布的集中程度,有效提升模型对全噪声区间的拟合能力。
模型框架



依托x-pred的核心理论优势,本文设计了极致极简、无领域偏见、完全自洽的JiT(Just image Transformers)像素扩散框架。该架构彻底摒弃扩散模型领域的一切专属优化与冗余设计,不使用VAE分词器、无需任何预训练权重、无对抗/感知/对齐等辅助损失、无卷积归纳偏置,纯粹基于原生ViT结构实现高分辨率图像生成,真正做到“回归基础、极简通用”。
基础架构设计
大尺寸图像补丁分块:打破传统像素扩散小补丁设计,采用超大非重叠像素补丁适配不同分辨率图像,大幅降低序列长度、减少计算冗余。具体配置:256×256分辨率采用16×16补丁(单补丁768维),512×512分辨率采用32×32补丁(单补丁3072维),1024×1024超高分辨率采用64×64补丁(单补丁12288维),统一序列长度、保证计算效率;
可降维补丁嵌入层:将原始高维像素补丁通过线性投影完成嵌入,搭配标准位置编码注入空间位置信息。创新性支持低秩瓶颈降维设计,可主动压缩嵌入维度、过滤高维冗余噪声信息,贴合流形学习的低维表达核心;
标准Transformer主干网络:完全沿用原生ViT编码器结构,无自定义卷积模块、无特殊跳跃连接,仅通过AdaLN-Zero归一化层注入扩散时间步信息与图像类别条件,实现条件生成;
原生像素还原输出:网络末端线性层直接输出对应补丁的干净像素预测值,严格执行x-pred范式,不做任何中间变换,从结构上保证模型直接学习降噪还原任务。
进阶优化组件
为进一步突破基础模型的性能上限,本文引入一系列通用跨领域Transformer优化组件,所有组件均无图像生成领域专属偏见,广泛适用于NLP、CV通用任务,可无缝迁移适配:SwiGLU激活函数提升非线性表达能力、RMSNorm归一化优化训练稳定性、RoPE旋转位置编码增强长距离空间依赖建模、qk-norm缓解注意力梯度爆炸问题、多上下文类别令牌强化类别条件约束,在不增加领域设计的前提下稳步提升生成性能。
核心算法流程
训练流程
基于logit-normal分布采样训练时间步,生成标准高斯噪声;
依托线性加噪公式,计算当前时间步的含噪图像样本;
将含噪样本与时间步输入JiT网络,直接预测干净像素图像;
通过x-pred转流速公式推导预测流速,与真实流速计算v-loss均方误差,反向传播更新网络参数,完成单步训练。
采样推理流程(Euler数值求解,轻量化高效推理)
初始化纯高斯噪声样本,作为生成起点;
逐时间步输入网络,预测干净图像并推导得到当前流速;
基于ODE迭代公式更新样本状态,逐步去除噪声、还原图像,直至完成全时间步迭代,输出最终生成结果。
实验数据


本文通过海量、分层的消融对照实验,从预测目标、网络容量、结构设计、分辨率尺度、训练超参等多个维度,系统性验证了x-pred范式的核心优势,推翻多项行业长期固有认知,得到一系列极具突破性的创新结论,彻底重构了高维扩散建模的底层逻辑:
预测目标是高维建模的决定性因素
低维小补丁场景:三者性能近似等价:在64×64小分辨率、低维补丁建模场景下,网络容量远大于输入维度,充足的拟合能力可弥补ϵ/v-pred的建模缺陷,三种预测范式的FID指标差异极小、性能基本持平。这也是过往研究误判三者完全等价、忽视预测目标本质差异的核心原因;
高维大补丁场景:性能两极分化:在256×256、512×512等高分辨率真实场景下,补丁维度大幅提升、远超常规网络隐藏层容量,ϵ-pred与v-pred范式因无法拟合全高维空间信息,出现灾难性性能失效,FID指标飙升至数百、生成图像完全失真混乱;唯有x-pred范式依托流形低维建模优势,始终保持稳定、高质量的生成效果;
超参调优无法弥补范式固有缺陷:论文验证,通过调整损失权重、偏移噪声等级、优化时间采样分布等主流调参手段,仅能小幅优化x-pred的生成效果,完全无法修复ϵ/v-pred在高维场景下的拟合崩溃问题。彻底证明:高维建模失效是预测目标本身的机理缺陷,而非训练超参问题。
网络容量无需匹配输入维度
业界长期存在固化认知:高维输入特征必须配备同等或更大维度的网络隐藏层,否则必然出现欠拟合、生成失效。
本文通过超大维度补丁实验彻底推翻该结论:JiT模型在单补丁3072维(512分辨率)、12288维(1024分辨率)的超高输入维度下,仅使用768维隐藏层的基础网络,依然可以实现稳定训练、高质量生成。核心机理在于x-pred仅学习低维流形有效信息,主动舍弃高维冗余噪声,实现了网络容量与输入维度的完全解耦。
瓶颈降维设计有益无害
论文针对补丁嵌入层开展低秩瓶颈消融实验,将原始768维高维嵌入特征,压缩至16/32/128/512维等不同低维瓶颈维度。实验结果表明:适度的维度瓶颈压缩不仅不会造成信息丢失、性能下降,反而能够小幅优化FID指标、提升生成画质。该结果完美印证经典流形学习理论:自然图像的有效语义信息是低维的,高维像素空间中的大部分维度均为无效冗余信息与噪声,瓶颈结构可主动过滤干扰、提纯有效特征。
模型具备极佳扩展性
JiT模型具备极强的尺度扩展性与分辨率泛化能力:从B(基础)、L(大)、H(超大)到G(巨型)模型尺度逐级放大,参数量与表达能力稳步提升,FID指标持续下降、生成细节与语义合理性显著优化。同时,得益于统一的大补丁分块策略,256分辨率与512分辨率模型的序列长度、计算量、参数量基本持平,彻底解决了传统高分辨率生成模型算力爆炸、成本倍增的行业痛点。最优的JiT-G/32模型在512×512分辨率ImageNet数据集上FID低至1.78,性能媲美一众依赖预训练的顶尖隐空间扩散模型。
重点实验数据


与像素空间扩散模型对比
相较于ADM-G、RIN、PixelFlow、PixNerd等经典像素空间扩散模型,JiT框架具备算力成本更低、模型参数更精简、生成性能更优异的三重优势。传统像素扩散模型普遍依赖超大参数量、密集卷积、复杂跳跃连接,算力开销极高;而JiT以极简原生Transformer为核心,无冗余结构、无预训练依赖,在大幅降低GFLOPs与训练成本的同时,量化指标与视觉效果全面超越同类型像素生成模型。
与隐空间扩散模型对比
相较于DiT、SiT、REPA等当前主流的高性能隐空间扩散模型,纯原生无预训练的JiT基础模型性能略有差距,核心原因是隐空间模型依托VAE压缩、自监督预训练、多辅助损失获得了更强的先验约束。但随着模型尺度放大,JiT-G等超大模型可追平甚至超越部分预训练隐空间模型。更关键的是,JiT具备完全自包含、无领域偏见、可跨领域迁移的核心优势,突破了隐空间模型依赖分词器、无法适配无分词器场景的固有局限。
跨分辨率泛化实验
论文开展跨分辨率迁移对照实验,充分验证大分辨率建模的信息优势:512分辨率训练的JiT模型,降采样适配256分辨率任务时,依然能够保持极佳的生成性能,无明显画质与指标损失;反之,256分辨率训练的模型上采样适配512高分辨率生成时,会严重丢失高频纹理、边缘细节等关键信息,FID显著劣化。该现象证明:高分辨率原生建模能够捕捉更完整、更丰富的图像底层与语义信息,具备更强的信息完整性与任务泛化性。
小小总结
来自 MIT 李天虹与何恺明的论文指出,现有扩散模型普遍预测噪声或流速度,但基于数据流形假设,噪声和流速分布在整个高维空间,对网络容量要求很高;相比之下直接预测干净图像(x‑pred)只需学习低维流形上的有效信息,即便容量有限的网络也可以完成任务,作者据此提出无需分词器、预训练与额外损失的极简 JiT(Just image Transformers)纯像素 Transformer 扩散模型,在 ImageNet 的 256/512/1024 分辨率上取得不错效果,证实 x‑pred 能够解决高维像素空间中/v‑pred 出现的灾难性退化,还发现适度瓶颈降维反而有益,为蛋白质、气象等难以设计分词器的原生高维数据提供了一套自包含的 “扩散 + Transformer” 通用范式。
原文链接:https://mp.weixin.qq.com/s/SbDd7DV-P3R2LZG7XRRZAQ
|