TransPolymer: a Transformer-based language model for polymer property predictions

Accurate and efficient prediction of polymer properties is essential for designing polymers for various applications, including polymer electrolytes, organic optoelectronics, and energy storage. Traditional methods rely on expensive and time-consuming experiments or simulations. Machine learning offers an alternative, but effective representation of polymers in a continuous vector space is crucial. Fingerprints have been used, but deep neural networks (DNNs) offer the advantage of directly learning expressive representations from data. Convolutional neural networks (CNNs) have been applied, but they struggle to consider molecular structure and interactions. Graph neural networks (GNNs) have shown promise, but they require explicit structural and conformational information, which can be computationally expensive to obtain. The varying degree of polymerization in polymers adds further complexity to their graph representation. Language models offer another approach, treating polymers as character sequences. This approach draws parallels between chemical sequences and natural language, suggesting that language models from computational linguistics could be adapted for polymer science. Recurrent neural networks (RNNs), such as LSTMs, have been explored, but they are limited by their reliance on previous hidden states and a tendency to lose information as the sequence length increases. Transformer models, known for their superior performance in natural language processing (NLP) tasks, offer a potential solution. Their self-attention mechanism allows them to capture relationships between tokens without relying on past hidden states. Several Transformer-based models have been successfully applied to property predictions of small organic molecules, following a pretrain-finetune pipeline. However, these models have not yet been extensively applied to polymers, which pose unique challenges due to their complex hierarchical structures, varying degrees of polymerization, and the often limited availability of well-labeled data. This paper addresses these challenges by introducing TransPolymer, a Transformer-based language model specifically designed for polymer property prediction.

Several studies have explored machine learning approaches for polymer property prediction. Rahman et al. utilized CNNs to predict mechanical properties of polymer-carbon nanotube surfaces, but these models faced challenges in capturing molecular structure and atomic interactions. GNNs, capable of learning representations from graphs, have yielded better results but still encounter difficulties with the varying degree of polymerization and the cost of obtaining necessary structural information. RNN-based models, treating polymers as character sequences, have also been investigated, with Simine et al. predicting spectra of conjugated polymers using LSTMs. Webb et al. applied LSTMs to predict polymer properties using coarse-grained polymer genomes, and Patel et al. extended this to copolymer systems. However, RNNs can struggle to capture long-range dependencies, which limits their effectiveness. The Transformer architecture, with its self-attention mechanism, offers advantages over RNNs in handling long-range dependencies. While Transformer models have been applied successfully to small molecule property prediction (e.g., SMILES-BERT, ChemBERTa), their application to polymers remains relatively unexplored.

TransPolymer utilizes a novel chemical-aware tokenization method to represent polymers as sequences of tokens. The repeating units of polymers are embedded using SMILES, and additional descriptors (degree of polymerization, polydispersity, chain conformation) are included. Copolymers are modeled by combining the SMILES of each repeating unit with their ratios and arrangements. Mixtures of polymers are represented by concatenating the sequences of each component and their descriptors. Each token represents an element, the value of a polymer descriptor, or a special separator. The Transformer encoder in TransPolymer is based on a RoBERTa architecture, composed of stacked self-attention and point-wise fully connected layers. The self-attention mechanism allows the model to capture relationships between tokens at different positions in a sequence. The model employs 6 hidden layers with 12 attention heads each. Hyperparameters are chosen based on ROBERTa settings and tuned based on model performance. To improve representation learning, the Transformer encoder is pretrained using Masked Language Modeling (MLM). In MLM, a percentage of tokens in the sequences are randomly masked, and the model is trained to predict these masked tokens based on the context. Approximately 5 million augmented unlabeled polymers from the PI1M database are used for pretraining. The pretrained model is then finetuned on ten datasets of polymers with various properties, using a multi-layer perceptron (MLP) regressor head. Data augmentation, involving the generation of non-canonical SMILES, is applied to enhance learning. Root mean square error (RMSE) and R² are used to evaluate model performance.

TransPolymer achieves state-of-the-art (SOTA) results on all ten benchmarks, surpassing baseline models such as Random Forest with ECFP fingerprints, LSTM, and unpretrained TransPolymer. The superior performance of TransPolymer is attributed to several factors: * **Masked Language Modeling (MLM) Pretraining:** Pretraining on a large unlabeled dataset significantly improves performance, enabling the model to learn generalizable features from polymer sequences. * **Finetuning:** Finetuning both the Transformer encoders and the regressor head is crucial for optimal performance, demonstrating that the model learns both generalizable and task-specific information. * **Data Augmentation:** Data augmentation, through non-canonical SMILES generation, enhances the model's ability to learn from the limited available data. * **Self-Attention Mechanism:** The visualization of attention scores reveals that TransPolymer successfully encodes chemical information about internal interactions within polymers and the factors influencing their properties. TransPolymer exhibits significant improvements over baseline models, particularly on datasets with noisy data or a limited number of data points. Ablation studies confirm the positive effects of each aspect of the model's design (pretraining size, finetuning strategy, data augmentation). The t-SNE visualization demonstrates that the chemical space covered by the pretrained model encompasses the downstream datasets, emphasizing the effectiveness of pretraining in representation learning. Furthermore, analysis of attention scores shows that TransPolymer effectively identifies key features within polymer sequences that influence properties.

TransPolymer's superior performance on diverse polymer property prediction tasks demonstrates the effectiveness of using Transformer-based models in polymer science. The findings address the challenges associated with limited data and complex polymer structures, showing how a carefully designed tokenization strategy combined with pretraining and data augmentation can significantly improve predictive accuracy. The model's ability to learn generalizable features from unlabeled data and adapt to specific tasks through finetuning highlights its potential as a valuable tool for accelerating materials discovery. The attention mechanism provides insights into the model's internal workings, revealing its ability to capture crucial chemical interactions and relationships.

TransPolymer represents a significant advance in polymer property prediction, showcasing the power of Transformer architectures and self-supervised learning. Its superior performance and generalizability suggest its potential for various applications in polymer design and discovery. Future work could explore the application of TransPolymer to larger and more diverse datasets, integrating it with active learning strategies for efficient exploration of chemical space, and further investigating the interpretability of the attention mechanism.

While TransPolymer demonstrates significant improvements, certain limitations exist. The performance is dependent on the quality and quantity of data used for both pretraining and finetuning. The interpretability of the attention mechanism, although providing some insights, remains a challenge, requiring further investigation. The model's computational cost could be a constraint for extremely large datasets or high-dimensional property predictions.