Frag2Seq
背景与动机
SBDD 的传统方法
高通量虚拟筛选和实验验证 —— 费钱费时费力
SBDD 对机器学习的挑战
要求模型捕捉 复杂的蛋白质-配体相互作用,同时 提高设计分子的药物相似性
相关工作
自回归模型和扩散模型
缺点:
- 只考虑 原子生成
- 扩散模型通常需要数千步才能生成,导致 生成效率低下
基于分子片段生成分子
缺点:
- 需要一个涉及多个神经网络的 复杂管道 来选择和连接片段
LMs 语言模型的优势
处理 大型数据集的效率高于基于扩散 的方法
利用 LMs 的上下文能力,基于分子片段生成,可以得到更真实的分子子结构和减少步骤以提高效率
从 大量生物和化学文本 中学习以完成各种 潜在任务,并生成分子和材料
LMs 面临的挑战
- 需要使 LMs 适应处理分子数据的 几何图结构
- 能否 准确模拟分子的物理和化学性质
- 如何在 LMs 中编码蛋白质上下文信息以有效 捕获蛋白质-配体相互作用,从而产生可以结合靶蛋白的分子
- 现有的 LMs 要么只考虑 2D 分子,要么需要考虑额外的设计以缓解对齐问题,在 SBDD 中没有现有的工作使用 LMs 直接生成 3D 配体
模型解析
1. 问题定义 (Problem Definition)
- 蛋白质口袋 (Protein Pocket): P={(si,bi)}i=1n
- si∈R3:第 i 个原子的 3D 笛卡尔坐标。
- bi:原子类型的 one-hot 向量。
- 配体 (Ligand): M={(vj,zj)}i=1m
- vj∈R3:第 j 个原子的 3D 笛卡尔坐标。
- zj:原子类型。
- 目标 (Objective): 学习一个条件生成模型 Model,从训练蛋白配对中捕获条件概率分布 p(M∣P)。
2. 3D 分子切片 (3D Molecular Slicing)
- 优点 (Advantages):
- 化学意义 (Chemical Significance): 片段通常对应功能基团(如苯环、羧基),生成的分子更“合理”。
- 效率更高 (Higher Efficiency): 生成步骤从原子数量级减少到片段数量级,大大加快生成速度。
- 结构更优 (Better Structure): 直接生成刚性片段,更好地保持片段内部的键长和键角,避免生成不合理的局部结构。
- 切片条件 (Slicing Criteria): 切断满足以下三个条件的可旋转化学键,以在保留重要功能基团完整性的前提下最大化分解分子:
- 键不在环中。
- 键类型为单键。
- 键上的起始原子和末端原子的度数大于1(避免打断如 -OH 或 -COOH 这样的重要官能团)。
- 切片结果 (Slicing Outcome): 对于分子 M,将其拆分为一组 3D 片段 {gi=(Zi,Vi)}i=1k,其中 Zi 表示原子类型矩阵,Vi 表示笛卡尔坐标。
3. 基于片段的三维分子词元化 (Fragment-based 3D Molecular Tokenization)
3.1. 基于三维图同构的原子排序 (Atom Reordering based on 3D Graph Isomorphism)
- 规范 SMILES (Canonical SMILES):
- 每个分子(片段)图使用规范 SMILES表示,其表示是唯一的。
- 这允许对分子数据进行高效索引和检索,确保原子重新排序的最大不变性。
- 定义 3.1. 3D 分子图的同构 (Definition 3.1. [3D Molecular Graph Isomorphism])
- 设 M1=(V1,Z1) 和 M2=(V2,Z2) 是两个 3D 分子图,其中 zi 是节点类型向量,vi 是分子 Mi 的节点坐标。设 ver(⋅) 表示顶点集合,attr(⋅) 表示节点属性,且不存在边。
- 设 M1≅M2 表示两个属性图是同构的。
- 如果存在一个双射 b:ver(M1)→ver(M2) 使得对于 M1 中的每个原子索引 i,有 ziM1=zb(i)M2,并且存在一个 3D 变换 τ∈SE(3) 使得 viM1=τ(vb(i)M2),则称两个 3D 分子 M1 和 M2 是 3D 同构的,记作 M1≅3DM2。
- 如果允许一个小的误差 ε 使得 ∣viM1−τ(vb(i)M2)∣≤ε,我们称这两个 3D 图是 ε-约束 3D 同构的。
- 引理 3.2. 3D 分子图同构的规范序 (Lemma 3.2. [Canonical Ordering for 3D Molecular Graph Isomorphism])
- 设 M1=(V1,Z1) 和 M2=(V2,Z2) 是两个遵循定义 3.1 的 3D 分子图。设 L:M→L 是一个函数,它将分子 M∈M(所有有限 3D 分子图的集合)映射到其规范序 L(M)∈L(所有可能规范序的集合),这个规范序是由规范 SMILES 产生的。那么以下等价关系成立:
L(M1)=L(M2)⇔M1≅3DM2
- 其中 M1≅3DM2 表示 M1 和 M2 是 3D 同构的。
- 设 M1=(V1,Z1) 和 M2=(V2,Z2) 是两个遵循定义 3.1 的 3D 分子图。设 L:M→L 是一个函数,它将分子 M∈M(所有有限 3D 分子图的集合)映射到其规范序 L(M)∈L(所有可能规范序的集合),这个规范序是由规范 SMILES 产生的。那么以下等价关系成立:
3.2. SE(3)-等变分子和片段的局部坐标系的构建 (Construction of SE(3)-Equivariant Local Coordinate Frames for Molecules and Fragments)
构建分子局部坐标系 (Molecule Local Coordinate Frame):
- 给定一个原子类型为 Z、原子坐标为 V 的 3D 分子 M 作为输入,构建其分子局部坐标系 m=(x,y,z)。
- 将切分的 3D 的片段按规范序进行排序(具体排序方式参考代码)。
- 坐标系基于规范顺序中的前三个非共线碎片中心构建。
- 设 ℓ1,ℓ2,ℓm 为这三个片段中心的索引,分子局部坐标系 m=(x,y,z) 计算如下:
x=normalize(vℓ2−vℓ1),y=normalize((vℓm−vℓ1)×x),z=x×y(1)
- normalize(⋅) 是归一化函数,将向量长度变为 1。
- v 表示世界坐标系中的坐标。
- 经过归一化后得到分子局部坐标系的基向量 m=(x,y,z)。
- 设 ℓ1,ℓ2,ℓm 为这三个片段中心的索引,分子局部坐标系 m=(x,y,z) 计算如下:
- 解释:
- 定原点 (Define Origin): 选择第一个片段的中心 vℓ1 作为这个局部坐标系的原点。
- 定 x 轴 (Define x-axis): 从原点 vℓ1 指向第二个片段中心 vℓ2 的方向,定义为 x 轴。
- 定 y 轴 (Define y-axis): 第三个片段中心 vℓm 与 x 轴定义了一个平面。y 轴垂直于这个平面(通过叉乘实现)。
- 定 z 轴 (Define z-axis): 根据右手定则,通过 x 和 y 的叉乘得到 z 轴。
构建分子局部坐标系 m 和世界坐标系 w 之间的变换,记作 m→w:
- 变换包含:
- 旋转矩阵 Rm→w∈R3×3,满足 ∣Rm→w∣=1。
- 平移向量 tm→w∈R3。
- 旋转矩阵构建 (Rotation Matrix Construction):
- 由于基向量 (x,y,z) 已经归一化且相互正交,我们可以直接将它们堆叠形成一个旋转矩阵:Rm→w=[xT,yT,zT]。
- 平移向量设置 (Translation Vector Setting):
- 平移向量可以设置为 tm→w=vℓ1,表示两个坐标系原点之间的位移。
- 变换包含:
构建片段局部坐标系 g 和变换 g→w:
- 构建过程 (Construction Process): 类似于分子局部坐标系的构建过程,使用片段中前三个非共线原子来构建片段局部坐标系,记作 g。
- 等价变换 (Equivariant Transformation): 片段局部坐标系也等价于输入分子的旋转和平移变换,可以获得旋转矩阵 Rg→w 和平移向量 tg→w,分别表示片段局部坐标系与世界坐标系之间的方向和位移。
- 在片段局部坐标系下,可以获得原子的局部坐标 Vgi∈Rq×3,其中 i 表示第 i 个片段,q 表示片段中的原子数量。
数据存储 (Data Storage):
- 将训练集中的每个切分片段保存到字典中。
- 键 (Key): 每个片段的规范 SMILES。
- 值 (Value): 原子类型和在相关片段局部坐标系下的原子局部坐标。
- 将训练集中的每个切分片段保存到字典中。
推导变换 g→m (Derivation of Transformation g→m):
- 将局部坐标系映射到世界坐标系:
Tm→w=[Rm→w0tm→w1],Tg→w=[Rg→w0tg→w1](2)
- 其中 T∈R4×4 表示齐次变换矩阵,0∈R3×1 是零向量。
- 使用坐标系变换的链式法则得到齐次变换:
Tg→m=Tm→w−1Tg→w=[Rm→wT0−Rm→wTtm→w1][Rg→w0tg→w1]=[Rm→wTRg→w0Rm→wT(tg→w−tm→w)1](3)
- 提取旋转矩阵和平移向量:
Rg→m=Rm→wTRg→w,tg→m=Rm→wT(tg→w−tm→w)(4)
- 保存片段局部坐标系原点到片段原子中心的位移 tg→c(g)(在将原子局部坐标从片段局部坐标系转换回世界坐标系时使用):
tg→c(g)=Vc(g)m−tg→m(5)
- 其中,c(g) 表示任何片段 g 的中心,Vc(g)m 是片段中心在分子局部坐标系下的坐标。
- 将局部坐标系映射到世界坐标系:
3.3. SE(3) 不变的分子片段局部表示 (SE(3)-Invariant Local Representation of Molecular Fragments): 获得片段的相对位置和相对朝向。
- 球面坐标表示 (Spherical Coordinate Representation): 利用函数 f(⋅) 将片段中心的坐标转换为分子局部坐标系下的球坐标。
- 对每个片段中心 ℓi(在世界坐标系下坐标为 vℓi),定义其球坐标:
diθiϕi=∥vℓi−vℓ1∥2,=arccos(dℓi(vℓi−vℓ1)⋅z),=atan2((vℓi−vℓ1)⋅y,(vℓi−vℓ1)⋅x)(6)
- 其中,(x,y,z) 是分子局部坐标系的基向量,vℓ1 是分子局部坐标系的原点(第一个片段中心)。
- 对每个片段中心 ℓi(在世界坐标系下坐标为 vℓi),定义其球坐标:
- 旋转向量表示 (Rotation Vector Representation): 利用函数 g(⋅,⋅)。
- 原因 (Reason): 欧几里德空间中的 3D 旋转只有三个自由度,旋转矩阵表示具有冗余信息(3x3矩阵有9个数字,但实际只需三个数字表示三个自由度),并且可能不必要地增加 LMs 的上下文长度。
- 已有旋转矩阵 Rg→m,提取其旋转角度 ψ 与单位旋转轴 a=(ax,ay,az),定义旋转向量:
m=ψa=(mx,my,mz)∈R3
- 紧凑度 (Compactness): 3个数值代替9个矩阵元素,同样具备 SE(3) 不变性。
- 逆映射 (Inverse Mapping): 从相对表示中重建原始分子。
- 使用坐标变换将原子坐标从片段局部坐标系转换到世界坐标系(片段 → 分子 → 世界):
tg→m=Vcm(G)−tg→c(G),Vm=VgRg→mT+tg→m,Vw=VmRm→wT+tm→w.(7)
- 其中:
- tg→m:片段局部坐标系到分子局部坐标系的平移向量。
- Vg:片段局部坐标系下的原子坐标。
- Vm:分子局部坐标系下的原子坐标。
- Vw:世界坐标系下的原子坐标。
- Rg→mT:片段局部坐标系到分子局部坐标系的旋转矩阵转置。
- Rm→wT:分子局部坐标系到世界坐标系的旋转矩阵转置。
- 使用坐标变换将原子坐标从片段局部坐标系转换到世界坐标系(片段 → 分子 → 世界):
- 引理 3.3. 3D 分子图同构的规范序 (Lemma 3.3. [Canonical Ordering for 3D Molecular Graph Isomorphism])
- 设 M=(V,Z) 是一个具有节点类型矩阵 Z 和节点坐标矩阵 V 的 3D 分子图。设 m 是基于 L(M) 中前三个非共线片段中心构建的 M 的等变局部坐标系,g 是基于 L(Gi) 中前三个非共线原子构建的任意片段 Gi 的等变局部坐标系。
- f(⋅) 是我们的函数,它将分子 M 的 3D 坐标矩阵 V 在分子局部坐标系 m 下映射到球面表示 S。g(⋅,⋅) 是将分子局部坐标系 m 和片段局部坐标系 g 映射到旋转向量 m 的函数。那么对于任意 3D 变换 τ∈SE(3),我们有:
f(V)=f(τ(V)),g(m,g)=g(τ(m,g))
- 给定球面表示 S=f(V) 和旋转向量 m=g(m,g),存在一个变换 τ∈SE(3) 使得:
f−1(S)=τ(V),g−1(m)=τ(m,g)
- 这个引理说明了所构建的表示具有 SE(3) 等变性质:
- 无论分子如何在 3D 空间中旋转或平移,其球面表示和旋转向量都保持不变。
- 从这些表示可以唯一地重构出原始的几何结构(经过某个 SE(3) 变换)。
3.4. Frag2Seq: 片段和几何感知词元化 (Frag2Seq: Fragment and Geometry-Aware Tokenization): 3D 分子片段到 1D 序列的可逆变换。
- 片段-位置向量 (Fragment-Position Vector):
- 给定一个具有 k 个片段的分子 M,我们将每个片段 gi 的片段向量表示为:
xi=[si,di,θi,ϕi,mxi,myi,mzi]
- 其中 si 是 gi 的规范 SMILES 字符串。
- 按照其 ℓ1,⋯,ℓk 连接成一个 1D 序列。
- 给定一个具有 k 个片段的分子 M,我们将每个片段 gi 的片段向量表示为:
- 序列化 (Serialization):
- 按照片段的规范序 ℓ1,…,ℓk 连接:
Frag2Seq(M)=concat(xℓ1∗,…,xℓk∗)
- 按照片段的规范序 ℓ1,…,ℓk 连接:
- 定理 3.4. 双射映射 (Theorem 3.4. [Bijective Mapping])
- 对任意两分子 M1,M2,有:
Frag2Seq(M1)=Frag2Seq(M2)⟺M1≅3DM2
- 当坐标值舍入化到小数点后 b 位时,仍保持 (10−b/2) 约束下等价。
- 对任意两分子 M1,M2,有:
4. 条件训练与生成 (Conditional Training and Generation)
- 训练目标 (Training Objective):
L(U;P)=i∑logpθ(ui∣ui−1,…,u1;P)(8)
- 其中:
- U={u1,...,un}: 整个分子的 token 序列(注意这里的 ui 是离散化后的 token,而不是原始的数值向量 xi)。
- P: 蛋白质口袋信息。
- pθ: 我们的语言模型,参数为 θ。
- 其中:
- 条件注入 (Conditional Injection):
- 使用了预训练的 ESM-IF1 模型来提取蛋白质的节点嵌入(特征)。
- 在每个注意力块中的多头自我注意之后,都在蛋白质节点嵌入和配体 token 嵌入之间添加交叉注意力:其中 Q 来自配体 token 嵌入,K 和 V 来自蛋白质节点嵌入。
- 生成过程 (Generation Process):
- 从一个起始 token 开始,模型自回归地一个一个预测 token,直到生成结束符。
- 然后,利用 Frag2Seq 的可逆性,将生成的 token 序列解码回 3D 分子结构。
