CurvLearn开源 | 阿里妈妈曲率学习框架详解

▐ 摘要

欧氏空间由于其直观的几何特性及简单有效的公式而被应用于几乎所有的深度学习模型。然而，近年来的研究发现，复杂结构数据（如社交网络、电商数据等）的建模精度往往会受到“平坦”的欧氏空间的制约，“弯曲”的曲率空间（如双曲空间和球形空间）由于具备更强的表征能力成为一种更佳的选择。

本文介绍的曲率学习框架（Curvature Learning Framework，以下简称 CurvLearn） 由阿里妈妈技术团队自研，旨在实现曲率空间下任意深度模型的训练与优化。基于 Tensorflow，CurvLearn 提供了丰富的流形实现，封装了完整的向量算子和高层 API，具备大规模分布式训练能力。目前该框架已在阿里广告关键业务场景落地，并取得了比欧氏空间更好的效果。

本文将对 CurvLearn 的相关背景、适用场景及实现逻辑做简要介绍，希望给大家带来一点启发与帮助，欢迎试用与交流讨论！

CurvLearn 开源地址：

https://github.com/alibaba/Curvature-Learning-Framework

▐ 框架背景

什么是曲率空间

曲率空间，也即微分流形，是三维欧氏空间中曲线和曲面概念的推广，呈现为局部平坦而整体“弯曲”的空间。

衡量空间“弯曲”程度的量即是曲率，曲率约接近于 ，空间越平坦，如欧氏空间曲率处处为 ；反之曲率绝对数值越大，空间畸变程度越高。根据空间各处曲率的不同，可以将曲率空间划分为常曲率空间及混合曲率空间两种。

常曲率空间指空间中各点截面曲率相等，即空间各处“弯曲”程度相同。根据曲率数值正负划分为双曲空间、欧氏空间及球面空间三种，各空间几何性质如下图所示。

图1：常曲率空间

混合曲率空间指空间各处曲率不同，整体呈现多样化的几何特性。实际上绝大部分常见的非欧空间都呈现出混合曲率特性，如下图展示了一个常见例子：环面空间。

图2：混合曲率空间示例：环面

曲率空间的特性

实际上不同空间暗涵了对数据分布的不同假设，一个常见例子是地图，在欧氏空间中计算真实城市距离会带来误差。因此不考虑真实数据分布而不假思索使用欧氏空间往往会带来模型精度上的损失。为此本文分别对不同空间适合的数据分布做出简要分析。

欧氏空间

• 欧氏空间各处均匀且平坦，具备各向同性及平移不变性，因此适合建模网格数据如图像等[1]。

图3：网格数据

双曲空间

• 双曲空间具备指数级体积容量 ，同存储量下样本间距更大，因此适合建模超大规模数据[2]。
• 双曲空间距离度规等价于幂律分布，互相导出，作为对数据分布的先验认知适合建模无尺度网络[3]。
• 双曲空间无损嵌入树结构，而欧氏空间无法避免表征损失，因此适合建模层次结构[4]。

图4：大规模数据/无尺度网络/层次结构

球面空间

• 球面空间距离度规等价于角度度规，具备旋转不变性，因此适合建模环状数据[4]。

图5：环状数据

混合曲率空间

• 混合曲率空间整体分布不均，各处呈现出不同的几何特性，因此适合建模多种结构特性并存的数据，如树环等[5]。

图6：复杂结构数据示例：树环

曲率空间的形式化定义

曲率空间 由流形 和度规 严格定义。 指定了空间可行域，度规 定义了各点处的坐标尺度，用来导出各点曲率 及测地线距离，欧氏空间下 即为单位矩阵 。对于流形上的每个点 ，切空间 为 点附近的一阶近似，属于欧氏空间。通过指数映射 可以将切空间上的点 映射至流形上：；与之可逆通过对数映射 可以将流形的点 映射至切空间上：。

常曲率空间拥有多种模型定义，如双曲空间包含洛伦兹模型（Lorentz Model）、克莱因模型（Klein Model）等。然而部分模型在曲率 时并不会收敛至欧氏空间，造成优化困难。幸运的是空间的球极平面投影（ Stereographic Projection）模型[4]可以避免这个问题，即庞加莱球模型（Poincare ́ball）以及投影球面模型（Projected Sphere）。这里给出三种空间模型的形式化定义：

