文献总结|使用强化学习和基于图的深度生成模型进行从头药物设计

Category 碎碎念

doi.org/10.1021/acs.jcim.2c00838

本文介绍于 2022 年发表在 Journal of Chemical Information and Modeling 上的一篇文章,文章原标题为 De Novo Drug Design Using Reinforcement Learning with Graph-Based Deep Generative Models,文章借鉴了 REINVENT 模型的强化学习策略,将其用于基于图的分子生成模型,得到了较好的结果。

引言

强化学习是机器学习中常用的一种策略,强化学习通过奖惩函数确定某环境下执行相应行为的优劣,从而做出决策。强化学习常常用来 fine-tune 深度生成模型,使其生成的分子逐渐向具有目标特征的分子靠近。

GraphINVENT 是基于图神经网络的分子生成模型,REINVENT 是使用强化学习和 RNN 基于序列的分子生成模型。文章从这两个模型中受到启发,将 GraphINVENT 基于图的分子生成模型与 REINVENT 的强化学习策略相结合,提出了本文的模型。

方法

n

文章中提出的模型主要由三个部分构成:

  1. 基于图的分子生成模型;
  2. 具有记忆力感知损失函数(memory-aware loss)的强化学习框架;
  3. 打分模型。

分子生成模型

分子生成模型来自于 GraphINVENT 模型,使用了门控图神经网络(GGNN)。GGNN 通过不断选取行为,从而修改输入的分子图。行为集合中包括 3 种可选的行为:添加原子、添加键和结束,将行为集合记作 \(\mathcal{A}\)。那么生成分子图 \(\mathcal{G}\) 的任务就可以描述为一个马尔可夫过程:代理模型依据一定的行为概率分布(action probability distributions, APD)不断选择行为,最终生成 \(\mathcal{G}\)。行为概率分布由当前的状态决定,也就是说,通过前一步骤添加原子或键决定下一步可能的行为。

每个步骤的 APD 可以记作 \(\mathrm{APD}_i\)\(f:\mathcal{G}_i\rightarrow \mathrm{APD}_i\),GGNN 模型就是将在 \(i\) 步骤生成的子图 \(\mathcal{G}_i\) 映射到\(\mathrm{APD}_i\) 的函数 \(f\),若用 \(a_i\) 表示每个步骤的行为,那么完整的图生成过程就可以表示为

$$\mathrm{G}_0\rightarrow a_0\sim\mathrm{APD}_0\rightarrow\mathcal{G}_1\rightarrow\cdots\rightarrow a_{n-1}\sim\mathrm{APD}_{n-1}\rightarrow\mathcal{G}_n$$

在 GGNN 训练中使用目标 \(\mathrm{APD}_\mathrm{t}\) 与预测 \(\mathrm{APD}_\mathrm{p}\) 的 KL 散度作为损失函数:

$$D_\mathrm{KL}(\mathrm{APD}_\mathrm{t}||\mathrm{APD}_\mathrm{p})=\sum_{a\in\mathcal{A}}\mathrm{APD}_\mathrm{t}(a)\log\frac{\mathrm{APD}_\mathrm{t}(a)}{\mathrm{APD}_\mathrm{p}(a)}$$

强化学习框架

强化学习的目的是基于先验策略 \(\pi_\mathrm{prior}\) 更新策略,使行为序列的预期打分增加,具体来说,就是在每个学习步骤中更新代理模型,保证其一直是最佳的代理模型。最佳代理模型提醒(the best agent reminder, BAR)损失函数定义为:

$$J(\boldsymbol{\theta})=(1-\alpha)J_\mathbb{A}(\mathbb{A},\mathbb{P};\boldsymbol{\theta})+\alpha J_\tilde{\mathbb{A}}(\mathbb{A},\tilde{\mathbb{A}};\boldsymbol{\theta})$$

其中 \(\alpha\in[0,1]\) 是设定的比例系数,\(\mathbb{P}\) 是先验模型,\(\mathbb{A}\)\(\tilde{\mathbb{A}}\) 分别是当前模型与最优模型(详见强化学习的策略梯度算法)。

打分函数

为了执行不同的优化任务,使用了不同的打分函数。

增加或减少分子的大小

为了改变生成分子的大小,提出了以下打分函数:

