Benchmarking Geneformer v1 vs v2 Bio Foundation Models

Before diving into the details, it’s helpful to revisit what bio foundation models (FMs) are and their role in biological research. Bio FMs are pretrained on massive biological datasets, such as DNA sequences or, in this case, RNA-seq data, which allows them to generate rich and informative embeddings.

These embeddings can then be applied to various downstream tasks, ranging from basic analyses like cell type annotation and batch integration to more complex applications in drug discovery, such as biomarker identification and target selection. If you’re looking for a deeper introduction to the concept, a great resource is available here.

‍

Architecture: The Old and the New

Now that we’ve covered the basics, let’s dive into the specifics. Before examining the results, it’s crucial to first understand the architectural and pretraining differences between Geneformer v1 and v2. Full details can be found in the Helical model card.

Geneformer v1 was pretrained on 29.9 million human single-cell RNA sequences from a wide range of tissues, all sourced from public datasets, available here. It comes in two variants, with 6-layer and 12-layer models, both using an input size of 2048 (genes per cell).

While this input size allows for some context awareness, it falls short in representing a full genome, which typically requires around 20,000 genes per cell (though this includes many genes with zero expression counts). To work around this, Geneformer v1 (as well as v2) employs rank value encodings to sort gene expressions, prioritizing those that better distinguish cell states. More details can be found in the full article.

Geneformer v2 was pretrained on a significantly larger corpus of biological data. The authors assembled a pre-training corpus of 103 million human cells, of which a filtered subset of 95 million cells — based on mutational burden — was used for pre-training. While the full corpus isn’t publicly available yet, a list of the datasets used can be found in the preprint.

One of the most notable improvements in v2 is the increased input size of 4096. This expanded input size captures 93% of the Genecorpus, meaning that 93% of the cells had detectable gene expression for fewer than 4096 genes. This extension enhances the model’s ability to capture biological variability while reducing the influence of zero-count genes.

Three variants of Geneformer v2 were introduced: a 12-layer and a 20-layer base model, along with a cancer-specific 12-layer model. The latter was tuned using continual learning on a separate corpus of 14 million cancerous cells. This model was domain-tuned to account for the variational complexity inherent in cancer-specific cellular landscapes, making it more effective at understanding the heterogeneity and distinct characteristics of cancer cell populations.

Another significant enhancement in Geneformer v2 is the introduction of a CLS token appended to all sequences. This addition allows the model to generate two types of embeddings: the traditional cell embedding, which averages the embeddings of each gene in the sequence (as in v1), and a CLS embedding, which outputs the embedding of the CLS token placed at the end of the sequence.

This means v2 can now accumulate and output meaningful information through the CLS token, providing a more comprehensive transcriptome representation.

It’s important to note that both versions of Geneformer are pretrained exclusively on human data. Therefore, the gene-mappings are limited to human RNA-seq data and hence cannot be easily applied to other species.

Both Geneformer v1 and v2 are available for immediate use through the Helical package. You can start using these models with just a few lines of code:

from helical.models.geneformer.model import Geneformer,GeneformerConfig
import anndata as ad

# Geneformer v2 (12-layer base model)
model_config = GeneformerConfig(model_name="gf-12L-95M-i4096", batch_size=10)
geneformer_v2 = Geneformer(model_config=model_config)

# You can use other model names in the config, such as:
# "gf-12L-30M-i2048" (Version 1.0)
# "gf-20L-95M-i4096" (Version 2.0, 20-layer model)
# "gf-12L-95M-i4096-CLcancer" (Version 2.0, Cancer-tuned)

# Example usage
ann_data = ad.read_h5ad("dataset.h5ad")
dataset = geneformer_v2.process_data(ann_data)
embeddings = geneformer_v2.get_embeddings(dataset)
print("Base model embeddings shape:", embeddings.shape)

For more details on how to use the Helical package, you can consult the full documentation here.

Now, without further ado, let’s dive into the benchmarking!

‍