这周我们简单介绍一个高效份子优化的方法。该工作由UIUC的Jimeng Sun组合MIT的Connor Coley组合作完成，对应的文章题目是Differentiable Scaffolding Tree for Molecule Optimization[1]，被2022年ICLR接受，主要的代码和数据发布在https://github.com/futianfan/DST。
内容：
思路：基于梯度的份子优化
份子的可微分骨架树
类梯度上升的优化算法
优化效果测试
由可微性失掉的可解释性
思路：基于梯度的份子优化
在药物发现中，份子优化，即找到具有理想性子的份子结构，是核心的一步。由于化学结构的复杂性，传统上我们只能依赖于一些启发式的组合优化算法，如遗传算法、树搜索等。最近深度进修的发展确实提供了一些新的思路，但目前大部分依赖图生成算法的份子优化都是通过图神经网络（graph neural network, GNN）来显式地生成一个份子，然后优化方针函数，通过反向传播梯度来更新GNN参数，使得网络生成具有优化性子的份子。例如在增强进修（reinforcement learning, RL）中方针函数是根据反馈（reward）定义的；在深度生成模型（deep generative model, DGM）中基于和方针份子的广义距离定义。但这些算法普遍优化能力不够强，尤其没有考虑oracle的成本，许多算法需要挪用数十万甚至百万次oracle才能失掉较好的结果，而这在实际份子设计过程中显然是不现实的。
我们知道，数值优化的核心就是在一个点估计其指向极值点的方向，而这一方向一般可以通过梯度估计。那我们是否可以估计一个份子的性子相对于结构的梯度，通过这个梯度估计方向进而优化一个份子？在Alan Aspuru-Guzik组的Deep Molecular Dreaming[2]一文中作家利用份子的字符串（SELFIES）表示实现了这一点：将份子看做每个位置字符的分布概率，通过一个一维卷积神经网络（convolutional neural network, CNN）进修其性子，失掉可微的性子预测器，进而估计份子的梯度，但是效果并不好。而在本文中，作家通过提出份子的可微分骨架树（Differentiable Scaffolding Tree, DST）这一概念，使得份子直接在结构层面上可微，进而构建了一个高效的优化算法。与其他生成模型不同，作家先预训练(pre-train)了一个以骨架树（ST）为输入的GNN来预测性子（标量）。然后在优化过程中固定GNN参数，每步迭代里首先构造份子相应的DST，通过前向传播用GNN来预测性子，然后优化方针性子并通过反向传播梯度来更新DST里的参数，进而优化份子结构。
目前主要的生成模型和DST的优化对比。
份子的可微分骨架树
首先我们明确本文关注从头份子优化（de novo molecule optimization），即以一个优化算法在一个隐式定义的小份子空间上找到性子较好的份子。而份子性子由一个Oracle给出，即给定一个份子，输入相对应的性子，可以看作一个黑盒函数（black box function of molecules），例如权衡一个份子的类药性的QED。
为了使份子在图结构层面上可微，首先需要扩展份子图的概念。另外为了避免生成环的中间步骤可能带来的不必要的麻烦，作家挑选在优化过程中用骨架树（scaffolding tree, ST）来表示份子，其节点定义为一个基本单位（substructure），包含了常见的原子和单环（详见原文附录Figure 5）。我们可以用节点的种别矩阵（node indicator matrix）和节点之间的链接矩阵（adjacency matrix）表示一个骨架树。其中种别矩阵每一行是一个one-hot向量，代表了该节点是哪一个基本单位，而链接矩阵中的每一个元素都是binary的数字，指示每一对节点之间是否毗连。
份子结构、骨架树（ST）、可微分骨架树（DST）示意。
构造可微分骨架树（DST）的核心是将节点的种别和毗连看做可进修的0到1的概率分布，而非0或1的binary code。其中为了实现毗连可微性的自洽，作家提出了与毗连等价的权重向量（node weight vector），通过一个节点的权重表示其存在与否，并通过权重构造毗连矩阵，将毗连与否的问题变成了该节点是否存在的问题：
从骨架树失掉可微分骨架树的方法如下所示：首先将份子结构抽象为骨架树, 然后将骨架树中每个结点毗连上一个拓展节点（expansion node）。每一个叶结点和拓展结点的权重和种别是可进修的（learnable）。其中节点种别是一个softmax的输入，保证和为1。结点权重是一个sigmoid的输入，保证在0-1之间。
类梯度上升的优化算法
为了进行优化，作家首先预先训练（pre-train）了一个GNN来做性子预测，其输入是一个份子的可微分骨架树，输入是方针的性子（标量），即。为了平衡不同权重节点的贡献，作家在GNN中使用了加权平均的read-out方法：
其中代表第回迭代之后的节点特征（node embedding）的第行，代表第个节点的权重，代表全毗连网络。
整个算法是一个迭代式优化。在单步迭代中，给定输入份子的DST，失掉了可微的性子预测替代后，可以通过任意梯度优化算法（文中使用了Adam）解
来失掉优化后的DST。在失掉优化后的DST后，根据其优化后的权重和种别，作家设计了如下三种在对应骨架树上的操作：（1）删除（SHRINK）叶结点权重小时，意味着该节点对性子提升没有帮助，或者说该节点的存在对性子提升有负面影响，所以在对应骨架树上删除该节点。（2）拓展（EXPAND）拓展节点权重大时，意味着该节点的存在对性子提升帮助是正向的，所以在对应骨架树上加上一个新的结点。该节点的种别也从对应softmax输入值中挑选。（3）替换（REPLACE）若一个节点权重改变不大，但种别改变较大时，不删除也不拓展，但是在优化后的分布中重新采样一个substructure（softmax输入中值比较大的）。
优化算法中的一步迭代。
根据更新后的DST采样其中一个操作，失掉对应的骨架树作为下一轮迭代的输入，如此我们迭代地优化份子。在每一轮迭代内DST的维数是固定的，由输入份子决定，因此一轮迭代只能失掉和原份子相差最多一个节点的份子。但在多轮迭代优化过程中，每一轮的输入都是上一轮的输入，因此相应的DST维度也会变化，使得生成份子的大小只受限于优化迭代次数。
在失掉优化后的骨架树后，作家枚举对应的所有可能的份子图，用oracle测量每一个份子。为了有更好的优化效果，并兼具输入份子的多样性，作家每次优化时同时优化多条轨迹，并在其中使用determinantal point process（DPP）来挑选保留的份子。即在每一部枚举出对应的多个可能的份子图之后，不是单纯根据性子好坏，而额外考虑了相似性矩阵的行列式，即挑选最大化下式的一批份子：
其中是subset 这批份子的性子分数的对角矩阵，而则是这批份子的相似性矩阵。可以注意到单纯依据的行列式挑选就是greedy的top-k挑选。而相似矩阵的行列式的最大化则鼓励subset内的diversity的增加（可以考虑一个的例子，对角线为1，非对角为彼此之间的相似度）。如此我们便失掉了一个完整的优化算法。
优化效果测试
作家首先权衡主要的优化效果，包含了单方针优化和多方针优化（同时优化多个性子）。为了能够有比较全面的对比，所有算法挑选了分数最高的100个份子权衡其：
新颖性（Nov）：生成份子不在训练集（如果有）中的比例；
多样性（Div）：生成份子的多样性，权衡对化学空间的探索能力；
方针性子的平均优化结果（APS）；
Oracle挪用的次数（#oracle）：我们关注有限的oracle挪用的情况下的表现，因为oracle挪用的数量是权衡一个算法效率的重要指标；
由于DST和一些其他算法的一部分oracle挪用可以离线完成（比如利用已有的标注数据），另一部分必须线上完成，所以#oracle是A+B的形式，A为线下挪用，B为线上挪用。
单方针优化结果。
多方针优化结果。
从实验结果上看，直接运用了梯度信息的DST优化效率最高，说明了由DST估计的梯度的可靠性。深度生成模型（LigGPT）由于其本身并不是一个迭代优化算法，表现并不好。增强进修方法（GCPN／MolDQN）虽然在无限oracle挪用的情况下能够失掉一些较好的结果，但是不出所料在限制oracle挪用的情况下表现相对不好。而以传统的组合优化方法为基础的算法（如GA+D，MARS）表现相对还是更好。
为了更系统地比较算法效率，作家测试了算法的oracle efficiency，即不同方法在不同oracle挪用次数下的优化性能。结果如下：
Oracle efficiency测试结果。其中DST-rand为DST的ablation study，即同样设定下每步随机挑选骨架树上的操作。
横轴为oracle挪用数量，纵轴为top-100个份子的平均性子（越高越好）。DST在三个任务上都取得了最好的效果，其他结论也和第一个实验类似。
由可微性失掉的可解释性
作家另外展示了由DST带来的份子性子的可解释性。通过观察当前DST对各个结点权重和梯度，我们可以看到不同的结点对性子提升的影响，进而分析不同基团、亚结构对性子的影响。
份子可解释性实例。
Therapeutics Data Commons（TDC）致力于推动机器进修和生命医学领域的融合，让更多的机器进修／生物医药研究者可以无门槛参与到实用且有价值的数据驱动模型开发中来。我们核心团队决定不定期suí biàn gē在本公众号分享领域相关的文章的解读，希望能对读者有所帮助，有任何反馈，或者有希望解读的文章与专题，欢迎在公众号后台留言！我们也欢迎相关专题的投稿，有兴趣的请关注！