RiT

一句话：RiT 系统地从几何角度比较了 pixel / SD-VAE / DINOv2 三个空间的流形性质（固有维度、有效秩、峰度、插值误差），发现 DINOv2 的边际统计特性（各向同性、近高斯）让 vanilla DiT + x-prediction 就可以实现 SOTA 生成——完全不需要 RAE 的 DDT wide head 或 RJF 的黎曼流匹配。

为什么值得读：

• 对于视觉编码器研究者：提供了一套评估「表示空间生成友好性」的分析框架（4-axis

  manifold comparison）

• 对于扩散模型研究者：证明 x-prediction 在合适表示空间中可以取代

  v-prediction + 架构改造

• 对于实践者：最少改动（只需改预测目标 + 噪声调度）就能把 DiT 从 FID

  2.27 带到 1.14

• 根本性启示：一个好的 encoder

  本身就决定了生成的上限，工程改造只能弥补 encoder 的缺陷

• arXiv: href="https://arxiv.org/abs/2605.21981">https://arxiv.org/abs/2605.21981
• PDF: href="https://arxiv.org/pdf/2605.21981">https://arxiv.org/pdf/2605.21981
• HTML: href="https://arxiv.org/html/2605.21981v1">https://arxiv.org/html/2605.21981v1
• 代码: href="https://github.com/lezhang7/RiT">https://github.com/lezhang7/RiT
• 作者主页:
  • Le Zhang: href="https://lezhang7.github.io/">https://lezhang7.github.io/
  • Aishwarya Agrawal: href="https://www.iro.umontreal.ca/~agrawal/">https://www.iro.umontreal.ca/~agrawal/

• HuggingFace Papers: href="https://huggingface.co/papers/2605.21981">https://huggingface.co/papers/2605.21981

相关资源

• RAE (核心前置): arxiv.org/abs/2510.11690 / href="https://rae-dit.github.io/">https://rae-dit.github.io/
• RJF (同期竞品): arxiv.org/abs/2602.10099 / href="https://github.com/amandpkr/RJF">https://github.com/amandpkr/RJF
• JiT (灵感来源): arxiv.org/abs/2511.13720 / github.com/LTH14/JiT
• REPA (互补方法): arxiv.org/abs/2410.06940 /

  github.com/sihyun-yu/REPA

• LightningDiT: arxiv.org/abs/2501.01423 /

  github.com/hustvl/LightningDiT

方法分类与发展脉络

论文将 DINOv2 特征空间做扩散的方法分为三类（按「如何处理径向歧义」分类）：

1. 架构方案 (RAE/DiT-DH): 用宽预测头

   (DDT head) 增加模型容量，让网络自己去处理径向歧义

2. 路径方案 (RJF, Kumar & Patel

   2026): 识别出几何干涉 (Geometric Interference)

   是根因，把欧几里得运输改为黎曼流匹配——让路径沿着流形测地线而非穿越球壳内部

3. 目标方案 (RiT, 本文): 将预测目标从

   v 改为 z₀

   (x-prediction)，回归目标天然在流形上，所以径向歧义在输出端被解决了

发展脉络

• 2023: DiT 在 SD-VAE latent space 上用 v-prediction 取得突破
• 2024: DINOv2 发布；人们开始尝试将其嵌入生成 pipeline
• 2025:
  • JiT: 证明 x-prediction 在像素空间有效，且无需 VAE
  • RAE: 将 DINOv2 编码器嵌为 diffusion tokenizer，但发现需要 DDT wide

    head

  • REPA: 用表示对齐加速 DiT 训练（互补方向）
  • LightningDiT: 架构现代化（SwiGLU, QK-norm, RoPE…）
• 2026:
  • RJF: 诊断出几何干涉，提出黎曼方案
  • RiT (本文): 证明 架构方案 + 路径方案 都不需要——x-prediction

    就够了

核心前置工作详情

为什么现有方法不够好：RAE 需要 19% 更多参数（839M vs 676M），RJF 需要黎曼几何的理论推导和 Jacobi 正则化。RiT 的论证是：如果底层表示空间足够好，这些额外复杂性就不再必要。

• 图 1：RiT 方法总览，紧贴方法路线对比，展示冻结 DINOv2 特征空间上的

  flow matching 训练路径。

Flow Matching

学习从噪声分布到数据分布的连续可逆速度场的生成范式，等价于连续归一化流的概率路径匹配。与扩散模型等价但参数化更灵活。

