Source-linked AI summary
BERTology of Molecular Property Prediction
Mohammad Mostafanejad, Paul Saxe, T. Daniel Crawford
TL;DR
Reported CLM performance on MPP benchmarks has been inconsistent, while scaling laws for encoder-only masked language models remain insufficiently established. The study analyzes hundreds of controlled experiments across dataset size, model size, tokenization, architecture, and standardization, finding that BERT performance is superior or similar to classical machine-learning models for all ADME endpoints and that larger models are more sample efficient and noise resilient.
Problem
MPP studies report inconsistent and contradictory CLM performance, and well-established scaling laws for encoder-only MLMs are absent.
Method
The study conducts hundreds of controlled experiments investigating how dataset size, model size, tokenization, architecture, and standardization affect CLM pre-training and fine-tuning.
Results
BERT is superior or similar to classical ML models for all ADME endpoints, while larger models are more sample efficient and noise resilient.
Takeaways & Limitations
The numerical trends favor the potential existence of a scaling law for encoder-only CLMs with an MLM objective.
Takeaways & Limitations
The reported trends are based on qualitative observations from small statistical samples of experiments, and confidence intervals can be overestimated.
Abstract
from arXiv · showhide
Chemical language models (CLMs) have emerged as promising competitors to popular classical machine learning models for molecular property prediction (MPP) tasks. However, an increasing number of studies have reported inconsistent and contradictory results for the performance of CLMs across various MPP benchmark tasks. In this study, we conduct and analyze hundreds of meticulously controlled experiments to systematically investigate the effects of various factors, such as dataset size, model size, and standardization, on the pre-training and fine-tuning performance of CLMs for MPP. In the absence of well-established scaling laws for encoder-only masked language models, our aim is to provide comprehensive numerical evidence and a deeper understanding of the underlying mechanisms affecting the performance of CLMs for MPP tasks, some of which appear to be entirely overlooked in the literature.
Results
Across controlled BERT experiments, pre-training performance depends on model and dataset size, standardization noise, and tokenization. Larger models generally perform better, use fewer samples efficiently, and show greater resilience to standardization noise.
- Model Initialization and Data Sampling Randomness: Randomness from data sampling and model initialization changes performance by approximately 1%, less than the variation caused by model size.The comparison uses independent runs with different random seeds.
- Model Initialization and Data Sampling Randomness: Increasing model size from Tiny-BERT to Base-BERT reduces validation pseudo-perplexity by more than 28%, from 1.5993 ± 0.0085 to 1.1455 ± 0.0013.The same model-size increase lowers average training and validation loss and increases validation accuracy and weighted-F1.
- Standardization Effects on Pre-training: Larger BERT models become more resilient to standardization noise, reducing its impact on V-wF1 and V-PPPL.Mixed standardization protocols can severely hamper masked-token prediction and degrade related validation metrics.
- The Effect of Tokenization on Pre-training: Using WordPiece or BPE tokenization, V-PPPL decreases as model size increases from Tiny-BERT to Base-BERT.With WordPiece (BPE), V-PPPL changes from 1.5978 ± 0.0138 (1.5435 ± 0.0332) to 1.1450 ± 0.0052 (1.1233 ± 0.0125).
- The Effect of Dataset and Model Sizes on Pre-training: Pre-training metrics improve as dataset size increases, while performance gains begin to plateau beyond a critical dataset-size threshold.Larger models also tend to be more sample efficient and achieve better performance with smaller dataset bins.
Model and Dataset Size Effects on Fine-tuning
Across the studied ADME tasks, larger pre-training datasets and models generally improved BERT fine-tuning performance, while standardization and benchmark quality materially affected interpretation. Base-BERT was superior or similar to classical ML models across all six endpoints, although some comparisons and endpoint evaluations have important limitations.
- Model size: Performance improved as model size increased from Tiny-BERT to Base-BERT for each pre-training dataset size.The reported pattern held across the examined downstream settings.
- Limitations: Solubility cross-validation results are difficult to assess fairly because of large standard deviations (≈0.68) and a strongly skewed experimental-data distribution.This limitation specifically concerns the solubility endpoint’s cross-validation results.
- Dataset size: Testing Pearson R and R2 increased, while MAE and RMSE decreased, as pre-training dataset size increased across HLM, hPPB, and solubility.These trends were reported for all BERT variants and were consistent with the corresponding figure-based testing analyses.
- Experimental factors: Random seed and tokenization choices had minor pre-training effects compared with model size and dataset size.The comparison covered model initialization and data sampling seeds, plus WordPiece versus BPE tokenization.
- Standardization: Standardization noise can severely degrade BERT pre-training performance, while larger models tend to be more resilient under some noise distributions.Resilience was especially reported when noise followed the “path of least destruction” and its distribution remained similar across training and validation splits.
Statistical Analysis of the Results
The analysis uses 95% confidence intervals to represent uncertainty, with critical values chosen from the Student’s t-distribution for each sample size. Because experiments often use only N = 3 or 5 runs, small-sample error bars require careful interpretation.
- 95% confidence intervals are reported as the sample mean plus or minus a Student’s t critical value multiplied by the standard error.The interval uses the sample mean, standard deviation, sample size, and the appropriate t-distribution critical value.
- The Student’s t critical value is selected at the 95% confidence level with N − 1 degrees of freedom.
- N = 3 or 5 is common because pre-training language models is computationally expensive, making the reported sample sizes small.
- For very small samples, the critical value can greatly exceed 1.96, substantially overestimating error bars and encouraging misinterpretation.The analysis therefore uses NIST’s Student’s t critical-value table rather than a fixed large-sample value.
- The study does not provide Cohen’s d or Hedges’ g for effect-size calculation and comparison across results.
Performance and Error Metrics
The study evaluates masked-language-model pre-training with pseudo-perplexity, accuracy, and weighted-F1, while fine-tuning BERT on standardized molecular data for six ADME regression endpoints. MLM reconstructs randomly masked tokens using their unmasked context.
- MLM trains BERT to reconstruct randomly masked input tokens from the surrounding unmasked sequence.In this setup, approximately 15% of tokens are masked, and only masked positions contribute directly to the loss.
- Pseudo-perplexity evaluates pre-training through the exponentiated negative pseudo-log-likelihood over masked tokens.
- Weighted-F1 accounts for token-frequency imbalance by weighting each token’s F1 score by its masked-position support.
- Accuracy measures whether the model selects the original token as its most probable vocabulary prediction.Vocabulary entries are treated as classes, with the unmasked reference token as the positive class.
- Fine-tuning uses the Biogen ADME dataset to predict six in vitro endpoints, including HLM, RLM, hPPB, rPPB, MDR1-MDCK ER, and solubility.The endpoint measurements include experimentally reported ADME properties such as intrinsic clearance, efflux ratio, protein binding, and solubility.
Standardization
The workflow compares PubChem and ChEMBL standardization pipelines and trains WordPiece and BPE tokenizers under matched vocabulary and sequence-length constraints. BERT is pre-trained with dynamic masked language modeling on PubChem SMILES, then fine-tuned and evaluated on ADME regression tasks.
- Two preprocessing pipelines are compared: PubChem’s OEChem-based workflow and ChEMBL’s RDKit-based workflow.
- WordPiece and BPE tokenizers are trained from scratch on standardized canonical isomeric SMILES and then compared experimentally.
- Both tokenizers use a vocabulary of 30,522 tokens, a maximum sequence length of 512, and a minimum token frequency of 2.
- BERT is pre-trained on PubChem standardized canonical isomeric SMILES with 15% dynamic masking and no next-sentence-prediction objective.Masking changes on the fly between batches, exposing the model to different masked tokens across epochs.
Extended Data:
The extended data documents experimental variability and examines how standardization noise and pre-training dataset size relate to BERT’s masked-language-model and downstream fine-tuning metrics. Reported figures use repeated runs and cross-validation procedures to organize these comparisons.
- Table 1 reports variations in masked-language-model performance under different model-weight initialization and experimental conditions.
- Figure 1 plots standardization noise against BERT’s pre-training validation loss and accuracy for masked language modeling.Panels a, c, and e show loss; panels b, d, and f show accuracy.
- Figure 1 averages each metric over three independent runs with different data-sampling and model-initialization seeds.
- Figures 2–4 plot Pearson R, R2, RMSE, and MAE against pre-training dataset-size bin index for HLM, hPPB, and solubility.Each fold uses a hyperparameter search over 50 models, selecting the best model by validation R2.
Supporting Information:
The paper is authored by Mohammad Mostafanejad, Paul Saxe, and T. Daniel Crawford, with affiliations at Virginia Tech and the Molecular Sciences Software Institute.
- Mohammad Mostafanejad, Paul Saxe, and T. Daniel Crawford are listed as the authors.
- The authors are affiliated with the Department of Chemistry at Virginia Tech.
- The authors are also affiliated with the Molecular Sciences Software Institute.
1 Computational Details
The experiments use distributed multi-GPU pre-training and single-GPU fine-tuning across heterogeneous computing nodes. Effective global batch size is controlled across configurations using node, GPU, micro-batch, and accumulation factors.
- Computing workflow: Pre-training uses distributed data parallelism across multiple GPU nodes, while fine-tuning and hyperparameter searches use one GPU per experiment.The computational workflow moves from distributed pre-training to single-GPU downstream experiments.
- Model-scale experiments: Tiny-BERT, Small-BERT, and Base-BERT models are pre-trained on different GPU configurations across Falcon compute nodes.The listed configurations include 16 A30 GPUs, 8 L40S GPUs, and 8 H200 or NVIDIA GPUs.
- Batch-size control: Effective global batch sizes of 1024 for pre-training and 512 for evaluation are adopted for all models.
- Batch-size control: Gradient accumulation is fixed at one step while GPU counts and memory capacities vary across distributed experiments.The protocol keeps the effective global batch size constant despite different hardware configurations.
- Batch-size control: The effective global batch size is computed as Beff = Nnodes × NGPU × BGPU × NGA.Nnodes is the number of nodes, NGPU is the number of GPUs per node, BGPU is the per-GPU batch size, and NGA is the number of accumulation steps.
2 The Effect of Standardization on Pre-training
This section describes molecular sanitization and database-specific standardization protocols, emphasizing that representation choices can alter token distributions and downstream CLM performance. ChEMBL and PubChem differ in handling parents, identifiers, tautomerism, and SMILES representations.
- Molecular sanitization: Invalid PubChem SMILES that trigger RDKit’s maximum BFS search error are excluded before standardization.The error is attributed to a highly symmetric fused ring.
- Molecular sanitization: RDKit sanitization processes molecules through ordered operations that clear properties, clean valences, update caches, assign aromaticity, and set related chemical features.Additional operations include ring symmetrization, Kekulization, radical assignment, conjugation, hybridization, chirality cleanup, and hydrogen adjustment.
- Representation choices: Tautomerism and aromaticity choices can change molecular properties, including structural similarity, biological activity, and topological or physicochemical features.The preferred tautomer may affect aromaticity depending on the aromaticity model.
- Database protocols: The ChEMBL pipeline combines validity checking, rule-based formatting, and parent-structure extraction that removes defined salts, solvents, and isotopic information.Its GetParent component excepts salts in organometallic compounds.
- Database protocols: PubChem uses canonical isomeric SMILES for unique structures, whereas ChEMBL does not impose canonical tautomeric forms.ChEMBL instead relies on medicinal chemists to select forms appropriate to target applications and notes tautomer interconversion under experimental conditions.
- Implications for CLMs: Different standardization pipelines can produce distinct vocabularies and token-frequency distributions, affecting tokenization, embeddings, and pre-training or fine-tuning outcomes.
3 The Effect of Dataset and Model Sizes on Pre-training
The study evaluates dataset-size and model-size effects on BERT pre-training by constructing six dataset bins and reporting their performance variations in tabulated results.
- Experimental design: Six dataset bins with different training-set sizes are created to study dataset-size effects on BERT pre-training.
- Experimental design: The first dataset bin contains 2,979,620 samples, and a correction factor expands the sixth bin to cover 80% of PubChem.The correction factor is b = 4 for k = 5 and zero otherwise.
- Reported measurements: Tables 1 and 2 report variations in BERT pre-training performance across dataset and model sizes.
- Reported measurements: The table notes identify the dataset-bin index, divergent or early-stopped models, excluded models, and whether higher or lower metric values are better.
4 Ablation Studies and Hyperparameter Tuning
The study evaluates BERT variants and classical machine-learning models on six Biogen ADME endpoints using standardized SMILES, fixed data splits, and cross-validated hyperparameter optimization. Performance is reported with multiple regression metrics for both cross-validation and held-out testing.
- The Biogen ADME dataset covers six in vitro endpoints, including HLM, RLM, hPPB, rPPB, MDR1-MDCK ER, and solubility.
- SMILES are standardized with PubChem’s protocol and tokenized using WordPiece before an 80:20% training-testing split within each property.
- The training data use 3-fold cross-validation for hyperparameter optimization, with grid search for classical models and Bayesian search for BERT models.
- Performance is assessed using R2, MSE, MAE, RMSE, and Pearson R, with cross-validation statistics reported as mean ± standard deviation.
- Tables 2–7 compare Tiny-, Small-, and Base-BERT with classical models across the six ADME labels, reporting cross-validation values and parenthesized test-set results.
- The supplementary tables define classical-model search spaces and optimized parameters, alongside Bayesian-optimized parameters for BERT variants.
5 Budget and Cost Estimations
The paper estimates computational costs for pre-training and evaluation, contrasting expensive BERT experimentation with much faster classical-model hyperparameter searches. It also reports concrete cloud-cost and runtime estimates for different BERT sizes.
- Small-BERT required about $2,700 and 5.3 days on eight NVIDIA L40S GPUs, while Tiny-BERT required about 2.3 days on 16 NVIDIA A30 GPUs.
- The estimated AWS cost for Tiny-BERT was around $2,300 using four g6e.12xlarge instances.
- A 3-fold cross-validation hyperparameter search using SVM can be completed in a few minutes on a personal multicore CPU computer.
- Pre-training 160 BERT models and subsequent fine-tuning took about 18 months on a shared university supercomputing cluster.
- The authors released code, scripts, data, and model artifacts to facilitate transparency, reproducibility, and further research.
6 BERT Architectural Details
The paper provides architectural details for the Tiny-, Small-, and Base-BERT model variants used in the experiments.
- Table 16 lists the architectural details of the Tiny-, Small-, and Base-BERT variants.