Derify/ModChemBERT-MLM

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:

pip install -U "transformers>=4.56.1,<5.0.0"

Load Model

from transformers import AutoModelForMaskedLM, AutoTokenizer

model_id = "Derify/ModChemBERT-MLM"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    dtype="float16",
    device_map="auto",
)

Fill-Mask Pipeline

from transformers import pipeline

fill = pipeline("fill-mask", model=model, tokenizer=tokenizer)
print(fill("c1ccccc1[MASK]"))

Architecture

• —Backbone: ModernBERT
• —Hidden size: 768
• —Intermediate size: 1152
• —Encoder Layers: 22
• —Attention heads: 12
• —Max sequence length: 256 tokens (MLM primarily trained with 128-token sequences)
• —Tokenizer: BPE tokenizer using MolFormer's vocab (2362 tokens)

Pooling (Classifier / Regressor Head)

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

<div align="center"> <img src="https://cdn-uploads.huggingface.co/production/uploads/656892962693fa22e18b5331/bxNbpgMkU8m60ypyEJoWQ.png" alt="ModChemBERT Training Pipeline" width="650"/> </div>

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.

Datasets

• —Pretraining: Derify/augmented_canonical_druglike_QED_Pfizer_15M (canonical_smiles column)
• —Domain Adaptive Pretraining (DAPT) & Task Adaptive Fine-tuning (TAFT): ADME (6 tasks) + AstraZeneca (4 tasks) datasets that are split using DA4MT's [3] Bemis-Murcko scaffold splitter (see domain-adaptation-molecular-transformers)
• —Benchmarking:
• —ChemBERTa-3 [7]
• —classification: BACE, BBBP, TOX21, HIV, SIDER, CLINTOX
• —regression: ESOL, FREESOLV, LIPO, BACE, CLEARANCE
• —Mswahili, et al. [8] proposed additional datasets for benchmarking chemical language models:
• —classification: Antimalarial [9], Cocrystal [10], COVID19 [11]
• —DAPT/TAFT stage regression datasets:
• —ADME [12]: admemicrosomstabh, admemicrosomstabr, admepermeability, admeppbh, admeppbr, admesolubility
• —AstraZeneca: astrazenecaCL, astrazenecaLogD74, astrazenecaPPB, astrazenecaSolubility

Benchmarking

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

<details><summary>Click to expand</summary>

ChemBERTa-3 Classification Datasets (ROC AUC - Higher is better)

ChemBERTa-3 Regression Datasets (RMSE - Lower is better)

Mswahili, et al. [8] Proposed 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.<br/> \* Published results from the ChemBERTa-3 [7] paper for optimized chemical language models using DeepChem scaffold splits.<br/> † AVG column shows the mean score across classification tasks.<br/> ‡ AVG column shows the mean scores across regression tasks without and with the clearance score.

</details>

Optimized ModChemBERT Hyperparameters

<details><summary>Click to expand</summary>

TAFT Datasets

Optimal parameters (per dataset) for the MLM + DAPT + TAFT OPT merged model:

Benchmarking Datasets

Optimal parameters (per dataset) for the MLM + DAPT + TAFT OPT merged model:

</details>

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

1. 1.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
2. 2.Behrendt, Maike, Stefan Sylvius Wagner, and Stefan Harmeling. "MaxPoolBERT: Enhancing BERT Classification via Layer-and Token-Wise Aggregation." arXiv preprint arXiv:2505.15696 (2025).
3. 3.Sultan, Afnan, et al. "Transformers for molecular property prediction: Domain adaptation efficiently improves performance." arXiv preprint arXiv:2503.03360 (2025).
4. 4.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).
5. 5.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).
6. 6.Grattafiori, Aaron, et al. "The llama 3 herd of models." arXiv preprint arXiv:2407.21783 (2024).
7. 7.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.
8. 8.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
9. 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
10. 10.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
11. 11.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
12. 12.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