流形上的学习：用表征编码器解锁标准扩散Transformer Learning on the Manifold: Unlocking Standard Diffusion Transformers with Representation Encoders

RJF方法的训练过程（Algorithm 1）包含以下步骤：（1）从数据集采样batch $x \sim D$，并归一化到单位球面 $x \leftarrow x/\|x\|$；（2）采样各向同性高斯噪声 $\epsilon \sim \mathcal{N}(0, I)$ 并归一化 $\epsilon \leftarrow \epsilon/\|\epsilon\|$；（3）使用Logit-Normal分布采样时间步 $t_{\text{raw}} \sim \text{LogitNormal}(\mu, \sigma)$，然后应用时间偏移 $t = \frac{s \cdot t_{\text{raw}}}{1+(s-1)t_{\text{raw}}}$；（4）计算测地线距离 $\Omega = \arccos(\langle \epsilon, x \rangle)$，使用SLERP插值 $x_t = \frac{\sin((1-t)\Omega)}{\sin(\Omega)} x + \frac{\sin(t\Omega)}{\sin(\Omega)} \epsilon$；（5）计算目标速度场 $u_t^M$（切空间向量）；（6）计算雅可比权重 $w_t = \text{sinc}^2((1-t)\Omega)$；（7）预测速度 $\hat{v} = v_\theta(x_t, t)$，投影到切空间 $\hat{v}_{\text{proj}} = \hat{v} - \langle \hat{v}, x_t \rangle x_t$；（8）计算正则化损失 $L = w_t \cdot \|\hat{v}_{\text{proj}} - u_t\|^2$。采样过程（Algorithm 2）使用测地线积分器：每步预测切空间速度，通过闭式三角旋转 $x_{t+\Delta t} = \cos(\|v\|\Delta t)x_t + \sin(\|v\|\Delta t)\frac{v}{\|v\|}$ 更新位置，确保轨迹沿大圆精确行进。

本文的技术新颖性体现在多个方面。首先，几何干扰的诊断是全新的——此前所有工作（包括RAE）都将其归因于容量不足，而本文通过精心设计的消融实验（掩蔽径向损失后窄模型也能收敛）给出了决定性证据。其次，将黎曼流匹配应用于表征编码器的生成建模是首创，此前RFM主要用于蛋白质骨架生成（SE(3)）、分子构象（环面）等科学领域。第三，雅可比正则化是理论上的重要贡献：标准RFM假设平坦度量，而本文通过求解超球面上的雅可比方程得到解析权重 $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$，这是几何感知的注意力机制——它自适应地降低数据端（$t=0$）的损失权重（聚焦效应抑制误差传播），提高噪声端（$t=1$）的权重（生成轨迹必须精确对齐流形）。最后，推理时使用指数映射积分器替代标准ODE求解器，保证离散化后轨迹仍在流形上，这是工程上的重要创新。

本文存在几个值得改进的弱点。第一，超球面假设过于理想化：实际中DINOv2等编码器的特征虽然因LayerNorm被约束在近似球面上，但并非严格球面，存在径向方差。可以考虑学习一个自适应的流形结构（如椭球面或更一般的黎曼流形），而非硬编码超球面假设。第二，雅可比权重函数 $\lambda(t, \Omega) = \text{sinc}^2((1-t)\Omega)$ 是基于解析推导的闭式解，但可能不是最优的。可以探索学习一个自适应的时间-距离依赖权重函数。第三，投影半径R作为超参数需要手动调优，可以考虑将其纳入训练过程，学习一个最优的输出缩放因子。第四，论文仅使用球面噪声投影（将高斯噪声归一化到单位球面），这种投影会改变噪声分布的各向同性，可能引入偏差。可以探索保测度的球面投影方法。第五，实验仅覆盖DiT系列架构，未验证在扩散模型之外的生成框架（如GAN、VAE）上的适用性。

基于本文的成果，未来研究可以从以下几个方向展开。首先，将黎曼流匹配推广到更一般的流形结构，如Stiefel流形（正交矩阵空间）、Grassmann流形（子空间空间）等，这可能对表征学习本身产生深远影响。其次，探索自适应几何学习：不预设流形类型，而是从数据中学习最优的黎曼度量，这可以通过可微分的度量张量实现。第三，将RJF应用于其他模态（如文本表征、音频表征），因为LayerNorm在Transformer中是普遍使用的，其他模态的表征空间可能也存在类似的超球面几何。第四，研究更高效的测地线积分方法，当前的指数映射积分器虽然精确但每步需要额外的三角函数计算。第五，将几何正则化思想应用于训练过程本身，如设计几何感知的学习率调度或优化器。第六，探索RJF与现有表征对齐方法（如REPA）的结合，可能产生协同效应。最后，研究在更高维度（如4K分辨率）和更大规模模型上的表现。

本文的复现条件较为良好。代码已在GitHub开源（https://github.com/amandpkr/RJF），这大大降低了复现门槛。实验基于公开的ImageNet-1K数据集和预训练的DINOv2-B、SigLIP、MAE编码器，这些资源均可公开获取。训练使用标准配置：Adam优化器（beta1=0.9, beta2=0.95）、学习率2e-4、batch size 1024、80 epochs、EMA衰减0.9995。这些超参数与LightingDiT一致，便于对比。然而，完整复现需要显著的计算资源：DiT-XL（677M参数）在ImageNet 256x256上训练80 epochs，使用全局batch size 1024，通常需要多块高端GPU。方法的核心实现（SLERP插值、雅可比权重计算、测地线积分器）并不复杂，但需要仔细处理切空间投影和数值稳定性（如 $\sin(\Omega)$ 接近零时的除法）。总体而言，复现难度中等，主要挑战在于计算资源而非算法复杂度。