[论文解读] Molecule Attention Transformer

Łukasz Maziarka, Tomasz Danel+4 more|arXiv (Cornell University)|Feb 19, 2020

Computational Drug Discovery Methods参考文献 44被引用 108

一句话总结

Molecule Attention Transformer (MAT) 增强 Transformer 自注意力以分子间距离和分子图结构，实现具有竞争力的分子性质预测与可解释的注意力；预训练 MAT 减少超参数调优并提升性能。

ABSTRACT

Designing a single neural network architecture that performs competitively across a range of molecule property prediction tasks remains largely an open challenge, and its solution may unlock a widespread use of deep learning in the drug discovery industry. To move towards this goal, we propose Molecule Attention Transformer (MAT). Our key innovation is to augment the attention mechanism in Transformer using inter-atomic distances and the molecular graph structure. Experiments show that MAT performs competitively on a diverse set of molecular prediction tasks. Most importantly, with a simple self-supervised pretraining, MAT requires tuning of only a few hyperparameter values to achieve state-of-the-art performance on downstream tasks. Finally, we show that attention weights learned by MAT are interpretable from the chemical point of view.

研究动机与目标

• 开发用于广泛分子性质预测任务的深度模型。
• 通过邻接信息和距离数据将化学结构信息并入 Transformer 的注意力机制。
• 证明简单的自监督预训练可以降低超参数调优并提升性能。
• 展示 MAT 学到的注意力权重在化学上具有可解释性。

提出的方法

• 用分子自注意力取代 Transformer 自注意力，通过受 λa、λd、λg 控制的加权和将注意力、距离和图信息融合在一起。
• 使用包含原子身份和局部特征的 26 维嵌入来表示原子。
• 利用 RDKit 从 3D 构象计算距离矩阵 D，并通过 g(D) 将其并入注意力（例如对行进行 softmax 或 exp(-d)）。
• 用一个虚拟节点扩增输入，以便在没有模式时允许模型跳过。
• 使用节点级掩码特征预测对 MAT 进行预训练，以稳定学习并减少超参数搜索。
• 在一个多样化的分子数据集上进行 scaffold 和随机划分的评估，并与基线图模型和指纹模型进行比较。

实验结果

研究问题

• RQ1 是否能够通过加入分子图结构和分子间距离的注意力机制，在多样的分子性质预测任务中实现具有竞争力的性能？
• RQ2 自监督预训练是否能在 MAT 中以最小的超参数调优实现较强的性能？
• RQ3 MAT 中的注意力权重是否能从化学角度进行解释？

主要发现

• MAT 在广泛数据集上的表现与浅层模型及若干基于图的方法相比具有竞争力甚至更优。
• 通过预训练，MAT 的表现优于非预训练变体且仅需调优少数超参数（学习率），显著降低搜索工作量。
• 预训练的 MAT 在所有数据集上的平均排名最好，仅使用 7 次超参数评估，而其他方法需要上百次。
• 消融研究表明 使用虚拟节点以及结合图和距离信息可提高各任务的稳定性和性能。
• MAT 的注意力头学习到可解释的化学模式，其中一些头专注于特定原子类型或局部子结构。