• 双曲空间以 表示，曲率 ，定义为，黎曼度量为 ；
• 欧氏空间以 表示，曲率 ，定义为 ，黎曼度规为 ；
• 球形空间以 表示，曲率 ，定义为 ，黎曼度量为 。

其中 为向量维度， 为单位矩阵，。

注意到上述三种模型的形式化定义存在相似之处，为了使得空间学习更为灵活，一种统一的非欧空间表示为 ，它是双曲空间、欧氏空间以及球形空间的一体化实现，分别对应 ， 以及 这三种情况。

混合曲率空间具备复杂多变的空间曲率，模型自由度极高。其中积空间（Product Space）具备可分解及独立优化的特性，通常作为混合曲率空间的典型模型[5]。积空间定义为 ，其中 代表笛卡尔积， 代表第 个空间的维度， 为第 个空间的曲率。

▐ 曲率学习框架

框架概览

曲率学习是一个开源的基于 Tensorflow 的非欧深度学习框架。它实现了多种非欧流形、非欧算子和黎曼优化器，基于与 Tensorflow 相似的底层接口，用户可以便捷的将模型从欧氏空间迁移至非欧空间。曲率学习整体架构如下图所示，自下而上分别包含流形（Manifolds）、算子（Operations）、模型（Models）、优化器（Optimizers）及应用（Applications）五大模块：

图7：曲率学习架构

Manifolds

非欧空间中向量操作需要满足流形约束，与欧氏空间不同，它们通常在陀螺矢量空间（Gyrovector Space）中进行定义[6]，其表达形式与流形特性相关联。框架实现了如下五大类基本流形：

• Euclidean：欧几里得空间，曲率恒零。
• Stereographic：使用球极平面投影的常曲率空间，曲率可以是任意实数值。
• PoincareBall：双曲几何中的庞加莱球模型，曲率恒负。
• ProjectedSphere：球面几何中的球极平面投影模型，曲率恒正。
• Product：多种流形笛卡尔积构成的积流形，曲率具备各向异性。

以球极平面投影模型 作为示例，框架实现了如下向量操作：

• 指数映射：将 的切空间向量 投影到对应流形上。

• 对数映射：将对应流形向量 投影到 的切空间上。

• 向量加法：流形向量 与 相加。

• 向量乘法：流形向量 扩充 倍。

• 平行移动：将 的切空间向量 移到 的切空间中发生的尺度变化。

• 测地线距离：流形上 和 的最短距离。

• Trigonometric函数：

Operations

模型中的向量操作远不止上述定义的几种。实际上诸如向量平均，向量拼接等不是良定义的，需要自行推导其形式，在满足合适的向量变换逻辑上，使得该操作尽可能高效易优化。通常而言，绝大部分非欧算子都遵循转移到切空间的运算流程：

• 向量 通过对数映射投影至切空间
• 切空间向量 通过欧氏算子进行变换
• 通过指数映射投影回非欧流形

由此框架导出了其余算子，如向量拼接，激活函数等。值得注意的是向量平均我们使用了爱因斯坦中值（Einstein Midpoint），因为相比于上述切空间变换，该方法精度能提高约200%[7]。

Models

基于非欧算子的堆叠组合可以自然的构建非欧模型，如非线性变换可以定义为：

另一种构建模型的方式为转移计算到切空间，即投影流形向量至切空间，在切空间中实现复杂模型变换，投影回流形。HGCN[8]实现的非欧图卷积神经网络正是一个典型的例子：

由于不同模型实现及侧重点不一致，这里不做进一步细节阐述。

Optimizers

由于底层向量定义在非欧流形上，直接优化向量需要使用黎曼优化器[9]。以 RSGD（Riemannian SGD）为例，基本思想为先计算欧氏空间下的梯度，接着投影至切空间，最后执行非欧空间下的梯度下降。其具体参数更新过程为：

其中 位于切空间，为 的黎曼梯度， 是欧氏梯度，通常有 ， 是学习率。出于性能方面的考虑，目前非欧优化更多的采用间接优化非欧向量的手段，即将底层向量定义在欧氏空间中，然后投影至非欧流形上，由此使用欧氏优化器间接优化非欧模型。

Applications

近年来，非欧表征学习得到了越来越多的关注，在理论方面也得到了持续的完善，在各个领域也有相应的应用，包括推荐系统[10]、图学习[11]、自然语言处理[12]等。在阿里妈妈搜索直通车场景中，框架在图学习召回和类目预测应用上进行了落地，并取得了不错的离线在线效果。

