MindSpore AI科学计算系列 | 图、流形与群等变网络

背景

空间是一个任意的流形(manifold)，而图(graph)是流形的离散形式。在上一篇科学计算文章中，我们从对称性的统一视角看待神经网络的普遍性质，并提出了具有对称性神经网络的设计蓝图。在这篇文章中，我们将从图和流形这两种空间形式来具体看神经网络的结构，以及如何添加额外的群等变性质。

1、图神经网络

一个图由顶点与边的集合组成。在图神经网络中，空间结构是图，在图的顶点上定义了特征场，同时每条边也可能具有相应的特征信息。如果边集为空，那么对应于点集，例如deep set。如果任意两点之间都有边，那么对应完全图，例如transformer。如果图的顶点与边具有稳定的周期性结构，那么对应格网，例如CNN。另外，边可以存在方向性，例如RNN。

图的对称性由顶点的置换对称群决定。为了保证置换对称的等变性或不变性，一般图神经网络使用某种汇聚机制将顶点的邻居信息进行汇聚，即所谓的信息传递机制。传递机制依复杂度不同大致可分为三类：（线性）卷积、注意力机制、信息传递（非线性卷积）。

图1：信息汇聚机制

在线性卷积中，通过对每一个邻居点之间乘一个可训练的参数，来实现特征汇聚。注意力机制是线性卷积的一个推广，它将汇聚的参数提升到一组可训练的参数。更一般的情况是使用非线性卷积，将当前顶点与邻居顶点以某种方式相结合，作为传递的信息汇聚到当前顶点进行更新。注意到，transformer所使用的汇聚方式即为注意力机制。这三类信息汇聚机制都可以看作某种图结构上的卷积。

2、群等变网络

当我们提前预知更多关于对称性信息时，我们可以通过在等变层增加更多的等变性约束，使得网络满足额外的等变性。例如，分子结构除了具有三维空间平移对称性外，还具有三维空间旋转对称性，因此具有空间旋转等变性的神经网络会具有更好的精度、更强的表达能力和可解释性，而且所需的数据量更少。

为了使网络具有额外对称群的等变性，主流方法可大致分为两类：一类是正则(regular)群卷积，将群对输入数据的变换信息扩展为额外的维度；另一类是可控(steerable)群卷积，通过群表示的性质，将特征场表示为一系列群等变基函数的组合。正则群卷积方法可以应用于更一般的群，而可控群卷积方法需要提前知道群表示的性质，但一般具有更好的精确性。我们曾经在《科学计算系列 | 等变神经网络与e3nn数学库》https://zhuanlan.zhihu.com/p/587704873这篇文章中介绍的等变网络即属于三维欧几里得群的可控群卷积方法。如果对群等变网络有进一步兴趣，可以参考https://uvagedl.github.io/相关的讲义与课程。

3、流形上神经网络

前面我们遇到的空间结构几乎都是离散的，或者可以进行简单的离散化。如果空间结构变成了一般的连续流形，那么情况会变得复杂得多。例如我们可以问下面几个问题：（1）如何从数学上描述流形上的特征场？（2）如何选取合适的坐标系？（3）为了进行卷积或者说信息汇聚，如何将流形上不同的点进行关联？（4）如何构建一个具有权值共享的卷积核？

这里面有些问题看起似乎很容易，但实际上会涉及到微分几何中的一些核心概念和思想，我们在这里只会简单介绍部分思路和基本概念，感兴趣的读者可以参考[2]进行系统性学习。

对于第一个问题，我们不妨来回顾一下流形的定义。流形是局部同构于欧几里得空间的拓扑空间。而描述流形上的特征场的数学结构则是纤维丛。类比于流形，纤维丛可以简单理解为局部同构于流形与纤维空间的直积，其中特征场所在的空间就是纤维空间。虽然纤维丛局部是平凡的，但整体上可能具有非平凡结构，如图2所示，一个莫比乌斯环可以看作由两个平凡的直积空间拼接而成，但宏观上具有非平凡的结构。

图2：莫比乌斯环

这同时引出了第二个问题，如何选择合适的坐标系，与整体流形的拓扑结构有紧密联系。对于欧几里得空间，我们可以简单地选取一个全局内禀坐标系，例如一个直角坐标系。而对于二维球面、莫比乌斯环这些流形，我们无法只选取单一一个全局内禀坐标系。在二维球面上，我们或许可以用经纬度来描述球面上的点，但在南北两极会出现奇点；在莫比乌斯环上，如果我们沿着环移动一圈，由于方向的改变，新的坐标会与旧的坐标发生矛盾。

为此，我们需要在流形上选取多组坐标系，在不同坐标系的交汇处，同一个点可以在不同坐标系中表达（平凡化），那么不同的表达应该是相容的，即可以通过一个规范变换进行转换，如图3所示。

图3：坐标系之间的相容性

这种不同坐标系之间的规范变换构成了群。一般情况规范变换群是一般线性群GL(n)，但当规范变换群选取更小的子群时，纤维丛会具有更多对称性的信息，形成所谓的G-结构（G-structure）。如图4所示，第一排的规范变换群是平凡的，意味着在底流形（二维平面或者扣去一点）上选取了单一一个全局坐标系，尽管不同点初的坐标系可能并不是规则排列的。在第二与第三排中，流形上的每一点都有多个坐标系，不同坐标系之间通过规范变换群G相联系，这里第二排的G是一个二元的反演群，第三排的G是SO(2)群。

图4：流形M与坐标系的规范变换群G

对于第三个问题，我们需要首先定义什么是平行移动（联络），进而将空间上不同的点拉回（pullback）到同一点上进行汇聚或者某种运算。对于第四个问题，为了进行流形上的卷积，我们还需要定义一个卷积核，使其满足坐标无关，并且在不同坐标表达下具有相同的形式（权值共享）。细致地讨论会涉及到更多微分几何背景知识，这里就不再展开，感兴趣读者推荐参考文献[2]。

在本系列的最后一篇文章中，我们将简单回顾一些几何深度学习的应用，以及未来在AI for Sicence领域的展望与设想。

参考文献

[1]Bronstein, Michael M., et al. "Geometric deep learning: Grids, groups, graphs, geodesics, and gauges." arXiv preprint arXiv:2104.13478 (2021).

[2]Weiler, Maurice, et al. "Coordinate Independent Convolutional Networks--Isometry and Gauge Equivariant Convolutions on Riemannian Manifolds." arXiv preprint arXiv:2106.06020 (2021).