metadata

datasets:
  - COCONUTDB
  - ChemBL34
language:
  - code
library_name: transformers
metrics:
  - perplexity
  - accuracy
pipeline_tag: fill-mask
tags:
  - fill-mask
  - chemistry
  - selfies
widget:
  - text: >-
      [C] [C] [=Branch1] [C] [MASK] [O] [C] [C] [N+1] [Branch1] [C] [C]
      [Branch1] [C] [C] [C]
    example_title: '[=O]'
  - text: >-
      [O-1] [P] [=Branch1] [C] [=O] [Branch1] [C] [MASK] [O] [P] [=Branch1] [C]
      [=O] [Branch1] [C] [O-1] [O-1] .[99Tc+4]
    example_title: '[O-1]'
model-index:
  - name: chemselfies-base-bertmlm
    results:
      - task:
          type: fill-mask
          name: Fill-Mask
        dataset:
          name: main-eval-uniform
          type: main-eval-uniform
        metrics:
          - type: perplexity
            value: 1.3978
            name: Perplexity
          - type: accuracy
            value: 0.8929
            name: MLM Accuracy
      - task:
          type: fill-mask
          name: Fill-Mask
        dataset:
          name: main-eval-varied
          type: main-eval-varied
        metrics:
          - type: perplexity
            value: 1.4759
            name: Perplexity
          - type: accuracy
            value: 0.876
            name: MLM Accuracy
license: cc-by-nc-sa-4.0

ChemFIE Base - A Lightweight Model Pre-trained on Molecular SELFIES

This model is a lightweight model pre-trained on SELFIES (Self-Referencing Embedded Strings) representations of molecules. It is trained on 2.7M unique and valid molecules taken from COCONUTDB and ChemBL34, with 7.3M total generated masked examples. It is a compact model with only 11M parameters while achieving decent performance:

On varied masking:
- Perplexity of 1.4759, MLM Accuracy of 87.60%
On uniform 15% masking:
- Perplexity of 1.3978, MLM Accuracy of 89.29%

The masking strategy for pretraining utilized dynamic masking approach with masking ratios ranging from 15% to 45% based on simple scoring to gauge molecular complexity.

Disclaimer: For Academic Purposes Only

The information and model provided is for academic purposes only. It is intended for educational and research use, and should not be used for any commercial or legal purposes. The author do not guarantee the accuracy, completeness, or reliability of the information.

Model Details
Usage
Bias and Limitations
Background
Training Details
Evaluation
Interpretability
Technical Specifications
Citation
Contact & Support My Work

Model Details

Model Description

Model type: Transformer (BertForMaskedLM)
Language: SELFIES
Maximum Sequence Length: 512 tokens
License: CC-BY-NC-SA 4.0
Training Dataset: COCONUTDB and ChemBL34
Resources for more information:
- Github Respository (coming soon)
- Detailed article (coming soon)

Usage

Intended Use

You can use this model for masked language modeling but it's mostly intended to be fine-tuned on a downstream task.

Direct Use

You can use this model directly with a pipeline for masked language modeling:

from transformers import pipeline

# text = "[C] [C] [=Branch1] [C] [=O] [O] [C] [C] [N+1] [Branch1] [C] [C] [Branch1] [C] [C] [C]"

text = "[C] [C] [=Branch1] [C] [MASK] [O] [C] [C] [N+1] [Branch1] [C] [C] [Branch1] [C] [C] [C]"
  
mask_filler = pipeline("fill-mask", "gbyuvd/chemselfies-base-bertmlm")

mask_filler(text, top_k=5)

"""
[{'score': 0.9974672794342041,
  'token': 8,
  'token_str': '[=O]',
  'sequence': '[C] [C] [=Branch1] [C] [=O] [O] [C] [C] [N+1] [Branch1] [C] [C] [Branch1] [C] [C] [C]'},
 {'score': 0.002122726757079363,
  'token': 34,
  'token_str': '[=S]',
  'sequence': '[C] [C] [=Branch1] [C] [=S] [O] [C] [C] [N+1] [Branch1] [C] [C] [Branch1] [C] [C] [C]'},
 {'score': 0.0002627855574246496,
  'token': 11,
  'token_str': '[=N]',
  'sequence': '[C] [C] [=Branch1] [C] [=N] [O] [C] [C] [N+1] [Branch1] [C] [C] [Branch1] [C] [C] [C]'},
 {'score': 8.700760372448713e-05,
  'token': 1,
  'token_str': '[C]',
  'sequence': '[C] [C] [=Branch1] [C] [C] [O] [C] [C] [N+1] [Branch1] [C] [C] [Branch1] [C] [C] [C]'},
 {'score': 2.8387064958224073e-05,
  'token': 2,
  'token_str': '[=C]',
  'sequence': '[C] [C] [=Branch1] [C] [=C] [O] [C] [C] [N+1] [Branch1] [C] [C] [Branch1] [C] [C] [C]'}]
"""

