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.

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):
    ...