$$\begin{equation} S_\mathrm{size}(A)= \begin{cases} 0, &\mathrm{if\ not\ \{PT,\ valid\ and\ unique\}}\\ 1-\frac{|n_\mathrm{nodes}-n^*_\mathrm{nodes}|}{\max_\mathrm{nodes}-n^*_\mathrm{nodes}}, &\mathrm{otherwise} \end{cases} \end{equation}$$

其中 \(n^*_\mathrm{nodes}\) 是分子中目标重原子数量,若要增加分子大小,设定高 \(n^*_\mathrm{nodes}\),若要减小分子大小,设定低\(n^*_\mathrm{nodes}\)\(\max_\mathrm{nodes}\) 是分子中最大节点数量,\(n_\mathrm{nodes}\) 是分子中重原子数量,\(A\) 指当前步骤的行为,\(\mathrm{PT}\) 指结束构造分子的行为。

增加分子类药性

类药性使用 RDKit 的 QED 进行评估,打分函数定义为

$$\begin{equation} S_\mathrm{QED}(A)= \begin{cases} 0, &\mathrm{if\ not\ \{PT,\ valid\ and\ unique\}}\\ \mathrm{QED}(\mathrm{Mol}(A)), &\mathrm{otherwise} \end{cases} \end{equation}$$

增加分子活性

使用支持向量机模型预测分子对于 DRD2 的活性,打分函数定义为

$$\begin{equation} S_\mathrm{QED}(A)= \begin{cases} 1, &\mathrm{if\ PT,\ valid,\ unique,}\\ &\quad \mathrm{QED}>0.5\ \mathrm{and\ activity} > 0.5\\ 0, &\mathrm{otherwise} \end{cases} \end{equation}$$

数据

训练模型的 DRD2 活性分子数据来自于 ChEMBL。

结果

不同打分函数

n

从使用不同打分函数生成的分子数据中可以看出,在训练过程中,生成分子的均分不断提高,合法分子和 PT 分子的比例也不断升高,说明打分函数都能够正确地奖惩生成分子,可以用于分子生成模型的训练。另外,选择减小分子尺寸或提高分子类药性或提高活性的打分函数后,独特分子的比例有所下降,这主要是因为指定的打分函数使模型进入了一个更小的化学空间。

n

文章分别从先验模型与单个 fine-tune 过的模型中选取 10000 个分子,在 10 个不同 fine-tune 过的模型中各选取 1000 个分子进行评估。针对平均 QED 这一指标,经过 fine-tune 的模型表现出了接近的数值且要高于先验模型。最重要的是,先验模型无法生成具有 DRD2 活性的分子,而在这一点上,经过 fine-tune 的模型生成的活性分子数量显著更高。最后,来自于 10 个不同模型的 10000 个分子中,活性分子与独特分子的占比都很高,说明 10 个模型间重叠程度小,将比单个模型更有潜力。

可合成性

文章使用 AiZynthFinder 分析分子的可合成性,AiZynthFinder 能通过循环树搜索的方式将分子通过化学反应分割为商用试剂分子。

n

可以看出,预训练模型生成的可合成分子的比例要低于训练集,可以认为 GraphINVENT 模型没能学习到可合成性这一特征。在经过 fine-tune 的模型中,除了增大分子尺寸的模型以外,都生成了较多可合成分子。主要原因是,一方面,减小分子的尺寸就是在提高分子的可合成性,另一方面,具有高类药性与 DRD2 活性的分子大多具有相似的化学结构或片段,这也使得该类分子具有较高可合成性。

讨论

模型在训练过程中并没有使用任何真实活性分子的数据,fine-tune 之前的模型无法生成具有真实活性的分子,而经过 fine-tune 的模型能够生成一部分真实具有活性的分子,表明在实际药物发现中尚未发现真实活性分子情况下,文章中所使用的模型具有应用潜力。

文章中使用记忆力机制能够使模型在训练中更加平滑地偏向于最优代理模型,也能够使模型不会忘记生成较优分子的行为序列,从而引导模型更好地完成生成分子任务。文章使用的 BAR 损失函数能够很好完成多种分子生成任务。

结论

文章将策略梯度强化学习应用于基于图的从头药物设计工具,并且使用该模型生成了具有目标特征的分子。在该模型中所使用的 BAR 损失函数能够用于模型训练,并提升了分子生成模型的性能。此外,模型中使用的 4 种打分函数也完成了 4 种生成具有相应特征分子的任务。