SELFormerMM: multimodal molecular representation learning via SELFIES, structure, text, and knowledge graph integration

Molecular representation learning is central to computational drug discovery, yet most models still rely on a single modality (e.g. sequences or graphs) and cannot easily unify complementary signals such as text and biological interaction networks. SELFormerMM addresses this gap as a unified multimodal framework: it integrates SELFIES encodings, 2D molecular graphs, natural-language descriptions, and knowledge-graph embeddings into a shared latent space, building on the transformer-based SELFormer. Self-supervised pre-training is carried out on a large-scale multimodal dataset of approximately 3 million molecules.

Overview of the SELFormerMM framework. (A) Contrastive multimodal pretraining integrates molecular information from four complementary modalities—SELFIES sequences, structural graphs, textual descriptions, and knowledge graph interactions—processed by modality-specific encoders and projection networks, then aligned in a shared representation space via contrastive learning with the SINCERE loss. (B) Downstream task finetuning adapts the pretrained multimodal backbone for molecular property prediction, using task-specific prediction heads for binary classification, multilabel classification, and regression.

Table of contents

• The Architecture of SELFormerMM
• Getting Started
• Dataset
• Generating SELFIES
• Generating modality embeddings
  • Graph (UniMol)
  • Text (Hugging Face encoder)
  • KG (DMGI + HeteroData)
• Producing multimodal embeddings
• Training and evaluating models
  • Pre-training
  • Fine-tuning
    • Fine-tuning from a released config
• Producing predictions with fine-tuned models
• Scoring fine-tuned models
• Citation
• License

The Architecture of SELFormerMM

SELFormerMM uses four modality-specific branches and a shared projection space that aligns modalities: the sequence path is [SELFormer]—a RoBERTa-based chemical language model. Text uses frozen SciBERT: mean-pooled 768d embeddings of natural language descriptions. Structure uses frozen Uni-Mol: 512d cls_repr from SMILES-derived 3D conformers. Knowledge uses a pretrained DMGI model on CROssBARv2-KG; relation-specific representations of compound nodes are mean-pooled to 128d per compound. Non-linear MLP projection heads (expansion layers, LayerNorm, ReLU, linear map to hidden size H = 768) map graph, text, and KG embeddings into RoBERTa’s space; zero vectors encode missing modalities. SINCERELoss implements multi-view, multi-positive supervised contrastive alignment (InfoNCE-style) so same-molecule views agree and different molecules separate. For downstream evaluation on MoleculeNet, a task-specific head receives the concatenation of the four embeddings (sequence [CLS] plus projected graph, text, and KG). This repository provides scripts for modality embeddings, pre-training, fine-tuning, and prediction.

Getting Started

We recommend Conda (or a fresh venv) with Python 3.10+. Install PyTorch for your CUDA/CPU setup, then:

pip install transformers torch selfies pandas numpy tqdm scikit-learn
# Graph embeddings
pip install rdkit unimol_tools
# Scaffold split (fine-tuning)
pip install chemprop
# KG embeddings (DMGI + PyG)
pip install torch-geometric safetensors

Clone or copy the repo and ensure the project root is on PYTHONPATH:

cd /path/to/SELFormerMM
export PYTHONPATH="${PYTHONPATH}:$(pwd)"

Dataset

Data and model download. Pretraining and fine-tuning datasets (including all raw modalities), single-modality embedding files (graph / text / kg), the pretrained SELFormerMM multimodal checkpoint, and fine-tuned task checkpoints are available from the Hugging Face dataset: HUBioDataLab/SELFormerMM.

1. SELFIES — From SMILES CSV: generate_selfies.py.
2. Graph embeddings — Same row order as your molecule table: generate_graph_embeddings.py → graph_embeddings.npy.
3. Text embeddings — Descriptions aligned by row: generate_text_embeddings.py → text_embeddings.npy.
4. KG embeddings — Compound nodes aligned by row with pre-trained DMGI: generate_kg_embeddings.py → kg_embeddings.npy.
5. Pre-train — Pretraining metadata CSV (selfies column) + three .npy files (same length): train_pretraining.py.
6. Fine-tune — One meta CSV (smiles, selfies, labels) + one NPZ bundling three modalities (graph, text, kg):

Generating SELFIES

python3 generate_selfies.py --smiles_dataset=data/input.csv --selfies_dataset=data/output_with_selfies.csv --smiles_column=smiles

