[论文解读] Geoopt: Riemannian Optimization in PyTorch

Max Kochurov, Rasul Karimov+1 more|arXiv (Cornell University)|May 6, 2020

Advanced Numerical Analysis Techniques参考文献 35被引用 49

一句话总结

Geoopt 是一个基于 PyTorch 的模块化 Riemannian 优化库，与 PyTorch 无缝集成，支持双曲几何，并为几何深度学习提供流形感知的优化器和层。

ABSTRACT

Geoopt is a research-oriented modular open-source package for Riemannian Optimization in PyTorch. The core of Geoopt is a standard Manifold interface that allows for the generic implementation of optimization algorithms. Geoopt supports basic Riemannian SGD as well as adaptive optimization algorithms. Geoopt also provides several algorithms and arithmetic methods for supported manifolds, which allow composing geometry-aware neural network layers that can be integrated with existing models.

研究动机与目标

• 为神经网络提供一个与 PyTorch 集成的、可直接使用的 Riemannian 优化接口。
• 支持多种流形和几何感知操作，以用于几何深度学习。
• 提供鲁棒、高效的优化算法，包括在流形上的自适应方法。
• 实现几何感知层的组合，并与现有的 PyTorch 模型兼容。

提出的方法

• 为在流形上的通用优化引入标准的 Manifold 接口。
• 实现符合 PyTorch 优化器的流形感知优化器（例如 Riemannian Adam、Riemannian SGD）。
• 通过在环境向量空间中的嵌入表示点和方向，并进行 egrad2rgrad 转换。
• 提供若干流形（Sphere、Stiefel、Birkhoff Polytope、Stereographic 超曲率模型、Lorentz）以及乘积/尺度变体。
• 通过投射回收、指数映射和切向量传输在乘积流形上实现高效更新。

实验结果

研究问题

• RQ1 如何将 Riemannian 优化作为可直接替换的组件集成到 PyTorch 生态系统中？
• RQ2 对深度学习应用最有用的流形和几何原语有哪些（例如双曲空间、单位约束、正交性）？
• RQ3 在流形上是否可以有效实现自适应和稀疏梯度优化器，以提升训练稳定性和性能？

主要发现

• Geoopt 提供了一个感知几何的优化框架，可以插入到 PyTorch 优化器和训练循环中。
• 它通过 Poincaré ball、Hyperboloid 和广义 κ-Stereographic 模型支持双曲深度学习，使几何感知的神经网络成为可能。
• 为流形参数提供自适应优化器（RiemannianAdam、SparseRiemannianAdam）和 SGD 变体。
• 乘积和尺度化流形允许将多种几何结合起来，并在黎曼优化中扩展标准流形。