Learning on the Manifold: Unlocking Standard Diffusion Transformers with Representation Encoders¶

会议: ECCV2026
论文: ECCV 原文
代码: https://github.com/amandpkr/RJF
领域: 图像生成
关键词: 表示编码器、球面潜空间、黎曼流匹配、Jacobi 正则化、扩散 Transformer

一句话总结¶

RJF 将表示编码器潜空间中的直线流匹配改为球面测地线流匹配，并按曲率导致的误差传播加权训练，使无需加宽的 131M 参数 DiT-B 在 ImageNet 256 类条件生成上达到 FID 4.95，加入引导后为 3.37，均训练 200 epoch。

研究背景与动机¶

潜空间扩散通常先用 VAE 压缩图像，再训练生成模型恢复其中的图像分布。VAE 以重建为目标，其潜变量未必像 DINOv2 这样的表示编码器一样具有强语义结构，因此 REPA 等方法会额外对齐生成模型与预训练特征。RAE 则直接用冻结的表示编码器替代 VAE 编码器，配合图像解码器，让生成模型直接学习语义特征分布。然而，把潜变量换成高维表示并不意味着标准 DiT 会顺利收敛：RAE 用增大模型宽度来处理困难，本文希望判断，真正需要增加的是表示能力，还是需要修改训练所要求模型完成的任务。

作者观察到 DINOv2-B 的 768 维特征具有非常集中的范数，近似分布在半径为 \(\sqrt{768}\) 的球壳上，而非填满整个欧氏空间。标准流匹配却用直线连接数据与噪声：即使两端都投到球面上，连线仍会穿过内部。结果是，网络既要学习语义方向，又要拟合训练路径人为引入的径向变化。单图像过拟合实验中，宽度 384 或 512 的模型在完整欧氏损失下失败，但移除径向误差后能学习角向部分。这支持“目标与数据几何不匹配”的解释，却还不足以证明大规模生成中的所有容量限制都不存在。

因此，本文并非简单提出“把特征归一化就好”，而是把路径、监督及采样放在同一个球面几何中考虑。噪声端点归一化只能保证起点合法，无法保证中间状态合法；而路径合法之后，不同时刻的速度误差仍可能因曲率而产生不同影响。核心 idea：用球面测地线消除离流形的训练路径，再用 Jacobi 场启发的权重校正误差传播，让标准 DiT 把容量用于学习语义方向，而非补偿人为引入的径向运动。

方法详解¶

整体框架¶

Riemannian Flow Matching with Jacobi Regularization，简称 RJF，保留标准扩散 Transformer 主干。训练图像经冻结的表示编码器得到特征，特征与高斯噪声都归一化到单位球面，再通过“球面测地线路径”构造中间状态和切向速度目标；“Jacobi 曲率加权”只影响训练损失。推理时从球面噪声出发，通过“测地线采样与解码”逐步生成特征，最后交给 RAE 解码器生成图像。

这里的类条件来自 ImageNet 类别，不是文本提示。原文把时间定义为数据端 \(t=0\)、噪声端 \(t=1\)，因此生成要从噪声端反向积分到数据端；不能照着训练插值方向把图像再次推向噪声。训练时真实图像提供监督，推理时不需要真实图像或编码器为待生成样本提供特征。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["训练图像与高斯噪声"] --> B["冻结编码器<br/>与单位球面归一化"]
    B --> C["球面测地线路径"]
    C -->|中间状态、时间与类别| D["标准 DiT<br/>预测速度"]
    C -.->|切向目标，仅训练| E["Jacobi 曲率加权"]
    D -.->|速度误差，仅训练| E
    N["推理：球面噪声与类别"] --> F["测地线采样与解码"]
    D -->|训练后速度场| F
    F --> G["生成图像"]

关键设计¶

1. 球面测地线路径：改变监督轨迹，而不只是改变端点分布

作者先把每个特征向量和采样的高斯噪声归一化，使两者处在同一单位球面上。随后用球面线性插值 SLERP 连接两端，沿大圆的短弧构造任意时刻的训练状态。这不是“直线插值完再归一化”的同义说法：SLERP 对角距离进行参数化，使路径和目标速度本身都由球面几何定义。模型每次接收的输入因而仍保持单位范数，不再看到因弦线穿过球心方向而产生的低范数中间状态。

目标速度由这条测地线路径对时间求导获得，其方向位于当前位置的切空间中。对单位球面，切向性意味着位置与速度正交：

\[ v_t \cdot x_t = 0. \]

