cellfm-800m

CellFM-800M

Model Description

CellFM is a large-scale foundation model pre-trained on transcriptomics of 100 million human cells using a retention-based architecture (MAE Autobin).

• Model Size: 800M parameters
• Pre-training Data: 100M human cells
• Architecture: Retention-based Transformer (MAE Autobin)
• Vocabulary: 24,072 genes
• Pre-training Task: Masked Autoencoding (MAE)

Model Details

• Source: biomed-AI/CellFM
• Original Framework: MindSpore
• Converted to: PyTorch (PerturbLab format)
• License: See original repository for details

Architecture Specifications

• Hidden Dimension: 1536
• Number of Layers: 40
• Number of Attention Heads: 48
• Dropout: 0.1
• Max Sequence Length: 2048 genes

Usage

Load Model

from perturblab.model.cellfm import CellFMModel

# Load pretrained model (automatically downloads if needed)
model = CellFMModel.from_pretrained('cellfm-800m')

# Or use short name
model = CellFMModel.from_pretrained('800m')

# Or from local path
model = CellFMModel.from_pretrained('./weights/cellfm-800m')

Generate Cell Embeddings

import scanpy as sc

# Load your data
adata = sc.read_h5ad('your_data.h5ad')

# Preprocess
adata = CellFMModel.prepare_data(adata)

# Get embeddings (use smaller batch size for 800M model)
embeddings = model.predict_embeddings(
    adata,
    batch_size=8,  # Smaller batch size for larger model
    return_cls_token=True,
)

# Access cell embeddings
cell_embeddings = embeddings['cell_embeddings']  # Shape: (n_cells, 1536)

Fine-tune for Classification

from perturblab.model.cellfm import CellFMModel, CellFMConfig

# Initialize model with classification head
config = CellFMConfig(
    model_name='800M',
    n_genes=24072,
    enc_dims=1536,
    enc_nlayers=40,
    enc_num_heads=48,
    num_cls=10,  # Number of cell types
)
model = CellFMModel(config, for_finetuning=True)

# Load pretrained weights
model.load_weights('./weights/cellfm-800m/model.pt')

# Get dataloaders
train_loader = model.get_dataloader(train_data, batch_size=4)['train']
val_loader = model.get_dataloader(val_data, batch_size=4)['train']

# Train
model.train_model(
    train_dataloader=train_loader,
    val_dataloader=val_loader,
    num_epochs=10,
    learning_rate=1e-4,
)

Perturbation Prediction

from perturblab.model.cellfm import CellFMPerturbationModel
from perturblab.data import PerturbationData

# Load perturbation data
data = PerturbationData.from_anndata(adata)
data.split_data(train=0.7, val=0.15, test=0.15)

# Initialize model
model = CellFMPerturbationModel.from_pretrained('cellfm-800m')

# Initialize perturbation head from dataset
model.init_perturbation_head_from_dataset(data)

# Train (use smaller batch size)
model.train_model(data, epochs=20, batch_size=4)

# Predict
predictions = model.predict_perturbation(data, split='test')

# Evaluate
metrics = model.evaluate(data, split='test')
print(f"Pearson correlation: {metrics['pearson']:.4f}")

Performance Notes

• Memory Requirements: ~3-4GB GPU memory for inference (batch_size=8)
• Recommended Batch Size: 4-8 for training, 8-16 for inference
• Inference Speed: ~2-3x slower than 80M model
• Loading Time: ~5-10 seconds

Model Architecture

• Encoder: Retention-based Transformer (MAE Autobin)
  • Auto-discretization embedding layer
  • 40 retention layers with 48 attention heads each
  • Hidden dimension: 1536
  • Layer normalization and residual connections
• Pre-training: Masked Autoencoding (MAE)
  • Masks 50% of genes
  • Reconstructs masked gene expression
• Output: Gene-level embeddings + CLS token (1536-dimensional)

Comparison with 80M Model

The 800M model provides significantly better representation quality due to its deeper architecture (40 layers vs 2 layers), at the cost of increased computational requirements.

Files

• config.json: Model configuration
• model.pt: Model weights (PyTorch state dict, ~3.0GB)
• README.md: This file
• .gitattributes: Git LFS configuration

Citation

If you use CellFM in your research, please cite:

@article{cellfm2024,
  title={CellFM: A Large-Scale Foundation Model for Single-Cell Transcriptomics},
  author={...},
  journal={...},
  year={2024}
}

References

• Original Repository: https://github.com/biomed-AI/CellFM
• PyTorch Version: https://github.com/biomed-AI/CellFM-torch
• Paper: [Link to paper when available]

Notes

• This model was converted from the original MindSpore checkpoint
• The gene vocabulary (24,072 genes) may differ from the 80M model (27,855 genes)
• For best results, ensure your data preprocessing matches the model's expected input format
• Use CellFMModel.prepare_data() to automatically preprocess your data