Rieoptax:JAX中的黎曼Optimization
Rieoptax: Riemannian Optimization in JAX
摘要:
我们介绍了Rieoptax,一个用于JAX中黎曼优化的开源Python库。我们发现,许多微分几何图元,如黎曼指数和对数映射,通常在Rieoptax中比Python中现有的框架更快,无论是在CPU还是GPU上。我们支持各种基本和高级随机最优化求解器,如黎曼随机梯度、随机方差缩减和自适应梯度法。建议的工具箱的一个显着特点是,我们也支持差分私有优化黎曼流形。
英文摘要:
We present Rieoptax, an open source Python library for Riemannian optimization in JAX. We show that many differential geometric primitives, such as Riemannian exponential and logarithm maps, are usually faster in Rieoptax than existing frameworks in Python, both on CPU and GPU. We support various range of basic and advanced stochastic optimization solvers like Riemannian stochastic gradient, stochastic variance reduction, and adaptive gradient methods. A distinguishing feature of the proposed toolbox is that we also support differentially private optimization on Riemannian manifolds.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。