Suiren-Base (1.8B)

Suiren-Base is a large-scale molecular domain foundation model developed by Golab (SAIS Physics Lab). With 1.8 billion parameters, it is designed to consolidate quantum chemical knowledge into a unified framework through massive self-supervised pre-training, enabling direct prediction of various quantum properties for small molecules.

Key Features

• Core Foundation Model: Serves as the backbone of the Suiren series, encoding broad chemical knowledge for predicting energy, forces, and generating high-quality atomic-level embeddings.
• Advanced Architecture: Utilizes a hybrid architecture combining EquiformerV2 with the Equivariant Spherical Transformer (EST). It captures deep inter-atomic interactions through high-order equivariant representations while significantly improving inference speed compared to traditional point-cloud networks.
• Self-Supervised Pre-training: Employs Equivariant Masked Position Prediction (EMPP), a task that forces the model to understand atomic interactions by reconstructing deleted atomic nodes and positions.
• Large-scale Dataset: Trained on the Full Qo2mol dataset (not been fully open-sourced), which contains about 100 million high-precision DFT calculation points, covering an extensive chemical space and various molecular sizes.

Suiren-Base model is the foundation model in Suiren family.

Usage

You can load the model using the provided API in the GitHub repository.

import torch
from suiren_models import ModelLoader

# Initialize loader
loader = ModelLoader(config_path='config_name.yml')
# example: loader = ModelLoader(config_path='suiren-base.yml')

# Load model architecture
model = loader.load_model()

# Load pre-trained weights
loader.load_weights(model, 'path/to/checkpoint')

# Load normalizer from config
loader.load_normalizer()

# Load normalizer from checkpoint (optional)
#  loader.load_normalizer('path/to/normalizer')

# Move model to device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()