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

We downloaded the canonicalized SMILES dataset utilized for Molformer [14], initially processing 33 million SMILES strings. After removing duplicates, we randomly selected 30 million entries for further analysis. Using DeepChem’s RDKitDescriptors Featurizer, we calculated descriptors for each SMILES string [15, 16]. This process revealed significant skewing in some property columns due to a few extreme values (e.g., ‘Kappa3’ values exceeding 1e+36, while most are below 1e+3). We addressed this by removing outliers and normalizing all properties using z-scores. Finally, we excluded any columns where the minimum and median values were identical, resulting in a curated initial training dataset of 30 million distinct compounds (30M rows) with 105 properties (106 columns: 1 identifier + 105 properties). To set the uniform training environment, we used the exact same training and validation data for all transformer models for each data size. Note that we shuffle the order of training data for each epoch.

All tokenizers for each transformer model are based on the tokenizer from ChemBERTa [7]. The authors have replaced the special tokens with the default settings from Hugging Face. The maximum sequence length for the tokenizers is 512. Note that the ChemLLaMA tokenizer adds the ‘eos’ tokens at the end of the sequence before padding, while the default LLaMA tokenizer does not.

The authors utilized three types of transformer-based models, representing encoder-based, decoder-based, and encoder-decoder-based architectures: RoBERTa, LLaMA, and BART, respectively. Default settings were used for each model, with adjustments made only to the special token IDs and vocabulary size, as illustrated in Figure 1. The models were trained with two parameter sizes, featuring 30 million (Medium) and 13 million (Small) parameters as calculated using the .num_param() method from the Hugging Face ‘transformers’ library. During the training, BART and LLaMA models calculated loss using the ‘eos’ token, whereas RoBERTa utilized the ‘bos’ token. We employed L1Loss, corresponding to Mean Absolute Error (MAE). The MTR model families discussed in this paper—ChemBERTa, ChemLLaMA, and ChemBART, see Table 2—correlate with the architectures of RoBERTa, LLaMA, and BART, respectively.

We configured both the training and validation phases with a batch size of 64. For optimization, we employed the AdamW optimizer, using its default settings. We also used the ‘Linear Warmup Cosine Annealing Learning Rate’ scheduler from Pytorch Lightning Bolts, which gradually ramped up the learning rate to a peak of 0.0001 by the end of the first epoch (starting from epoch 0) and sustained this rate across a total of 7 epochs. Our MTR models were trained on four NVIDIA A30 GPUs using a Distributed Data Parallelism (DDP) strategy, facilitated by the Pytorch Lightning framework.

We selected six benchmarking tasks from DeepChem to evaluate the performance of our transformer models after fine-tuning. These tasks, detailed in Table 3, are representative of common benchmark challenges in cheminformatics. For regression analyses, we utilized the ‘bace_regression,’ ‘delaney,’ and ‘lipo’ datasets. For classification tasks, we selected ‘bace_classification,’ ‘hiv,’ and ‘tox21_sr_p53.’ Consistency in preprocessing was maintained by using the same tokenizers as those applied in the pre-training phase of our models.

In the fine-tuning process, we utilized the 18 pre-trained models from our MTR series. These models were initially frozen to preserve the learned features and then appended with linear layers incorporating the Gelu activation function, as depicted in Figure 1. For the loss functions, the fine-tuned models employed Mean Absolute Error (MAE, or L1 Loss) for regression tasks and Binary Cross-Entropy with Logits (BCE with Logits Loss) for classification tasks.

During the fine-tuning of our models, we allocated one GPU per model due to the relatively smaller size of the fine-tuning datasets compared to those used for pre-training. Consequently, we did not employ parallel processing techniques and instead adhered to the default training settings. We retained the same learning rate scheduler as in the pre-training phase but increased the peak learning rate from 0.0001 to 0.01 for both regression and classification tasks to better suit the scale of the datasets. All other training configurations were maintained as in the initial setup. To ensure the reliability of our results and reduce potential biases from single-model evaluations, we repeated the fine-tuning process five times. This approach enabled us to compile and analyze statistics across multiple iterations, enhancing the robustness of our findings.

During the training of the MTR models, we observed that BART and RoBERTa-based models exhibited similar learning rate behaviors across all model sizes. However, ChemBART’s learning rate decreased more rapidly than ChemBERTa’s at the outset, across all sizes, as shown in Figure 2 (data for 20 million parameters displayed, with other sizes presenting similar trends). Regarding validation loss, the smaller-sized models of ChemBART and ChemBERTa displayed comparable loss values across various data sizes and epochs. However, for the medium-sized models, ChemBART consistently showed lower validation losses in all cases, as illustrated in Figure 3. ChemLLaMA achieved the lowest validation loss values, irrespective of epoch or model size.

