GeneMamba: An Efficient and Effective Foundation Model on Single Cell Data

1. El Problema

El análisis de datos de secuenciación de ARN de células individuales (scRNA-seq) ha revolucionado la biología, permitiendo estudiar la heterogeneidad celular a alta resolución. Sin embargo, el procesamiento computacional de estos datos presenta desafíos significativos:

• Alta dimensionalidad y escasez: Los datos son extremadamente dispersos y ruidosos.
• Limitaciones de las Arquitecturas Transformer: Los modelos basados en Transformers (como scBERT, scGPT, scFoundation), aunque potentes, sufren de una complejidad computacional cuadrática ($O(N^2)$) debido al mecanismo de auto-atención. Esto los hace ineficientes y poco escalables para secuencias largas, características de los transcriptomas de células individuales.
• Dependencias de largo alcance: Los Transformers a menudo tienen dificultades para capturar dependencias de largo alcance esenciales para modelar interacciones regulatorias genéticas y transiciones de estado celular.

2. Metodología: GeneMamba

GeneMamba es un modelo fundamental (foundation model) diseñado específicamente para scRNA-seq que supera las limitaciones de los Transformers mediante el uso de Modelos de Espacio de Estados (SSM), específicamente la arquitectura Mamba, adaptada bidireccionalmente.

A. Preprocesamiento y Tokenización

• Tokenización Basada en Rangos: A diferencia de métodos que usan discretización por "bins" o proyecciones de valores continuos, GeneMamba utiliza una estrategia basada en rangos (similar a Geneformer).
• Proceso:
  1. Se normalizan los valores de expresión genética para cada célula.
  2. Los genes se ordenan en cada célula de mayor a menor expresión.
  3. Se seleccionan los genes con los rangos más altos (top 2048 o 4096) como tokens de entrada. Esto prioriza los genes que distinguen estados celulares sobre los genes de "housekeeping" universalmente altos.

B. Arquitectura Bi-Mamba

El núcleo del modelo es el módulo Bi-Mamba, una extensión bidireccional del modelo Mamba:

• Procesamiento Bidireccional: A diferencia de Mamba estándar (unidireccional), Bi-Mamba procesa la secuencia en dos direcciones: hacia adelante y hacia atrás (reversa).
• Mecanismo de Puerta (Gating): Las salidas de los dos flujos se combinan mediante un mecanismo de puerta aprendible (sigmoid) que pondera dinámicamente la información contextual anterior y posterior.
• Ventaja: Esto permite capturar dependencias contextuales completas (regulaciones genéticas que dependen de genes tanto "antes" como "después" en la secuencia ordenada) con una complejidad computacional lineal ($O(N)$), lo que reduce drásticamente el uso de memoria y tiempo de cómputo.

C. Objetivos de Entrenamiento (Pre-entrenamiento)

El modelo se entrena con una función de pérdida compuesta por dos componentes:

1. Pérdida de Predicción del Siguiente Gen ($L_{lang}$): Similar al entrenamiento de lenguaje natural, predice el siguiente token (gen) basado en el contexto previo dentro de la misma célula.
2. Pérdida de Vías Biológicas (Pathway Loss, $L_{pathway}$): Utiliza una función de pérdida InfoNCE (Contrastive) para forzar que los genes que pertenecen a la misma vía biológica tengan representaciones (embeddings) similares, mientras que los genes de vías diferentes se separen. Esto inyecta conocimiento biológico directo en el modelo.

3. Contribuciones Clave

1. Escalabilidad y Eficiencia: GeneMamba es capaz de procesar conjuntos de datos masivos (más de 50 millones de células) con un costo computacional significativamente menor que los modelos basados en Transformers, gracias a la arquitectura SSM lineal.
2. Arquitectura Bi-Mamba: Introduce un diseño bidireccional eficiente para datos de scRNA-seq, mejorando la capacidad de capturar dependencias globales sin sacrificar la velocidad.
3. Interpretabilidad Biológica: Mediante la tokenización basada en rangos y la pérdida de vías, el modelo logra una mayor interpretabilidad biológica y una mejor organización funcional de los embeddings genéticos.
4. Validación Exhaustiva: Se demuestra la eficacia del modelo en múltiples tareas downstream, superando a modelos de vanguardia como GeneFormer, scGPT y scFoundation.

4. Resultados Experimentales

El modelo fue pre-entrenado en ~30 millones de células únicas (filtradas de CELLxGENE) y evaluado en varias tareas:

• Integración de Múltiples Lotes (Multi-batch Integration):
  • GeneMamba superó a los baselines en la corrección de efectos de lote (batch effects) mientras preservaba la variación biológica.
  • Logró puntuaciones superiores en métricas como Avg-bio (preservación biológica) y Avg-batch (corrección de lote) en conjuntos de datos como PBMC12k, COVID-19 y Perirhinal Cortex.
• Anotación de Tipos Celulares:
  • En tareas de clasificación supervisada, GeneMamba alcanzó la mayor precisión y puntuación Macro-F1 en conjuntos de datos como hPancreas (97.13% de precisión) y Myeloid_b, demostrando su capacidad para distinguir tipos celulares sutiles.
• Reconstrucción de Rangos Genéticos:
  • En experimentos de reconstrucción, GeneMamba mostró una fidelidad superior al reconstruir el orden de expresión de los genes.
  • Superó a GeneFormer y a una versión unidireccional de Mamba (GeneMamba_U) en métricas como BLEU y Spearman, confirmando que el contexto bidireccional es crucial para entender la estructura de los datos.
• Análisis de Correlación Genética:
  • El modelo demostró una mejor capacidad para distinguir pares de genes positivos (misma vía) de negativos, manteniendo la topología de las redes de interacción genética.

5. Significado e Impacto

GeneMamba representa un cambio de paradigma en el modelado de datos de células individuales:

• Viabilidad de Modelos Grandes: Demuestra que es posible entrenar modelos fundamentales a escala masiva (decenas de millones de células) que antes eran prohibitivos debido a las limitaciones de memoria de los Transformers.
• Eficiencia de Recursos: Reduce el tiempo de entrenamiento y el consumo de energía, haciendo que el análisis de alto rendimiento sea más accesible.
• Fundamento para la Investigación Futura: Proporciona una base robusta y escalable para futuras aplicaciones en biología y biomedicina, incluyendo el descubrimiento de fármacos, el modelado de enfermedades y la comprensión de la regulación génica, al tiempo que ofrece una explicabilidad biológica superior gracias a su diseño informado por la biología.

En resumen, GeneMamba combina la eficiencia computacional de los Modelos de Espacio de Estados con la profundidad contextual bidireccional y el conocimiento biológico, estableciendo un nuevo estándar para los modelos fundamentales en la transcriptómica de células individuales.