跳到正文
北京时间
原文
HuggingFace Daily Papers(社区热门论文)·· 2026-05-21精选AI 评分70

RiT:在表示空间中使用原生扩散变换器已足够

RiT: Vanilla Diffusion Transformers Suffice in Representation Space

AI 导读

本研究探讨预训练表示空间在流匹配学习中的优势。比较像素、SD-VAE与DINOv2特征后发现,尽管像素与DINOv2的内在维度相近,但DINOv2在几何统计特性(如有效秩、协方差条件等)上表现更优,使回归过程更稳定。基于此,我们提出了表示图像变换器(RiT),它使用冻结的DINOv2特征,通过x-prediction目标训练一个原生扩散变换器。在ImageNet 256×256生成任务上,RiT性能优于参数量更多的DiT^DH-XL模型,且生成的常微分方程仅需少量步骤即可高效求解。

推荐理由

这篇论文没发明新架构,但通过剖析DINOv2特征的统计属性,证明简单结构在表示空间也能做出SOTA,对做图像生成的人来说是个省钱省参数的好思路。

正文 · AI 翻译
摘要

基于预测的流匹配——即回归干净数据点而非环境速度——已知能够在像素空间中有效利用低维流形结构[18]。我们探究一个预训练表示空间,在包含具有可比内在维度的低维数据流形的同时,是否能为流匹配学习提供更有利的分布。通过沿四个几何轴比较像素、SD-VAE 和 DINOv2 特征,我们发现像素和 DINOv2 具有几乎相同的内在维度(均为 ),但 DINOv2 表现出更高的有效秩、更好的协方差条件性、更低的超额峰度以及更低的流形上插值误差;SD-VAE 潜变量始终处于中间水平,这表明优势源于表示学习目标而非单纯的压缩。这些统计特性使得流匹配回归具有良好的条件性,并消除了先前 DINOv2 扩散方法所使用的专用预测头或黎曼传输的需求。我们提出了表示图像 Transformer(RiT):一种在冻结的 DINOv2 特征上通过预测训练的普通扩散 Transformer,仅增加了维度感知噪声调度和联合 [CLS]-块建模。在 ImageNet 上,RiT 在无引导条件下达到 FID 1.45,在无分类器引导条件下达到 FID 1.14,以更少的参数(676M vs. 839M)超越了 DiT-XL。由此产生的 ODE 在粗离散化下可高效求解:在无分类器引导下,Heun 步数 已达到 FID 2.0,步数 达到 1.25,无需蒸馏或一致性训练。代码见 https://github.com/lezhang7/RiT。

1 引言

流匹配 [19, 7] 学习一个速度场,该速度场沿线性路径将高斯噪声输送到数据。当数据集中在低维流形附近时,预测——将网络参数化为输出干净数据点而非环境空间速度——将回归目标置于该流形上,正如 JiT [18] 在像素空间中所证明的那样。一个自然的问题是,预训练的表征空间在包含一个内在维度相当的数据流形的同时,是否能为学习流匹配速度场提供更有利的分布。

沿四个几何轴比较像素、SD-VAE [23] 和 DINOv2 [21] 特征,我们发现像素和 DINOv2 具有几乎相同的内在维度(均为 ),但相对于 嵌入该流形的方式不同。像素流形是各向异性的,具有强非高斯性的逐坐标边缘分布,并且存在穿越低密度区域的线性弦。DINOv2 特征表现出近各向同性的方差、近高斯性的逐坐标边缘分布 [38],以及近似在流形上的线性插值。这些是边缘性质,而非联合性质:DINOv2 特征仍然集中在 维流形上,但每个坐标向 的传输路径短且条件良好。

第 2 节量化了这些差距:与像素相比,DINOv2 实现了更高的有效秩、更好的协方差条件数、更低的超额峰度以及更低的流形上插值误差。SD-VAE 的潜在表示始终介于两者之间,表明这一优势源于表征学习目标而非单纯的压缩。这些分布优势与 DINOv2 在流形外中间状态处的特定病理现象并存:逐 token 的 LayerNorm 固定了 ,因此线性流匹配路径会穿越编码器从未输出的环境区域,并且在此类 处的 目标会获得一个较大的径向分量。

针对这种径向模糊性,目前主流的应对方案是架构层面的。RAE [44] 采用了一种专门的宽预测头(DDT [37]),置于 -预测之上,并配合一个 ViT 解码器,将 DINOv2 特征映射回像素。同期工作 [16] 将这种现象称为几何干扰,并用基于范数集中球面上的黎曼流匹配替代了欧几里得传输。这两种修改都增加了架构或传输路径的复杂性。

我们采用了一种目标端的替代方案:-预测。在这种参数化方式下,网络回归的是 ,它天然位于数据流形上,因此径向模糊性在网络的输出端(其目标位于流形上)得到解决,而非在其输入端( 仍处于流形外)。这种重参数化本身并非新事物 [18];其在此处的有效性源于与 DINOv2 的各向同性逐坐标方差和近似高斯边缘分布的结合,这使得去噪回归的条件足够良好,以至于可以用一个普通的 DiT 来处理。我们将这种结合实例化为表征图像 Transformer(RiT)(第 3 节):一个在表征空间中通过 -预测流匹配训练的普通扩散 Transformer,并辅以维度感知的噪声调度和联合 [CLS]-patch 建模。正如 DiT 在 SD-VAE 潜空间上运行,RiT 在由冻结的编码器-解码器提供的表征空间上运行;我们使用了 RAE [44] 的冻结 DINOv2 编码器和 ViT 解码器。因此,RiT 直接对高维 DINOv2 特征分布进行建模,而无需为生成任务调整编码器。在 ImageNet 上,RiT 在无引导条件下达到了 FID 1.45,在无分类器引导条件下达到了 FID 1.14,以更少的参数超越了 DiT-XL。由此产生的 ODE 在少量 Heun 步骤内收敛,实现了 5 步 FID 2.0 和 10 步 FID 1.25(有引导),且无需知识蒸馏或一致性训练(第 4.3 节)。

2 流匹配的表征空间几何

Refer to caption
(a) PCA 谱
Refer to caption
(b) 优化条件
Refer to caption
(c) 流形上的插值
图 1:像素空间、SD-VAE 和 DINOv2 的流形分析。(a) PCA 谱:累计方差(上图)和对数尺度下的各分量方差(下图);衰减越平缓表示分布越均匀。(b) 传输路径上的条件数;DINOv2 在 处的条件数始终优于像素空间。(c) 插值重建 MSE;像素空间始终偏离流形,而 DINOv2 在整个过程中保持接近流形。

流形假设——即数据集中在低维曲面上——无论采用何种表示方式都成立。不同表示方式的差异在于,该流形相对于 的位置是否有利,从而决定传输路径是否较短,以及 ODE 是否能在少量步骤内高效求解。我们从四个互补维度刻画每种表示方式——内在维度(流形的真实自由度)、有效秩(方差在各方向上的均匀程度)、边际高斯性(每个坐标与 的相似度)以及流形上的线性插值(线性弦是否靠近数据)——每个维度都预测了使流匹配变得更容易或更困难的不同机制。我们在三个空间上对 10,000 张 ImageNet 图像进行了量化:(i) 原始像素(, ),(ii) DINOv2-Base 特征(, ),以及 (iii) SD-VAE 潜变量(, )——即潜扩散模型 [23] 所使用的预训练 VAE。像素空间和 DINOv2 具有相同的环境维度,从而可以直接进行几何比较;加入 SD-VAE 则可以将表示学习训练(以 DINOv2 的自监督学习为代表)的效果与通用压缩的效果分离开来。

内在维度是流形的真实自由度——即剥离环境冗余后,描述数据所需的独立方向数量。两个具有可比内在维度的空间,其流形的基础复杂度也相当;因此,流形匹配学习难度的任何差异,必然源于流形的定位方式,而非其规模大小。我们采用 TwoNN 估计器 [8],该方法在局部均匀性假设下,通过第二近邻与第一近邻距离之比进行最大似然估计(附录 C)。对 10 个独立的 5000 点子样本进行自助法重采样,得到像素空间为 ,DINOv2 空间为 ——两者几乎相同,1 维度的差距完全在组合估计标准差范围内。因此,两个空间共享本质上相同的基础流形维度;DINOv2 的优势,通过排除法可知,在于该流形相对于噪声的嵌入方式。

