The Role of Model Architecture and Scale in Predicting Molecular Properties: Insights from Fine-Tuning RoBERTa, BART, and LLaMA

TL;DR

This work presents a uniform framework to compare fine-tuning of large language models for SMILES-based molecular property prediction, evaluating RoBERTa, BART, and LLaMA across 18 configurations and six DeepChem benchmarks. By pre-training multi-task regression models and then fine-tuning them under consistent conditions, the study disentangles architecture- and scale-related effects, showing that LLaMA-based models often achieve the lowest validation losses and that model size substantially influences performance. The results reveal that absolute validation loss is not a definitive predictor of downstream performance for fine-tuning, with medium-sized ChemLLaMA models frequently offering robust, cross-task results, while ChemBART excels in regression tasks. Overall, the findings provide practical guidance for selecting LLM architectures and data scales in cheminformatics, highlighting the value of large-scale LLaMA models with ample pre-training data for both regression and classification molecular-property tasks.

Abstract

This study introduces a systematic framework to compare the efficacy of Large Language Models (LLMs) for fine-tuning across various cheminformatics tasks. Employing a uniform training methodology, we assessed three well-known models-RoBERTa, BART, and LLaMA-on their ability to predict molecular properties using the Simplified Molecular Input Line Entry System (SMILES) as a universal molecular representation format. Our comparative analysis involved pre-training 18 configurations of these models, with varying parameter sizes and dataset scales, followed by fine-tuning them on six benchmarking tasks from DeepChem. We maintained consistent training environments across models to ensure reliable comparisons. This approach allowed us to assess the influence of model type, size, and training dataset size on model performance. Specifically, we found that LLaMA-based models generally offered the lowest validation loss, suggesting their superior adaptability across tasks and scales. However, we observed that absolute validation loss is not a definitive indicator of model performance - contradicts previous research - at least for fine-tuning tasks: instead, model size plays a crucial role. Through rigorous replication and validation, involving multiple training and fine-tuning cycles, our study not only delineates the strengths and limitations of each model type but also provides a robust methodology for selecting the most suitable LLM for specific cheminformatics applications. This research underscores the importance of considering model architecture and dataset characteristics in deploying AI for molecular property prediction, paving the way for more informed and effective utilization of AI in drug discovery and related fields.

Figures (3)

• Figure 1: The Architecture of Multi-Task Regression model and Fine-Tuned model.
• Figure 2: Training Loss vs. Steps for Each Model Type by Model Size with a Data Size of 20 millions. The loss trends for 10 million and 20 million datasets are similar.
• Figure 3: Validation Loss vs. Steps for Each Model Type Across All Model Sizes and Data Sizes (10, 20, and 30 million).