• --smiles_dataset — Input CSV path.
• --selfies_dataset — Output CSV (adds column selfies).
• --smiles_column — SMILES column name (default smiles).
• --on_error — keep | empty | raise for encoder errors.

Generating modality embeddings

You can download precomputed embeddings here, or generate embeddings for your own dataset using the following scripts.

Graph (UniMol)

python3 generate_graph_embeddings.py --input_csv=data/pretraining_datasets/pretraining_dataset_meta.csv --output_npy=data/pretraining_datasets/graph_embeddings.npy --smiles_column=smiles --batch_size=512 --embedding_dim=512 --normalize=1

Optional: --output_csv, --id_column, --gpu_ids=0,1, --use_gpu=0 (CPU), --batch_timeout_seconds, --min_batch_size.

Text (Hugging Face encoder)

python3 generate_text_embeddings.py --input_csv=data/pretraining_datasets/pretraining_dataset_meta.csv --output_npy=data/pretraining_datasets/text_embeddings.npy --text_column=Description --model_name=allenai/scibert_scivocab_uncased --max_length=512 --normalize=1

KG (DMGI + HeteroData)

python3 generate_kg_embeddings.py --checkpoint_path=data/models/DMGI/dmgi_model.pt --heterodata_path=data/pretraining_datasets/selformermm_kg_heterodata.pt --output_npy=data/pretraining_datasets/kg_embeddings_new.npy --align_meta_csv=data/pretraining_datasets/pretraining_dataset_meta.csv --align_meta_column=chembl_id --node_type=Compound --out_channels=128 --normalize=1

Producing multimodal embeddings

This script loads the pretrained SELFormerMM. Pass the single modality .npy files (graph / text / KG), aligned row-wise with the pretraining metadata CSV.

python3 produce_multimodal_embeddings.py \
  --selfies_csv=data/pretraining_datasets/pretraining_dataset_meta.csv \
  --selfies_column=selfies \
  --pretrained_multimodal_dir=data/models/SELFormerMM \
  --graph_embs=data/pretraining_datasets/graph_embeddings.npy \
  --text_embs=data/pretraining_datasets/text_embeddings.npy \
  --kg_embs=data/pretraining_datasets/kg_embeddings.npy \
  --output_npy=data/multimodal_embeddings.npy \
  --output_mode=concat \
  --batch_size=32 \
  --max_len=512

• --output_mode — concat: (N, 4*D); stacked: (4N, D).
• --output_csv / --id_column — Optional CSV with ID + embedding columns.

Training and evaluating models

Pre-training

Generate the SELFIES CSV and single modality embeddings, or download ready-to-use files from here, then train the SELFormerMM multimodal backbone using the following script.

python3 train_pretraining.py \
  --selfies_csv=data/pretraining_datasets/pretraining_dataset_meta.csv \
  --selfies_column=selfies \
  --graph_embs=data/pretraining_datasets/graph_embeddings.npy \
  --text_embs=data/pretraining_datasets/text_embeddings.npy \
  --kg_embs=data/pretraining_datasets/kg_embeddings.npy \
  --model_path=HUBioDataLab/SELFormer \
  --tokenizer_path=HUBioDataLab/SELFormer \
  --save_dir=data/models/SELFormerMM \
  --batch_size=40 \
  --max_len=512 \
  --epochs=267 \
  --lr=2e-5 \
  --val_frac=0.1 \
  --save_every=1 \
  --save_embeddings=data/final_pretrain_embeddings.npy

• --selfies_csv — CSV with SELFIES column (required).
• --graph_embs / --text_embs / --kg_embs — Optional .npy; omitted modalities are zero-filled.
• --model_path — HF model id/local directory of unimodal SELFormer.
• --tokenizer_path — HF tokenizer id/local directory of unimodal SELFormer.
• --save_dir — Checkpoints under epoch_XXX/ plus final pytorch_model.bin + config + tokenizer.
• --save_embeddings — Optional memmap .npy for all rows after training.

Fine-tuning

Use the pretrained SELFormerMM multimodal backbone to fine-tune on your own dataset using the following script. Inputs are a task CSV (<task>.csv) with selfies, smiles, and label column(s) and (<task>_embs.npz) file that contains the single modality embeddings. Ready-to-use datasets can be found here.