指标 像素 SD-VAE DINOv2
均值 0.923 0.467 0.114
中位数 0.958 0.228 0.083
0.0% 74.2% 98.7%
70.6% 86.7% 99.7%
媒体内容 · 前往原文查看
表 1:边缘高斯性。三个表示空间的超额峰度。

[Uncaptioned image]

媒体内容 · 前往原文查看
图 2:峰度分布。DINOv2 的边缘分布紧密集中在 (高斯分布)附近;SD-VAE 处于中间水平;像素空间则严重偏离。

[Uncaptioned image]

媒体内容 · 前往原文查看
图 3:跨类别插值。每组的上行:像素空间混合(出现鬼影伪影)。下行:在 DINOv2 表示空间中进行插值,然后通过 RAE 解码器解码回像素(实现平滑的语义过渡)。

有效秩量化了方差在主方向上的均匀分布程度。当所有方差集中于单一方向时(即高维空间中的细针状),有效秩等于1;当方差完全均匀分布时(即各向同性球体),有效秩等于环境维度。由于流匹配源本身是各向同性的,数据侧更高的有效秩意味着从噪声到数据的传输路径更短、更均匀。具体而言,基于归一化的PCA特征值[24],图1(a)绘制了各分量方差(对数尺度)及其累积值:像素空间前50个分量捕获了总方差的XX%,而DINOv2则为YY%。像素空间、SD-VAE和DINOv2的有效秩分别为A、B和C——像素空间与DINOv2之间存在D倍的差距。DINOv2的逐token LayerNorm通过构造进一步固定了方差,因此这种高有效秩将特征集中在一个近似各向同性的球壳附近,该球壳半径为R,包含d维数据流形[38, 16]。

优化条件反映了训练过程中不同方差方向能否被并行学习:条件良好的回归沿所有方向以可比速率收敛,而条件不良的回归则会过拟合高方差方向,同时忽视低方差方向。时间t处的流匹配本质上就是这样一种回归,其有效协方差在t=0时的纯噪声与t=1时的干净数据之间进行插值;因此,不良条件会传播到后期调度训练中。具体而言,在局部高斯近似下,Ahamed等人[1]证明回归协方差为Σ_reg;我们采用标准条件数κ作为诊断指标。图1(b)绘制了κ随t的变化:两个空间均从t=0附近(κ≈1,此时Σ≈I)开始,并随t→1单调增长至κ→∞。在t=0.9处(代表后期调度精细数据拟合),像素空间达到κ≈10^4,而DINOv2保持在κ≈10^2——存在两个数量级的差距,使得所有方差分量都能以可比速率被学习。与各向同性噪声相同的分布邻近性还收紧了后验分布,缩小了逐对速度目标的不可约方差——这是促成第4节中更快收敛的另一个独特机制。

边际高斯性衡量每个独立坐标的一维分布与高斯分布的接近程度。源分布在每个轴向上都是高斯的,因此数据侧每个维度上越接近高斯的边际分布,就能使该维度从噪声到数据的传输路径越短、行为越规整。我们使用每个维度的超额峰度,该值对于高斯分布为零,对于比高斯分布更重(易出现离群点)的尾部为正,对于更轻的尾部为负。如表1和图2所示,DINOv2的维度明显更接近高斯分布:98.7%的维度满足条件(相比之下SD-VAE为74.2%,像素为0%),其中位数低于像素,也低于SD-VAE。这仅反映了边际行为;下面的插值实验将探究联合几何结构。

流形上的线性插值。前三个轴总结了每个方向上的方差;最后一个轴探究联合几何结构。流匹配沿着直线路径传输样本,因此如果数据点之间的线性弦本身偏离了流形,中间状态也会偏离,导致速度目标定义不清。跨类别图像插值使这一点具体化:像素插值会产生鬼影伪影,这是路径穿过低密度空洞的特征,而DINOv2插值则产生平滑的语义过渡(图3)。我们通过往返重建误差来量化这一点(完整流程见附录C):每个中间帧——无论是通过像素混合获得,还是通过在DINOv2空间中进行线性插值后经RAE解码获得——都经过相同的DINOv2编码器-RAE解码器流水线[44],与输入之间的MSE衡量了偏离流形的距离。由于两种条件都经过相同的流水线,编码器-解码器的重建偏差是共享的;剩余的差距则单独表明该帧位于图像流形之上还是之外。像素帧产生的误差高于DINOv2帧(对比);图1(c)显示DINOv2在整个过程中始终接近流形,而像素则始终偏离流形。

摘要。Pixel 和 DINOv2 的内在维度几乎相同(均为 ),但 DINOv2 更适合流匹配学习:其有效秩更高、协方差条件更好、超额峰度更低、流形内插误差更小;SD-VAE 始终处于中间水平,这表明优势源于表征学习目标,而非单纯的压缩。这些特性预示着,DDT 头部、黎曼传输和更宽的主干网络并非获得竞争性性能所必需——我们在第 3-4 节中通过一个 vanilla DiT 和 -预测验证了这一预测。

3 RiT:用于表征空间扩散的 Vanilla DiT

Refer to caption
图 4:RiT 架构。冻结的 RAE 编码器/解码器(灰色)包围着一个通过 -预测训练的 vanilla DiT;[CLS] token 和 patch token 共享自注意力机制,并设有独立的 和 头部。

在第 2 节的几何特性指导下,我们实例化了表征图像 Transformer(RiT)(图 4):一个 vanilla DiT 主干网络(SwiGLU [26]、RMSNorm [43]、2D-RoPE [32]、QK-归一化 [11],以及遵循 JiT [18] 的上下文类别 token),采用针对 DINOv2 特征定制的训练方案。我们复用 RAE [44] 中冻结的带寄存器的 DINOv2 编码器 [21] 和 ViT 解码器,在像素与特征之间进行转换。编码器输出 patch token 和一个 [CLS] token ;两者经过投影、拼接并共同参与注意力计算,由独立的线性头部预测 和 。完整细节见附录 D。

流匹配预备知识。流匹配 [19, 7] 学习一个速度场,该场将噪声沿直线路径 传输至数据,其中 为纯噪声, 为干净数据;该路径的时间导数为 。标准的 -预测目标训练一个以时间步 和类别 为条件的网络,以回归该速度:

(1)

生成过程通过欧拉或休恩求解器,将学习到的常微分方程从 积分到 。

3.1 标准化特征上的 -预测

逐元素标准化。DINOv2 的逐 token LayerNorm 将每个 token 内部固定,但保留了跨数据集的逐通道方差异质性(跨通道分布,§4)。因此,在扩散之前,我们使用训练集上预先计算的统计量,将 patch token 和 [CLS] token 均标准化为零均值和逐元素单位方差:(类似地对于 ),这种对角预条件器 [1] 降低了数据协方差的条件数,并放松了 LayerNorm 对原始 DINOv2 特征施加的近似恒定范数约束 [16]。在解码前应用逆变换。此后,我们用 表示标准化后的特征。我们发现这一步是前提条件而非优化:在原始 DINOv2 特征上训练会完全发散(表 3)。

-预测。如 §2 所述,DINOv2 的范数集中意味着线性流匹配路径会经过编码器从未输出的环境区域;在此类流形外 上,-目标会获得一个与数据流形正交的径向分量——即 Kumar 和 Patel [16] 诊断出的几何干扰现象,他们通过使用 SLERP 路径的黎曼流匹配 [2] 解决了该问题。在 -预测下,网络必须拟合这个径向分量,因此将容量花费在范数方向而非切向(沿流形)方向上。我们通过改变输出参数化,更简单地解决了同一问题。

将 设为标准化后的 DINOv2 特征,-预测 [18] 改为直接输出 ,预测速度为 ;将其代入 -预测损失 (1) 可得

(2)
媒体内容 · 前往原文查看
表 2:-预测 与 -预测。FID-50K,ImageNet ,Heun 50 步,无引导。
80 轮 200 轮 400 轮
-预测 3.17 2.11 1.86
-预测 2.63 1.89 1.70

