Mol-JEPA

Mol-JEPA

Mol-JEPA is a multi-modal molecular foundation model trained on 5 million molecules spanning various modalities. It is based on a joint embedding predictive architecture and trained to predict missing modalities using the remaining ones. The predicted embeddings can be used for molecular machine learning downstream tasks.

<img src="overview.png" alt="Mol-JEPA" width="400">

Requirements

pip install torch torch-geometric transformers safetensors rdkit molfeat numpy

Embedding generation

The model is compatible with the HuggingFace AutoModel API. Because it ships custom modeling code, load it with trust_remote_code=True:

from transformers import AutoModel

model = AutoModel.from_pretrained("Flogrammer/Mol-JEPA", trust_remote_code=True)
model.eval()

Then a list of smiles can be provided and three types of outputs are returned: the predicted modality embeddings, the global CLS token and the latent modality embeddings from the second last layer. The result is a ModelOutput, so the fields are accessible both by name and by index.

smiles_list = ["Cn1cnc2n(C)c(=O)n(C)c(=O)c12", "CC(=O)Oc1ccccc1C(=O)O"]
out = model(smiles_list)

print("Predicted embeddings shape:", out.predictions.shape) # batch size, modalities, embedding dimension
>>> Predicted embeddings shape: torch.Size([2, 12, 512])

print("CLS token shape:", out.cls.shape) # batch size, embedding dimension
>>> CLS token shape: torch.Size([2, 512])

print("Latent embeddings shape:", out.embeddings.shape) # batch size, modalities, latent dimension
>>> Latent embeddings shape: torch.Size([2, 13, 512])

# Index access (predictions, cls, embeddings) is also supported:
predictions, cls, embeddings = out[0], out[1], out[2]

For larger datasets it is recommended to call the model in chunks.

CLS embedding downstream prediction

The CLS embeddings can be directly used as feature matrix with shape [n_samples, features]. We recommend to use TabICL or any other foundation model, as this led to the best performance in our experiments.

from tabicl import TabICLRegressor

# Take features from above
cls_features = out.cls

# Do some splitting
train_idx, test_index = ...

model = TabICLRegressor(random_state=123)
model.fit(cls_features[train_idx], y[train_idx])
test_pred = model.predict(cls_features[test_idx])

Multimodal embeddings downstream prediction

While the CLS token can be a powerful summary, it might miss modality-specific details. Therefore, another possibility is to use all predicted modality embeddings and combine them in the downstream model. We recommend to use a transformer module as follows:

class TransformerProbe(nn.Module):
    def __init__(self, n_tokens, token_dim, n_layers=2, n_heads=4, dropout=0.1):
        super().__init__()
        self.n_tokens = n_tokens
        self.token_dim = token_dim
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=token_dim,
            nhead=n_heads,
            dim_feedforward=token_dim * 4,
            dropout=dropout,
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(
            encoder_layer, num_layers=n_layers, enable_nested_tensor=False
        )
        self.head = nn.Linear(token_dim, 1)

    def forward(self, x):
        # x: (B, n_tokens * token_dim) -> (B, n_tokens, token_dim)
        x = x.view(-1, self.n_tokens, self.token_dim)
        x = self.transformer(x)
        # Mean pool over tokens, then predict
        x = x.mean(dim=1)
        return self.head(x)


# Extract the number of modalities
n_tokens = out.predictions.shape[-1]

# Build model
hidden_dim = 512
model = TransformerProbe(n_tokens=n_tokens, token_dim=hidden_dim)

# Optimize
epochs = 10
for _ in range(epochs):
    ...