Regression Transformer: Concurrent sequence regression and generation for molecular language modeling

Jannis Born; Matteo Manica

Source-linked AI summary

Regression Transformer: Concurrent sequence regression and generation for molecular language modeling

Jannis Born, Matteo Manica

arXiv:2202.01338v3Feb 1, 2022cs.LGcs.AIcs.CLq-bio.BM

TL;DR

Molecular and protein generative models lack an inductive bias for continuous properties, motivating a model that unifies prediction with controllable generation. The Regression Transformer treats regression as conditional sequence modeling and uses one multitask model for both tasks. It matches or surpasses conventional regression models and specialized generative approaches across the reported biochemical evaluations.

Problem
Molecular and protein generative models lack an inductive bias that reflects continuous properties, while prediction and conditional generation remain separated.
Method
The Regression Transformer tokenizes numerical and textual sequences and alternates training objectives so one model performs property prediction and conditional sequence generation.
Results
The RT matches or surpasses conventional regression models across several property-prediction benchmarks and specialized conditional generative models in constrained molecular generation.
Takeaways & Limitations
The RT supports property-driven, substructure-constrained exploration of molecular and protein space while bridging predictive and generative tasks.
Takeaways & Limitations
The authors identify fine-grained regression as an area for improvement and note that poor property prediction can cause undesired effects in self-consistency training.

Abstract

from arXiv · show

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

1 Introduction

The Regression Transformer reframes regression as conditional sequence modeling, allowing one multitask language model to perform property prediction and property-driven generation across scientific domains.

Moving the [MASK] location lets the same model predict numerical tokens from text or generate text conditioned on numerical and textual tokens.This supports regression and conditional generation without finetuning task-specific heads.
Existing approaches typically separate property prediction from conditional molecular generation, leaving both tasks weakly entangled.This separation persists across GANs, reinforcement learning, VAEs, graph neural networks, flows, and diffusion models.
The RT learns the joint distribution of numerical and textual tokens instead of using a designated regression head.It uses tokenized numbers and cross-entropy loss, with numerical encodings supplying an inductive bias for decimal-order and semantic proximity.
The paper evaluates RT across small molecules, proteins, chemical reactions, and natural text using predictive and generative tasks.Small-molecule experiments include synthetic drug-likeness data and three MoleculeNet property-prediction datasets.

2 Results

Across molecular, protein, and reaction tasks, the Regression Transformer (RT) achieved competitive property prediction while also supporting conditional generation from continuous property prompts. Alternating objectives and self-consistency improved generation, enabling local design and constrained optimization across chemical and protein sequences.

Small-molecule modeling: 0.017 MAE: RT with alternating objectives outperformed k-NN and SMILES-BERT on synthetic QED prediction.The reported baselines were 0.054 for k-NN and 0.020 for SMILES-BERT.
Conditional molecular generation: Novel conditional molecules exceeded 99% novelty while following continuous property prompts, with Grover externally validating generated-molecule properties.The RT generated molecules conditioned on QED and molecular properties including aquatic solubility.
MoleculeNet regression: Across ESOL, FreeSolv, and Lipophilicity, RT outperformed XGBoost and was on par with or mildly inferior to a conventionally fine-tuned XLNet regressor.BERT-based approaches remained stronger on these benchmarks, while RT matched XLNet on Lipophilicity and stayed within standard deviation on ESOL and FreeSolv.
Conditional molecular generation: 614% and 103% average-improvement advantages over Junction-Tree-VAE and GCPN, respectively, were reported for constrained pLogP optimization across similarity settings.The RT required no inference-time optimization and also predicted pLogP with Pearson correlation 0.92.
Protein language modeling: Protein experiments combined strong prediction with generation: RT reached Spearman ρ > 0.994 for Boman’s index, ρ = 0.84 for synthetic peptide generation, and ρ = 0.44 for TAPE stability adaptation.On TAPE fluorescence, RT outperformed all reported methods in Spearman correlation but struggled with intra-mode precision.
Chemical reaction modeling: On reaction yields, RT outperformed fingerprint- and quantum-mechanics-based methods and matched or nearly matched Yield-BERT, while precursor reconstruction reached 98.2% for aryl-halides.For Buchwald–Hartwig and Suzuki–Miyaura datasets, the table reports RT correlations of 0.939±0.01 and 0.81±0.02, respectively.

3 Discussion

The Regression Transformer bridges property prediction and conditional generation in one multitask model, performing strongly across molecular and protein tasks. Its current scope includes single-property experiments, while finer-grained regression and reaction modeling remain areas for future work.

Core contribution: The RT bridges property prediction and conditional generation in a single multitask Transformer.The paper characterizes this as a flexible model that excels at both previously disjoint tasks.
Property prediction: The RT learns continuous molecular properties from small datasets, surpasses conventional regression models on several benchmarks, and sometimes competes with regression-trained Transformers.These results use tokenized numerical properties without ratio-scale information.
Conditional generation: Across chemical and protein tasks, the RT generates novel sequences that seemingly adhere to primed continuous properties.The paper identifies property-driven, substructure-constrained molecular or protein design as a potential application.
Scope: The reported experiments examined singular properties, although the RT also naturally scales to multiproperty prediction.Pretrained multiproperty models are available through the software described by the paper.
Future directions: Future work should improve fine-grained regression and intensify reaction modeling.The authors note that the RT effectively generalizes forward reaction and retrosynthesis models, while pointing to a fine-grained regression failure mode.

Software and Data

