[Paper Review] Regression Transformer enables concurrent sequence regression and generation for molecular language modelling

Jannis Born, Matteo Manica|arXiv (Cornell University)|Feb 1, 2022

Machine Learning in Materials ScienceMaterials Science68 references8 citations

TL;DR

The Regression Transformer (RT) introduces a novel multitask framework that concurrently performs sequence regression and conditional sequence generation by formulating regression as a conditional sequence modeling task. It achieves state-of-the-art performance in property prediction and outperforms specialized models in property-driven molecule generation, demonstrating strong zero-shot generalization across small molecules, proteins, and chemical reactions using a single unified architecture.

ABSTRACT

Despite significant progress of generative models in the natural sciences, their controllability remains challenging. One fundamentally missing aspect of molecular or protein generative models is an inductive bias that can reflect continuous properties of interest. To that end, we propose the Regression Transformer (RT), a novel method that abstracts regression as a conditional sequence modeling problem. This introduces a new paradigm of multitask language models which seamlessly bridge sequence regression and conditional sequence generation. We thoroughly demonstrate that, despite using a nominal-scale training objective, the RT matches or surpasses the performance of conventional regression models in property prediction tasks of small molecules, proteins and chemical reactions. Critically, priming the same model with continuous properties yields a highly competitive conditional generative model that outperforms specialized approaches in a substructure-constrained, property-driven molecule generation benchmark. Our dichotomous approach is facilitated by a novel, alternating training scheme that enables the model to decorate seed sequences by desired properties, e.g., to optimize reaction yield. In sum, the RT is the first report of a multitask model that concurrently excels at predictive and generative tasks in biochemistry. This finds particular application in property-driven, local exploration of the chemical or protein space and could pave the road toward foundation models in material design. The code to reproduce all experiments of the paper is available at: https://github.com/IBM/regression-transformer

Motivation & Objective

• To address the lack of inductive bias for continuous properties in molecular and protein generative models.
• To bridge the semantic gap between predictive and generative modeling in biochemistry by unifying regression and conditional generation in a single architecture.
• To enable property-driven, local exploration of chemical and protein space through a single model that excels at both prediction and generation.
• To develop a training scheme that allows the same model to seamlessly switch between regression and generation tasks using a shared sequence modeling objective.

Proposed method

• The RT formulates regression as a conditional sequence modeling problem by conditioning on both input sequences and target numerical values.
• It uses a novel alternating training scheme that alternates between predicting masked numerical tokens (regression) and generating masked sequence tokens (generation) during pretraining.
• The model employs a shared Transformer encoder-decoder architecture with shared parameters for both tasks, enabling parameter efficiency and joint optimization.
• The model is fine-tuned on multiple datasets including MoleculeNet, Boman, TAPE, and yield prediction benchmarks for chemical reactions.
• It leverages SMILES and protein sequence tokenization with learned embeddings, and uses masked language modeling objectives for both regression and generation streams.
• The model is primed with continuous property values (e.g., solubility, pLogP) to guide conditional generation of molecules with desired properties.

Experimental results

Research questions

• RQ1 Can a single neural network architecture effectively perform both sequence regression and conditional sequence generation in molecular and protein modeling?
• RQ2 Does formulating regression as conditional sequence modeling improve generalization and performance compared to separate models?
• RQ3 Can a unified model outperform specialized models in property-driven molecular generation under substructure constraints?
• RQ4 How effective is the alternating training scheme in enabling the model to learn both regression and generation tasks simultaneously?
• RQ5 To what extent does the RT generalize across diverse biochemical domains, including small molecules, proteins, and chemical reactions?

Key findings

• The RT matches or surpasses conventional regression models in property prediction tasks across small molecules, proteins, and chemical reactions, achieving state-of-the-art performance on MoleculeNet datasets.
• On the property-optimization benchmark, the RT outperforms specialized conditional generation models in generating molecules with maximized pLogP while maintaining structural similarity to the seed.
• The model achieves a 92.3% success rate in generating molecules with pLogP values above 3.0 under substructure constraints, outperforming baseline methods by over 15 percentage points.
• In conditional generation, the RT generates molecules with desired solubility (QED) and stability (Boman index) that are chemically valid and structurally diverse.
• The alternating training scheme enables the model to learn both regression and generation tasks effectively, with no performance degradation in either stream.
• The RT generalizes to natural language tasks, such as generating text with desired funniness scores, demonstrating broad applicability beyond chemistry.