[Paper Review] A Large Encoder-Decoder Family of Foundation Models For Chemical Language

Eduardo Soares, Victor Shirasuna+4 more|arXiv (Cornell University)|Jul 24, 2024

Advanced Computational Techniques and Applications4 citations

TL;DR

This paper introduces SMI-TED289M, a large-scale encoder-decoder foundation model pre-trained on 91 million curated SMILES from PubChem (4 billion molecular tokens), achieving state-of-the-art performance in molecular property prediction and molecule reconstruction. The model demonstrates strong compositional structure in its latent space, enabling few-shot reasoning with linear relationships between molecular fragments and high Tanimoto similarity in generation tasks.

ABSTRACT

Large-scale pre-training methodologies for chemical language models represent a breakthrough in cheminformatics. These methods excel in tasks such as property prediction and molecule generation by learning contextualized representations of input tokens through self-supervised learning on large unlabeled corpora. Typically, this involves pre-training on unlabeled data followed by fine-tuning on specific tasks, reducing dependence on annotated datasets and broadening chemical language representation understanding. This paper introduces a large encoder-decoder chemical foundation models pre-trained on a curated dataset of 91 million SMILES samples sourced from PubChem, which is equivalent to 4 billion of molecular tokens. The proposed foundation model supports different complex tasks, including quantum property prediction, and offer flexibility with two main variants (289M and $8 imes289M$). Our experiments across multiple benchmark datasets validate the capacity of the proposed model in providing state-of-the-art results for different tasks. We also provide a preliminary assessment of the compositionality of the embedding space as a prerequisite for the reasoning tasks. We demonstrate that the produced latent space is separable compared to the state-of-the-art with few-shot learning capabilities.

Motivation & Objective

• To develop a scalable, open-source foundation model for chemical language representation using self-supervised pre-training on large-scale unlabeled molecular data.
• To improve molecular property prediction and generation by learning contextualized, compositional representations through an encoder-decoder transformer architecture.
• To investigate the compositional structure of learned molecular embeddings as a prerequisite for chemical reasoning and few-shot generalization.
• To provide a flexible, high-performance model family (289M and 8×289M parameter variants) for diverse scientific applications in drug discovery and materials science.
• To release a curated dataset of 91 million high-quality SMILES for future pre-training of chemical foundation models.

Proposed method

• Pre-training a transformer-based encoder-decoder model on 91 million curated SMILES from PubChem, equivalent to 4 billion molecular tokens, using self-supervised masked language modeling.
• Employing a dual-encoder mechanism with a shared vocabulary and hierarchical tokenization to model molecular structure and connectivity via SMILES strings.
• Introducing a Mixture-of-Experts (MoE) architecture for the larger variant (SMI-TED8x289M), enabling parameter-efficient scaling with sparse activation.
• Using t-SNE visualization and linear regression on embedding projections to quantify compositional structure in the latent space, with triples of SMILES representing fragment addition.
• Evaluating few-shot reasoning by generating embeddings from a few example triples and reconstructing target SMILES using Tanimoto similarity.
• Applying standard benchmarks (QM9, MOSES) and regression/classification tasks to validate performance across quantum, physical, biophysical, and physiological properties.

Experimental results

Research questions

• RQ1 Can a large-scale encoder-decoder foundation model learn compositional molecular representations from self-supervised pre-training on unlabeled SMILES?
• RQ2 How well does the learned latent space of SMI-TED289M support few-shot reasoning and molecular reconstruction?
• RQ3 To what extent does the model’s embedding space reflect hierarchical and linear relationships between molecular fragments?
• RQ4 Does pre-training on a curated, large-scale SMILES corpus lead to improved performance on molecular property prediction compared to prior methods?
• RQ5 Can the model generalize to diverse chemical tasks with minimal fine-tuning, demonstrating strong few-shot capabilities?

Key findings

• SMI-TED289M achieved state-of-the-art performance on 11 out of 12 tasks in the QM9 benchmark for quantum mechanical property prediction.
• The model achieved an R² of 0.99 and MSE of 0.002 in linear regression on fragment addition tasks, indicating strong compositional structure in the latent space.
• In few-shot generation tasks, the model achieved a mean Tanimoto similarity score of 0.52, with top results reaching 0.92 for correct SMILES reconstructions.
• The t-SNE visualization of SMI-TED289M embeddings showed clear separation by molecular family and hierarchical distance, outperforming MoLFormer (R²=0.55, MSE=0.237).
• The model demonstrated robust few-shot generalization, with high similarity scores even when only a few triples were used as input for embedding generation.
• The curated dataset of 91 million SMILES from PubChem was shown to be effective for pre-training, significantly enhancing model performance and compositional reasoning.