ModChemBERT: ModernBERT as a Chemical Language Model
ModChemBERT is a ModernBERT-based chemical language model (CLM), trained on SMILES strings for masked language modeling (MLM) and downstream molecular property prediction (classification & regression).
Usage
Install the transformers library starting from v4.56.1:
Kallergis et al. [1] demonstrated that the CLM embedding method prior to the prediction head was the strongest contributor to downstream performance among evaluated hyperparameters.
Behrendt et al. [2] noted that the last few layers contain task-specific information and that pooling methods leveraging information from multiple layers can enhance model performance. Their results further demonstrated that the max_seq_mha pooling method was particularly effective in low-data regimes, which is often the case for molecular property prediction tasks.
Multiple pooling strategies are supported by ModChemBERT to explore their impact on downstream performance:
cls: Last layer [CLS]
mean: Mean over last hidden layer
max_cls: Max over last k layers of [CLS]
cls_mha: MHA with [CLS] as query
max_seq_mha: MHA with max pooled sequence as KV and max pooled [CLS] as query
sum_mean: Sum over all layers then mean tokens
sum_sum: Sum over all layers then sum tokens
mean_mean: Mean over all layers then mean tokens
mean_sum: Mean over all layers then sum tokens
max_seq_mean: Max over last k layers then mean tokens
Note: ModChemBERT’s max_seq_mha differs from MaxPoolBERT [2]. MaxPoolBERT uses PyTorch nn.MultiheadAttention, whereas ModChemBERT's ModChemBertPoolingAttention adapts ModernBERT’s ModernBertAttention.
On ChemBERTa-3 benchmarks this variant produced stronger validation metrics and avoided the training instabilities (sporadic zero / NaN losses and gradient norms) seen with nn.MultiheadAttention. Training instability with ModernBERT has been reported in the past (discussion 1 and discussion 2).
Training Pipeline
ModChemBERT Training Pipeline
Rationale for MTR Stage
Following Sultan et al. [3], multi-task regression (physicochemical properties) biases the latent space toward ADME-related representations prior to narrow TAFT specialization. Sultan et al. observed that MLM + DAPT (MTR) outperforms MLM-only, MTR-only, and MTR + DAPT (MTR).
Checkpoint Averaging Motivation
Inspired by ModernBERT [4], JaColBERTv2.5 [5], and Llama 3.1 [6], where results show that model merging can enhance generalization or performance while mitigating overfitting to any single fine-tune or annealing checkpoint.
Benchmarks were conducted using the ChemBERTa-3 framework. DeepChem scaffold splits were utilized for all datasets, with the exception of the Antimalarial dataset, which employed a random split. Each task was trained for 100 epochs, with results averaged across 3 random seeds.
The complete hyperparameter configurations for these benchmarks are available here: ChemBERTa3 configs
Evaluation Methodology
Classification Metric: ROC AUC
Regression Metric: RMSE
Aggregation: Mean ± standard deviation of the triplicate results.
Input Constraints: SMILES truncated / filtered to ≤200 tokens, following ChemBERTa-3's recommendation.
Results
Click to expand
ChemBERTa-3 Classification Datasets (ROC AUC - Higher is better)
ADME/AstraZeneca Regression Datasets (RMSE - Lower is better)
Hyperparameter optimization for the TAFT stage appears to induce overfitting, as the MLM + DAPT + TAFT OPT model shows slightly degraded performance on the ADME/AstraZeneca datasets compared to the MLM + DAPT + TAFT model.
The MLM + DAPT + TAFT model, a merge of unoptimized TAFT checkpoints trained with max_seq_mean pooling, achieved the best overall performance across the ADME/AstraZeneca datasets.
Bold indicates the best result in the column; italic indicates the best result among ModChemBERT checkpoints.
* Published results from the ChemBERTa-3 [7] paper for optimized chemical language models using DeepChem scaffold splits.
† AVG column shows the mean score across classification tasks.
‡ AVG column shows the mean scores across regression tasks without and with the clearance score.
Optimized ModChemBERT Hyperparameters
Click to expand
TAFT Datasets
Optimal parameters (per dataset) for the MLM + DAPT + TAFT OPT merged model:
Dataset
Learning Rate
Batch Size
Warmup Ratio
Classifier Pooling
Last k Layers
adme_microsom_stab_h
3e-5
8
0.0
max_seq_mean
5
adme_microsom_stab_r
3e-5
16
0.2
max_cls
3
adme_permeability
3e-5
8
0.0
max_cls
3
adme_ppb_h
1e-5
32
0.1
max_seq_mean
5
adme_ppb_r
1e-5
32
0.0
sum_mean
N/A
adme_solubility
3e-5
32
0.0
sum_mean
N/A
astrazeneca_CL
3e-5
8
0.1
max_seq_mha
3
astrazeneca_LogD74
1e-5
8
0.0
max_seq_mean
5
astrazeneca_PPB
1e-5
32
0.0
max_cls
3
astrazeneca_Solubility
1e-5
32
0.0
max_seq_mean
5
Benchmarking Datasets
Optimal parameters (per dataset) for the MLM + DAPT + TAFT OPT merged model:
Dataset
Batch Size
Classifier Pooling
Last k Layers
Pooling Attention Dropout
Classifier Dropout
Embedding Dropout
bace_classification
32
max_seq_mha
3
0.0
0.0
0.0
bbbp
64
max_cls
3
0.1
0.0
0.0
clintox
32
max_seq_mha
5
0.1
0.0
0.0
hiv
32
max_seq_mha
3
0.0
0.0
0.0
sider
32
mean
N/A
0.1
0.0
0.1
tox21
32
max_seq_mha
5
0.1
0.0
0.0
base_regression
32
max_seq_mha
5
0.1
0.0
0.0
clearance
32
max_seq_mha
5
0.1
0.0
0.0
esol
64
sum_mean
N/A
0.1
0.0
0.1
freesolv
32
max_seq_mha
5
0.1
0.0
0.0
lipo
32
max_seq_mha
3
0.1
0.1
0.1
antimalarial
16
max_seq_mha
3
0.1
0.1
0.1
cocrystal
16
max_cls
3
0.1
0.0
0.1
covid19
16
sum_mean
N/A
0.1
0.0
0.1
Intended Use
Primary: Research and development for molecular property prediction, experimentation with pooling strategies, and as a foundational model for downstream applications.
Appropriate for: Binary / multi-class classification (e.g., toxicity, activity) and single-task or multi-task regression (e.g., solubility, clearance) after fine-tuning.
Not intended for generating novel molecules.
Limitations
Out-of-domain performance may degrade for: very long (>128 token) SMILES, inorganic / organometallic compounds, polymers, or charged / enumerated tautomers are not well represented in training.
No guarantee of synthesizability, safety, or biological efficacy.
Ethical Considerations & Responsible Use
Potential biases arise from training corpora skewed to drug-like space.
Do not deploy in clinical or regulatory settings without rigorous, domain-specific validation.
Hardware
Training and experiments were performed on 2 NVIDIA RTX 3090 GPUs.
Citation
If you use ModChemBERT in your research, please cite the checkpoint and the following:
@software{cortes-2025-modchembert,
author = {Emmanuel Cortes},
title = {ModChemBERT: ModernBERT as a Chemical Language Model},
year = {2025},
publisher = {GitHub},
howpublished = {GitHub repository},
url = {https://github.com/emapco/ModChemBERT}
}
References
Kallergis, G., Asgari, E., Empting, M. et al. Domain adaptable language modeling of chemical compounds identifies potent pathoblockers for Pseudomonas aeruginosa. Commun Chem 8, 114 (2025). https://doi.org/10.1038/s42004-025-01484-4
Behrendt, Maike, Stefan Sylvius Wagner, and Stefan Harmeling. "MaxPoolBERT: Enhancing BERT Classification via Layer-and Token-Wise Aggregation." arXiv preprint arXiv:2505.15696 (2025).
Sultan, Afnan, et al. "Transformers for molecular property prediction: Domain adaptation efficiently improves performance." arXiv preprint arXiv:2503.03360 (2025).
Warner, Benjamin, et al. "Smarter, better, faster, longer: A modern bidirectional encoder for fast, memory efficient, and long context finetuning and inference." arXiv preprint arXiv:2412.13663 (2024).
Clavié, Benjamin. "JaColBERTv2.5: Optimising Multi-Vector Retrievers to Create State-of-the-Art Japanese Retrievers with Constrained Resources." arXiv preprint arXiv:2407.20750 (2024).
Grattafiori, Aaron, et al. "The llama 3 herd of models." arXiv preprint arXiv:2407.21783 (2024).
Singh R, Barsainyan AA, Irfan R, Amorin CJ, He S, Davis T, et al. ChemBERTa-3: An Open Source Training Framework for Chemical Foundation Models. ChemRxiv. 2025; doi:10.26434/chemrxiv-2025-4glrl-v2 This content is a preprint and has not been peer-reviewed.
Mswahili, M.E., Hwang, J., Rajapakse, J.C. et al. Positional embeddings and zero-shot learning using BERT for molecular-property prediction. J Cheminform 17, 17 (2025). https://doi.org/10.1186/s13321-025-00959-9
Mswahili, M.E.; Ndomba, G.E.; Jo, K.; Jeong, Y.-S. Graph Neural Network and BERT Model for Antimalarial Drug Predictions Using Plasmodium Potential Targets. Applied Sciences, 2024, 14(4), 1472. https://doi.org/10.3390/app14041472
Mswahili, M.E.; Lee, M.-J.; Martin, G.L.; Kim, J.; Kim, P.; Choi, G.J.; Jeong, Y.-S. Cocrystal Prediction Using Machine Learning Models and Descriptors. Applied Sciences, 2021, 11, 1323. https://doi.org/10.3390/app11031323
Harigua-Souiai, E.; Heinhane, M.M.; Abdelkrim, Y.Z.; Souiai, O.; Abdeljaoued-Tej, I.; Guizani, I. Deep Learning Algorithms Achieved Satisfactory Predictions When Trained on a Novel Collection of Anticoronavirus Molecules. Frontiers in Genetics, 2021, 12:744170. https://doi.org/10.3389/fgene.2021.744170
Cheng Fang, Ye Wang, Richard Grater, Sudarshan Kapadnis, Cheryl Black, Patrick Trapa, and Simone Sciabola. "Prospective Validation of Machine Learning Algorithms for Absorption, Distribution, Metabolism, and Excretion Prediction: An Industrial Perspective" Journal of Chemical Information and Modeling 2023 63 (11), 3263-3274 https://doi.org/10.1021/acs.jcim.3c00160
