文献总结|使用上下文增强的分子表示提升少样本药物发现的效果

Category 碎碎念

OpenReview

本文介绍于 2023 年发表在 ICLR 2023 上的一篇文章,文章原标题为 Context-enriched molecule representations improve few-shot drug discovery,文章介绍了一种可以用于药物发现的少样本学习模型 MHNfs,MHNfs 通过 Hopfield 网络用上下文数据集少样本的强化分子表示,提升了分子性质预测的准确度。

深度学习已经成为了药物发现中的重要工具,但目前大部分深度学习方法都是通过大训练集获得分子信息。药物发现中的深度学习方法通常需要大量的生物试验数据,这在实际的药物研发过程中很难获取。少样本学习解决了药物发现中有效数据较少的问题,少样本学习主要有 3 种方法:

  1. 基于数据增强的方法(Data-augmentation based approaches):变换已有数据达到增加数据量的目的。
  2. 基于词嵌入与最近邻的方法(Embedding-based and nearest neighbour approaches learn approaches):学习词嵌入的空间,从已有数据邻近位置取得新数据(相似分子)。
  3. 基于优化和微调的方法(Optimization-based or fine-tuning methods):将大规模的预训练模型放在已有数据上微调,使其迁移到新的化学空间。

文章提出了一种新的 MHNfs 模型用于少样本的药物发现,模型使用联想记忆来提取原始数据中的共现和协变结构从而强化其分子表示,在少样本数据集 FS-Mol 上达到了最佳效果。

方法

原理

药物发现中所使用的模型 \(g(\boldsymbol{m})\) 用于在给定分子表示 \(\boldsymbol{m}\in\mathcal{M}\) 的情况下预测分子性质或活性 \(\hat{y}\)。深度学习模型中的分子编码器将分子的一些低级表示(如 SMILES、分子图等)映射为模型空间的表示 \(f^\mathrm{ME}:\mathcal{M}\rightarrow\mathbb{R}^d\),再通过后续计算给出分子性质。

在少样本的情况下,只有分子的小数据集 \(\{\boldsymbol{x}_1,\cdots,\boldsymbol{x}_N\}\) 与对应分子是否具有活性的数据 \(\boldsymbol{y}=\{y_1,\cdots,y_N\}\)。这里将数据集 \(\{(\boldsymbol{x}_n,y_n)\}_{n=1}^N\) 称为 support set,少样本学习就是要正确预测不在 support set 中 \(\boldsymbol{x}\) 所对应的 \(y\)

文章中的模型分为 3 个模块:

$$ \begin{align} \text{context module: }&\quad&\boldsymbol{m}'&=f^\mathrm{CM}(\boldsymbol{m},\boldsymbol{C})\\ &\quad&\boldsymbol{X}'&=f^\mathrm{CM}(\boldsymbol{X},\boldsymbol{C})\\ \text{cross-attention module: }&\quad&[\boldsymbol{m}'',\boldsymbol{X}'']&=f^\mathrm{CAM}([\boldsymbol{m}',\boldsymbol{X}'])\\ \text{similarity module: }&\quad&\hat{y}&=f^\mathrm{SM}(\boldsymbol{m}'',\boldsymbol{X}'',\boldsymbol{y}) \end{align} $$

\(\boldsymbol{m}\in\mathbb{R}^d\) 是分子的词嵌入表示,\(\boldsymbol{X}\in\mathbb{R}^{d\times N}\) 是 support set 中分子的词嵌入表示,\(\boldsymbol{C}\in\mathbb{R}^{d\times M}\) 是另一个更大的分子数据集(context set)中分子的词嵌入表示。

\(f^\mathrm{CM}\) 交换 \((\boldsymbol{m},\boldsymbol{C})\) 间与 \((\boldsymbol{X},\boldsymbol{C})\) 间的上下文信息,得到强化的表示 \(\boldsymbol{m}'\)\(\boldsymbol{X}'\)。拼合两个增强的表示,\(f^\mathrm{CAM}\) 计算两者间注意力,得到进一步增强的 \(\boldsymbol{m}''\)\(\boldsymbol{X}''\),最后结合二者的信息进行预测。上面的过程可以描述成

