[论文解读] BARTSmiles: Generative Masked Language Models for Molecular Representations

Gayane Chilingaryan, Hovhannes Tamoyan+7 more|arXiv (Cornell University)|Nov 29, 2022

Machine Learning in Materials Science被引用 15

一句话总结

BARTSmiles 提出了一种大规模、自监督的生成式掩码语言模型，用于分子表征，其在包含17亿个分子的ZINC20数据集上进行预训练，使用的计算资源比以往方法多一个数量级。该模型在11项分类、回归和生成任务中均达到最先进性能，可解释性分析表明，积分梯度法突出了已知的毒性药效团结构，且关键神经元与下游任务高度相关。

ABSTRACT

We discover a robust self-supervised strategy tailored towards molecular representations for generative masked language models through a series of tailored, in-depth ablations. Using this pre-training strategy, we train BARTSmiles, a BART-like model with an order of magnitude more compute than previous self-supervised molecular representations. In-depth evaluations show that BARTSmiles consistently outperforms other self-supervised representations across classification, regression, and generation tasks setting a new state-of-the-art on 11 tasks. We then quantitatively show that when applied to the molecular domain, the BART objective learns representations that implicitly encode our downstream tasks of interest. For example, by selecting seven neurons from a frozen BARTSmiles, we can obtain a model having performance within two percentage points of the full fine-tuned model on task Clintox. Lastly, we show that standard attribution interpretability methods, when applied to BARTSmiles, highlight certain substructures that chemists use to explain specific properties of molecules. The code and the pretrained model are publicly available.

研究动机与目标

• 开发一种稳健、可扩展的自监督预训练策略，专用于分子表征的掩码语言建模。
• 使用远超以往模型的计算资源，训练一个大规模生成式分子语言模型——BARTSmiles。
• 在包括分类、回归和生成在内的多样化下游任务上评估BARTSmiles，以确立最先进性能。
• 通过神经元层面分析和可解释性方法，探究预训练表征是否隐式编码了与任务相关的信息。
• 通过对比归因图（如积分梯度）与已知的化学结构警示，验证模型的可解释性。

提出的方法

• 在完整的ZINC20数据集（17亿个分子）上，使用SMILES字符串对类BART的Transformer模型进行预训练，目标为掩码语言建模。
• 采用稳健的消融驱动预训练策略，以优化分子表征学习的超参数和架构选择。
• 在11项多样化下游任务上微调冻结的BARTSmiles模型，包括毒性预测（ClinTox）、溶解度（ESOL）和反应预测。
• 应用积分梯度法进行特征归因，以识别对模型预测有贡献的原子/结构片段。
• 通过识别激活值与下游任务性能高度相关的单个神经元，开展神经元层面分析。
• 通过线性探测冻结神经元（7–15个神经元）实现超过全微调模型90%的性能，表明任务特定表征学习能力。

实验结果

研究问题

• RQ1 在SMILES字符串上训练的大规模、自监督掩码语言模型，能否学习到可泛化的分子表征，并在多样化任务中超越先前方法？
• RQ2 BARTSmiles中的预训练表征在多大程度上隐式编码了与毒性或溶解度等下游任务相关的信息？
• RQ3 如积分梯度等可解释性方法是否能突出已知与分子性质相关的化学有意义结构片段（如结构警示）？
• RQ4 BARTSmiles中少量冻结神经元能否实现接近最先进水平的下游任务性能，表明其具备高效的表征学习能力？
• RQ5 模型的注意力机制和归因行为是否与药物化学中的专家知识（如致突变性结构警示）保持一致？

主要发现

• BARTSmiles在所评估的11项分子性质预测、生成和回归任务中全部达到最先进性能。
• 由于其规模和稳健的预训练策略，该模型在所有任务中，尤其是在低数据场景下，显著优于先前的自监督分子表征方法。
• 积分梯度归因图始终突出已知的致突变性结构警示（如硝基基团和环氧基团），证实其与专家知识的一致性。
• 在ClinTox任务中，仅7–15个冻结神经元的线性组合即可实现超过全微调模型90%的性能，表明存在强大的任务特定表征学习。
• 在ESOL回归任务中，模型将极性官能团（如-OH、C=O）识别为影响溶解度的关键因素，而对烃链的归因呈负向，与化学直觉一致。
• 观察到积分梯度分数存在系统性偏差——首字母标记的token始终获得更高的归因分数——通过均值减法处理后，显著提升了可解释性保真度。