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

1. Problématique

L'analyse des données de séquençage ARN à cellule unique (scRNA-seq) a révolutionné la biologie en permettant l'étude de l'hétérogénéité cellulaire. Cependant, l'analyse computationnelle de ces données se heurte à des défis majeurs :

• Dimensionnalité et Sparsité : Les données sont de très haute dimension (des milliers de gènes) et extrêmement clairsemées (beaucoup de zéros).
• Limites des Transformers : Les architectures basées sur les Transformers (comme scBERT, scGPT, Geneformer), bien que performantes, souffrent d'une complexité computationnelle quadratique ($O(N^2)$) due au mécanisme d'attention auto-attention. Cela les rend peu évolutifs pour les séquences longues typiques des transcriptomes cellulaires.
• Dépendances à long terme : Les Transformers ont souvent du mal à capturer efficacement les dépendances à long terme, essentielles pour modéliser les interactions régulatrices entre gènes et les transitions d'états cellulaires.
• Coût : L'entraînement de modèles fondationnels sur des millions de cellules est extrêmement coûteux en temps et en ressources matérielles.

2. Méthodologie

Le papier propose GeneMamba, un modèle fondationnel basé sur les Modèles d'Espace d'État (SSM - State Space Models), spécifiquement une architecture Bi-Mamba (Bidirectional Mamba).

A. Prétraitement et Tokenisation

• Données d'entrée : Le modèle traite une matrice d'expression génique où les lignes sont les cellules et les colonnes les gènes.
• Tokenisation par Rang (Rank-based Discretization) : Contrairement aux méthodes basées sur des "bins" (intervalles de valeurs) ou des projections continues, GeneMamba utilise une approche basée sur le rang.
  • Les valeurs d'expression sont normalisées (correction de la profondeur de séquençage et médiane des valeurs non nulles).
  • Les gènes sont ensuite classés par ordre décroissant d'expression au sein de chaque cellule.
  • Cette approche est robuste aux effets de lot (batch effects) et au bruit, tout en mettant l'accent sur les gènes discriminants plutôt que sur les gènes d'entretien (housekeeping) universellement hauts.
  • Seuls les top 2 048 ou 4 096 gènes sont sélectionnés comme séquence d'entrée.

B. Architecture Bi-Mamba

Le cœur du modèle est le module Bi-Mamba, conçu pour capturer le contexte bidirectionnel (amont et aval) tout en conservant une complexité linéaire.

• Principe SSM : Les SSM modélisent les séquences via un état latent qui évolue dans le temps, permettant une modélisation efficace des dépendances à long terme.
• Bidirectionnalité : Contrairement au Mamba standard (unidirectionnel), GeneMamba traite la séquence dans les deux sens :
  1. Passage Avant : Traitement séquentiel normal.
  2. Passage Arrière : Traitement de la séquence inversée.
  3. Fusion : Les sorties des deux directions sont combinées via un mécanisme de porte (gating mechanism) apprenable. Cela permet au modèle d'intégrer dynamiquement les contextes amont et aval pour chaque token (gène).
• Efficacité : L'architecture offre une complexité computationnelle et mémoire linéaire par rapport à la longueur de la séquence, contrairement au quadratique des Transformers.

C. Objectifs de Pré-entraînement

Le modèle est pré-entraîné avec une fonction de perte composite combinant deux objectifs biologiques :

1. Prédiction du prochain gène ($L_{lang}$) : Une tâche de modélisation de langage (Next Token Prediction) où le modèle prédit le gène suivant dans la séquence de rangs, apprenant ainsi les relations séquentielles.
2. Perte de Voie Génique ($L_{pathway}$) : Une perte de type InfoNCE (Contrastive) qui force la similarité des embeddings des gènes appartenant à la même voie biologique (pathway) et les éloigne des gènes d'autres voies. Cela ancre la structure biologique dans l'espace latent.

3. Contributions Clés

1. Nouvelle Architecture Évolutif : Introduction de GeneMamba, le premier modèle fondationnel pour le scRNA-seq basé sur l'architecture Mamba bidirectionnelle, capable de traiter plus de 50 millions de cellules avec une efficacité accrue.
2. Tokenisation Biologiquement Informée : Adoption d'une tokenisation par rang qui s'aligne naturellement sur les processus biologiques (interactions régulatrices) et réduit la sensibilité au bruit technique.
3. Apprentissage Multi-Objectifs : Intégration de connaissances biologiques (voies géniques) directement dans la fonction de perte pour améliorer l'interprétabilité et la cohérence fonctionnelle des embeddings.
4. Validation à Grande Échelle : Démonstration de la scalabilité et de la performance sur des tâches variées (intégration de lots, annotation, reconstruction) surpassant les modèles basés sur les Transformers.

4. Résultats Expérimentaux

Le modèle a été pré-entraîné sur 29,8 millions de cellules (données CELLxGENE) et évalué sur plusieurs tâches de downstream :

• Intégration Multi-Lots (Multi-batch Integration) :
  • GeneMamba surpasse Harmony, scGPT et GeneFormer.
  • Il parvient à corriger les effets de lot tout en préservant la variation biologique (meilleur score Avg-bio), résolvant le compromis classique entre correction technique et conservation biologique.
• Annotation des Types Cellulaires :
  • Performances compétitives, atteignant la meilleure précision (0.9713) et Macro-F1 (0.7710) sur le dataset hPancreas.
  • Excellente robustesse sur des datasets déséquilibrés (Myeloid_b).
• Reconstruction du Rang Génique :
  • GeneMamba démontre une capacité supérieure à reconstruire l'ordre des gènes par rapport aux modèles unidirectionnels (GeneMamba_U) et aux Transformers (GeneFormer).
  • Métriques clés : Distance de Levenshtein très faible (6 vs 23 pour GeneFormer) et score BLEU élevé (0.987), indiquant une fidélité exceptionnelle à la structure d'entrée.
• Analyse de Corrélation Génique :
  • Le modèle distingue mieux les paires de gènes positives (même voie) des paires négatives que les modèles de base, prouvant sa capacité à capturer la topologie des réseaux de régulation.
• Efficacité Computationnelle :
  • Réduction drastique du temps d'entraînement et des FLOPs par rapport aux Transformers. Par exemple, pour une séquence de 2048, GeneMamba est environ 42% plus rapide que le Transformer équivalent.

5. Signification et Impact

GeneMamba représente une avancée majeure pour la bio-informatique et la biologie computationnelle :

• Passage à l'échelle (Scalability) : Il rend possible l'entraînement de modèles fondationnels sur des ensembles de données massifs (des dizaines de millions de cellules) qui étaient auparavant inaccessibles aux Transformers en raison de contraintes de mémoire et de temps.
• Interprétabilité Biologique : En combinant la dynamique d'état bidirectionnelle avec des pertes basées sur les voies biologiques, le modèle offre non seulement de meilleures performances prédictives, mais aussi des représentations plus interprétables et biologiquement pertinentes.
• Nouveau Standard : Il propose une alternative viable et plus efficace aux architectures Transformer pour l'analyse du transcriptome à cellule unique, ouvrant la voie à de futures découvertes en médecine personnalisée, découverte de médicaments et compréhension des maladies complexes.

En résumé, GeneMamba unifie la profondeur contextuelle des modèles de langage modernes avec l'efficacité computationnelle des SSM, créant une colonne vertébrale robuste pour l'avenir de l'omique à cellule unique.