这等价于在重新加权后的 -预测损失(附录 B)。因此,-形式和 -形式作为损失函数是等价的,但它们给网络带来了不同的学习问题,因为输出参数化决定了实际拟合的是哪个函数。在 -预测下,网络必须拟合 :一个依赖于流形外 的目标,在 接近 时发散,并跨越整个环境空间。在 -预测下,网络拟合 :一个通过构造位于低维数据流形上且不显式依赖于 的目标。流形外的弦在网络输入处持续存在;在输出处,目标被限制在数据流形内,而非跨越整个环境空间。这种流形定位特性本身并非 DINOv2 所独有 [18];这里独特之处在于它与 DINOv2 的各向同性逐坐标方差和近高斯边缘分布相结合,这两者共同约束了目标并平滑了其对 的依赖关系(§2),因此一个普通的 DiT 就足够了。表 2 在经验上证实了这一点:在相同的架构、编码器和噪声调度下,-预测始终优于 -预测。

3.2 联合 CLS–Patch 建模

在表征空间中操作的一个独特优势是可以直接访问 [CLS] token——这是一个全局语义摘要,编码了与局部块内容互补的类别、布局和外观信息。在基于 VAE 特征的标准潜在扩散模型中,这种全局 token 并非潜在表征本身的组成部分;而在表征空间中,它是内在固有的。我们在同一扩散过程中将 [CLS] 与块联合建模:对其进行投影,前置到块序列中,并参与双向自注意力机制,将空间证据聚合到全局上下文中,再将优化后的引导信息广播回局部 token。一个独立的线性头产生 [CLS] 预测 ,从而产生一个辅助的 -预测损失(以速度形式书写,以与公式 1 对称;等价于附录 B 中相同重新加权后的 ),以及总目标函数 。在训练期间,[CLS] 噪声与块噪声独立采样,以避免方差坍缩([CLS] 是单个向量,而块噪声是 )。在推理时,我们在(可选独立的)引导尺度下,使用 Heun + 无分类器引导联合推进 [CLS] 和块;在所有报告的实验中,我们对两者使用相同的尺度,但该机制允许解耦。仅解码块 token,而丢弃 [CLS]。我们还在初始化时通过(耦合噪声)将两个噪声流耦合——这在收敛时带来了微小但一致的改进(§4.3)。

3.3 维度感知噪声调度

Refer to caption
图 5:时间采样(上)和每个 token 的信噪比(下)。

的信噪比为 ,但每个 token 的有效信噪比随每个 token 的维度缩放 [13]:对于一个 维的 token,噪声幅度随 增长,而信号保持在单位尺度,因此更高维度的 token 需要更低的 (更多噪声)才能达到相同的相对破坏程度。一个 DINOv2-Small token 具有 ,即 的每像素维度,因此像素空间的调度会在噪声状态下训练不足。遵循 RAE [44] 和 SD3 [7],我们对 应用维度相关的时间偏移,并设置 。这将中位数 从 推至 (中位数信噪比降低 ,图 5);§4 表明,这弥补了与像素空间 logit-正态基线相比的 FID 差距(800 个 epoch 时从 3.17 降至 1.44)。

4 实验

实验设置。RiT-XL 拥有 28 层,隐藏维度 1152,16 个注意力头,FFN 扩展比为 ,总计 6.76 亿参数。我们在 ImageNet 上使用冻结的 DINOv2-Small 编码器()和预训练的 RAE 解码器 [44] 进行训练,使用 8 块 H200 GPU(每轮 12 分钟)。我们遵循 RAE [44] 的方法,采用类别平衡采样(每类 50 张图像)进行 FID-50K 评估。完整超参数见附录 D。

图 6:ImageNet 上的收敛性比较。FID-50K 与训练轮次对比。
Refer to caption
媒体内容 · 前往原文查看
表 3:消融研究。默认设置:-预测(表 2)、逐元素标准化、时间偏移调度、、DINOv2-S。每一行替换一个因素。“”表示训练发散。

200 轮 400 轮 800 轮
默认设置(我们的) 1.83 1.67 1.44
预处理
无标准化† 387.4 362.1 344.8
噪声调度
Logit-normal 4.45 4.11 3.17
分类损失权重
无分类损失 1.89 1.70 1.63
1.86 1.64 1.50
编码器
DINOv2-Base 2.20 1.78 1.56

4.1 收敛性与效率

图 6 比较了 RiT-XL 与基线模型。与 RAE-XL (DINOv2-S) [44](一个使用与 RiT-XL 相同编码器、解码器和参数量(6.76 亿)的 -预测 DiT-XL,用于隔离第 3 节的设计选择)相比,RiT-XL 在每个轮次都领先,并在 800 轮时达到 FID 1.45(优于 1.87)。在 100 轮时,RiT-XL 已经匹配了更大的 RAE-XL (DINOv2-B) 基线在 720 轮时的表现(加速比);在 200 轮时,它匹配了 RAE-XL (DINOv2-S) 在 800 轮时的表现(加速比)。RiT 还在 20 到 200 轮内超越了表示对齐方法 REPA [41] 和 REG [39] 在 800 轮时的 FID。并行的 RJF [16] 通过在范数集中球面上的黎曼流匹配来解决相同的 DINOv2 径向模糊问题;在图 6 所示的匹配的 80 轮预算下,RiT-XL 达到了 FID 2.48(DINOv2-S),而 RJF 为 3.62(DINOv2-B)。

4.2 每个配方选择都是必要的

除非另有说明,消融实验使用完整配方,每次改变一个因素;我们报告在 Heun 50 步、无引导条件下的 ImageNet FID-50K。

逐元素标准化。原始的 DINOv2 特征具有异质的每通道方差(范围跨通道)。在原始特征上训练会发散:损失振荡,并且 FID 在整个训练过程中保持在随机初始化水平()。

噪声调度。时间偏移在原始 JiT logit-normal 调度上缩小了 FID 差距(800 轮时从 3.17 降至 1.44),这证实了当每个 token 的维度增长时(此处与像素维度相比),将训练密度重新分配至更高噪声区域至关重要。

CLS token。在没有 CLS 建模的情况下,FID 停滞在 1.63;加入后达到 1.44。注意力可视化(附录 I)显示,[CLS] 在早期层聚合粗略的场景线索,在中间层整合对象-上下文关系,并在后期层将精炼后的引导信息广播回去。

编码器大小。尽管特征维度减半,DINOv2-Small 在 800 轮时始终优于 DINOv2-Base(1.44 对比 1.56)。模型容量在此并非瓶颈:DINOv2-B 在相同条件下具有两倍的潜在维度,因此去噪器必须回归更高维度的目标,却未获得相应的底层结构增益。因此,第 2 节中对 DINOv2-Base 的分析是一种保守描述——主实验中使用的小型流形至少同样有利,且回归任务更为简单。

4.3 高效的 ODE 收敛实现少步生成

第 2 节中建立的四个几何特性——高有效秩、良好条件化的协方差、近似高斯的边缘分布以及流形上的线性插值——共同预测,在 DINOv2 特征上,噪声到数据的 ODE 应能以少量 Heun 步高效求解。我们通过测量像素空间中的 Heun 求解器截断误差直接验证了这一点,并表明在严格的采样预算下,这能转化为数量级的增益——无需任何蒸馏或一致性训练。

媒体内容 · 前往原文查看
表 4:采样调度消融实验。六种 ODE 时间离散化调度和 Heun 步数下的 FID-50K,分别在有和无需分类器引导的条件下。每个单元格报告独立/耦合噪声 FID。RiT-XL 基于 DINOv2-S,800 轮。
无引导(CFG,Heun 步数) 有引导(CFG,Heun 步数)
调度 5 10 25 50 5 10 25 50
均匀 12.8 / 12.7 6.88 / 6.76 2.34 / 2.29 1.61 / 1.58 10.8 / 10.7 6.19 / 6.12 1.93 / 1.90 1.30 / 1.28
EDM 2.37 / 2.34 1.61 / 1.58 1.49 / 1.47 1.47 / 1.46 2.01 / 1.99 1.33 / 1.32 1.17 / 1.15 1.16 / 1.14
余弦 5.96 / 5.86 2.16 / 2.12 1.57 / 1.56 1.48 / 1.45 5.69 / 5.63 1.91 / 1.88 1.29 / 1.28 1.19 / 1.18
Power-2 2.41 / 2.39 1.74 / 1.72 1.50 / 1.48 1.47 / 1.44 1.99 / 1.98 1.48 / 1.46 1.18 / 1.16 1.16 / 1.15
Log-SNR 4.56 / 4.51 1.96 / 1.95 1.51 / 1.50 1.46 / 1.44 3.79 / 3.78 1.73 / 1.74 1.23 / 1.22 1.18 / 1.17
时间偏移 2.44 / 2.38 1.59 / 1.58 1.47 / 1.45 1.46 / 1.44 1.99 / 1.99 1.27 / 1.25 1.15 / 1.14 1.15 / 1.14
图 7:在匹配的 NFE 下的少步 FID。