在接下来的文章中以图学习召回为例，我们会详细描述如何使用曲率学习框架提高表征效果，欢迎持续关注。

▐ 曲率学习训练技巧

无论是在业界，还是我们自己的实践经验，都发现曲率模型训练可能存在着数值不稳定性的问题，主要包括如下两个原因：

• 节点Embedding容易超出曲率空间定义域。例如双曲空间的定义域为一个有限半径的球体，在训练的过程中，节点的 Embedding 容易靠近空间边缘，由于数值精度问题容易导致向量操作不合法。
• 梯度爆炸或弥散。由于曲率空间引入了复杂的计算逻辑，例如大量的 以及 函数等，位于这些函数的极值区域较难被优化。

为了解决这些问题，我们经过大量尝试总结了一些比较有效的训练优化手段：

• 参数初始化。曲率空间训练对参数初始化比较敏感，我们发现维持任意节点初始距离为较小的常值对提高训练稳定性很有帮助。一种可行的方案是使用小方差的正态分布初始化向量，同时使用 Xavier 初始化模型参数。
• 曲率 Embedding 正则。为了让节点 Embedding 尽量远离空间边缘等极值区域，直观的想法是通过正则项约束向量分布。作为 的扩展，正则项被定义为：

• 学习率及梯度处理。模型分布式优化往往涉及到梯度累加的过程，使用Warm-up及梯度裁剪有助于在训练初期维持梯度数值稳定性。

▐ 曲率ANN检索工具

由于曲率空间距离不是简单的内积、Cosine 或 L2，传统的 ANN 检索工具已经不再适用，为此我们开发了支持一系列非欧空间距离的 ANN 检索工具，包括双曲距离、球形距离以及混合空间距离，为了弥补距离公式复杂度提高带来的性能损失，工程上做了一系列的优化措施，包括分布式、数据及指令并行化，使得检索性能大为提高，相较于欧氏检索性能损耗被缩小至一个可以接受的范围内。目前我们正在进行工具的完善，计划于近期在 Github 进行开源发布。

▐ 展望

曲率空间作为一个全新领域，可以视作欧氏空间的扩展，天然适合建模复杂网络，在大规模场景下具有很大的应用潜力。

曲率空间独立于模型，理论上可以推广到任意的模型结构。越来越多的学术工作着眼于探索新的非欧流形表征及特性应用，并取得了亮眼的效果。

希望我们的曲率学习框架能为非欧空间在工业场景的落地添砖加瓦，铺平道路。

CurvLearn 开源地址：

https://github.com/alibaba/Curvature-Learning-Framework

参考文献

[1] Bronstein, Michael M. et al. “Geometric Deep Learning: Going beyond Euclidean data.” IEEE Signal Processing Magazine 34 (2017): 18-42. 

[2] Nickel, Maximilian and Douwe Kiela. “Poincaré Embeddings for Learning Hierarchical Representations.” NIPS (2017). 

[3] Krioukov, Dmitri V. et al. “Hyperbolic Geometry of Complex Networks.” Physical review. E, Statistical, nonlinear, and soft matter physics 82 3 Pt 2 (2010): 036106. 

[4] Bachmann, Gregor et al. “Constant Curvature Graph Convolutional Networks.” ICML (2020). 

[5] Gu, Albert et al. “Learning Mixed-Curvature Representations in Product Spaces.” ICLR (2019). 

[6] Ganea, Octavian-Eugen et al. “Hyperbolic Neural Networks.” NIPS (2018). 

[7] Lou, Aaron et al. “Differentiating through the Fréchet Mean.” ICML (2020). 

[8] Chami, Ines et al. “Hyperbolic Graph Convolutional Neural Networks.” NIPS (2019). 

[9] Bécigneul, Gary and Octavian-Eugen Ganea. “Riemannian Adaptive Optimization Methods.” ICLR (2019). 

[10] Vinh Tran, Lucas et al. “HyperML: A Boosting Metric Learning Approach in Hyperbolic Space for Recommender Systems.” WSDM (2020). 

[11] Zhu, Shichao et al. “Graph Geometry Interaction Learning.” NIPS (2020). 

[12] Chen, Boli et al. “Probing BERT in Hyperbolic Spaces.” ICLR (2021).

END

欢迎关注「阿里妈妈技术」，了解更多~

疯狂暗示↓↓↓↓↓↓↓