The paper releases code, datasets, and ready-to-use RT implementations through public repositories and GT4SD. These resources support reproduction, custom training, inference, and access to single- and multiproperty models.

Code: The complete experiment codebase is publicly available in the IBM Regression Transformer repository.The repository is intended to facilitate reproduction of all experiments.
Data: Public data sources cover MoleculeNet, molecular optimization, protein language modeling, and reaction-yield experiments.The paper lists separate URLs for each dataset group.
Software: GT4SD provides ready-to-use RT pipelines for training, finetuning, and inference, including QED, ESOL, protein-stability, and multiproperty models.A demonstration notebook is also distributed through the GT4SD repository.

4 Methods

The Regression Transformer reformulates regression and conditional generation as sequence-modeling tasks over mixed numerical and textual tokens. Its XLNet-based architecture uses numerical encodings and alternating objectives to predict properties, generate conditioned sequences, and evaluate property-driven molecular and reaction tasks.

Architecture: The RT uses an XLNet backbone to combine autoregressive modeling with bidirectional context for filling arbitrary masked positions.This supports generation when discontiguous molecular regions are masked while retaining access to the full sequence.
Tokenization: Numbers are split into digit-and-decimal-place tokens, while molecules use SELFIES tokens and proteins use amino-acid tokens.For example, 12.3 becomes [1_1, 2_0, ., 3_-1], preserving decimal order.
Numerical encodings: Numerical encodings add a proximity-aware inductive bias to numerical tokens and are summed with word embeddings and XLNet relative positional encodings.Their pairwise distances are symmetric and decay monotonically with the represented float value.
Training objectives: The alternating training scheme uses one objective for property prediction and another for conditional text generation.Property prediction masks numerical tokens, whereas conditional generation preserves property tokens and reconstructs masked textual tokens.
Training objectives: The self-consistency formulation can reward generated molecules that differ from training samples while adhering to the target property, but may cause undesired side effects when property prediction is poor.The scheme uses greedy-decoded predictions during optimization.
Evaluation: Conditional generation is evaluated with rank-based metrics for scaffold decoration and with top-3 accuracy and molecular similarity for reaction reconstruction.The property-optimization benchmark maximizes a property while maintaining a fixed similarity constraint.

A1 Training and evaluation procedure

Experiments extend an XLNet backbone with numerical tokenization, a numerical-embedding layer, and custom training objectives.

Training and evaluation procedure: All experiments use XLNet with the proposed tokenization scheme, a 16-dimensional numerical-embedding layer, and custom training objectives.The implementation expands the HuggingFace XLNet backbone rather than replacing it.

Tokenization

The RT workflow combines molecular text, property values, numerical encodings, and alternating objectives in one model. Evaluation uses greedy decoding for property prediction and beam search for conditional generation.

Tokenization: The workflow represents molecular strings and multiple property tags as mixed text-and-number sequences.The figure indicates that the model scales to multiple property tags.
Tokenization: Numerical values are tokenized into ordered character-level tokens, with a separator distinguishing numerical and textual tokens.The tokenization preserves each digit’s decimal order.
Tokenization: The RT combines numerical encodings with learned embeddings and relative positional encodings before alternating between prediction and generation objectives.Greedy decoding is used for property prediction, while beam search is used for conditional sequence generation.

A1.1.1 QED dataset.

On the synthetic QED dataset, training begins with vanilla permutation language modeling and then alternates specialized property-prediction and text-generation objectives.

A1.1.1 QED dataset.: Models are first trained with vanilla permutation language modeling until validation perplexity saturates, then refined with alternating objectives.The objectives alternate every 50 steps during refinement.
A1.1.1 QED dataset.: The QED experiments ablate the self-consistency loss by setting α to 0 and 1.Both SELFIES and SMILES vocabularies are evaluated.

A1.1.2 MoleculeNet dataset.

MoleculeNet regression values were normalized to [0, 1] and rounded to three decimal places, with models warm-started from QED initialization.

MoleculeNet models were warm-started using QED initialization and trained for 50k steps with batch size 4 and early stopping.

A1.1.3 Property-optimization benchmark

The appendix evaluates the Regression Transformer across molecular, protein, and reaction settings, including property prediction, conditional generation, and constrained scaffold modification. Results show strong multitask behavior alongside specific trade-offs and failure modes.

Property-optimization benchmark: RT ablations consistently outperformed JT-VAE and GCPN by a wide margin in mean pLogP improvement on constrained property optimization.The benchmark compares float-based numerical encodings and self-consistency-loss settings.
Training-scheme ablations: Alternating optimization of property and text tokens was highly effective for both regression and generation, while self-consistency improved generation but slightly reduced regression accuracy.The reported trade-off was attributed to self-evaluation using decoded sequences paired with original property values.
Protein conditional generation: More masked amino-acid residues improved conditional-generation Spearman ρ but reduced similarity to the seed sequence.
Scaffold hopping: Scaffold hopping masks only non-scaffold tokens, making the task more constrained and often limiting the attainable drug-likeness range.The setup used SMILES models because scaffolds could not be determined trivially in SELFIES.
Scaffold hopping: The RT cannot fill zero or multiple tokens at one mask location, which can prevent the best action when the scaffold already matches or exceeds the requested QED.In the example, a scaffold with QED 0.87 was paired with a 0.86 primer.
Protein property prediction: On fluorescence data, token-level classification caused collapse toward a common numerical pattern, while the RT improved mode assignment but struggled with fine-grained regression.The failure arose near a repeated dark-mode value, and Gaussian noise was added as a non-algorithmic remedy.