Refer to caption
Refer to caption
图 8:DINOv2 ODE 在少数 Heun 步内收敛。像素空间截断误差与步数的关系(128 条轨迹的平均值;每个模型与自身的 对比)。RiT 的衰减()对比 JiT 的 ;后期 斜率对比 ,虚线表示 Heun 渐近线。

像素空间截断误差测量。对于每个模型,我们从匹配对中为每个空间生成 128 条轨迹,在 加上一个参考值 下运行 Heun 采样,将每个端点解码到像素空间(RiT 通过冻结的 RAE 解码器),并测量 Frobenius 距离 。由于每个模型都与自身的 步参考值进行比较,该指标将 ODE 收敛速度与任一端点的绝对质量分离开来:一个恰好收敛到较差不动点的模型不会因此被错误地奖励。JiT 使用其自身论文中报告的调度。

RiT 的截断误差衰减比 JiT 更陡。图 8 显示了平均值。RiT 的截断误差从 下降到 (),而 JiT 仅下降到 ()。在 (RiT 的默认值)时,RiT 与其 步端点的 Frobenius 距离在 以内——这与表 4 的发现一致,即 RiT 的 FID 已经匹配了收敛的 。相同 下的 JiT 仍处于 Frobenius 距离 ,其 FID 也远未收敛(图 7:JiT-H 的 -NFE FID 为 )。在二阶 Heun 误差界 下,衰减率的差距转化为 DINOv2 边际流相应更小的有效曲率。这是第 2 节中量化的四个几何属性的经验对应:更高的有效秩缩短了传输距离,高斯性匹配了源分布,更紧的后验降低了速度-目标方差,流形上的插值使其保持在定义明确的区域——每个属性都独立地预测了一个更平滑、更易于积分的速度场。

少步生成。图7在匹配的NFE条件下,单独展示了表示空间的作用。在10 NFE时,像素空间的JiT-H得到FID 26.2,DINOv2空间的DiT-XL得到3.29,而RiT-XL达到2.38——相比像素空间提升了一个数量级,并且明显优于配备了DDT的DINOv2基线;在20 NFE时,这一排序同样成立(JiT-H为,DiT-XL为,RiT-XL为)。在RiT自身内部(表4),采用时间偏移调度的5步Heun方法在无引导下达到FID 2.44,有引导下达到1.99;10步达到1.59和1.25;而25步已匹配完全收敛——超越了表5中大多数先前的VAE潜空间基线,这些基线通常需要更多采样步数。

采样调度消融实验。表4比较了六种ODE时间离散化调度(正式定义见附录G)。在步数时,所有非均匀调度都收敛到相近的FID(无引导1.43–1.45,有引导1.14–1.15),证实了ODE已被良好近似。在5步时,三种集中式调度(EDM、power-2、time-shift:FID)相比均匀间隔调度(12.7)有显著提升,因为它们将更多评估步分配给了速度场变化最剧烈的高噪声区域。耦合噪声(§3.2,表4中的独立/耦合单元)在不同调度和步数下使FID均匀变化;我们在主要结果中采用了它,因为它能带来微小但一致的增益,但并未将其视为核心组件。

从几何到经验。第2节中分离出的三种机制直接对应RiT的性能提升。(i) 更好的协方差条件(对比)使得所有方差方向都能以可比速率进行训练。(ii) 更紧致的后验分布缩小了不可约的目标方差,并有助于加速收敛(图6)。(iii) 高有效秩(更短的传输路径)、近似高斯的边缘分布(更平滑的源数据插值)以及流形上的插值点(无空洞穿越)带来了低有效曲率,这通过更快的截断误差衰减(图8)和少步数机制(图7)得到了验证。像素级潜变量在这三种机制上均失败,SD-VAE则至少在其中两种上失败,因此这些机制在经验上是独立的;它们在DINOv2上的共同实现解释了为什么一个普通的M DiT在收敛性和严格采样预算下均能超越配备M DDT的基线模型。

Refer to caption
图9:在ImageNet上精选的RiT-XL样本,覆盖了ImageNet的各类别。
媒体内容 · 前往原文查看
表5:在ImageNet上的类别条件图像生成。使用普通的DiT-XL骨干网络(6.76亿参数)且仅需25步Heun采样,RiT取得了最佳FID。值得注意的是,在基于表征的方法中,RiT使用了最小的DINOv2变体(DINOv2-S);DiT-XL使用DINOv2-B,而FAE [10] 微调了一个DINOv2-G编码器,将其特征压缩为潜变量用于生成。
方法 自编码器 训练轮数 参数量 无引导 有引导
FID IS 精确率 召回率 FID IS 精确率 召回率
像素扩散
ADM [6] – 400 5.54亿 10.94 101.0 0.69 0.63 3.94 215.8 0.83 0.53
PixelFlow-XL [3] – 320 6.77亿 – – – – 1.98 282.1 0.81 0.60
PixNerd-XL [36] – 320 7亿 – – – – 1.93 298.0 0.80 0.60
PixelDiT-XL [42] – 320 7.97亿 – – – – 1.61 292.7 0.78 0.64
JiT-G [18] – 600 20亿 – – – – 1.82 292.6 0.79 0.62
潜变量扩散
DiT-XL [22] SD-VAE 1400 6.75亿 9.62 121.5 0.67 0.67 2.27 278.2 0.83 0.57
SiT-XL [20] SD-VAE 1400 6.75亿 8.61 131.7 0.68 0.67 2.06 270.3 0.82 0.59
MaskDiT [45] SD-VAE 1600 6.75亿 5.69 177.9 0.74 0.60 2.28 276.6 0.80 0.61
MDTv2-XL [9] SD-VAE 1080 6.75亿 – – – – 1.58 314.7 0.79 0.65
REPA-XL [41] SD-VAE 800 6.75亿 5.78 158.3 0.70 0.67 1.29 306.3 0.79 0.64
LightningDiT [40] VA-VAE 800 6.75亿 2.17 205.6 0.77 0.65 1.35 295.3 0.79 0.65
DDT-XL [37] SD-VAE 400 6.75亿 6.27 154.7 0.68 0.69 1.26 310.6 0.79 0.65
SVG-XL [27] SVGTok 1400 6.75亿 3.36 181.2 – – 1.92 264.9 – –
REG-XL [39] SD-VAE 800 6.75亿 1.80 - - - 1.36 299.4 0.77 0.66
RAE-XL [44] RAE-DINOv2-S 800 6.76亿 1.87 209.7 0.80 0.63 1.41 309.4 0.80 0.63
RAE-XL[44] RAE-DINOv2-B 800 839M 1.51 242.9 0.79 0.63 1.16 261.0 0.77 0.67
FAE-XL [10] FAE-DINOv2-G 800 675M 1.48 239.8 0.81 0.63 1.29 268.0 0.80 0.64
RiT-XL(我们的方法) RAE-DINOv2-S 800 676M 1.45 231.6 0.81 0.62 1.14 299.7 0.80 0.63

4.4 与先前方法的比较

表5报告了完整的比较结果。RiT的无引导FID为1.45,已经超越了所有无表示编码器方法的引导FID(PixelDiT-XL 1.61,JiT-G 1.82,MDTv2-XL 1.58,DiT-XL 2.27,SiT-XL 2.06)——结构良好的语义空间如此忠实地捕捉了数据分布,使得CFG变得远非必要。在基于表示的方法中,RiT取得了最佳的无引导FID(1.45,对比DiT-XL 1.51、FAE-DINOv2-G 1.48、RAE-XL 1.87)和最佳的引导FID(1.14,对比DiT-XL 1.28、FAE 1.29、REPA-XL 1.29)。在无引导层面,RiT相对于FAE的优势很微弱(1.45对比1.48),但这是在显著更简单的设置下取得的:FAE使用了最大的DINOv2变体(DINOv2-G),并联合微调其编码器以用于生成,而RiT使用了最小的变体(DINOv2-S),且编码器完全冻结——这表明第2节中讨论的几何优势已经可以直接开箱即用,无需编码器的协同适应。更根本的是,FAE和RiT解决的是不同的问题:FAE通过调整编码器以产生更利于扩散的潜在空间(将DINOv2-G的特征压缩为生成潜在表示)来简化生成,而RiT则直接着手对现有的高维表示分布进行建模,而不改变编码器。因此,这两项贡献在很大程度上是正交的——编码器侧的适应(FAE)和去噪器侧的方案(RiT)原则上可以组合使用。

