Entropic RNOT：打破流形束缚，实现高效几何最优传输

TL;DR

传统的流形最优传输（OT）往往在“计算效率”与“几何忠实度”之间艰难权衡。本文提出的 Entropic RNOT 框架通过引入熵正则化和神经半对偶参数化，不仅提升了在大规模数据上的扩展性，还首次实现了流形上 OT 的摊销推理（Amortization）。在蛋白质-配体对接实验中，它展现了极强的现实应用潜力，显著优化了刚体姿态精度。

背景定位：为什么要关注流形 OT？

在机器学习中，数据往往并非分布在平坦的欧几里得空间，而是存在于带曲率的流形上。例如：

• S2 (球面)：地理坐标或全景图像；
• SO(3) (旋转群)：物体的姿态控制；
• SPD(3) (正定矩阵流形)：扩散张量成像或协方差描述子；
• SE(3)：刚体在 3D 空间中的运动。

如果强行使用欧几里得距离来度量这些数据，会产生严重的几何畸变。

核心动机：现有方法的痛点

1. 离散 Sinkhorn 的局限：传统的离散算法内存消耗随样本量 $N$ 的平方增长（$O(N^2)$），且无法对未见样本进行快速推理。
2. 确定性映射的盲目性：常规神经传输尝试学习一个硬性的映射 $y=T(x)$，但在正则化环境下，真实的传输目标往往是一个耦合分布。

核心贡献：Entropic RNOT 的数学直觉

Entropic RNOT 的逻辑非常优雅。它不再直接学习映射，而是学习 Schrödinger 势函数。

1. 神经半对偶参数化 (Neural Semidual)

作者利用“拉回机制（Pullback）”，将欧几里得空间的神经网络映射到流形上，定义目标侧势函数 $g_{h}eta$。通过极大化下述目标函数： $${\mathcal{J}}_{\epsilon }(g)=\int gdu+\int {\mathcal{T}}_u^{\epsilon }gd\mu$$ 其中 ${\mathcal{T}}_u^{\epsilon }$ 是软 c-变换，它能捕捉到流形的测地线距离成本。

2. 两种内在代理映射 (Intrinsic Surrogates)

当我们需要从概率耦合中提取一个点坐标时，Entropic RNOT 提供了两个工具：

• 重心投影 (Barycentric Projection)：适用于负曲率流形（如 Hyperbolic Space），具有良好的理论收敛保证。
• 热平滑模式提取 (Heat-smoothed Mode Finding)：利用流形热核对离散原子进行平滑，即使目标分布是离散的，也能通过梯度上升找到最优传输点。

图 1：算法在不同几何体（球面、旋转群、正定矩阵、双曲面）上的表现对比

实验战绩：速度与精度的双重飞跃

1. 扩展性测试

相较于传统的离散流形 Sinkhorn 算法，Entropic RNOT 的训练时间和内存占用在样本量增加时保持恒定（受 batch size 驱动），推理吞吐量呈线性扩展。

图 2：随样本规模 $N$ 增加时的计算开销对比

2. 实战：蛋白质-配体对接 (SE(3) Pose Refinement)

这是本文最令人兴奋的应用。在 CrossDocked2020 数据集上，Entropic RNOT 将原始对接工具 GNINA 产生的候选位姿进行“降噪”和“精修”。

• 关键数据：Top-1 RMSD 从 11.24 Å 骤降至 3.47 Å。
• 直观效果：算法成功将许多错误的偏离位姿修正到了晶体结构真实的结合口袋中。

图 3：Entropic RNOT (蓝色) 成功将不准确的初始位姿引导至与晶体结构 (绿色) 重合的位置

深度洞察

Entropic RNOT 的成功在于它承认了“不确定性”。通过熵正则化，它不是在找一个绝对的 $x\to y$ 映射，而是在学习流形上的统计规律。其利用热核进行平滑处理的思路，巧妙地避开了黎曼几何中非凸性的坑。

局限性展望：目前该研究主要集中在固定熵参数 $\epsilon$ 的情况，对于 $\epsilon \to 0$（即逼近硬最优传输）的过程仍需进一步探索。同时，在生物分子建模中，如何引入受体口袋的动态背景信息将是一个重要的研究方向。

总结

Entropic RNOT 为处理流形结构数据提供了一套可扩展、有理论保障且在工程上易于实现的闭环方案。对于从事几何深度学习和计算生物学的研究者来说，这篇论文提供了极其宝贵的工具箱。