韩国科学技术院提出HI-Mol模型，仅使用训练集2%的数据即可实现分子生成

寻找新分子一直是化学中一个基本而关键的问题，因为它与实现重要应用(如药物发现和材料设计)有着密切的关系。基于生成模型最近在其他领域(例如图像和视频)在学习高维数据分布方面的突破的启发，一些研究已经考虑使用大型分子数据集训练深度生成模型来学习分子分布。尽管分子生成方法在基于大规模数据集来寻找化学上有效的新分子的任务上表现出优秀的性能，然而，现有的生成模型在小样本数据集的分子生成任务上仍面临挑战。

2024年5月2日，韩国科学技术院Seojin Kim等人在人工智能顶级会议ICML2024上发表文章Data-Efficient Molecular Generation with Hierarchical Textual Inversion。

作者提出了一种新的分子生成方法，分子生成层次文本反转(Molecular generation with Hierarchical textual Inversion, HI-Mol)。HI-Mol的灵感来自于层次信息的重要性，例如，在理解分子分布时，粗粒度和细粒度的特征。作者提出在深度学习领域采用最新的文本反转技术的基础上，使用多层次嵌入来反映这种层次特征，从而实现高效的分子生成。实验证明了HI-Mol相比于现有方法更优秀的性能，仅使用训练集2%的数据即可实现高效的分子生成。

文本反转方法是一种简单而强大的小样本生成技术，它将给定样本中的共同概念作为语义嵌入空间中的单个标记来学习。受其成功的激励，HI-Mol的目标是学习分子作为文本标记的常见化学概念，并将它们用于数据高效生成的目标。然而，与图像或语言数据不同，具有相似语义的分子通常具有完全不同的结构，这使得很难简单地将共同概念作为单个文本标记来学习。HI-Mol的贡献在于通过在框架中采用分子特异性先验(即分子的层次结构)来解决这一挑战，从而通过文本反转技术实现数据高效的分子生成，如图1所示。

图1 HI-Mol结构图

在深度学习领域，最近对于文本到图像生成的研究提出了文本反转方法，其目的是学习一个共同的概念，并将其用于概念嵌入(或个性化)生成。为了实现这一点，文本反转优化了新令牌（token）的单个文本嵌入，使用冻结的预训练文本到图像模型实现图像之间共享学习，然后使用给定的具有与用于训练的样本完全相同目标的图像来优化这个令牌嵌入。

对于分子数据而言，基于分子可以分层聚类的化学先验的启发，HI-Mol引入了多层次嵌入，以捕获小样本分子中的粗粒度和细粒度特征，其中每个分子的多层次标记的选择性分配允许纳入分子的分层特征。具体来说，HI-Mol首先提出了一种分子特异性文本反转来学习分子的期望分布。与原始的文本反转只使用一个共享令牌[S*]不同，HI-Mol还使用了“分层”令牌[I*]和[D*]，其中I表示对分子中的某些簇（即相邻的若干原子或官能团的集合）作为令牌，称为中间（intermediate）标记，D表示对分子中的某些原子作为令牌，称为细节（detailed）标记。因此，HI-Mol能够学习从原子级别，到官能团（或者簇）级别，再到整个分子级别的寓意特征。

经过训练，每个分子x被表示为文本令牌[S*][I*][D*]，并以无监督的方式进行训练，因此不需要关于每个分子的特定信息。实际上，可以通过这个简单的选择方案学习到一些信息丰富的簇级特征，尽管这一方法没有输入任何关于给定分子数据的先前化学知识。“多层次”标记设计对于分子的成功反转尤为重要，因为分子与现有文本反转方法中通常使用的图像具有不同的性质。传统文本反演中的图像输入在视觉上是相似的，例如，同一只狗的不同姿势的图片，而分子通常具有完全不同的结构，即使它们具有共同的化学概念，或者属于同一类别。这种差异使得很难将一类分子作为一个简单的单个标记来学习。

因此，作者通过结合化学文献强调分子数据可以分层聚类的原理，利用预训练文本到分子模型中的语法信息划分各个分子的簇，同时，通过分层文本反转框架从学习分布中采样，包括利用在文本反转中获得的学习到的分子层次信息，即中间标记和细节标记，从分子的目标分布中采样。例如，通过分子i（表示为[S*][Ici*][Di*]）和分子j（表示为[S*][Icj*][Dj*]）的加权λ([S*][Ici*][Di*])+(1-λ)([S*][Icj*][Dj*])采样生成新的分子。