Background

Three weeks ago, I had an idea to train a sentence transformer based on chemical "language" which so far I looked up back then, had not yet existed. While trying to do so, I found this wonderful and human-readable new molecular representation called SELFIES - developed by Aspuru-Guzik group. I found this representation fascinating and worth to explore, due to its robustness and at least so far proven to be versatile and easier to train a model using it. For more information on SELFIES, you could read this blogpost or check out their github.

My initial attempt focused on training a sentence transformer based on SELFIES, with the goal of enabling rapid molecule similarity search and clustering. This approach potentially offers advantages over traditional fingerprinting algorithms like MACCS, as the embeddings are context-aware. I decided to fine-tune a relatively lightweight NLP-trained miniLM model by Nils Reimers, as I was unsure about training from scratch and didn't even know about pre-training at that time.

The next challenges were how to properly make molecule pairs that is diverse yet informative, and how to label them. After tackling those, I trained the model on a dataset built from natural compounds taken from COCONUTDB. After some initial training, I pushed the model to Hugging Face to get some feedback. Gladly, Tom Aarsen provided valuable suggestions, including training a custom tokenizer, exploring Matryoshka embeddings, and considering training from scratch. The attempt to implement Aarsen's suggestions, specifically in training from scratch is the main goal of this project as well as a first experience for me.

Lastly before going into the details, it's important to note that this is the result of a hands-on learning project, and as such - beside my insufficient knowledge - it may not meet rigorous scientific standards. Like any learning journey, it's messy and I myself constrained by financial, computational, and time limitations. I've had to make compromises, such as conducting incomplete experiments and chunking datasets. However, I am more than happy to receive any feedback, so that I can improve both myself and future models/projects. A more detailed article discussing this project in details is coming soon.

Training Details

Training Data

Data Sources

The dataset combines two sources of molecular data:

Natural compounds from COCONUTDB (Sorokina et al., 2021)
Bioactive compounds from ChemBL34 (Zdrazil et al., 2023)

Data Preparation

Fetching: Canonical SMILES were extracted from both databases.
De-duplication:
- Each dataset was de-duplicated internally.
- The combined dataset ("All") was further de-duplicated to ensure unique entries.
Validity Check and Conversion: A dual validity check was performed using RDKit and by converting them into SELFIES

Filtering and Chunking

Filtering by Lipinski's Rule of Five or its subsets (e.g., Mw < 500 and LogP < 5) was omitted to maintain broader coverage for potential future expansion to organic and inorganic molecules such in PubCHEM and ZINC20.
The dataset was chunked into 13 parts, each containing 203,458 molecules, to accommodate the 6-hour time limit on Paperspace's Gradient.
Any leftover data was randomly distributed across the 13 chunks to ensure even distribution.

Validation Set

10% of each chunk was set aside for validation.
These validation sets were combined into a main test set, totaling 810,108 examples.

Dataset	Number of Valid Unique Molecules	Generated Training Examples
Chunk I	207,727	560,859
Chunk II	207,727	560,859
Chunk III	207,727	560,859
Chunk IV	207,727	560,859
Chunk V	207,727	560,859
Chunk VI	207,727	560,859
Chunk VII	207,727	560,859
Chunk VIII	207,727	560,859
Chunk IX	207,727	560,859
Chunk X	207,727	560,859
Chunk XII	207,727	560,859
Chunk XI	207,727	560,859
Chunk XIII	207,738	560,889
Total	2,700,462	7,291,197

Training Procedure

Tokenizer Setup

The tokenizer is a combination of my own pretrained tokenizer on the merged COCONUTDB+ChemBL34 SELFIES dataset with vocabularies from zpn's word-level tokenizer trained on PubChem. This approach was chosen to ensure comprehensive coverage while maintaining relevance to biological compounds. The tokenizer was modified to suit the BertTokenizer format, using whitespace to split input tokens.

When using or fine-tuning this model, it's crucial to separate each SELFIES token with a whitespace. For example:

[C] [N] [C] [C] [C] [C@H1] [Ring1] [Branch1] [C] [=C] [N] [=C] [C] [=C] [Ring1] [=Branch1]

To ensure coverage, the tokenizer underwent evaluation to cover all tokens in the training data. Unrecognized tokens were identified and incorporated into the tokenizer. Additionally, my previous pre-training issues, such as improper tokenization of dot symbol prefixes in complex molecules (e.g., ".[Cl]"), were addressed and resolved.