x-prediction

在流匹配/扩散中，网络输出直接预测 clean data point z₀ 而非噪声 ε 或速度

v。速度通过 v=(z₀-z_t)/(1-t) 隐式计算。

v-prediction

网络输出直接预测速度场 v，是流匹配的标准参数化方式。

有效秩 (Effective Rank)

衡量协方差矩阵的方差分布均匀程度。erank=1 表示全集中在单一方向，=D

表示完美各向同性。

TwoNN 估计器

基于二近邻距离比的固有维度 MLE 估计方法 (Facco et al.,

2017)。对高维数据鲁棒。

超额峰度 (Excess Kurtosis)

κ = μ₄/σ⁴ - 3。κ=0 是高斯分布，正值为厚尾，负值为薄尾。

Geometric Interference

Kumar & Patel (2026) 诊断的现象：DINOv2 特征空间的 Euclidean

流匹配路径会穿越流形内部的低密度区域（因为 LayerNorm

使特征集中在球壳上），导致 v-prediction 学习困难。

DDT Head

RAE 中使用的宽预测头 (Deep, Direct Token

head)，通过扩大最后一层的宽度来增加模型容量，处理高维特征空间的径向歧义。

LayerNorm 球壳

DINOv2 的 per-token LayerNorm 固定

||z||²=d，使特征集中在一个近似各向同性的球壳表面。这对生成有利（有效秩高）也有弊（线性路径穿越球壳内部）。

• 图 2：PCA 谱与有效秩分析，放在几何概念解释附近，说明 representation

  space 的方差分布。

1. 为什么相同的固有维度 (d̂≈33)

   会产生截然不同的生成效果？几何上的「嵌入方式」到底意味着什么？

2. x-prediction 和 v-prediction 的区别只是重参数化——为什么在 DINOv2

   空间区别这么大，但在像素空间就不够？

3. 如果编码器越小越好（DINOv2-Small >

   DINOv2-Base），那是否存在一个「最优特征维度」的 scaling law？

4. 这三种方案（架构 / 路径 /

   目标）是正交的可叠加设计，还是互相替代的？如果叠在一起会更好吗？

5. 这组几何分析能否反过来指导 tokenizer

   设计——我们能否训练一个显式优化有效秩和边际高斯性的编码器？

论文要解决的具体问题

给定一个预训练表示空间（DINOv2），其数据流形的固有维度（≈33）和像素空间几乎一样，但在此空间做 flow matching 时，standard DiT 使用 v-prediction 会失败（发散或 FID 很差），需要针对性架构改造。为什么？有没有更简单的方案？

Insight: 问题的根源不在架构容量，在几何性质

之前的直觉是「DINOv2 特征空间维度太高（384D 或 768D），标准 head 容量不够 → 需要宽 head」。但本文通过 4-axis 几何分析证明：根因不是容量，是几何。

具体来说：

• 图 3：协方差条件数对比，放在几何诊断结论附近，作为 latent space 可生成性的判断依据。

• 像素空间的流形是「歪的」：方差集中在少数方向（有效秩=45），大部分方向几乎无信号。从 N(0,I) 到数据的运输路径不对称——有些方向要走很远，有些几乎不动。x-prediction 在像素空间有效，是因为回归目标在歪流形上也比 v-prediction 好。但 DiT baseline 依然很弱（FID > 9）。

• DINOv2 的流形是「正圆的」：方差均匀分布在所有方向（有效秩=327），每个坐标的边际分布近乎高斯（κ≈0.083）。这意味着从 N(0,I) 到数据的每条运输路径长度大致相等且很短。但这带来一个新问题——LayerNorm 固定 ||z||=√d 使特征在球壳上，线性路径 z_t = t·z₀+(1-t)·ε 要穿过球壳内部（低密度区域）。所以 v-prediction 的径向分量巨大。

• x-prediction 解决了这个问题：它让网络输出 ẑ₀（在流形/球壳上），而不是 v̂（在环境空间）。损失的对比仍然在速度空间计算，但回归目标始终在流形上，所以网络只需要学习「从噪声到流形表面的投影」——这个东西在 DINOv2 空间中很简单。

用类比理解

想象你要把散落在房间内的积木（噪声）捡起来放到一个球形的盘子里（数据流形）。

• v-prediction

  告诉你：「从你现在的位置，以某个速度和方向走向盘子」。但如果你的位置在盘子内部，这个方向向量可能非常大（径向歧义）。