作者利用QM9数据集进行实验，以显示HI-Mol的数据效率。该数据集包含超过10万个分子，因此已成为评估大规模分子生成框架的流行基准。在这里，作者整个QM9训练集的2%或10%子集来训练HI-Mol，而其他基线方法是在完整的训练集上训练的。对比指标如下：

Frechet ChemNet距离(FCD)：使用预先训练的ChemNet测量源和目标分布之间距离的度量。

邻域子图成对距离核(NSPDK)：另一个度量源分布和目标分布之间差距的度量，基于使用基于图的分子表示的算法计算。

有效性(Valid.)：生成的分子具有有效化学结构的比例。

唯一性(Unique.)：生成的分子的多样性，基于从生成模型中获得的不同样本占有效分子总数的比例。

新颖性（Novelty）：未包含在训练集中的有效分子的比例。

作者将HI-Mol与一些具有代表性的方法进行了比较。对比方法包括一系列基于图（graph）或SMILES的方法，以及一系列使用或未使用分子语法（簇）信息的方法。如表1所示，将FCD，NSPDK，有效性，唯一性和新颖性作为对比指标，前两个指标越低越好，后三个指标越高越好。仅使用2%的数据的HI-Mol在FCD, NSPDK和有效性这三个指标上已超越现有模型，并具有较高的唯一性和新颖性（大于75%），使用10%数据的HI-Mol总体表现更好。

表1 与其他方法对比

作者设计了消融实验来验证模型设计的有效性。为了验证HI-Mol框架中每个组件的有效性，作者比较了一些组件被排除在整体框架之外的结果。具体来说，作者比较了以下设置的生成性能:(1)不使用反转技术，(2)使用共享标记[S*]，(3)使用共享标记[S*]和细节标记[D*]，(4)使用所有三种类型的令牌，以及(5)应用额外的语法信息。首先，(1)不使用反转技术和(2)单个共享令牌的反转没有表现出合理的性能，即它们仅达到0.4%的有效性。在(3)和(4)中，通过学习分子中的底层特征，在反转框架中引入底层令牌，显著提高了生成质量。最后，(5)HI-Mol进一步提高了生成性能。

表2 消融实验

作者还进行了案例分析。为了验证中间标记[I*]设计的效果，作者在图2中通过提供分配给相同中间标记的分子组来可视化聚类结果。如图所示，分子根据其共同的亚结构，如长碳链或磺酰基苯基，被很好地分组。这表明，这种簇级别层次语义的学习在分子生成中确实是有益的。

图2 案例分析

作者还验证了HI-Mol生成的分子可以用来提高分类器的性能，用于分子性质预测。在这里，作者从MoleculeNet基线数据集中收集了三个数据集(HIV, BBBP和BACE)的活性和非活性分子，分别训练活性分子和非活性分子的分子生成模型，然后从模型中生成分子。在表3中报告了ΔROC-AUC的值，这是通过将生成的分子额外添加到原始的小样本训练数据（每一类的训练集中只包括16个或32个样本）时，ROC-AUC分数的改善来计算的，越高越好。结果表明HI-Mol在三个数据集上始终表现出优势，证明了HI-Mol即使在分子数量有限的情况下，也能学习到每个分子性质预测任务的共同概念，即活跃性和非活跃性。在标签信息难以实现的小样本预测的实际场景中，HI-Mol确实在改进分类器方面发挥了重要作用。

表3 分子性质预测

本文提出了HI-Mol，利用分子专用的文本反转实现数据高效的分子生成框架。具体而言，作者提出了一种分子专用的文本反转方案来捕获分子的层次信息。与在给定训练数据上优化单个共享标记的视觉域文本反转不同，作者为分子文本反转设计了多级标记。该方法广泛地利用了给定分子的信息，有效地缓解了分子数量有限的问题。HI-Mol仅使用训练集2%的数据即可实现高效的分子生成，超越了现有的基线方法，有望在高效生成分子方面开启尚未被探索但至关重要的研究方向。这项工作将促进分子生成的研究，并可以加快许多重要生成任务的发展，例如新药物的发现与设计。

参考文献：

Kim et al. Data-Efficient Molecular Generation with Hierarchical Textual Inversion. ICML. 2024