这个条件说明，局部运动改变的是方向，而不是向球内缩或向球外扩。原文用预测速度与测地线目标速度的平方误差训练 DiT，球面诱导度量允许在环境空间中计算这一误差。需要区分“目标没有径向速度”和“任意网络输出都自动切向”：前者由路径保证，后者还涉及输出约束或实现细节，正文并未充分展开；不应仅凭目标的切向性就宣称预测误差的径向分量恒为零。

2. Jacobi 曲率加权：让速度误差的重要性随几何传播而变

测地线路径解决了“在哪里学”，却没有完全解决“各处误差是否一样重要”。在球面这样的正曲率空间里，邻近测地线可以逐渐聚拢，某个时刻的速度扰动对后续位置的影响并不等同于平直空间中的线性累积。Jacobi 场描述邻近测地线因扰动产生的分离，作者据此设计与时间、两端夹角相关的权重，对流匹配误差重新加权，而不是再加一个独立的语义教师或特征对齐网络。

按原文的时间约定，这个权重降低数据端附近误差的相对贡献，更重视噪声端附近的速度准确性。直觉上，采样从噪声端启动，这一阶段决定最初朝哪个语义方向移动；曲率导致的聚拢又会影响后续扰动的实际作用。正文式 (13) 将权重写为 sinc 平方形式，式 (14) 用其加权速度误差；本地 PDF 文本中的公式存在缺括号与符号粘连，因此这里保留机制与方向，不把修复后的表达式冒充作者精确公式。消融显示该步骤是路径修正后的精炼，而不是主要收益来源。

3. 测地线采样与解码：训练在球面上，生成也沿球面走

如果训练使用球面路径，推理却用普通欧氏 Euler 步直接把切向速度加到当前位置，有限步长仍会让状态离开球面。因此作者用指数映射积分：将当前位置沿预测切向方向旋转，旋转角由速度大小和步长共同决定。对一个给定的切向量，这个更新落在大圆上；它保持单步几何约束，但不意味着用有限步数精确解出了整个随时间变化的神经 ODE。作者还用旋转与归一化处理数值漂移，实验使用 50 步采样。

采样后的单位方向不能不加处理地当作原始编码器特征输入解码器。作者分析了最终特征的投影半径，发现恢复 DINOv2-B 原始范数约 27.7 时 FID 为 7.79，而半径约 45 时可达 6.77。这表明解码器对幅值敏感，“语义主要在角向”并不等于“解码对范数不敏感”。范数调整发生在解码阶段，与训练中沿单位球面学习方向是两个不同环节；这也是复现时不能漏掉的设置，且最佳半径并非由前面的球面假设直接推出。

一个完整示例¶

以 ImageNet 的一个类别和一张 \(256\times256\) 训练图像为例，DINOv2-B 为图像产生一组 768 维特征向量。下面只看其中一个向量，实际生成器共同处理整个空间 token 序列，而不是独立生成每个 token。该向量先被归一化，同时抽取一个高斯向量并归一化；类别标签则作为 DiT 的条件。

为了直观看清直线路径的问题，可以考虑两个近似正交的单位端点。这是几何示意计算，不是新增的实验结果：

\[ \left\|\frac{x+\epsilon}{2}\right\| \approx \frac{1}{\sqrt{2}} \approx 0.707. \]

也就是说，直线中点明显落到球内，而 SLERP 中点仍在球面上。训练随机抽一个时刻，让 DiT 根据该球面状态预测对应切向速度，再用 Jacobi 权重衡量误差。推理则不再使用这张真实图像：从一组新的球面噪声 token 出发，给定类别，按 50 步测地线更新产生方向，调整输出特征幅值后由 RAE 解码器还原图像。

损失函数 / 训练策略¶

最终目标是经过几何权重调整的速度匹配误差，不是图像像素重建损失，也不是额外的 DINO 特征蒸馏损失。表示编码器被冻结，本文生成训练使用 RAE 解码器；文中报告的 131M、459M、677M 是对应生成主干的参数规模，不能直接理解为包含编码器、解码器与引导模型在内的总系统成本。

默认训练遵循 LightningDiT 配置：ImageNet-1K、\(256\times256\) 分辨率、80 epoch、全局 batch size 1024。优化器是 Adam，学习率 \(2\times10^{-4}\)，两个动量系数为 0.9 和 0.95；梯度范数裁剪为 1.0，EMA 衰减为 0.9995。作者还沿用 RAE 的维度相关噪声调度偏移，设置为 \(n=4096\)，它不等同于上述 DINO 特征的 768 维。