$$ \begin{align} &\underset{\textsf{symbolic or}\atop\textsf{low-level repr.}}{m}\overset{f^\mathrm{ME}}{\longrightarrow}\underset{\textsf{molecule}\atop\textsf{embedding}}{\boldsymbol{m}}\overset{f^\mathrm{CM}}{\longrightarrow}\underset{\textsf{context}\atop\textsf{repr.}}{\boldsymbol{m}'}\overset{f^\mathrm{CAM}}{\longrightarrow}\underset{\textsf{similarity}\atop\textsf{repr.}}{\boldsymbol{m}''}\\ &\underset{\textsf{symbolic or}\atop\textsf{low-level repr.}}{x_n}\overset{f^\mathrm{ME}}{\longrightarrow}\underset{\textsf{molecule}\atop\textsf{embedding}}{\boldsymbol{x}_n}\overset{f^\mathrm{CM}}{\longrightarrow}\underset{\textsf{context}\atop\textsf{repr.}}{\boldsymbol{x}_n'}\overset{f^\mathrm{CAM}}{\longrightarrow}\underset{\textsf{similarity}\atop\textsf{repr.}}{\boldsymbol{x}_n''} \end{align} $$

模型

n

MHNfs 由 Transformer 中的 encoder 部分构建,具有与 Transformer 类似的结构与工作方式。

模型中的上下文模块由现代 Hopfield 网络(Modern Hopfield Network, MHN)实现:

$$ \mathrm{Hopfield}(\boldsymbol{\Xi},\boldsymbol{C}):=(\boldsymbol{W}_E\boldsymbol{C})\mathrm{Softmax}\left(\beta(\boldsymbol{W}_C\boldsymbol{C})^\top(\boldsymbol{W}_\Xi\boldsymbol{\Xi})\right) $$
$$ \boldsymbol{m}'=\mathrm{Hopfield(\boldsymbol{m},\boldsymbol{C})},\quad\boldsymbol{X}'=\mathrm{Hopfield}(\boldsymbol{X},\boldsymbol{C}) $$

MHN 能够计算两个输入间的注意力,最后更新得到的分子表示就具有参考分子集 \(\boldsymbol{C}\) 中的联想记忆。

交叉注意力模块替换了原来 Transformer 中的多头注意力机制,但功能仍然类似,用于记算输入分子 \(\boldsymbol{m}'\) 与 support set \(\boldsymbol{X}'\) 之间的注意力,再次更新分子表示:

$$[\boldsymbol{m}'',\boldsymbol{X}'']=\mathrm{Hopfield}([\boldsymbol{m}',\boldsymbol{X}'],[\boldsymbol{m}',\boldsymbol{X}'])$$

在最后的相似性模块中,模型计算输入分子 \(\boldsymbol{m}''\) 与 support set \(\boldsymbol{X}''\) 中每个分子 \(\boldsymbol{x}_n''\) 之间的相似性 \(k(\boldsymbol{m}'',\boldsymbol{x}_n'')\),并使用所有相似性的加权平均表示输入分子,用该表示计算输入分子的性质:

$$\hat{y}=\mathrm{Sigmoid}\left(\tau^{-1}\frac 1N\sum_{n=1}^Ny_n'k(\boldsymbol{m}'',\boldsymbol{x}_n'')\right)$$

文章这么做的理由是,考虑现实中的情况,当药物化学家对某系列化合物只有有限的活性数据(support set)而要预测(同一靶点或类似结构的)一化合物(query molecule)的活性时,化学家会将该化合物与手头已有数据的化合物对比,再在化合物库(context set)中对比,综合考虑各项因素后得出判断。模型所做的 MHN 计算以及平均相似性,就是简化了的上述过程,文章认为这样的设计有助于模型模仿化学家的思考方式。

数据

文章使用用于少样本药物发现的标准数据集 FS-Mol 作为模型的数据集。该数据集中的分子来自于 ChEMBL 27,其中定义了 4938 个训练任务,40 个验证任务与 157 个测试任务,平均每个任务下只有 94 个数据点。

文章使用 ECFPs 分子指纹与 RDKit 描述符来作为初始的分子表示。

实验

n

实验结果如上表所示,文章对比了各个模型在 FS-Mol 测试集上的 ΔAUC-PR,除了 ADKF-IFT 在 Hydol. 与 Oxid. 小部分任务上优于 MHNfs,其他模型的结果都不如 MHNfs,而且 MHNfs 在全部任务的整体结果上优于其他全部模型,所以文章认为 MHNfs 在 FS-Mol 测试集实现了目前药物发现少样本学习的最优性能。

结论

文章提出了一种可以用于药物发现的少样本学习模型 MHNfs,MHNfs 参考了现实中化学家面对少样本数据时的策略,通过设想的一种上下文增强的方式更新了输入模型的分子表示,使其具有更多大数据集中的背景信息。在实验中,测试结果表示这种增强的分子表示确实提高了模型预测的准确率,MHNfs 也在该任务上达到了最优的性能。