• x-prediction

  告诉你：「直接告诉我盘子的表面在哪里」。一旦你知道了表面位置，速度自然就推导出来了。

在像素空间中，这个「盘子」是歪斜的椭圆形，所以即使知道表面位置，也很难准确回归。但 DINOv2 的「盘子」是正圆的，所以回归很简单。

完整 Pipeline

1. 编码：输入图像 x → RAE 的 frozen DINOv2-Small encoder → DINOv2

   feature z₀ (384×16×16)

2. 训练噪声注入：采样 t ~ truncated logit-normal(μ=0, σ=1) + time

   shift s≈4.9 → ε ~ N(0, I) → z_t = t·z₀ + (1-t)·ε

3. 模型主体：RiT-XL (28 层 DiT block: SwiGLU FFN, QK-normed

   attention, VisionRoPE, adaLN) 接收 z_t (16×16 patch grid) + t + class

   label y + 独立的 [CLS] noise token

4. 预测：输出 ẑ₀ (clean patches) + clŝ₀ (clean CLS token)
5. 损失：L = MSE(v̂_patch, v_patch) + 0.2·MSE(v̂_cls, v_cls)
6. 解码：生成完成后，z₀ → RAE ViT decoder → 重建像素 x̂

数学公式

Flow matching 速度场 (真实值):

线性插值路径: z_t = t ⋅ z₀ + (1 − t) ⋅ ϵ,  ϵ ∼ 𝒩(0, I)

x-prediction 重参数化:

损失函数: ℒ = MSE(v̂_patch, v_patch) + 0.2 ⋅ MSE(v̂_cls, v_cls) 其中 v 的每个分量做 clamp_min(1-t, ε_t=0.05) 防止除以零。

维度感知噪声调度（时间移位）:

模型架构细节

每个 DiT block:

1. adaLN modulation: timestep embedding + class embedding → 求和后经

   SiLU-Linear 投影为每层 scale/shift

2. Multi-head self-attention with QK-normalization (RMSNorm on Q, K

   before attention) + VisionRoPE ([CLS]/register tokens 不参与)

3. SwiGLU FFN: FFN(x) = (SiLU(xW₁) ⊙ xW₃)W₂

32 个 learnable in-context tokens（RiT-L 在 layer 8 注入），从 class embedding + learnable positional embeddings 初始化，在所有后续层参与 self-attention。

与已有方法技术对比

实现技巧

• Epsilon clamp ε_t=0.05 同时在训练和采样中使用，防止 (1-t) 接近 0

  时的数值问题

• EMA 双轨跟踪 0.9999/0.9996——推测用于早期/后期不同稳定性需求
• 独立的 [CLS] 预测 head：[CLS] token 和 patch tokens 共享 backbone

  但独立投影

• 仅中间 50% 层开启 dropout（但主配置设为 0.0）——设计保留了灵活性

数据集：ImageNet 256×256（标准 ImageFolder 结构）

超参数表：

计算成本：

• 硬件：8× NVIDIA H200
• 吞吐：~12 分钟/epoch → 总 ~160 小时 (≈6.7 天)
• RiT-XL 参数量：676M
• 推理成本极低：5-25 Heun steps

主实验 (ImageNet 256×256)

消融实验关键发现：

1. x-prediction vs v-prediction on DINOv2 space：v-prediction

   baseline 比 x-prediction 差约 2 FID（原文未精确给出数值）

2. 编码器大小：DINOv2-Small (d=384) FID 1.44 > DINOv2-Base

   (d=768) FID 1.56 → 低维特征更容易建模

3. 维度感知噪声调度：Time-shift / EDM / Power-2 在 ≤10 步时远优于

   Uniform

4. Few-step 生成能力（本文亮点）：

无需蒸馏或一致性训练，ODE 本身就容易求解。

1. 耦合 vs 独立噪声：耦合噪声略好，但两者差距不大
2. 训练收敛：RiT 在 ~600 epochs (FID 1.54) 开始收敛，继续改善到 800

   epochs (FID 1.45)

0. 仓库定位与阅读状态

• 仓库链接：https://github.com/lezhang7/RiT
• 代码与论文关系：官方实现，完整包含训练、评估、采样代码
• commit：latest（主分支）
• 是否已执行 github-repo-read：否（本次仅从 README 获取概述信息）
• 代码完整性：完整（训练脚本、评估脚本、模型定义、checkpoint

  下载）