Generating Dynamic Masked Sequence

The key method in this project is the implementation of a dynamic masking rate based on molecular complexity. I think we can heuristically infer a molecule's complexity based on the syntactic characteristics of SELFIES. Simpler tokens will have only one character, such as "[N]" (l = 1; ignoring the brackets), while more complex ones would be ".[N+1]" (l = 4). Relatively rare atoms compared to the CHONS, like [Na] (l = 2), and ionized metals like [Fe+3] (l = 4), also vary in complexity. To normalize and infer the density of many character tokens, we can sum of all tokens length ratio with the molecule's length. I will refer to this simple score as the "complexity score" hereafter. We can then normalize it and use it to determine a variable masking probability ranging from 15% to 45%. Additionally, we can employ three different masking strategies to introduce further variability. This approach aims to create a more challenging and diverse training dataset while getting the most out of it, potentially leading to a more robust and generalizable model for molecular representation learning. Each SELFIES string's complexity is calculated based on the logarithm of the sum of token ratios with the sequence length.

1. Complexity Score Calculation

The raw complexity score is calculated using the formula:

$Sc = \log\left[\sum\left(\frac{l_{\text{token}}}{n_{\text{tokens}}}\right)\right]$

Example outputs:

Sentence A:
Tokens: ['[C]', '[C]', '[=Branch1]', '[C]', '[=O]', '[O]', '[C]']
Sum of token lengths: 29
Number of tokens: 7
Raw complexity score: 1.4214

==================================================

Sentence B:
Tokens: ['[C]', '[N+1]', '[Branch1]', '[C]', '[C]', '[Branch1]', '[C]', '[C]', '[C]']
Sum of token lengths: 41
Number of tokens: 9
Raw complexity score: 1.5163

2. Normalization

The raw score is then normalized to a range of 0-1 using predefined minimum (1.39) and maximum (1.69) normalization values which determined from dataset's score distributions:

$Sc_{norm} = max(0, min(1, (Sc - min_{norm}) / (max_{norm} - min_{norm})))$

3. Mapping to Masking Probability

I decided to use quadratic mapping with 0.3 steps, ensuring smooth masking probability adjustment in range between 15% to 45% with more complex molecules having a higher masking probability:

$P_{\text{mask}} = 0.15 + 0.3 * (Sc_{norm})^2$

4. Multi-Strategy Masking

Three different masking strategies are employed for each SELFIES string:

Main Strategy:
- 80% chance to mask the token
- 10% chance to keep the original token
- 10% chance to replace with a random token
Alternative Strategy 1:
- 70% chance to mask the token
- 15% chance to keep the original token
- 15% chance to replace with a random token
Alternative Strategy 2:
- 60% chance to mask the token
- 20% chance to keep the original token
- 20% chance to replace with a random token

5. Data Augmentation

Each SELFIES string is processed three times, once with each masking strategy.
This hopefully triples the effective dataset size and introduces variability in the masking patterns.

6. Masking Process

Tokens are randomly selected for masking based on the calculated masking probability.
Special tokens ([CLS] and [SEP]) are never masked.
The number of tokens to be masked is determined by the masking probability and the length of the SELFIES string.

This methodology aims to create a diverse and challenging dataset for masked language modeling of molecular structures, adapting the masking intensity to the complexity of each molecule and employing multiple masking strategies to improve model robustness and generalization. Also, beside masking differently based on complexity scores, the on-the-fly data generation might ensure that each run and batches - the data are masked differently. But additional and further confirmation of this is needed.

Training Hyperparameters

Batch size = 128
Num of Epoch = 1
Total steps on all chunks = 56,966
Training time on each chunk = 03h:24m / ~205 mins

I am using Ranger21 optimizer with these settings:

Core optimizer = madgrad
Learning rate of 2e-05

num_epochs of training = ** 1 epochs **

using AdaBelief for variance computation
Warm-up: linear warmup, over 964 iterations (0.22)

Lookahead active, merging every 5 steps, with blend factor of 0.5
Norm Loss active, factor = 0.0001
Stable weight decay of 0.01
Gradient Centralization = On

Adaptive Gradient Clipping = True
    clipping value of 0.01
    steps for clipping = 0.001

I turned off the warm down, since in prior experiments it led to instability of losses in my case. For more information about Ranger21, you could check out this repository.

Evaluation

Dataset: main-eval
Number of test examples: 810,108

Varied Masking Test

Chunk	Avg Loss	Perplexity	MLM Accuracy
I-IV	0.4547	1.5758	0.851
V-VIII	0.4224	1.5257	0.864
IX-XIII	0.3893	1.4759	0.876