RiT使用的组件规模统一更小。去噪器有M个参数(比DiT-XL的M个参数更小,无DDT头);编码器是DINOv2-S(最小的DINOv2变体),而DiT-XL使用的是DINOv2-B,FAE使用的是DINOv2-G。RiT还在具有竞争力的召回率下取得了最高的无引导精确率,这证实了当表示分布有利于流匹配时,一个现成的DINOv2编码器和一个基础骨干网络就已足够。

5 相关工作

面向图像的扩散与流匹配。扩散模型与流匹配[12, 30, 19, 7]是现代图像生成的基础。潜在扩散模型[23]通过VAE压缩图像;DiT/SiT[22, 20]用Transformer架构取代了U-Net。一个关键的设计选择是预测目标——即 、 或 。JiT[18]表明,通过将目标置于低维数据流形上,-预测在像素空间中显著优于其他方案。我们的工作扩展了这一见解,刻画了预训练表示空间相对于 如何以不同方式嵌入该流形,并证明DINOv2特别适合使用标准主干网络进行 -预测。

利用表示进行生成。VA-VAE、EQ-VAE和Diffusability通过塑造自编码器训练来生成适合扩散的潜在表示[40, 15, 28]。REPA[41]增加了一个仅在训练期间激活的DINOv2对齐损失。REG[39]通过将DINOv2 [CLS]嵌入到SD-VAE轨迹中来弥合这种训练-推理差距;RiT的不同之处在于它原生运行在DINOv2空间中,其中[CLS]是固有的。REPA和REG均未分析流形几何结构,且两者都在SD-VAE潜在表示上保留了 -预测。RAE[44]用DINOv2编码器替换了VAE,但采用了 -预测,并且需要DDT头来处理病态速度场;并行的RJF[16]则使用黎曼流匹配,在范数集中球面上采用SLERP路径。我们证明,使用逐元素标准化进行 -预测,足以用标准DiT对DINOv2特征进行建模——无需架构修改,也无需黎曼重构。

少步生成与蒸馏。渐进式蒸馏[25]、一致性模型[31, 29]和修正流[19]通过训练专用的少步学生模型或拉直教师模型的轨迹来降低采样成本。这些方法与RiT的贡献是正交的:RiT表明,基础模型本身在几何友好的表示空间中已经能够达到具有竞争力的少步FID分数,无需任何蒸馏或一致性损失,并且仍然是此类方法的天然教师模型。

走向理解与生成的统一。统一的视觉模型[5, 33, 4, 34, 35]通常为感知(CLIP/DINOv2)和合成(SD-VAE)任务保留独立的编码器,并在生成侧附加特定于任务的架构组件。RiT 在 DINOv2 空间中具备竞争力的生成能力表明了一种更简洁的替代方案:单一的语义表示和单一的普通 Transformer 主干——无需 DDT 头、无需黎曼重构、无需表示对齐损失——即可服务于这两类任务。在匹配编码器下的训练加速(§4)以及具有竞争力的少步采样(5 步 Heun 采样 FID 为 2.0,10 步为 1.25)进一步降低了将生成头附加到现有感知栈上的实际成本,使得基于 DINOv2 空间的 RiT 成为统一流水线的有前景基础,在这些流水线中,相同的特征可驱动分类、检索和合成。

6 结论

我们提出了 RiT,这是一个在冻结的 DINOv2 特征上使用 -预测训练的普通 DiT 模型。它在 ImageNet 上,使用比 DiT-XL 更少的去噪器参数(6.76 亿 vs 8.39 亿)和最小的 DINOv2 变体(DINOv2-S),实现了无引导 FID 1.45 和有引导 FID 1.14 的成绩;并且无需蒸馏或一致性训练即可支持少步生成(引导 FID:5 步 Heun 采样为 2.0,10 步为 1.25)。第 2 节的分析表明,当特征分布满足我们识别出的四个几何轴——高有效秩、条件良好的协方差、近似高斯边缘分布以及流形上的线性插值——时,表示空间扩散在架构上会变得更简单。我们的结果论证了,只要表示的分布几何特性已经有利,就应该优先采用目标侧的重构(-预测),而非架构侧(DDT 头)或传输侧(黎曼流匹配)的解决方案。

致谢

我们感谢 Mila IDT 团队及其技术支持人员对 Mila 计算集群的维护。我们也感谢 NVIDIA 以计算资源形式提供的物质支持。在本项目期间,Aishwarya Agrawal 获得了加拿大 CIFAR AI 讲席教授奖项的支持。