1. 代码组织总览

从 README 看，顶层包括：

• scripts/ — train.sh, eval.sh, download_assets.sh
• stats/ — normalization stats
• fid_stats/ — Inception FID reference stats
• requirements.txt — 依赖清单

模型主体直接继承自 JiT 和 RAE 代码库。RAE 的 frozen DINOv2 encoder + ViT decoder 从 HuggingFace 拉取。

2-8 => 代码仓库存在但本次未展开 github-repo-read，相关实现细节待核实。

总体判断：仓库运维良好，train.sh 和 eval.sh 提供完整的一次性运行入口，支持 OUTPUT_DIR/IMAGENET_PATH 等环境变量覆盖路径。

定量总结：

• RiT-XL 以 676M 参数在 ImageNet 256×256 上达到 FID 1.45 (no CFG) /

  1.14 (CFG)

• 比 RAE (DiT-DH-XL) 少 19% 参数（676M vs 839M），FID 更好
• 5 步达到 FID 2.0，10 步达到 FID 1.25——无需蒸馏
• 编码器越小越好：DINOv2-S (384D) 优于 DINOv2-B (768D)

失败案例与局限性：

• DINOv2-Base

  编码器效果更差，原因未完全解释（仅说高维特征更难建模）

• v-prediction 在 DINOv2 空间到底差多少？原文只说 "~2 FID

  worse"，未精确给出

• 只在 ImageNet 256×256 上验证，未扩展至更大分辨率或文本条件
• 依赖 RAE 的编码器-解码器 pipeline，迁移至其他编码器需要重新验证
• 为什么 weight decay=0.0？原文未解释（但常见于其他扩散模型）

论文自己的结论：DINOv2 特征空间的统计特性（近各向同性、边际近高斯、插值在流形上）使 x-prediction flow matching 回归良态，这消除了对专门预测头（如 DDT）或黎曼运输（如 RJF）的需求。vanilla DiT + x-prediction 就够了。

局限性（来自原文 + 补充分析）：

1. 仅在 ImageNet 256×256 验证，未测试分辨率缩放或文本条件
2. 对 DINOv2-Base 效果下降的原因分析不够深入
3. x-prediction vs v-prediction 的消融数据不够完整（只说 ~2 FID

   worse）

4. 依赖于 RAE 提供的冻结编码器-解码器，不是端到端方案

未来方向：

• 探索 x-prediction 在其他表示空间（SigLIP, CLIP, MAE,

  i-JEPA）的普适性

• 将几何分析框架用于指导 tokenizer / 编码器设计
• RiT + REPA 的联合训练（表示对齐 +

  正确预测目标，可能进一步加速）

• 扩展到文本条件生成和更大分辨率

对我们的视觉编码器研究的具体启发：

1. 4-axis geometry analysis 可以直接复用：评估一个 visual tokenizer 时，不仅要看重建指标（PSNR/FID/LPIPS），还应该检查这四项：

   • 有效秩（方差分布均匀性）
   • 边际峰度（每坐标的高斯性）
   • 协方差条件数（数值稳定性）
   • 线性插值重建误差（流形连续性）

   这些指标可以预测 tokenizer 的「生成友好性」。

2. SSL 预训练目标影响编码器的隐式几何：DINOv2 的 iBot / self-distillation 目标天然产生各向同性 + 边际高斯分布，这是生成友好的。SD-VAE 的重建损失（L1 + LPIPS + adversarial）产生中间结果。这说明「生成 tokenizer」和「表示 tokenizer」的优化目标会导致不同的隐式几何——这可能是统一表示和生成的关键设计空间。

3. x-prediction 是最小入侵式生成方案：不改变架构，只改变预测目标。这降低了编码器设计者的负担——他们只需要关注编码器本身的质量，不需要为生成侧的特殊需求定制架构。

4. scaling law 新视角：编码器不是越大越好（DINOv2-Small > DINOv2-Base），可能存在一个「最优特征维度 vs 生成性能」的 trade-off curve。这值得系统研究。

可做的后续实验：

• 对 SigLIP / CLIP / MAE / i-JEPA 特征做同样的 4-axis 几何分析
• 尝试在 FAE feature 或自定义 tokenizer 上直接跑 x-prediction
• 设计 tokenizer

  训练目标时加入几何正则化（如最大化有效秩、最小化边际峰度）