arXiv:2605.04255stat.MLcs.LG2026-05

提出统一框架,让机器学习在弯曲空间中更准地做数据传输

Entropic Riemannian Neural Optimal Transport

论文配图:Entropic Riemannian Neural Optimal Transport
图 1 · 摘自论文原文
  • 用神经网络学出曲面内嵌的最优传输映射
  • 在球面、旋转群等空间上性能超越传统方法
  • 适合处理蛋白对接等需精准姿态匹配的任务

许多机器学习问题涉及定义在球面、旋转群、双曲空间及一般黎曼流形上的数据,欧氏几何会扭曲距离与平均值,导致最优传输(OT)失真。现有流形OT方法多采用可泛化的离线映射,而熵正则化提升了离散OT的可扩展性,但二者优势长期分离。本文提出熵正则黎曼神经最优传输(Entropic RNOT),将内在熵正则OT与流形上的可泛化离线评估统一起来。方法通过神经拉回参数化学习单个目标侧的薛定谔势能,恢复诱导吉布斯耦合,并利用条件分布构建内在传输代理:在Cartan-Hadamard流形上为巴氏投影,在随机完备流形上为热平滑条件代理,后者可将可能原子型目标分布变为绝对连续分布。对固定正则化参数ε>0,证明该假设类在强概率度量下可恢复熵正则最优耦合。由此,巴氏代理在L²中收敛,热平滑代理在固定热时间下稳定且热时间趋零时渐近无偏。理论保证适用于紧支撑数据在可能非紧流形上的情形。实验表明,本方法在S²、SO(3)、SPD(3)、SE(3)和H²基准上表现优于或匹敌欧氏、切空间与对数欧氏基线,相对于离散流形Sinkhorn更具可扩展性;在蛋白质-配体对接任务中,无需重训练或逐例优化即可精修姿态。

原文摘要 · Abstract (English)

Many machine learning problems involve data supported on curved spaces such as spheres, rotation groups, hyperbolic spaces, and general Riemannian manifolds, where Euclidean geometry can distort distances, averages, and the resulting optimal transport (OT) problem. Existing manifold OT methods have pursued amortized out-of-sample maps, while entropic regularization has made discrete OT more scalable, but these advantages have remained largely disjoint. We propose Entropic Riemannian Neural Optimal Transport (Entropic RNOT), a unified framework that combines intrinsic entropic OT with amortized out-of-sample evaluation on Riemannian manifolds. Our method learns a single target-side Schrödinger potential through a neural pullback parameterization, recovers the induced Gibbs coupling, and uses the resulting conditional laws to construct intrinsic transport surrogates. These include barycentric projections on Cartan-Hadamard manifolds and heat-smoothed conditional surrogates on stochastically complete manifolds, the latter turning possibly atomic target laws into absolutely continuous ones. For fixed regularization $\varepsilon>0$, we prove that the proposed hypothesis class recovers the entropic optimal coupling in strong probabilistic metrics. As consequences, barycentric surrogates converge in $L^2$, while heat-smoothed surrogates are stable at fixed heat time and asymptotically unbiased as the heat time vanishes. The guarantees hold for compactly supported data on possibly noncompact manifolds. Empirically, our method matches or improves over Euclidean, tangent-space, and log-Euclidean baselines on benchmarks over $\mathbb{S}^2$, $\mathrm{SO}(3)$, $\mathrm{SPD}(3)$, $\mathrm{SE}(3)$, and $\mathbb{H}^2$, scales favorably relative to discrete manifold Sinkhorn, and in a protein-ligand docking application, refines poses on $\mathrm{SE}(3)$ without retraining or per-instance optimization.

最优传输黎曼流形神经网络蛋白对接

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。