[Paper Review] BERT Learns (and Teaches) Chemistry

Josh Payne, M. Sammer Srouji+2 more|arXiv (Cornell University)|Jul 11, 2020

Machine Learning in Materials ScienceMaterials Science34 references21 citations

TL;DR

This paper proposes using BERT with SMILES string representations of molecules to automatically identify chemically relevant substructures—such as functional groups—through attention mechanism analysis, without relying on domain-specific heuristics. It demonstrates that BERT's attention heads can highlight property-impacting substructures, and fine-tuning BERT on solubility, drug-likeness, and synthesis accessibility tasks yields interpretable, chemically meaningful attention patterns, particularly for solubility prediction.

ABSTRACT

Modern computational organic chemistry is becoming increasingly data-driven. There remain a large number of important unsolved problems in this area such as product prediction given reactants, drug discovery, and metric-optimized molecule synthesis, but efforts to solve these problems using machine learning have also increased in recent years. In this work, we propose the use of attention to study functional groups and other property-impacting molecular substructures from a data-driven perspective, using a transformer-based model (BERT) on datasets of string representations of molecules and analyzing the behavior of its attention heads. We then apply the representations of functional groups and atoms learned by the model to tackle problems of toxicity, solubility, drug-likeness, and synthesis accessibility on smaller datasets using the learned representations as features for graph convolution and attention models on the graph structure of molecules, as well as fine-tuning of BERT. Finally, we propose the use of attention visualization as a helpful tool for chemistry practitioners and students to quickly identify important substructures in various chemical properties.

Motivation & Objective

• To identify chemically significant molecular substructures in a data-driven way without relying on handcrafted rules or domain expertise.
• To explore whether BERT's attention mechanisms can learn to highlight functional groups that influence molecular properties such as solubility, toxicity, and drug-likeness.
• To evaluate the utility of BERT's learned representations as features for graph neural networks and for fine-tuning on downstream chemical regression tasks.
• To develop attention visualization as an interpretable tool for chemists to identify key substructures in molecular properties.
• To investigate whether pretraining on large-scale SMILES data enables BERT to learn chemically meaningful representations beyond standard atom featurization.

Proposed method

• Fine-tune BERT on a large dataset of SMILES strings representing molecules to learn contextual representations of atoms and substructures.
• Analyze attention heads across BERT layers to identify which atoms or substructures receive high attention during prediction tasks.
• Use BERT-learned representations as input features for graph convolutional and self-attention models on molecular graphs to predict solubility, drug-likeness, and synthesis accessibility.
• Fine-tune BERT directly on regression tasks (e.g., QED, SAS, solubility) using the ZINC15 dataset to assess performance and attention behavior.
• Visualize attention matrices to interpret which molecular fragments are most influential for specific chemical properties.
• Propose a contrastive loss to encourage BERT to learn consistent embeddings across different SMILES representations of the same molecule, improving structural awareness.

Experimental results

Research questions

• RQ1 Can BERT’s attention mechanisms automatically identify chemically relevant substructures in molecules without prior domain knowledge?
• RQ2 Do BERT-learned representations improve performance on molecular property prediction tasks such as solubility, drug-likeness, and synthesis accessibility?
• RQ3 How do attention patterns in BERT change when fine-tuned on different chemical regression tasks?
• RQ4 Can attention visualization serve as a practical interpretability tool for chemists to identify key functional groups in molecular properties?
• RQ5 Does pretraining on large-scale SMILES data lead to more robust and chemically meaningful representations than standard atom featurization?

Key findings

• BERT’s attention heads, particularly after fine-tuning on solubility, showed intuitive attention on known property-impacting substructures, suggesting meaningful pattern recognition.
• Fine-tuning BERT on regression tasks led to better performance than training from random weights, indicating that pretraining provides useful inductive bias.
• Attention patterns across different tasks were largely similar, except for notable differences in layers 10 and 11 for the solubility task, indicating task-specific attention learning.
• The model’s learned embeddings were not superior to canonical atom featurization in performance, suggesting limitations in capturing full molecular graph information.
• Despite limited performance gains, attention visualization revealed interpretable, chemically plausible substructure highlights, supporting its use as an interpretability tool.
• The authors observed that rare atoms not in BERT’s vocabulary had minimal impact, indicating robustness to vocabulary limitations.