小模型另报告 200 epoch 的延长训练。需要引导时采用 Autoguidance，强度为 1.6，而不是把所有结果默认理解成无分类器引导。原文未在此展开引导辅助模型的额外资源消耗，因此“无需加宽”是主干架构层面的结论，不宜直接换算成端到端成本节省比例。

实验关键数据¶

主实验¶

任务为 ImageNet-1K \(256\times256\) 类条件生成，以 50,000 张生成图像评估。FID 衡量真实与生成图像的特征分布差异，越低越好；正文没有清楚标明 FID 参考统计使用的具体划分。下表取原文表 1，同一规模均为 DINOv2-B 特征、80 epoch、无引导，参数规模在每行配对中一致。

生成主干	参数	EFM FID	RJF FID	FID 绝对降低
DiT-B	131M	24.21	6.77	17.44
DiT-L	459M	6.31	4.21	2.10
DiT-XL	677M	4.28	3.62	0.66

表 2 中，XL 的无引导结果还包括 IS 186.2、Precision 0.82、Recall 0.52；IS 越高通常表示分类可辨性与类别多样性更好，Precision/Recall 分别刻画生成分布的质量与覆盖。相同 80 epoch 表块中，REPA-E 的无引导 FID 为 3.46，REG 为 3.40，均低于 RJF 的 3.62。因此本文优势不能表述为所有方法上的最佳 FID；同时 0.52 的 Recall 提醒我们不能仅凭高 Precision 宣称覆盖最好。

消融实验¶

下表对应原文表 3 及第 4.3 节的文字说明，均为标准 DiT-B/1 与 DINOv2-B，评估条件同上；除最后一行外均无引导。SN 仅把噪声及潜变量投到球面，RFM 改用测地线，RJF 再加入 Jacobi 加权。

配置	Epoch	FID	解释
欧氏流匹配 EFM	80	24.32	直线路径基线
EFM + SN	80	21.99	端点合法，路径仍穿球内
RFM	80	7.06	测地线路径带来主要改善
RJF	80	6.77	Jacobi 加权再降低 0.29
RJF	200	4.95	延长训练后的无引导结果
RJF + Autoguidance	200	3.37	引导强度 1.6

原文存在数值口径不一致：表 1 的 B 模型 EFM 是 24.21，表 3 是 24.32；第 4.2 节称 XL 相对欧氏基线改善 1.19，但表 1 的 4.28 与 3.62 相减为 0.66。本笔记分别保留各表数值，差值按对应表计算，不将两组基线合并，也不沿用错误差值。

关键发现¶

最大收益来自改变路径：同一消融中 SN 到 RFM 的 FID 从 21.99 降到 7.06，远大于 RFM 到 RJF 的 0.29 改善。
几何收益没有在 XL 上消失，但绝对幅度比 B 小；“标准模型完全无法收敛”不能泛化到所有规模，因为 XL 的 EFM 本身已达 4.28。
第 5.2 节报告 SigLIP 的基线/RJF FID 为 130.21/10.39，MAE 为 50.48/19.82，支持跨编码器有效，但并不意味着不同表示最终质量相同。

亮点与洞察¶

将“高维特征难生成”具体化为训练路径的几何冲突，给出可操作的诊断方式。先拆径向与角向误差，比直接把失败归因于模型太小更有解释力。
从端点归一化到测地线路径再到曲率加权，消融能分辨每一步的作用。最可迁移的思路是检查中间状态，而不只检查真实数据和噪声的边缘分布。

局限与展望¶

作者明确受计算资源限制，主要在 80 epoch 预算下比较；小模型补到 200 epoch，没有建立与长训练强方法完全等预算的全面优势。
本笔记判断：单图过拟合支持几何干扰解释，却不能排除大规模、多模态分布中的容量瓶颈；近似球壳也不证明真实语义分布均匀覆盖整个球面。
本笔记判断：切向输出处理、解码半径选择和引导额外成本需要更完整的复现说明。可优先检验逐层归一化后的实际几何，而非假定任意 LayerNorm 输出都严格满足同一球面模型。
本笔记判断：验证范围主要是 ImageNet 256 图像生成，尚不能直接外推到文本到图像、视频或其他分辨率；多随机种子和训练耗时报告也有助于评估收益稳定性。

评分¶

新颖性: 4/5。几何诊断与表示空间生成的结合有价值，基础黎曼流匹配并非首创。
实验充分度: 3/5。覆盖尺度、编码器和组件消融，但训练预算、复现细节与部分数字口径限制结论。
写作质量: 3/5。核心直觉易理解，不过 SOTA 措辞、差值及几何保证表述需要更谨慎。
价值: 4/5。对不愿扩大主干宽度的表示空间生成研究提供了清晰、可检验的方向。