Uniform 15% Masking Test (80%:10%:10%)

Chunk	Avg Loss	Perplexity	MLM Accuracy
XII	0.3349	1.3978	0.8929

Interpretability

Attention Head Visualization

(coming soon)

Neural Stacks Visualization

(coming soon)

Attributions in Determining Masked Tokens

(coming soon)

Technical Specifications

Model Architecture and Objective

Layers: 8
Attention Heads: 4
Hidden Size: 320
Intermediate Size: 1280 (4H)
Attention Type: SDPA

Compute Infrastructure

Hardware

Platform: Paperspace's Gradients

Compute: Free-P5000 (16 GB GPU, 30 GB RAM, 8 vCPU)

Software

Python: 3.9.13
Transformers: 4.42.4
PyTorch: 2.3.1+cu121
Accelerate: 0.32.0
Datasets: 2.20.0
Tokenizers: 0.19.1
Ranger21: 0.0.1
Selfies: 2.1.2
RDKit: 2024.3.3

Citation

If you find this project useful in your research and wish to cite it, please use the following BibTex entry:

@software{chemfie_basebertmlm,
  author = {GP Bayu},
  title = {{ChemFIE Base}: Pretraining A Lightweight BERT-like model on Molecular SELFIES},
  url = {https://huggingface.co/gbyuvd/chemselfies-base-bertmlm},
  version = {1.0},
  year = {2024},
}

References

`@article{zdrazil2023chembl,
  title={The ChEMBL Database in 2023: a drug discovery platform spanning multiple bioactivity data types and time periods},
  author={Zdrazil, Barbara and Felix, Eloy and Hunter, Fiona and Manners, Emma J and Blackshaw, James and Corbett, Sybilla and de Veij, Marleen and Ioannidis, Harris and Lopez, David Mendez and Mosquera, Juan F and Magarinos, Maria Paula and Bosc, Nicolas and Arcila, Ricardo and Kizil{\"o}ren, Tevfik and Gaulton, Anna and Bento, A Patr{\'i}cia and Adasme, Melissa F and Monecke, Peter and Landrum, Gregory A and Leach, Andrew R},
  journal={Nucleic Acids Research},
  year={2023},
  volume={gkad1004},
  doi={10.1093/nar/gkad1004}
}

@misc{chembl34,
  title={ChemBL34},
  year={2023},
  doi={10.6019/CHEMBL.database.34}
}

@article{sorokina2021coconut,
  title={COCONUT online: Collection of Open Natural Products database},
  author={Sorokina, Maria and Merseburger, Peter and Rajan, Kohulan and Yirik, Mehmet Aziz and Steinbeck, Christoph},
  journal={Journal of Cheminformatics},
  volume={13},
  number={1},
  pages={2},
  year={2021},
  doi={10.1186/s13321-020-00478-9}
}

@article{krenn2020selfies,
  title={Self-referencing embedded strings (SELFIES): A 100\% robust molecular string representation},
  author={Krenn, Mario and H{\"a}se, Florian and Nigam, AkshatKumar and Friederich, Pascal and Aspuru-Guzik, Alan},
  journal={Machine Learning: Science and Technology},
  volume={1},
  number={4},
  pages={045024},
  year={2020},
  doi={10.1088/2632-2153/aba947}
}

Contact & Support My Work

G Bayu ([email protected])

This project has been quiet a journey for me, I’ve dedicated hours on this and I would like to improve myself, this model, and future projects. However, financial and computational constraints can be challenging.

If you find my work valuable and would like to support my journey, please consider suppoting me here. Your support will help me cover costs for computational resources, data acquisition, and further development of this project. Any amount, big or small, is greatly appreciated and will enable me to continue learning and explore more.

Thank you for checking out this model, I am more than happy to receive any feedback, so that I can improve myself and the future model/projects I will be working on.

gbyuvd
/

chemselfies-base-bertmlm

ChemFIE Base - A Lightweight Model Pre-trained on Molecular SELFIES

Disclaimer: For Academic Purposes Only

Table of Contents

Model Details

Model Description

Usage

Intended Use

Direct Use

Background

Training Details

Training Data

Data Sources

Data Preparation

Filtering and Chunking

Validation Set

Training Procedure

Tokenizer Setup

Generating Dynamic Masked Sequence

Training Hyperparameters

Evaluation

Varied Masking Test

Uniform 15% Masking Test (80%:10%:10%)

Interpretability

Attention Head Visualization

Neural Stacks Visualization

Attributions in Determining Masked Tokens

Technical Specifications

Model Architecture and Objective

Compute Infrastructure

Hardware

Software

Citation

References

Contact & Support My Work