From Tokens to Cells: Foundation Models for Single-Cell Biology - Akram Baharlouei, Altos Labs
Akram Baharlouei from Altos Labs discusses the engineering hurdles in developing foundation models for single-cell biology. The talk covers the importance of single-cell analysis for cellular rejuvenation and drug discovery, the complexities and challenges of single-cell data (particularly RNA-seq), and a comparative analysis of current foundation model approaches, highlighting the limitations of transformer-based models and the potential of flow matching techniques like PrimeFlow.
Akram Baharlouei, a Machine Learning Engineer at Altos Labs, presents an engineering perspective on building foundation models for single-cell biology, emphasizing the goal of restoring cell health and resilience to combat age-related diseases.
Why Single-Cell Biology Matters
Single-cell biology offers a granular understanding of life, moving beyond bulk analysis to individual cellular states. A prime example is Shinya Yamanaka’s discovery of Yamanaka factors in 2006, which demonstrated the ability to reprogram aged skin cells into embryonic stem cell-like states. This breakthrough, recognized with a Nobel Prize, opened avenues for regenerative medicine, allowing for tissue regeneration and partial reprogramming to only change cellular age, aiming to restore youthful functions and potentially treat age-related diseases with therapies like mRNA medicine. Efforts like the Human Cell Atlas project aim to map every human cell, contributing to the long-term vision of virtual cell, tissue, and human models, or “digital twins.” This holistic understanding is critical for overcoming the paradoxical “Moore’s Law in reverse” observed in drug development, where drug output declines despite technological advancements. AI-driven virtual models hold the promise of significantly reducing the 10-year, billion-dollar drug development pipeline.
Measuring Single-Cell Data: Modalities and Challenges
Studying a single cell reveals immense complexity. Various modalities capture different aspects:
- Genome: DNA sequencing, largely consistent across cells.
- RNA-seq (Transcriptomics): Measures gene expression profiles (active genes and their quantities) at a given time. This is a 20,000-scale matrix per cell. It’s the most widely used modality for foundation models due to easier measurement and scalability (e.g., PCR technologies) leading to datasets of tens of millions to billions of cells.
- Proteomics: Measures proteins, which perform most bodily functions. However, it’s very hard and low-throughput to measure.
- Morphology: Captures cell shape and structure, useful for spatial dimension analysis through imaging technologies.
Despite the abundance of RNA-seq data, it presents significant challenges for foundation models:
- Heterogeneity and Noise: Even identical cells measured at the same time can yield different results due to inherent biological processes (e.g., cell cycles, bursts of activity) and technical variations between labs and machines.
- Snapshot Nature: Current sequencing methods capture only a “snapshot” of a dynamic cellular “movie,” making it difficult to infer continuous processes.
These data characteristics highlight a need for not just more data, but also higher quality and more realistic measurement technologies.
Foundation Models for Single-Cell Data
Considering the high-dimensional, multi-state nature of single-cell data, some speculate that quantum computing might be a natural fit, but current efforts focus on state-of-the-art AI.
Transformer-Based Models
Models like SCGPT and Informer adapt BERT-style architectures from natural language processing:
- Representation: Cells are treated like “sentences,” and genes as “tokens.”
- Training: They mask certain genes and predict their counts, using bidirectional attention to understand gene relationships.
- Applications: Used for tasks like cell type prediction and perturbation response modeling.
- Limitations: These models often compress high-dimensional single-cell data into a latent vector, leading to loss of information. Benchmarking studies (e.g., NeurIPS papers on multimodal data and perturbation response) show that despite high computational cost, these complex models sometimes perform on par with or even worse than simpler linear models in the single-cell domain, unlike their success in language and imaging.
Flow Matching-Based Models
Models such as PrimeFlow represent a newer approach:
- Approach: Instead of compressing data, they aim to predict and match the distribution of the data, often starting from Gaussian noise.
- Performance: These models appear to perform better than autoregressive or autoencoder-based models, as they capture the underlying data distribution more effectively rather than just predicting means.
- Example: Visualizations show PrimeFlow effectively matching ground truth distributions, while other models might only approximate the mean.
Key Takeaways
- Single-Cell Importance: It’s crucial for cellular rejuvenation, realizing “digital human” concepts, and accelerating drug development.
- Data Quality: While RNA-seq provides scalable data, advancements are needed in measuring other modalities and, critically, in improving the quality and realism of data acquisition to reflect biological processes more accurately, rather than just scaling current noisy methods.
- Model Direction: Flow matching models currently show more promise for single-cell data by focusing on distribution matching. Achieving generalizability and scaling these models effectively for unseen data will require massive improvements in both the quantity and quality of biological data.
References: The speaker highlighted recent work including “SCGENE scope,” “perturbation” (both from NeurIPS last year), and “PrimeFlow.”