参考文献

  • Ahamed 等人 [2026] Shadab Ahamed、Eshed Gal、Simon Ghyselincks、Md Shahriar Rahim Siddiqui、Moshe Eliasof 和 Eldad Haber。预条件得分与流匹配。arXiv 预印本 arXiv:2603.02337,2026 年。
  • Chen 和 Lipman [2023] Ricky TQ Chen 和 Yaron Lipman。一般几何上的流匹配。arXiv 预印本 arXiv:2302.03660,2023 年。
  • Chen 等人 [2025a] Shoufa Chen、Chongjian Ge、Shilong Zhang、Peize Sun 和 Ping Luo。PixelFlow:基于流的像素空间生成模型。arXiv 预印本 arXiv:2504.07963,2025a。
  • Chen 等人 [2025b] Xiaokang Chen、Zhiyu Wu、Xingchao Liu、Zizheng Pan、Wen Liu、Zhenda Xie、Xingkai Yu 和 Chong Ruan。Janus-Pro:通过数据和模型扩展实现统一的多模态理解与生成。arXiv 预印本 arXiv:2501.17811,2025b。
  • Deng 等人 [2025] Chaorui Deng、Deyao Zhu、Kunchang Li、Chenhui Gou、Feng Li、Zeyu Wang、Shu Zhong、Weihao Yu、Xiaonan Nie、Ziang Song 等。统一多模态预训练中的涌现特性。arXiv 预印本 arXiv:2505.14683,2025 年。
  • Dhariwal 和 Nichol [2021] Prafulla Dhariwal 和 Alexander Nichol。扩散模型在图像合成上击败 GAN。收录于 NeurIPS,2021 年。
  • Esser 等人 [2024] Patrick Esser、Sumith Kulal、Andreas Blattmann、Rahim Entezari、Jonas Müller、Harry Saini、Yam Levi、Dominik Lorenz、Axel Sauer、Frederic Boesel 等。扩展修正流 Transformer 以实现高分辨率图像合成。收录于第四十一届国际机器学习大会,2024 年。
  • Facco 等人 [2017] Elena Facco、Maria d’Errico、Alex Rodriguez 和 Alessandro Laio。通过最小邻域信息估计数据集的本质维度。《科学报告》,7(1):12140,2017 年。
  • Gao 等人 [2023] Shanghua Gao、Pan Zhou、Ming-Ming Cheng 和 Shuicheng Yan。MDTv2:掩码扩散 Transformer 是一种强大的图像合成器。arXiv 预印本 arXiv:2303.14389,2023 年。
  • Gao 等人 [2025] Yuan Gao、Chen Chen、Tianrong Chen 和 Jiatao Gu。一层足矣:适配预训练视觉编码器用于图像生成。arXiv 预印本 arXiv:2512.07829,2025 年。
  • Henry 等人 [2020] Alex Henry、Prudhvi Raj Dachapally、Shubham Shantaram Pawar 和 Yuxuan Chen。Transformer 架构的查询-键归一化。发表于《计算语言学协会发现:EMNLP 2020》,第 4246–4253 页,2020 年。
  • Ho 等人 [2020] Jonathan Ho、Ajay Jain 和 Pieter Abbeel。去噪扩散概率模型。《神经信息处理系统进展》,33:6840–6851,2020 年。
  • Hoogeboom 等人 [2023] Emiel Hoogeboom、Jonathan Heek 和 Tim Salimans。简单扩散:面向高分辨率图像的端到端扩散。发表于《国际机器学习大会》,第 13213–13232 页。PMLR,2023 年。
  • Karras 等人 [2022] Tero Karras、Miika Aittala、Timo Aila 和 Samuli Laine。阐明基于扩散的生成模型的设计空间。《神经信息处理系统进展》,35:26565–26577,2022 年。
  • Kouzelis 等人 [2025] Theodoros Kouzelis、Ioannis Kakogeorgiou、Spyros Gidaris 和 Nikos Komodakis。Eq-VAE:用于改进生成式图像建模的等变正则化潜在空间。arXiv 预印本 arXiv:2502.09509,2025 年。
  • Kumar 和 Patel [2026] Amandeep Kumar 和 Vishal M Patel。在流形上学习:利用表示编码器解锁标准扩散 Transformer。arXiv 预印本 arXiv:2602.10099,2026 年。
  • Levina 和 Bickel [2004] Elizaveta Levina 和 Peter Bickel。本征维度的最大似然估计。《神经信息处理系统进展》,17,2004 年。
  • Li 和 He [2025] Tianhong Li 和 Kaiming He。回归本源:让去噪生成模型去噪。arXiv 预印本 arXiv:2511.13720,2025 年。
  • Liu 等人 [2022] Xingchao Liu、Chengyue Gong 和 Qiang Liu。流直且快:学习使用修正流生成和传输数据。arXiv 预印本 arXiv:2209.03003,2022 年。
  • Ma 等人 [2024] Nanye Ma、Mark Goldstein、Michael S Albergo、Nicholas M Boffi、Eric Vanden-Eijnden 和 Saining Xie。Sit:使用可扩展插值 Transformer 探索基于流和扩散的生成模型。发表于《欧洲计算机视觉会议》,第 23–40 页。Springer,2024 年。
  • Oquab 等人 [2024] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby 等。Dinov2:无监督学习鲁棒视觉特征。Transactions on Machine Learning Research,2024 年。
  • Peebles 和 Xie [2023] William Peebles 与 Saining Xie。基于 Transformer 的可扩展扩散模型,2023 年。URL https://arxiv.org/abs/2212.09748。
  • Rombach 等人 [2022] Robin Rombach、Andreas Blattmann、Dominik Lorenz、Patrick Esser 和 Björn Ommer。基于潜在扩散模型的高分辨率图像合成。收录于 IEEE/CVF 计算机视觉与模式识别会议论文集,第 10684–10695 页,2022 年。
  • Roy 和 Vetterli [2007] Olivier Roy 与 Martin Vetterli。有效秩:一种有效维度的度量。收录于 2007 年第 15 届欧洲信号处理会议,第 606–610 页。IEEE,2007 年。
  • Salimans 和 Ho [2022] Tim Salimans 与 Jonathan Ho。用于扩散模型快速采样的渐进式蒸馏。arXiv 预印本 arXiv:2202.00512,2022 年。
  • Shazeer [2020] Noam Shazeer。GLU 变体改进 Transformer。arXiv 预印本 arXiv:2002.05202,2020 年。
  • Shi 等人 [2025] Minglei Shi、Haolin Wang、Wenzhao Zheng、Ziyang Yuan、Xiaoshi Wu、Xintao Wang、Pengfei Wan、Jie Zhou 和 Jiwen Lu。无需变分自编码器的潜在扩散模型。arXiv 预印本 arXiv:2510.15301,2025 年。
  • Skorokhodov 等人 [2025] Ivan Skorokhodov、Sharath Girish、Benran Hu、Willi Menapace、Yanyu Li、Rameen Abdal、Sergey Tulyakov 和 Aliaksandr Siarohin。改进自编码器的可扩散性。arXiv 预印本 arXiv:2502.14831,2025 年。
  • Song 和 Dhariwal [2023] Yang Song 与 Prafulla Dhariwal。训练一致性模型的改进技术。arXiv 预印本 arXiv:2310.14189,2023 年。
  • Song 等人 [2020] Yang Song、Jascha Sohl-Dickstein、Diederik P Kingma、Abhishek Kumar、Stefano Ermon 和 Ben Poole。基于随机微分方程的分数生成建模。arXiv 预印本 arXiv:2011.13456,2020 年。
  • Song 等人 [2023] Yang Song、Prafulla Dhariwal、Mark Chen 和 Ilya Sutskever。一致性模型。2023 年。
  • Su 等人 [2024] 苏剑林、Murtadha Ahmed、陆宇、潘胜锋、文博、刘云峰。Roformer:基于旋转位置编码的增强型 Transformer 架构。Neurocomputing,568:127063,2024。
  • Team [2024] Chameleon 团队。Chameleon:混合模态早期融合基础模型。arXiv 预印本 arXiv:2405.09818,2024。
  • Tong 等人 [2025] 童胜邦、David Fan、李嘉晨、熊云阳、陈鑫磊、Koustuv Sinha、Michael Rabbat、Yann LeCun、谢赛宁、刘壮。Metamorph:通过指令微调实现多模态理解与生成。收录于 IEEE/CVF 国际计算机视觉大会论文集,第 17001–17012 页,2025。
  • Tong 等人 [2026] 童胜邦、郑博阳、王子腾、唐炳达、马南野、Ellis Brown、杨继涵、Rob Fergus、Yann LeCun、谢赛宁。利用表示自编码器扩展文本到图像的扩散 Transformer。arXiv 预印本 arXiv:2601.16208,2026。
  • Wang 等人 [2025a] 王帅、高子腾、朱晨辉、黄伟林、王利民。Pixnerd:像素神经场扩散。arXiv 预印本 arXiv:2507.23268,2025a。
  • Wang 等人 [2025b] 王帅、田智、黄伟林、王利民。DDT:解耦扩散 Transformer,2025b。
  • Wang 和 Isola [2020] 王同舟、Phillip Isola。通过超球面上的对齐与均匀性理解对比表示学习。收录于国际机器学习大会,第 9929–9939 页。PMLR,2020。
  • Wu 等人 [2025] 吴戈、张申、石瑞静、高尚华、陈振远、王磊、陈兆伟、高鸿程、唐瑶、杨健等。面向生成的表示纠缠:训练扩散 Transformer 比想象中简单得多。arXiv 预印本 arXiv:2507.01467,2025。
  • Yao 等人 [2025] 姚景峰、杨彬、王兴刚。重建 vs. 生成:化解潜在扩散模型中的优化困境。收录于计算机视觉与模式识别大会论文集,第 15703–15712 页,2025。
  • Yu 等人 [2024] Sihyun Yu、Sangkyung Kwak、Huiwon Jang、Jongheon Jeong、Jonathan Huang、Jinwoo Shin 和 Saining Xie。《面向生成的表征对齐:训练扩散 Transformer 比想象中更简单》。arXiv 预印本 arXiv:2410.06940,2024 年。
  • Yu 等人 [2025] Yongsheng Yu、Wei Xiong、Weili Nie、Yichen Sheng、Shiqiu Liu、Anima Anandkumar 和 Arash Vahdat。《PixelDiT:面向图像生成的像素扩散 Transformer》。arXiv 预印本 arXiv:2511.20645,2025 年。
  • Zhang 和 Sennrich [2019] Biao Zhang 和 Rico Sennrich。《均方根层归一化》。Advances in Neural Information Processing Systems,第 32 卷,2019 年。
  • Zheng 等人 [2025] Boyang Zheng、Nanye Ma、Shengbang Tong 和 Saining Xie。《基于表征自编码器的扩散 Transformer》。arXiv 预印本 arXiv:2510.11690,2025 年。
  • Zheng 等人 [2023] Hongkai Zheng、Weili Nie、Arash Vahdat 和 Anima Anandkumar。《使用掩码 Transformer 快速训练扩散模型》。TMLR,2023 年。

附录 A 局限性

DINOv2 编码器偏差。RiT 继承了冻结的 DINOv2 编码器的归纳偏差。DINOv2 的自监督学习目标更强调语义内容而非光度细节,且先前的研究已观察到其在精细纹理、薄结构和小物体上的特征分辨率较弱。这些偏差会直接传播到 RiT 所能生成的内容中,因为 RAE 解码器是在相同的特征上运行的。联合编码器微调(如 FAE [10] 中所示)可以缓解这一问题,但代价是放弃我们所倡导的更简单的冻结编码器方案。

类别条件与分辨率。所有 RiT 结果均基于 ImageNet 数据集上的类别条件设置。我们尚未评估文生图生成、更高分辨率(例如 或 ),或非图像模态。第 2 节的几何分析是在 DINOv2-Base 和 DINOv2-Small 尺度下的 ImageNet 图像上进行的;相同的四个轴在更大的模型/数据尺度下或在文本条件设置下是否仍然成立,有待未来工作研究。