Additionally, we confirmed that larger MTR models tend to reach saturation more quickly than their smaller counterparts, aligning with previous findings [17]. Interestingly, as Figure 3 illustrates, ChemLLaMA models of the same size exhibited lower validation losses with smaller datasets compared to larger ones throughout the gradient steps. Nonetheless, models trained with larger datasets ultimately achieved lower validation losses as the number of gradient steps increased.

We assess the performance of various finely-tuned models across multiple tasks using the following categories: Model Type, Model Type & Model Size, Model Type & Data Size, and Model Size & Data Size. Our evaluation process is as follows:

1. 1.

   We compute the average RMSE (for regression tasks) or AUC (for classification tasks) for each group: Model Type, Model Type & Model Size, Model Type & Data Size, and Model Size & Data Size.

2. 2.

   We establish the optimal performance — the lowest average RMSE in regression tasks or the highest average AUC in classification tasks — as the benchmark for each group.

3. 3.

   For each individual model run, we measure the deviation of the RMSE or AUC from the established benchmark. These deviations are signed to reflect that a lower RMSE or higher AUC is preferable, aligning with performance standards for both regression and classification tasks.

4. 4.

   We then aggregate these deviations, grouped by the aforementioned categories, into an Error Sum (ES) for each task.

5. 5.

   We calculate the Total Error Sum (TES) by summing all the ES values, providing a comprehensive performance metric across tasks. Additionally, we determine the Standard Deviation (STD) based on the deviations from the benchmark identified in Step 3.

This structured approach enables a detailed model performance analysis relative to established standards across different configurations.

Table 5 displays the average best metrics (RMSE/AUC) for our finely-tuned models, grouped by model type. For example, the average of the best metric for the BACE benchmark task is calculated from 30 RMSE values, derived from combinations of 2 model sizes, 3 data sizes, and 5 runs. This yields three comparative averages—one for each model type. Among them, ChemBART records the lowest average RMSE at 0.896. Average metrics for other tasks are calculated similarly. Table 5 splits the results by model size.

When focusing solely on the transformer’s Model Type and disregarding the size, ChemBART outperforms the other models in regression tasks but ranks lowest in classification tasks. Whereas, ChemLLaMA exhibits the worst performance (highest TES) in regression tasks across both sizes, as shown in Table 7. However, ChemLLaMA medium has the best overall performance (lowest TES values) among all Model Types for medium-sized models in both regression and classification tasks, as highlighted in Table 7. These tables categorize the best metric sets from each dataset, highlighting differences across model configurations.

The baseline values for each dataset are taken from Table 5. Additionally, the best metrics are organized by MTMS from the original best metric set for each dataset.

As previously mentioned, ChemLLaMA with a ’Medium’ model size demonstrates the best performance across all regression and classification tasks, even though ChemLLaMA ‘Small’ has the worst performance in regression tasks among all MTMS combinations. The ’Medium’ size ChemBART also performs impressively in regression tasks.

Furthermore, our analysis reveals that the performance of each Model Type generally scales with Model Size across all task types, with the notable exception of ChemBERTa in regression tasks.

When the average metrics and TES are calculated from the best model and grouped by Model Type and Data Size (MTDS), as presented in Tables 9 and 11, ChemBART, utilizing 10M of MTR data, excels in regression tasks compared to other groups. For classification tasks, ChemLLaMA with 30M of data achieves the lowest TES value. ChemBERTa with 30M of data also displays strong performance in classification, alongside other groups with similarly low TES values. Additionally, performance for ChemBERTa in regression tasks and for ChemLLaMA in classification tasks shows a clear proportional relationship to data size.

Considering Model Size and Model Type, as detailed in Tables 9 and 11, a Medium Model Size with 10M Data Size yields the best (lowest) TES value in regression tasks. Generally, larger Model Sizes outperform smaller ones across all Data Sizes. In classification tasks, Small Model Sizes with 30M Data Size achieve the lowest TES. However, with the exception of the 30M data scenarios, Medium Model Sizes tend to exhibit lower errors than smaller sizes. Notably, in regression tasks, Small Model Sizes demonstrate a non-proportional relationship between TES and Data Size. Conversely, in classification tasks, there is a consistent proportional relationship between performance and Data Size, regardless of Model Size.