A context-aware multi-modal molecular foundation model that jointly learns from 1D (SELFIES), 2D (graphs), and 3D (conformer ensembles) representations for molecular property prediction.

Model Description

MolFM-Lite is a multi-modal molecular foundation model designed for molecular property prediction tasks in drug discovery. It achieves state-of-the-art results on MoleculeNet benchmarks by:

Multi-modal encoding: Learning from 1D sequences (SELFIES), 2D graphs (GIN), and 3D structures (SchNet) simultaneously
Conformer ensemble attention: Modeling molecular flexibility using multiple conformers weighted by Boltzmann statistics
Context conditioning: Accounting for experimental variables using FiLM layers
Cross-modal fusion: Using cross-attention to let modalities inform each other

Model Architecture

Input Representations
        |
+-------+-------+-------+
|       |       |       |
1D      2D      3D Conformers
|       |       |
Transformer  GIN   SchNet
|       |       |
|       |   Conformer
|       |   Attention
|       |       |
+-------+-------+
        |
  Cross-Modal Fusion
        |
  Context Conditioning (FiLM)
        |
  Prediction Head + Uncertainty

Performance (State-of-the-Art Results)

Dataset	Task	Metric	MolFM-Lite	Previous SOTA	Improvement
BBBP	Blood-Brain Barrier	AUC	0.956	0.894	+6.9%
BACE	Beta-secretase Inhibition	AUC	0.902	0.878	+2.7%
Tox21	Toxicity (12 tasks)	AUC	0.848	0.795	+6.7%
Lipophilicity	Solubility	RMSE	0.570	0.631	-9.7%

Note: For RMSE, lower is better. For AUC, higher is better.

Available Checkpoints

model_bbbp.pt - Fine-tuned on BBBP (Blood-Brain Barrier Penetration)
model_bace.pt - Fine-tuned on BACE (Beta-secretase Inhibition)
model_tox21.pt - Fine-tuned on Tox21 (Toxicity prediction, 12 tasks)
model_lipophilicity.pt - Fine-tuned on Lipophilicity (Solubility prediction)

Usage

python

import torch
from src.models.molfm import MolFMLite
from src.data.preprocessing import MoleculePreprocessor

# Load model
model = MolFMLite(
    hidden_dim=256,
    hidden_dim_3d=128,
    num_layers_1d=4,
    num_layers_2d=4,
)

# Load checkpoint
checkpoint = torch.load("model_bbbp.pt", map_location="cpu")
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

# Process molecule
preprocessor = MoleculePreprocessor()
features = preprocessor.process_molecule("CC(=O)OC1=CC=CC=C1C(=O)O")  # Aspirin

# Predict
with torch.no_grad():
    prediction = model(features)
    print(f"BBB Penetration Probability: {prediction.item():.3f}")

Training Details

Framework: PyTorch 2.0
Hardware: AWS SageMaker ml.g4dn.xlarge (1x NVIDIA T4 GPU)
Training Cost: Under $20 total (no pretraining required)
Training Time: ~6 GPU hours for all 4 datasets
Optimizer: AdamW with learning rate 5e-5
Early Stopping: Patience of 15 epochs

Model Hyperparameters

Parameter	Value
Hidden Dimension (1D/2D)	256
Hidden Dimension (3D)	128
Transformer Layers	4
GIN Layers	4
SchNet Interactions	3
Conformers per Molecule	5
Attention Heads	8

Limitations

Requires RDKit for molecular preprocessing
3D conformer generation can be slow for large molecules
Performance may vary for molecules significantly different from training distribution

Citation

bibtex

@article{shah2026molfm,
  title={MolFM-Lite: A Multi-Modal Molecular Foundation Model with Context-Aware Predictions},
  author={Shah, Syed Omer},
  journal={GitHub},
  year={2026},
  url={https://github.com/Syedomershah99/molfm-lite}
}

Author

Syed Omer Shah

Email: syedomer@buffalo.edu
GitHub: @Syedomershah99
Affiliation: University at Buffalo

License

MIT License

OmerShah/MolFM-Lite

MolFM-Lite: Multi-Modal Molecular Foundation Model