1. 几何深度学习的基本概念
深度学习已经被应用于各种科学问题中(AI4Science),包括分子性质预测、药物合成与设计、动力学模拟等等。深度学习相较于传统方法(基于计算、模拟、数据库搜索)来说,推理速度快,具有学习能力。但是,科学问题的标注数据非常昂贵,想要学习到有用、符合常识的信息,我们需要将几何先验加入到模型的设计中来降低参数的搜索空间。具体来说,自然界的一些性质(例如能量、力、电势)关于变换具有不变性(invariance)或等变性(equivariance)。如果x具有性质y,,那么如果对于任意对x的空间变换满足,则该性质具有不变性;如果,则具有等变性。
2. 论文介绍
2.1 等变图神经网络在抗体设计的应用 (Oral)
Motivation
本文研究的问题是抗体设计。在医学中,相较于小分子药物,抗体药物的靶向性强,副作用小,在临床上的应用更加广泛,尤其是癌症、免疫疾病的治疗。抗体是Y形状的蛋白质,由重链和轻链组成。抗体的前端与抗原结合,称为CDR片段,不同抗体的CDR片段不同,而其余部分都是相同的。因此抗体设计的关键在于CDR片段的设计。之前的方法存在三个缺陷:没有考虑所有的上下文信息(轻链,重链和抗原);缺少3D几何结构的信息;基于自回归的生成方式比较低效。
数据建模
本文以每个氨基酸作为图中的一个节点,节点的表示由两部分组成:氨基酸的类型作为特征,backbone原子(,,,)的空间位置共同构成空间特征。相距在某个阈值内的氨基酸之间在图中用边相连。因此,CDR设计的任务就转化成了预测CDR氨基酸的节点类型的任务。
模型:MEAN
模型分为两个部分。内部的消息传递在轻链、重链、抗体各自的内部进行更新,使用EGNN(等变图神经网络,被称为几何深度学习界的GCN)作为网络结构。外部的消息传递在三部分之间进行,使用注意力机制。
模型的输出包括两部分:氨基酸的embedding和空间坐标。模型的训练有T步,第步将和经过1-6描述的encoder后,得到更新后的和,使用对CDR部分的边进行重构,再将和作为CDR部分的图特征,构成第步的图作为新的输入。模型的训练loss由序列预测和结构预测两部分组成。
实验
本文进行了三种实验。第一个是CDR的生成式预测,数据集为Structural Antibody Database,分别评估CDR-H1,CDR-H2,CDR-H3的生成准确率;第二个是带有抗体的CDR生成,数据集为RAbD;第三个是评估抗体与抗原结合的亲和性,在SKEMPI V2.0数据集上,用结合的能量作为指标。
简评
本文的模型设计并不复杂,但值得学习的点有很多:在几何GNN领域考虑了一个新的问题;对抗体设计问题的理解比较透彻,因此对问题的建模也考虑到了各方面的信息;模型训练有一些trick可以参考,比如空间坐标的设计、生成过程的迭代模式。
2.2 等变Transformer (Spotlight)
Motivation
一方面,现在有很多等变图神经网络的研究;另一方面,图的Transformer也是图大模型的热点研究之一。因此,很有必要实现一个等变的图Transformer。存在的挑战就是,如何让图Transformer具有等变性/不变性。
数据建模
本文将3D信息加入模型的方式是不可约特征(irreducible feature,irreps feature)。详细关于irreps的定义参见原论文,我们只说结论:节点的特征是由L组feature组成的,第0组是标量,剩下的是向量。这些向量是3D信息的球面谐波(Spherical Harmonics)得到的。每组feature都有多个channel,因此irreps feature之间运算需要使用张量运算。
模型
Equiformer实际上就是把Transformer的各部分替换成等变的操作(如(a)所示),并且作用在irreps feature上。
实验
本文做了三个数据集的实验,分别是QM9(分子性质预测),MD17(分子动力学),OC20(能量预测)。模型在三个数据集上都取得了SOTA。
简评
本文创新性很好,提出了如何做几何GNN,并且也取得了很好的实验效果。本文的数学部分值得深入理解,包括球面谐波、张量乘积等等。这些数学表达是对3D信息进行Encode的常用方式,尤其是分子领域。
2.3 蛋白质表示学习的几何GNN预训练
Motivation
本文研究的问题是蛋白质功能预测。蛋白质的功能指的是蛋白质发挥的作用,包括催化、组建细胞、运送信号分子等等。蛋白质有序列信息和结构信息。之前的蛋白质功能预测的工作主要集中在直接从氨基酸序列预测蛋白质功能,例如CNN、LSTM、Transformer,也有一些蛋白质语言模型的工作例如ProtBERT。而最近,AlphaFold,RosettaFold等模型能够从蛋白质序列得到其结构信息,再加上能够利用空间信息的几何GNN,这启发我们是否能够使用蛋白质结构信息预测其功能,因为蛋白质功能与结构密切相关。
数据建模
模型:GEAR
Encoder基于RGCN,并且有可学习边embedding的变体。整个Encoder是E(3)不变的,并不具有等变性。
模型是对比学习的框架,采用了两步的增广策略:第一步是基于序列或空间距离进行子图的裁剪,第二步是我们常见的拓扑增广。
实验
下游任务共有四个:Enzyme Commission (EC) number prediction,预测蛋白质作为酶的类别;Gene Ontology (GO) term prediction,预测蛋白质在基因本体论中的类别;Fold classification,预测蛋白质折叠的类别;Reaction classification,预测蛋白质参与反应的类别。这些类别都是根据蛋白质功能归类的。共有两种实验类别。第一种是单独评测GEAR这种GNN设计,与其他GNN进行比较;第二种是评测预训练模型,使用AlphaFold2的公开数据集进行预训练,再与其他蛋白质预训练模型比较。两种实验都取得了比较好的效果。
简评
本文第一次尝试进行蛋白质任务的GNN预训练方法,考虑了蛋白质的结构信息。但本文对于结构信息的考虑比较单一,并没有使用等变GNN,3D信息仅仅用于定义图结构,网络完全是E(3)不变的,表达能力有限。
长按下图并点击“识别图中二维码”
即可关注北邮 GAMMA Lab 公众号
原文始发于微信公众号(北邮 GAMMA Lab):专题解读 | ICLR'23的几何深度学习
- 左青龙
- 微信扫一扫
- 右白虎
- 微信扫一扫
评论