Fine-tuning from a released config

Every released task ships the exact hyperparameters that produced its checkpoints, so you do not have to re-derive the flags. There are two forms:

• data/models/finetuned/<task>/finetuning_optimized_config.json in this repository — the task's hyperparameters with no seed pinned, so --seed is yours to choose (42 if you leave it out).
• models/finetuned/<task>/seed_<n>/finetuning_optimized_config.json in the dataset repository — the same values plus the seed of the checkpoint sitting next to it. Use these to reproduce a published checkpoint, since the seed also fixes the split.

A config records the dataset paths, the split settings (use_scaffold, train_frac/val_frac/test_frac) and the optimization settings (batch size, epochs, learning rates, weight decay). Run fine-tuning directly from one:

python3 train_finetuning.py \
  --config=data/models/finetuned/bbbp/finetuning_optimized_config.json \
  --save_dir=runs/bbbp

Any flag on the command line overrides the config, which is how you change one setting while keeping the rest fixed:

python3 train_finetuning.py \
  --config=data/models/finetuned/bbbp/finetuning_optimized_config.json \
  --seed=123 --save_dir=runs/bbbp_seed123

The paths inside a config are relative to this repository's root, with the dataset mounted at data/. Download it there and everything resolves:

huggingface-cli download HUBioDataLab/SELFormerMM --repo-type dataset \
  --include "models/SELFormerMM/*" "finetuning_datasets/classification/bbbp/*" \
  --local-dir ./data

Training writes two checkpoints into --save_dir: best_model.pt, the epoch with the best validation metric, and model.pt, the final epoch. best_metrics.json records which epoch was selected and its validation and test metrics. Pass --save_best_val=0 to skip the former.

Binary classification example:

python3 train_finetuning.py \
  --dataset_meta_csv=data/finetuning_datasets/classification/bbbp/bbbp.csv \
  --dataset_embs_npz=data/finetuning_datasets/classification/bbbp/bbbp_embs.npz \
  --model_path=HUBioDataLab/SELFormer \
  --tokenizer_path=HUBioDataLab/SELFormer \
  --pretrained_multimodal_dir=data/models/SELFormerMM \
  --task_type=binary \
  --label_column=p_np \
  --num_labels=2 \
  --use_scaffold=1 \
  --batch_size=32 \
  --max_len=128 \
  --epochs=25 \
  --backbone_lr=5e-5 \
  --head_lr=1e-5 \
  --weight_decay=0.1 \
  --checkpoint_every=25 \
  --save_dir=data/models/finetuned/bbbp

Multi-label example:

python3 train_finetuning.py \
  --dataset_meta_csv=data/finetuning_datasets/classification/sider/sider.csv \
  --dataset_embs_npz=data/finetuning_datasets/classification/sider/sider_embs.npz \
  --model_path=HUBioDataLab/SELFormer \
  --tokenizer_path=HUBioDataLab/SELFormer \
  --pretrained_multimodal_dir=data/models/SELFormerMM \
  --task_type=multilabel \
  --num_labels=0 \
  --use_scaffold=0 \
  --batch_size=8 \
  --max_len=128 \
  --epochs=25 \
  --backbone_lr=5e-5 \
  --head_lr=5e-6 \
  --weight_decay=0.001 \
  --checkpoint_every=25 \
  --save_dir=data/models/finetuned/sider

Regression example:

python3 train_finetuning.py \
  --dataset_meta_csv=data/finetuning_datasets/regression/freesolv/freesolv.csv \
  --dataset_embs_npz=data/finetuning_datasets/regression/freesolv/freesolv_embs.npz \
  --model_path=HUBioDataLab/SELFormer \
  --tokenizer_path=HUBioDataLab/SELFormer \
  --pretrained_multimodal_dir=data/models/SELFormerMM \
  --task_type=regression \
  --label_column=freesolv \
  --num_labels=1 \
  --use_scaffold=0 \
  --batch_size=4 \
  --max_len=128 \
  --epochs=25 \
  --backbone_lr=5e-5 \
  --head_lr=2e-4 \
  --weight_decay=0.001 \
  --checkpoint_every=25 \
  --save_dir=data/models/finetuned/freesolv

Useful flags: --config=finetune.json (inject defaults), --save_split_csvs=dir, --train_frac / --val_frac / --test_frac, --test_eval_every, --device=cuda.