局部高斯假设在分析中成立。协方差条件诊断(第2节)基于局部高斯近似;其他三个几何轴对假设要求较低,但仍报告聚合标量,无法排除流形上那些有利性质失效的对抗性区域。流匹配结果从实证角度佐证了几何论断,但并未确立每个轴对观察到的效率提升而言是单独必要的。

附录B 速度损失与重加权ε预测损失的等价性

我们证明,采用ε预测参数化的速度MSE损失(式2)等价于重加权的ε预测损失。给定前向过程,目标速度为:

(3)

其中第二个等式由代入得到。在ε预测下,网络输出为,预测速度为。将两者代入速度损失:

(4)

因此,速度损失等于由重加权的ε预测损失,该权重在(接近干净数据)时增大损失。因此,这两个损失作为泛函是等价的;区别在于网络的参数化方式,它决定了实际拟合的函数(§3.1):ε预测使成为网络的直接输出,因此其回归目标始终位于数据流形上,而v预测要求网络产生环境速度,该速度依赖于流形外区域,并在时发散。

附录C 流形分析细节

本附录为第2节中使用的流形分析指标提供正式定义和实现细节。所有实验均在10,000张随机采样的ImageNet训练图像上进行。

PCA谱与有效秩。

我们在每个空间的展平特征向量上拟合包含512个主成分的PCA。令表示样本协方差矩阵的特征值,表示归一化特征值。有效秩[24]定义为

(5)

当所有方差集中在单一方向(最大各向异性)时,该值等于1;当方差完全均匀(最大各向同性)时,该值等于。对于流匹配,有效秩越高意味着数据分布越接近各向同性高斯源,从而需要复杂度更低的速场。

内在维度(TwoNN)。

TwoNN 估计器 [8] 通过第一近邻与第二近邻距离的比值来估计内在维度。对于每个点 ,设 和 分别为其到最近邻和第二近邻的距离,且 。在数据在 维流形上局部均匀的假设下,MLE 估计量为

(6)

其中 是满足 的有效样本数量。我们对 5000 个点进行子采样,并分块计算成对欧氏距离以控制内存。

在 处的鲁棒性。

在如此高的环境维度下,近邻距离会趋于集中,任何单次运行的内在维度估计都可能存在噪声。我们通过对 10 个独立的 5000 点子样本进行 TwoNN 自助法(bootstrap)计算,报告样本均值与标准差:

空间 (均值 标准差,10 次自助法)
像素
DINOv2

像素与 DINOv2 之间的维度差距远低于合并标准差(, ),因此这两个估计在统计上无法区分。已知 MLE 估计器 [17] 的较大 变体在高环境维度下存在向上偏差 [8],且不具备这种收敛特性;因此我们依赖 TwoNN 作为主要估计器。DINOv2 的优势不在于流形维度,而在于第 2 节中其他三个轴所表征的全局几何结构。

边缘高斯性(超额峰度)。

对于每个维度 ,超额峰度为

(7)

其中 和 分别是每个维度的均值和标准差。高斯分布的 ;正值表示重尾,负值表示轻尾。我们报告所有维度上 的中位数作为标量汇总。

流形内插值评分。

对于线性插值路径上的每个中间帧,我们通过统一流程下的重建误差来衡量其保持在自然图像流形上的程度:图像 编码(DINOv2) 解码 MSE 与输入图像对比。该流程对像素空间和 DINOv2 空间的插值帧均相同地应用:

  • •

    像素插值:中间帧经过编码 解码流程。重影伪影(脱离流形)会产生高 MSE,因为编码器会将其投影到最近的有效表示上。

  • •

    DINOv2 插值:中间表示首先被解码为像素空间帧,然后通过相同的编解码流水线。

通过相同的流水线测量两者,比较是无偏的:唯一的区别在于该帧是由像素混合还是 DINOv2 潜在空间插值生成的。我们对 100 个同类对进行平均,每个对包含 11 个插值步骤。

合理性检查:对 DINOv2 插值结果进行编码器往返测试。

为了排除一种简单的解释,即 DINOv2 插值流水线在结构上具有接近零的重建误差,我们还测量了 和 在特征空间中的接近程度。如果编码器只是将任意输入重新投影到其最近的合法表示,那么像素混合与 DINOv2 之间的 MSE 差距可能会被人为地放大。我们报告了 与其重新编码对应项之间的平均余弦相似度,该相似度在整个插值路径中保持较高水平;因此,图 1(c) 中的差距反映了像素混合的流形外位置,而非编码器处理每个输入时的基线不对称性。

附录 D 架构与超参数

D.1 模型架构

RiT 使用现代化的 DiT 主干网络(遵循 JiT [18]/LightningDiT [40]),在 DINOv2 特征的空间网格上运行。我们的主要实验使用 DINOv2-Small ();我们还在编码器消融实验中报告了 DINOv2-Base () 的结果。表 6 总结了模型变体。

媒体内容 · 前往原文查看
表 6:RiT 模型变体。所有模型均使用 patch 大小为 1、SwiGLU FFN(MLP 比例为 4)、QK-归一化和 VisionRoPE。输入维度取决于编码器:对于 DINOv2-Small(主要实验)为 ,对于 DINOv2-Base 为 。
模型 层数 隐藏维度 注意力头数 FFN 维度 参数量
RiT-L 24 1024 16 4096 458M
RiT-XL 28 1152 16 4608 676M

每个 DiT 块包含:

  1. 1.

    adaLN 调制:时间步和类别嵌入被求和 (),并通过共享的 SiLU–Linear 层投影到每层的缩放/偏移参数。

  2. 2.

    多头自注意力,带有 QK-归一化(在注意力计算前对 Q 和 K 进行 RMSNorm)和用于二维空间位置编码的 VisionRoPE。[CLS] 和 register token 被排除在 RoPE 之外。

  3. 3.

    SwiGLU FFN:。

最后一层使用 adaLN 调制的 RMSNorm,随后通过线性投影输出通道(DINOv2-Small 为 384,DINOv2-Base 为 768)。一个独立的线性头预测 [CLS] token。代码支持仅在中间 50% 的层中应用注意力丢弃和投影丢弃;我们的主要 RiT-XL 训练配置将两者的比率均设为 0。

遵循 JiT [18],我们在中间层(RiT-L 为第 8 层)注入 32 个可学习的上下文 token。这些 token 从类别嵌入初始化,并添加了可学习的位置嵌入,参与所有后续层的自注意力,并在最终投影之前被丢弃。它们在不修改核心 DiT 模块的情况下,为类别条件生成提供了额外的容量。

D.2 训练超参数

媒体内容 · 前往原文查看
表 7:RiT-XL 在 ImageNet 上的训练超参数
超参数 取值
硬件 8 块 NVIDIA H200 GPU
吞吐量 每个 epoch 12 分钟
优化器 AdamW (β₁, β₂)
基础学习率 1e-4(按 batch size / 256 缩放)
学习率调度 常数(预热后)
预热 epoch 数 5
权重衰减 0.0
梯度裁剪 1.0(最大范数)
总 epoch 数 800
批量大小 1536(每 GPU 8 × 192)
EMA 衰减 0.9999 / 0.9996(双跟踪)
标签丢弃 0.1
注意力 / 投影丢弃 0.0(主要配置)
噪声调度 截断对数正态分布 (μ, σ)
时间偏移 s = 1.0(对于 σ = 1.0)
噪声尺度 1.0
Epsilon 裁剪 0.05
CLS 损失权重 0.2

D.3 采样超参数

媒体内容 · 前往原文查看
表 8:采样超参数
超参数 取值
ODE 求解器 Heun(二阶)
步数 25
CFG 尺度(patch) 3.7
CFG 尺度(CLS) 3.7
CFG 区间
Epsilon 裁剪 0.05
生成精度 FP32
EMA 模型 0.9999

D.4 伪代码

媒体内容 · 前往原文查看
训练

⬇

t = sample_logit_normal(B, shift=s)

eps = randn_like(z0) * sigma

zt = t * z0 + (1 - t) * eps

eps_cls = randn_like(z_cls) * sigma

zt_cls = t * z_cls + (1 - t) * eps_cls

z0_hat, cls_hat = dit(zt, t, y, zt_cls)

