是关于几何图神经网络综述论文 A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications 的一个简单笔记。
# 数学基础
文章中用到了一些常见的群,包括:
- E(d): 欧几里得群 / 保距变换群 / 刚体变换群
- T(d): 平移变换群
- O(d): 正交变换群
- SO(d): 特殊正交变换群,即旋转群
- SE(d): 特殊欧几里得群,即平移旋转群
- Lie Group: 李群
- SN: n 个元素的置换群
# 群表示
一个群 G 的表示 (representation) 实际上就是一个从群到矩阵空间的映射
ρ(g):G→GL(V)
其中 V 是一个线性空间,则 GL(V) 是该线性空间上变换对应的矩阵表示。
群表示将群作用具体刻画为矩阵计算,为群作用的研究提供了方便。
# 不变性和等变性
对于群 G, 若函数 ϕ:X→Y 是 G- 不变的 (G-invariant), 当且仅当其满足
ϕ(g⋅x)=ϕ(x),∀g∈G.
相应地,对于群 G, 若函数 ϕ:X→Y 是 G- 等变的 (G-equivariant), 当且仅当其满足
ϕ(g⋅x)=g⋅ϕ(x),∀g∈G.
如果采用群表示的语言描述,则等变性可以表示为
ϕ(ρX(g)x)=ρY(g)ϕ(x),∀g∈G,
其中 ρX,ρY 是输入和输出的群表示。
# 数据结构
# 模型
# 不变图神经网络 (Invariant GNN)
# SchNet
SchNet 是 2018 年发表于 The Journal of Chemical Physics 期刊的一篇文章,17 年底就挂在了 Arxiv 上。
该网络在 MPNN 的消息计算过程中做文章,将消息改变为仅依赖于相对距离 rij=∥xi−xj∥ 的量,从而可以实现 E(3) 不变。
该网络的每一层(采用原文献符号)可以表示为
xil+1=(Xl∗Wl)i=j=1∑natomsxjl∘Wl(rj−ri),
其中 i 是原子标号,x∈RF 是特征向量,W:R3→RF 是将三维坐标映射到特征维度的全连接神经网络,natoms 为原子数。
# DimeNet
DimeNet 是 ICLR 2020 上提出的一种不变图神经网络。作者注意到,对于一个环己烷 C6H12 和两个环丙烷 2C3H6, 仅依靠相对距离信息是几乎无法区分的(两种碳碳键之间的键长差距极小),且各原子的邻域相同。因此,引入便于观察的角度信息是自然的想法。于是,相比于 SchNet 只采用相对距离作为消息,DimeNet 首次在消息中引入了角度信息。
基于原文符号,其消息计算过程可以表示为
mji(l+1)=fupdate⎝⎛mji(l),k∈Nj∖{i}∑fint(mkj(l),eRBF(ji),aCBF(kj,ji))⎠⎞
也就是说,这里的消息计算除了依赖该边先前消息 mji(l) 之外,还依赖了二阶消息 mkj(l), 以及 eRBF(ji),aCBF(kj,ji) 两项。其中 eRBF(ji) 是原子间距离 dji 的 RBF 表示,也就是依赖于相对距离的一项信息。aCBF(kj,ji) 是边 kj,ji 夹角的衡量,是基于圆贝塞尔函数 (circular Bessel function) 表示的。
径向基函数 (radial basis function, RBF) 是一种输出仅依赖于输入坐标与某个中心点之间距离的函数。常见的一种 RBF 是高斯函数,即
ϕ(r)=exp(−2σ2r2).
因此,该方法的创新点是首先引入了夹角作为消息的一部分,用来进行原子间相对关系的表示。但是此时的消息计算需要依赖二阶消息,因此计算复杂度是一个限制模型性能的重要因素。
# GemNet
GemNet 是 DimeNet 的进一步扩展,在消息传递中采用了二面角信息,该信息采用球贝塞尔函数 (spherical Bessel function) 表示。
mji′=σmsg⎝⎛mji,k∈Ni∖{j}∑l∈Nk∖{i,j}∑σint(mlk,eRBF(lk),eCBF(ikl),eSBF(jikl))⎠⎞.
# LieConv
LieConv 还没完全看懂。
# 等变图神经网络
等变 GNN 是满足
ϕ(g⋅G)=g⋅ϕ(G),∀g∈E(3).
消息计算和更新过程具体定义为
message passing:update:mij,mij=ϕmsg(hi,hj,xi,xj,eij),hi′,xi′=ϕupd(hi,{(mij,mij)}j∈Ni).
一般来说,等变 GNN 可以分成两类,分别是基于标量化的模型 (scalarization-based models) 和高度可操纵模型 (high-degree steerable models).
# 基于标量化的模型
# EGNN
EGNN 是最经典的基于标量化的等变模型。
基于原文符号,EGNN 可以表示为
mij=ϕe(hil,hjl,∥∥∥xil−xjl∥∥∥2,aij),xil+1=xil+Cj=i∑(xil−xjl)ϕx(mij),mi=j=i∑mij,hil+1=ϕh(hil,mi).
实际上它是
ϕmsg:ϕupd:mij=σ1(hi,hj,∥xi−xj∥2,eij),mij=(xi−xj)σ2(mij),hi′=σ3(hi,∑j∈Nimij),xi′=xi+γj∈Ni∑mij.
在 EGNN 中,还提到了将速度信息引入模型的方法,但速度信息的引入方式更像 Polyak 的动量梯度下降算法,即
vl+1=ϕv(hil)viinit+Cj=i∑(xil−xjl)ϕx(mij),xl+1=xl+vl+1.
这显然是可以改进的。
# GMN
GMN 是在 EGNN 基础上,引入更多通道,从而提高了等变模型的表征能力。对于一个节点 i, 其输入的几何特征从 xi 改为多通道的 Vi∈R3×C. 可以表示为
mij=σ1⎝⎛hi,hj,∥Vij⊤Vij∥FVij⊤Vij,eij⎠⎞,Mij=Vijσ2(mij).
若取 C=2, 则可以令 Vi=(xi,x˙i). 然后同时对两个特征进行学习。可以看到,此时的 xi 与 x˙i 地位等价,不存在从属关系。这样的模型可学习自由度更高,从而有更大潜力得到更优的模拟效果。
# PAINN
PAINN 全称 Polarizable Atom Interaction Neural Network, 即可极化原子相互作用神经网络。因此,它描述的是极性分子的分子结构,即分子中不同原子之间的相互作用。
# 局部框架 (Local Frame)
为了应对旋转中产生的对称性问题,可以构建旋转等变的局部框架,从而将几何信息投影为不变表示。例如 ClofNet 采用类似物理学中换用质心系的思想实现平移不变。即计算 xic=xi−xc, 其中 xc=N1∑i=1Nxi, 然后将 xic 作为输入。
# 高度可操纵模型
# 数学基础
# 勒让德方程和勒让德多项式
连带勒让德方程 (Associated Legendre polynomials) 形式如下:
dxd((1−x2)dxdy)+(l(l+1)−1−x2m2)y=0,l=0,1,2,….
它是由球坐标系下求解拉普拉斯方程 ∇2u=0 导出的。其通解为
Plm(x)=(−)m(1−x2)m/2dxmdmPl(x),m=0,1,2,…,
其中 Pl(x) 是勒让德多项式 (Legendre polynomials).
当连带勒让德方程参数 m=0 时,转化为下式:
dxd((1−x2)dxdy)+l(l+1)y=0,l=0,1,2,…,
称为勒让德方程 (Legendre equation). 通解
Pl(x)=2l1s=0∑⌊l/2⌋s!(l−s)!(l−2s)!(−)s(2l−2s)!xl−2s,x∈[−1,1]
称为勒让德多项式 (Legendre polynomial), 记作 Pl(x). 也可以用超几何函数表示为
Pl(x)=2n(n!)2(2n)!xnF(−2n,21−n,21−n,x−2).
# 球谐函数
球谐函数可以看作单位球面上的一组 Fourier 基,它天然借用了 Legendre 多项式表示。
球谐函数定义为
Ylm(θ,ϕ):=(−)m4π(2l+1)(l+m)!(l−m)!Plm(cosθ)eimϕ,l=0,1,2,…,−l≤m≤l.
其中含有两个参数(参考 ChatGPT):
- l 称为度数 (degree) 或角动量量子数,它决定了球谐函数的总角动量大小。l 的取值范围是 0,1,2,…, 通常被称为阶数。较大的 l 对应更高的角动量,表现为球谐函数在球面上的更多变化和波峰波谷。
- m 称为阶数 (order) 或磁量子数,它表示角动量在空间中的投影。对于给定的 l, m∈{−l,−l+1,...,0,...,l−1,l}. 球谐函数描述了球面上在不同方向上的分布,而 m 决定了在球面上的旋转对称性。具体而言,m 的正负号表示旋转方向,而 ∣m∣ 决定了旋转的角度。
# 函数的球谐展开
基本的球谐函数是 Ym(l)(θ,ϕ):S2→C 的映射,其中 θ 是极角,ϕ 是方位角,l,m∈Z,l>0,−l≤m≤l. 将一系列 m 组合为一个新的向量
Y(l)=(Yl−l(θ,ϕ),Yl−l+1(θ,ϕ),…,Yll(θ,ϕ))⊤
那么,type-l 函数的球谐展开可以表示为
f(x)=m=−l∑lcmYlm(x)=c⊤Y(l)(x).
一般地,如果函数是多个 type 特征混合而成的。例如此处考虑其包括 type-l1 和 type-l2 特征,那么函数的球谐展开可以表示为
f(x)=c1⊤Y(l1)(x)+c2⊤Y(l2)(x)=(c1⊤,c2⊤)(Y(l1)(x),Y(l2)(x)).
关于实球谐函数的具体计算,可以参考 Table of spherical harmonics - Wikipedia.
# Wigner-D 矩阵与旋转等变
Wigner-D 矩阵 (Wigner-D matrix) 是旋转群 SO(3) 的一种不可约表示:
ρ(g):=D(l)(g)∈R(2l+1)×(2l+1),g∈SO(3),
假设 (ρ,V) 是群 G 的一个表示,若对 ∀g∈G, V 的子空间 U 是线性变换 ρ(g) 的不变子空间,则称 U 是 G- 不变子空间 (invariant subspace). 若 U 的补空间 U′ 也是 G - 不变子空间,则称 (ρ,V) 是完全可约的 (completely reducible), 并称 U′ 是 U 在 V 中的 G- 不变补空间 (invariant complement).
若 V 的 G - 不变子空间仅有 {0} 和 V 本身,则称 (ρ,V) 是不可约的 (irreducible).
Wigner-D 矩阵的作用是在球谐函数表示下实现旋转等变性。假设一个旋转算子为 Rg, 则
Ym(l)(Rgx)=m′=−l∑lDm′m(l)(g)Ym′(l)(x).
写成向量,那就是
Y(l)(Rgx)=D(l)(g)Y(l)(x),
其中 D(l)(g)∈R(2l+1)×(2l+1). 结合上面的分析,具有两个 type 的函数的球谐展开加入旋转等变可以表示为
f(Rgx)=(c1⊤,c2⊤)(Y(l1)(Rgx),Y(l2)(Rgx))=(c1⊤,c2⊤)(D(l1)(g)Y(l1)(x),D(l2)(g)Y(l2)(x))=(c1c2)⊤(D(l1)(g)00D(l2)(g))(Y(l1)(x)Y(l2)(x))
是一个超级 cool 的二次型,可以将空间按照 type 划分为不同的不变子空间。
# 张量积与 CG 张量乘
前面一直在说 "type", 但终究还是没有给出一个严谨的定义。
的确还没完全想明白。
总之现在采用比较玄学的方式说明下面的知识。
注意到,一个位矢 x=(x1,x2,x3)⊤ 是 type-1 的。对本身作张量积,得到
x⊗x=(x1,x2,x3)⊗(x1,x2,x3)=(x12,x1x2,x1x3,x2x1,x22,x2x3,x3x1,x3x2,x32)⊤.
然后把里面这个 R9 的向量里面的不同维度线性组合一下,或者说加一个神经网络层,就可以得到
CG 张量乘 (Clebsch-Gordan tensor product), 则是把两个不同 type 类型特征进行克罗内克积,然后接入一个可学习的参数矩阵,筛选上述的对应 type 特征的过程。可以形式化表示为
V(l)=(V(l1)⊗cgWV(l2)),
这里要求 ∣l1−l2∣≤l≤l1+l2,l∈N. 这个式子展开之后就是
vm,c(l)=c1=1∑C1c2=1∑C2wc1,c2,cm1=−l1∑l1m2=−l2∑l2Q(l1,m1)(l2,m2)(l,m)vm1,c1(l1)vm2,c2(l2),
这里的 Q(l1,m1)(l2,m2)(l,m) 是 CG 系数,实际上是一个类似示性函数的东西,当且仅当 ∣l1−l2∣≤l≤l1+l2 时该系数不为 0.
对于 multi-type 的情况,就用空心字母表示张量,于是有
V(L)=V(L1)⊗cgWV(L2).
# TFN
TFN 全称 Tensor Field Networks, 即张量场网络。可以表示为
Mij(L)=Y(L)(∥xij∥xij)⊗cgWVj(L).
# SEGNN
SEGNN 即 Steerable E(3) Equivariant GNN, 可操控的 E(3) 等变图神经网络。
111
# 几何图的表征能力
# 几何图上的 WL-Test
WL-test 是古早的图同构判定算法,它与 GCN 类似,通过多次传递消息然后判断各节点的哈希组成集合是否相同,公式表示为
ci(t)=HASH(ci(t−1),{{cj(t−1)∣j∈Ni}}).
该算法复杂度是 O(∣V∣⋅∣E∣) 的。
GWL 框架 (Geometric Weisfeiler-Leman framework) 提出了 WL-test 的几何版本。
注意到,原本的 WL-test 是采用节点为中心的、局部的更新过程进行图同构判断,GWL-test 仍然按照此想法,但是从邻居节点周围收集的信息同时包括不变的拓扑信息和等变的几何信息,可以公式表述为
gi(t)=HASH((ci(t−1),gi(t−1)),{{(ci(t−1),gj(t−1),xij)∣j∈Ni}})ci(t)=I−HASH(t)(gi(t))
这里的 I−HASH(t) 只是强调它在每一此迭代过程中都是一个单射 (injection), 并无深意。
# 应用
# 原子级任务 (Particles)
# 物理动力学模拟 (Physics dynamics simulation)
这里的典型例子是多体模拟,此外还可以建模流体、刚体、人体等多种结构。抽象为数学形式,它实际上是一个初值问题,基于 t 时刻初始图 G(t) 对 k 时间步后的系统进行预测:
X(t+k)=ϕθ(G(t)).
此外,输入也可以更改为过去一段时间的信息。
# 小分子级任务 (Small Molecules)
在小分子中,取原子作为几何图节点,几何图连边对应化学键,则分子构成几何图。分子中相对距离是重要的输入特征,还可以用来通过 KNN 构建邻居节点。
# 分子性质预测 (Molecular roperty prediction)
分子性质预测是最基本的任务,可以公式表述如下:
y,y=ϕθ(G).
由于可以同时进行多任务预测,且可以对标量和矢量进行预测,因此左侧是 y,y.
# 分子动力学模拟 (Molecular dynamics simulation)
与物理动力学模拟类似,分子动力学模拟也可以公式表述为
X(t+k)=ϕθ(G(t)).
但分子动力学与的输入几何图建模不同,导致可以施加的操作也略有差别。我们可以对分子中的每个节点施加力 F∈R3, 也可以对分子本身的能量 E∈R 进行建模。
# 分子生成 (Molecular generation)
分子生成是人工智能药物发现的重要任务,通过对已有分子结构对应几何图进行概率学习,预测可能存在的新分子结构。
分子生成的方式一般包括两种:
- 构象生成 (conformation generation): 基于图的 2D 拓扑表示生成 3D 空间结构,该过程主要依赖几何图实现相对距离预测和坐标重建。
- 从头生成 (de novo generation): 可以使用自回归模型,也可以基于扩散模型生成。
# 分子预训练 (Molecular pretraining)
用 ϕθ(G) 代表表示模型,用 y^(G) 代表基于几何图结构的伪标签,那么自监督训练过程就是
θ=θargminL(y^(G),ϕθ(G)).
# 蛋白质任务 (Proteins)
注意到蛋白质是由单个或多个氨基酸残基长链盘曲折叠组成的生物大分子,因此几何图构建包括两种方式
- 将全原子作为节点;
- 将各氨基酸残基视为节点,以 α - 碳位置作为空间坐标。
蛋白质是多模态的结构,可以关注其一维氨基酸序列,也可以关注三维空间结构。
# 蛋白质性质预测 (Protein property prediction)
按照粒度分类,蛋白质性质预测任务可以分为三类:
- 蛋白质级别预测:希望基于蛋白质的结构,给出蛋白质功能的预测,或是对蛋白质输出一个评分。其可以包括如下任务:
别骂了还没看懂 bg.
- 残基级别预测:任务是捕获单个氨基酸之间的结构依赖性。
- 原子级别预测:主要任务是发现蛋白质上原子是否存在潜在配体结合位点。
# 蛋白质生成 (Protein generation)
# 蛋白质预训练 (Protein pretraining)