Learning on the Manifold: Unlocking Standard Diffusion Transformers withRepresentation Encoders

该论文深入探讨了标准扩散Transformer (Diffusion Transformers, DiT) 在处理预训练表示编码器（如DINOv2）生成的高维特征空间时遇到的收敛失败问题。先前的研究（如RAE）将此归因于模型容量瓶颈，并提出了通过“宽度扩展”（width scaling）增加DiT模型宽度来解决。然而，本文通过严谨的几何分析，指出这种失败并非源于容量不足，而是由于一个根本性的几何冲突，即“几何干涉”（Geometric Interference）。
核心问题与几何分析：
论文指出，DINOv2等表示编码器生成的特征向量并非均匀分布在环境欧几里得空间中，而是严格限制在一个固定半径的超球面上 ()，所有的语义信息都编码在角度分量中，而径向分量几乎为常数（由于LayerNorm等操作）。标准欧几里得流匹配 (Euclidean Flow Matching, EFM) 算法在构建从数据到高斯噪声的条件概率路径时，采用线性插值方式。这种线性路径在超球面上形成一条“弦”，穿过超球面的低密度内部区域（off-manifold区域），而非沿着流形表面。这意味着模型被迫在表示空间未定义的区域学习速度场，同时耗费其有限的建模能力去最小化与超球面几何结构冲突的径向误差。论文通过消融实验证明，当忽略径向误差，仅优化角度分量时，即使是较窄的DiT模型也能快速收敛，这有力反驳了“容量瓶颈”假说，并强调了“几何干涉”才是症结所在。先前的“宽度扩展”解决方案，本质上是赋予模型足够多的参数来“记忆”通过这些无效区域的、不适定的径向向量场，而非从根本上解决问题。
核心方法：黎曼流匹配与雅可比正则化 (Riemannian Flow Matching with Jacobi Regularization, RJF)：
为解决上述几何冲突，论文提出了RJF方法，旨在直接在超球面流形上定义生成过程：

黎曼流匹配 (Riemannian Flow Matching, RFM)：

测地线概率路径： RJF将欧几里得线性插值替换为球面线性插值 (Spherical Linear Interpolation, SLERP)。给定数据和噪声，测地线路径通过SLERP定义为：
​​​​​​​
其中是和之间的测地距离（角度）。SLERP确保了始终保持在超球面上 (，从而消除了径向分量的塌缩问题，并确保生成过程遵循流形的内在几何结构。
切空间速度场： 这种约束意味着速度向量必须严格位于处的切空间中，即。目标黎曼速度场通过对测地线路径求导得到：
​​​​​​​ ​​​​​​​
通过优化预测模型去逼近，模型仅学习语义转换（角度变化），而避免在径向分量上浪费容量，从而有效解决了“几何干涉”。损失函数定义为：
​​​​​​​
测地线（指数映射）积分： 在采样阶段，RJF采用测地线（指数映射）积分（Geodesic (Exponential Map) Integration）来更新，确保轨迹沿大圆弧精确移动，保持在流形表面。更新公式为：
​​​​​​​ ​​​​​​​ ​​​​​​​
其中是预测的切向速度。
雅可比正则化 (Jacobi Regularization)：

尽管RFM确保了路径在流形上，但标准速度匹配目标仍然是“几何无感知”的，它均匀地处误差。在正曲率流形（如超球面）上，测地线会聚焦，导致速度误差的传播是非线性的。为了最大化生成保真度，RJF通过引入雅可比场（Jacobi Fields）来修正损失函数。
雅可比场量化了速度扰动引起的测地线分离。通过求解超球面的雅可比方程，论文推导出一个几何权重因子：

这个权重因子在（数据附近）处对误差进行降权，因为测地线聚焦会缓解扰动，而在（噪声附近）处优先处理误差，因为生成轨迹必须精确对齐特征流形。最终的雅可比正则化目标函数为：
​​​​​​​
通过优化这个曲率校正后的目标，模型能更有效地学习高维潜在空间。

实验结果与贡献：

收敛性与效率： RJF使标准DiT-B架构（131M参数）在ImageNet 256x256图像生成任务上有效收敛，无引导FID达到4.95（200 epoch），有引导FID达到3.37，而现有方法在此设置下难以收敛。在大规模DiT-XL模型上，RJF在80个epoch内实现了3.62的FID，优于标准流匹配基线（4.28）和基于VAE的DiT模型（4.29）。
消除宽度扩展需求： 论文证明，通过尊重潜在空间的内在几何结构，标准DiT架构无需计算成本高昂的“宽度扩展”即可实现高性能生成。
泛化能力： RJF在不同DiT架构（如LightingDiT、DDT、DiT DH）和不同表示编码器（DINOv2、SigLIP、MAE）上均表现出一致的性能提升，表明其方法的普遍性。
消融研究： 详细的消融研究证实了RFM和Jacobi正则化各自对性能提升的关键作用。仅将噪声投影到球面上带来的改善微乎其微，而采用RFM则能显著提升性能，再结合Jacobi正则化则能进一步优化。
半径对解码器敏感性： 论文还发现，在推理阶段调整投影半径可以进一步提升生成质量，表明RAE解码器对特征幅值敏感。

总而言之，该论文通过深入的几何洞察，揭示了将标准扩散Transformer应用于表示编码器特征空间时收敛失败的根本原因，并提出了一种创新的几何对齐方法RJF，成功解锁了现有DiT架构的潜力，实现了高效且高质量的生成。