v = (z0 - zt) / (1 - t).clamp_min(eps_t)

v_hat = (z0_hat - zt) / (1 - t).clamp_min(eps_t)

v_cls = (z_cls - zt_cls) / (1 - t).clamp_min(eps_t)

v_cls_hat = (cls_hat - zt_cls) / (1 - t).clamp_min(eps_t)

loss = mse(v_hat, v) + lam * mse(v_cls_hat, v_cls)

loss.backward()

媒体内容 · 前往原文查看
采样(Euler)

⬇

z = randn(B, C, H, W) * sigma

z_cls = randn(B, C) * sigma

dt = 1.0 / K

for i in range(K):

t = i / K

z0_hat, cls_hat = dit(z, t, y, z_cls)

v = (z0_hat - z) / (1 - t).clamp_min(eps_t)

v_cls = (cls_hat - z_cls) / (1 - t).clamp_min(eps_t)

z = z + dt * v

z_cls = z_cls + dt * v_cls

x = rae_decoder(z)

媒体内容 · 前往原文查看
图 10:RiT 训练(左)和采样(右)的 PyTorch 风格伪代码。采样代码为清晰起见展示了欧拉方法;在所有主要实验中,我们使用二阶 Heun 求解器,该求解器额外对当前步的速度和预测的下一步进行平均。与标准流匹配的关键区别:-预测和联合 CLS 建模。

附录 E 编码器规模消融实验

编码器规模消融实验见第 2 节(正文),完整消融实验结果见表 3。DINOv2-Small()尽管特征维度减半,但始终优于 DINOv2-Base(),在 800 个 epoch 时达到了 FID 1.44 对比 1.56。较低维度的潜在空间(对比)更易于去噪器建模,同时 DINOv2-Small 仍保留了足够的语义信息(不同编码器规模下 TwoNN 本征维度相当)。

附录 F 未经筛选的样本网格图

图 11 展示了 RiT-XL(DINOv2-Small 编码器,760 个训练 epoch)使用 Heun 采样器(100 步,无分类器引导尺度 3.7)生成的未经筛选的样本。这些样本涵盖了多种 ImageNet 类别,包括动物、食物、风景、交通工具和植物,展示了 RiT 在广泛语义类别中生成高保真、多样化图像的能力。

Refer to caption
图 11:ImageNet 上未经筛选的 RiT-XL 样本。涵盖不同类别的 28 个样本:金刚鹦鹉、水母、火烈鸟、帝企鹅、金毛寻回犬、西伯利亚哈士奇、北极狐、狮子、君主斑蝶、小熊猫、大熊猫、气球、航天飞机、冰淇淋、芝士汉堡、披萨、悬崖、珊瑚礁、火山和雏菊。使用 Heun 100 步、CFG 尺度 3.7 生成。

附录 G 采样调度分析

表 9 列出了本工作中评估的六种 ODE 时间离散化调度。每种调度将一个归一化的步索引()映射到一个时间步,其中是纯噪声,是干净数据。

媒体内容 · 前往原文查看
表 9:采样调度定义。是 Heun 步数,。
调度 公式
均匀
余弦
对数信噪比均匀
EDM [14]
幂次-2
时间偏移

图 12 可视化了这些调度函数及其对生成质量的影响。面板 (a) 展示了调度函数:均匀分布将步骤均匀分配,而 EDM、power-2 和 time-shift 则将步骤集中在(高噪声端)附近;余弦和对数信噪比在两端点附近都更密集。面板 (b) 和 (c) 展示了相应的 FID 随 Heun 步数变化的函数。将更多评估分配给高噪声区域(速度场变化最剧烈的区域)的调度,在低步数时实现了显著更好的 FID,而所有非均匀调度在步数较多时趋于收敛。

Refer to caption
图 12:采样调度对比。(a) 步数的调度函数。(b, c) 无引导和有引导下的 FID-50K 与 Heun 步数关系图。耦合噪声,RiT-XL 在 DINOv2-S 上,800 个 epoch。
媒体内容 · 前往原文查看
表 10:完整采样调度消融实验(包括 2 步)。ODE 时间离散化调度和 Heun 步数的 FID-50K。每个单元格:独立噪声 / 耦合噪声。RiT-XL 在 DINOv2-S 上,800 个 epoch。
无引导(CFG,Heun 步数) 有引导(CFG,Heun 步数)
调度 2 5 10 25 50 125 2 5 10 25 50 125
均匀 23.67 / 23.75 12.84 / 12.72 6.88 / 6.76 2.34 / 2.29 1.61 / 1.58 1.47 / 1.45 20.67 / 20.63 10.80 / 10.68 6.19 / 6.12 1.93 / 1.90 1.30 / 1.28 1.16 / 1.15
EDM 11.78 / 11.70 2.37 / 2.34 1.61 / 1.58 1.49 / 1.47 1.47 / 1.46 1.47 / 1.45 10.43 / 10.34 2.01 / 1.99 1.33 / 1.32 1.17 / 1.15 1.16 / 1.14 1.14 / 1.13
余弦 23.67 / 23.75 5.96 / 5.86 2.16 / 2.12 1.57 / 1.56 1.48 / 1.45 1.45 / 1.44 20.67 / 20.63 5.69 / 5.63 1.91 / 1.88 1.29 / 1.28 1.19 / 1.18 1.15 / 1.14
Power-2 11.16 / 11.09 2.41 / 2.39 1.74 / 1.72 1.50 / 1.48 1.47 / 1.44 1.45 / 1.43 9.44 / 9.37 1.99 / 1.98 1.48 / 1.46 1.18 / 1.16 1.16 / 1.15 1.15 / 1.14
对数信噪比 23.67 / 23.75 4.56 / 4.51 1.96 / 1.95 1.51 / 1.50 1.46 / 1.44 1.45 / 1.44 20.67 / 20.63 3.79 / 3.78 1.73 / 1.74 1.23 / 1.22 1.18 / 1.17 1.15 / 1.14
Time-shift 14.05 / 14.03 2.44 / 2.38 1.59 / 1.58 1.47 / 1.45 1.46 / 1.44 1.45 / 1.43 8.59 / 8.64 1.99 / 1.99 1.27 / 1.25 1.15 / 1.14 1.15 / 1.14 1.15 / 1.14

附录 H 随机样本

图 13 展示了来自 RiT-XL 的 192 个随机生成样本(每类 8 个,共 24 个类别),未经任何筛选或精心挑选。每一行对应一个 ImageNet 类别。该模型在所有类别中均能持续生成高质量且多样化的样本。

Refer to caption
图 13:ImageNet 上随机(未经筛选)的 RiT-XL 样本。每一行展示单个类别的 8 个独立生成样本。展示的 24 个类别为:金刚鹦鹉、金毛寻回犬、西伯利亚哈士奇、北极狐、狮子、君主斑蝶、大熊猫、气球、冰淇淋、芝士汉堡、披萨、悬崖、珊瑚礁、火山、雏菊、火烈鸟、帝企鹅、水母、水獭、小熊猫、猎豹、航天飞机、喷泉以及蠵龟。使用 Heun 100 步、CFG 尺度 3.7 生成。

附录 I CLS–图像块注意力分析

图 14 可视化了 [CLS] 与图像块 token 之间跨层和跨时间步的双向注意力。我们观察到一种清晰的阶段性通信模式:

CLS→图像块(左侧)。在早期层中,[CLS] 广泛关注显著的前景区域,聚合粗略的物体线索。在中间层,其注意力扩展到上下文/背景区域,形成全局场景摘要。在后期层,[CLS] 重新聚焦于语义关键细节(例如头部和眼睛),这些细节对结构一致性和感知真实感有强烈影响。

图像块→CLS(右侧)。图像块 token 在更深层中越来越多地查询 [CLS],在语义重要区域依赖最强,而低信息量的背景图像块对其依赖较少。这些观察表明 [CLS] 充当了全局消息枢纽:它收集分布式的证据,整合物体-上下文关系,并将精炼后的全局指导广播回图像块 token,从而改善物体-背景分离和最终生成质量。

Refer to caption
图 14:RiT 中逐层的 CLS–图像块通信。左侧:[CLS] 到图像块的注意力从粗略的场景聚合过渡到语义显著区域。右侧:图像块到 [CLS] 的注意力显示出全局信息交换,随后是聚焦的精炼过程。

来源:HuggingFace Daily Papers(社区热门论文) · arxiv.org