Producing predictions with fine-tuned models

This script loads the fine-tuned SELFormerMM multimodal backbone and produces predictions on a new dataset using the following script. Inputs are a task CSV (<task>.csv) with selfies, smiles, and label column(s) and (<task>_embs.npz) file that contains the single modality embeddings. Ready-to-use datasets and our finetuned models can be found here. Three seeds are published per task, so --model_dir points at one seed's directory (<task>/seed_42, seed_123 or seed_3407).

Binary classification example:

python3 predict.py \
  --model_dir=data/models/finetuned/bbbp/seed_42 \
  --input_meta_csv=data/finetuning_datasets/classification/bbbp/bbbp.csv \
  --input_embs_npz=data/finetuning_datasets/classification/bbbp/bbbp_embs.npz \
  --output_csv=data/finetuning_datasets/classification/bbbp/bbbp_predictions.csv \
  --task_type=binary \
  --num_labels=2 \
  --batch_size=16 \
  --max_len=512 \
  --output_id_column=smiles

Multi-label example:

python3 predict.py \
  --model_dir=data/models/finetuned/sider/seed_42 \
  --input_meta_csv=data/finetuning_datasets/multilabel/sider/sider.csv \
  --input_embs_npz=data/finetuning_datasets/multilabel/sider/sider_embs.npz \
  --output_csv=data/finetuning_datasets/multilabel/sider/sider_predictions.csv \
  --task_type=multilabel \
  --batch_size=16 \
  --max_len=512 \
  --output_id_column=smiles

Regression example:

python3 predict.py \
  --model_dir=data/models/finetuned/freesolv/seed_42 \
  --input_meta_csv=data/finetuning_datasets/regression/freesolv/freesolv.csv \
  --input_embs_npz=data/finetuning_datasets/regression/freesolv/freesolv_embs.npz \
  --output_csv=data/finetuning_datasets/regression/freesolv/freesolv_predictions.csv \
  --task_type=regression \
  --num_labels=1 \
  --batch_size=16 \
  --max_len=512 \
  --output_id_column=smiles

Scoring fine-tuned models

evaluate.py rebuilds the exact train/val/test split a checkpoint was trained with, scores the model on it, and prints the metrics — ROC-AUC and PRC-AUC for classification, RMSE/MSE/MAE for regression.

Three seeds (42, 123, 3407) are published per task. Point --model_dir at a task directory to score all of them and report the mean ± standard deviation, which is the form the paper's table is stated in:

python3 evaluate.py \
  --model_dir=data/models/finetuned/bbbp \
  --config=data/models/finetuned/bbbp/finetuning_optimized_config.json \
  --device=cuda

[evaluate] 3 seeds: 42=0.9484  123=0.9264  3407=0.9605
[evaluate] roc-auc = 0.9451 +/- 0.0141  (mean +/- std over seeds)

To score one checkpoint, pass that seed's config and leave --model_dir out:

python3 evaluate.py \
  --config=data/models/finetuned/bbbp/seed_123/finetuning_optimized_config.json \
  --device=cuda

• --model_dir — Task directory (scores every seed_<n> subdirectory) or a single checkpoint directory. Defaults to the directory holding --config.
• --model_file — Weights file inside the checkpoint directory; model.pt by default, best_model.pt for the validation-selected epoch of your own run.
• --split — test (default), val or train.
• --output_json — Write the per-seed results and the aggregate to a file.
• --device — Pass cuda explicitly. Without it the script falls back to CPU when a GPU cannot be initialized, which on the larger tasks turns minutes into hours.

Citation

Please refer to our article for more information. If you use the SELFormerMM method or the datasets provided in this repository, please cite this paper:

@article{ulusoy2026selformermm,
  title   = {SELFormerMM: multimodal molecular representation learning via SELFIES,
             structure, text, and knowledge graph integration},
  author  = {Ulusoy, Erva and Bostanc{\i}, {\c{S}}evval and Deniz, Bora Engin and Do{\u{g}}an, Tunca},
  journal = {Bioinformatics},
  volume  = {42},
  number  = {Supplement\_2},
  pages   = {btag451},
  year    = {2026},
  doi     = {10.1093/bioinformatics/btag451}
}

License

Copyright (C) 2026 HUBioDataLab

This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License along with this program. If not, see http://www.gnu.org/licenses/.