[Paper] 在流形上学习：利用表示编码器解锁标准 Diffusion Transformers
Source: arXiv - 2602.10099v1
概述
论文 “Learning on the Manifold: Unlocking Standard Diffusion Transformers with Representation Encoders” 阐明了传统扩散‑Transformer 模型在直接从高级表征编码器（如 CLIP、DINO）生成数据时为何会遇到困难。作者指出，问题并非模型容量不足，而是欧氏扩散动力学与编码器特征所在的超球面流形之间存在 几何不匹配。通过重新设计扩散过程以遵循底层黎曼几何，他们使得一个标准的 131 M 参数 Diffusion Transformer（DiT‑B）能够成功训练，并达到最先进的图像合成质量。
关键贡献
- 识别几何干扰: 证明欧几里得流匹配（Euclidean flow‑matching）会迫使概率质量穿过编码器超球面特征空间的低密度内部，导致训练崩溃。
- 黎曼流匹配与Jacobi正则化 (RJF): 引入一种遵循特征流形测地线并补偿曲率引起误差传播的扩散公式。
- 无需宽度扩展: 表明标准 DiT‑B 架构（131 M 参数）在不使用此前被认为必需的宽度增加技巧的情况下即可收敛。
- 实证验证: 在 ImageNet‑256 上实现 3.37 的 FID，突破了先前扩散‑Transformer 方法会发散的局限。
- 开源发布: 提供干净的 PyTorch 实现（https://github.com/amandpkr/RJF），便于可重复性研究和后续工作。
方法论
-
Problem Setup – 作者从一个 representation encoder 开始，该编码器将图像映射到高维超球面上的点（例如，归一化的 CLIP embeddings）。传统的扩散模型定义了一个 Euclidean stochastic differential equation (SDE)，在环境空间中在噪声和数据之间进行插值。
-
Geometric Analysis – 通过可视化编码器特征的密度，他们观察到大部分质量位于球面表面，而 Euclidean diffusion trajectory 大部分时间位于球体内部，那里没有真实数据。这种 “geometric interference” 导致梯度信号差，训练失败。
-
Riemannian Flow Matching – 与 Euclidean SDE 不同，他们在流形上构建了一个 Riemannian flow：
- 扩散路径遵循 geodesics（球面上的最短路径）。
- 速度场通过 Riemannian optimal transport 定义，确保概率质量在整个训练过程中保持在流形上。
-
Jacobi Regularization – 曲率在沿 geodesics 积分时会放大小误差。作者借用了微分几何中的 Jacobi equation 来正则化学习到的向量场，使流在曲率引起的漂移下保持稳定。
-
Training Pipeline – RJF loss 替代了 vanilla Diffusion Transformer (DiT‑B) 中的标准 flow‑matching loss。无需架构更改、额外层或宽度扩展。
结果与发现
| 模型（Params） | 训练设置 | FID（ImageNet‑256） | 备注 |
|---|---|---|---|
| DiT‑B (131 M) + Euclidean Flow | 标准 | 未收敛 | 由于几何干扰导致崩溃 |
| DiT‑B (131 M) + RJF (proposed) | 相同超参数 | 3.37 | 匹配或超过先前宽度扩展基线 |
| DiT‑L (large, 300 M) + Euclidean Flow (baseline) | 更宽的模型 | ~3.5 | 需要 >2× 参数才能获得可比质量 |
关键要点
- 几何对齐 是主要瓶颈，而非原始容量。
- RJF 在 相同 模型规模下恢复了稳定训练，计算和内存消耗约比宽度扩展的替代方案减半。
- 定性样本显示纹理更锐利、伪影更少，尤其是在编码器流形曲率较高的区域。
实际意义
- Cost‑Effective High‑Fidelity Generation: 开发者现在可以在普通 GPU 上部署 diffusion‑transformer 流水线，而无需增大模型规模，使大规模图像合成更具成本效益。
- Plug‑and‑Play with Existing Encoders: RJF 可与任何归一化表示编码器（CLIP、DINO、SimCLR）配合使用，打开了基于语义嵌入的 条件 生成的大门，无需重新训练编码器。
- Better Integration in Multi‑Modal Systems: 由于生成过程遵循编码器的几何结构，下游任务（例如文本到图像、风格迁移）使用相同嵌入时会更加一致。
- Reduced Training Instability: 团队可以避免在扩大宽度或添加临时技巧的反复试验；RJF 损失是对标准 diffusion 损失的直接替换。
- Potential for Other Manifolds: 同样的黎曼流匹配思路可以适用于图嵌入、双曲空间或任何已知几何结构的潜在空间，扩大其在视觉之外的影响。
限制与未来工作
- 流形假设: RJF 假设编码器的输出位于一个行为良好的超球面上。产生未归一化或高度各向异性嵌入的编码器可能需要额外的预处理。
- Jacobi 正则化的计算开销: 虽然模型规模保持不变，计算曲率感知正则化会给训练时间增加一个适度的常数因子。
- 对超高分辨率的可扩展性: 实验仅限于 256×256 图像；扩展到 1024×1024 或视频生成可能会出现新的几何挑战。
- 更广泛的流形类型: 未来工作可以探索自适应流形学习（与扩散过程共同学习度量），或将 RJF 应用于非球面流形，例如用于多模态嵌入的乘积流形。
如果你有兴趣自行尝试 RJF，作者在 GitHub 上提供了可直接运行的实现和预训练检查点。该方法为现有的扩散‑Transformer 流水线提供了一个简洁的、几何感知的升级，将理论洞见转化为实际的性能提升。
作者
- Amandeep Kumar
- Vishal M. Patel
论文信息
- arXiv ID: 2602.10099v1
- 分类: cs.LG, cs.CV
- 发布: 2026年2月10日
- PDF: 下载 PDF