HEIMDALL: Disentangling tokenizer design for robust transfer in single-cell foundation models

Ce papier présente HEIMDALL, un cadre unifié qui démontre que la conception des tokenizers, et non l'architecture du modèle, est déterminante pour assurer la robustesse et la transférabilité des modèles de fondation en biologie cellulaire face aux décalages de distribution.

L'essentiel

🧬 Le Problème : Comment parler aux robots du vivant ?

Imaginez que vous voulez enseigner à un robot (une intelligence artificielle) à comprendre le corps humain. Pour cela, vous lui donnez des données sur des milliards de cellules. Chaque cellule est comme une petite usine qui produit des milliers de produits chimiques différents (les gènes).

Le problème, c'est que les cellules n'ont pas de "liste de courses" standardisée.

• Dans un livre, les mots sont toujours dans le même ordre.
• Dans une image, les pixels sont toujours alignés en grille.
• Mais dans une cellule, les gènes sont un sac de billes en vrac, sans ordre précis, avec des quantités variables.

Pour que le robot comprenne, il faut transformer ces données brutes en une "phrase" qu'il peut lire. C'est ce qu'on appelle le tokeniseur (ou "tokeniseur"). C'est comme un traducteur qui convertit le langage cellulaire en langage robotique.

Jusqu'à présent, chaque équipe de chercheurs inventait son propre traducteur, au hasard, comme si on essayait de cuisiner un gâteau sans recette précise. Certains gagnaient, d'autres perdaient, et personne ne savait vraiment pourquoi.

🔍 La Solution : HEIMDALL, le "Démonteur de Moteurs"

Les auteurs de cet article ont créé un outil appelé HEIMDALL. Imaginez HEIMDALL comme un atelier de mécanique ultra-perfectionné où l'on peut démonter n'importe quel moteur de voiture (ici, n'importe quel modèle d'IA) pour voir exactement quelle pièce fait fonctionner le véhicule.

Au lieu de comparer deux voitures entières (qui ont des pneus, des moteurs et des carrosseries différents), HEIMDALL permet de comparer uniquement le système d'injection de carburant (le tokeniseur), tout en gardant le reste identique.

HEIMDALL décompose le traducteur en trois pièces maîtresses :

1. L'Identité (FG) : Comment le robot reconnaît le nom du gène ? (Est-ce qu'il lit la séquence d'ADN ? Ou utilise-t-il une liste de correspondance ?)
2. L'Expression (FE) : Comment le robot comprend la quantité ? (Est-ce qu'il voit un chiffre précis, ou le classe-t-il dans des catégories comme "peu", "moyen", "beaucoup" ?)
3. L'Ordre (FC) : Dans quel ordre le robot lit-il la liste des gènes ? (Par ordre alphabétique ? Par ordre de quantité ? Au hasard ?)

🌍 Les Découvertes : Quand tout va bien, et quand ça dérape

Les chercheurs ont testé ces traducteurs dans quatre situations différentes, comme un test de conduite sur différents terrains :

1. Sur la même route (Données identiques)
Si le robot s'entraîne sur des cellules de foie et est testé sur d'autres cellules de foie, peu importe le traducteur utilisé. Tout le monde arrive à destination. C'est comme conduire sur une autoroute lisse : n'importe quelle voiture va bien.

2. Sur une nouvelle route (Changement de tissu)
Si le robot s'entraîne sur le foie mais doit conduire dans le cerveau, les choses changent. Là, le choix du traducteur devient crucial. Certains traducteurs perdent le nord, d'autres arrivent à destination.

3. Dans un autre pays (Changement d'espèce)
C'est le test le plus dur : entraîner le robot sur des cellules humaines et le tester sur des cellules de souris.

• La découverte clé : La plupart des traducteurs échouent lamentablement car ils ne comprennent pas que "Gène A" chez l'homme est le cousin de "Gène B" chez la souris.
• Le gagnant : Un traducteur spécial (UCE) qui lit directement la "recette" chimique (la séquence d'ADN) du gène, plutôt que de se fier à un nom. C'est comme si, au lieu de lire "Pomme" (qui s'écrit différemment en français et en espagnol), le robot lisait la forme et le goût du fruit. Il comprend que c'est le même fruit, peu importe la langue.

4. Avec une carte incomplète (Changement de gènes)
Parfois, on utilise des appareils de mesure qui ne voient que 100 gènes au lieu de 20 000.

• La découverte : Le traducteur qui utilise une "mémoire de groupe" (savoir quels gènes travaillent souvent ensemble) fonctionne beaucoup mieux. C'est comme deviner le mot manquant dans une phrase en regardant le contexte, même si le mot n'est pas écrit.

💡 La Leçon Principale : Pas de solution unique, mais des principes clairs

L'article nous apprend une chose fondamentale : il n'existe pas de "traducteur universel parfait" qui fonctionne dans tous les cas.

Cependant, HEIMDALL a identifié les trois axes sur lesquels il faut jouer pour réussir :

1. L'identité du gène : Il faut utiliser des informations biologiques profondes (comme la séquence d'ADN) pour que le robot comprenne les liens entre les espèces.
2. L'expression : Il faut coder intelligemment les quantités (pas juste des zéros et des uns).
3. L'ordre : L'ordre dans lequel on présente les gènes au robot compte énormément. Les trier par "importance" (quantité) aide souvent le robot à mieux comprendre.

🚀 Conclusion : Vers des robots plus robustes

En résumé, HEIMDALL est une boîte à outils qui permet aux scientifiques de ne plus deviner, mais de construire des intelligences artificielles pour la biologie de manière rationnelle.

Au lieu de dire "J'ai choisi ce modèle parce qu'il est à la mode", les chercheurs peuvent maintenant dire : "J'ai choisi ce traducteur parce que sa pièce 'Identité' est conçue pour gérer les changements d'espèces, et sa pièce 'Ordre' est optimisée pour les tissus".

C'est une étape majeure pour rendre les outils d'IA fiables, non seulement pour la recherche, mais aussi pour la médecine future, où l'on voudra peut-être utiliser un modèle entraîné sur des souris pour aider des patients humains, ou analyser des tissus rares avec des équipements limités. HEIMDALL est la boussole qui permet de ne pas se perdre dans ce voyage.

Résumé technique

1. Problématique

Les modèles de fondation pour l'ARN de séquençage à cellule unique (scRNA-seq), ou scFMs, émergent comme des outils puissants pour l'analyse biologique. Cependant, leur performance est souvent incohérente, en particulier lors de tâches de transfert (par exemple, d'un tissu à un autre, d'une espèce à une autre, ou avec des panels de gènes différents).

Le problème central identifié par les auteurs réside dans le tokenization (la manière dont le profil d'expression génique d'une cellule est converti en entrées de modèle). Contrairement au texte ou aux images, les données de cellule unique n'ont pas de schéma de tokenisation canonique. Les choix de conception actuels dans les scFMs sont souvent :

• Heuristiques et entrelacés : Il est difficile de distinguer si une performance provient de l'architecture du modèle, des données d'entraînement ou de la représentation des entrées.
• Manque d'évaluation systématique : Les benchmarks existants comparent des modèles pré-entraînés complets, ce qui empêche d'attribuer les différences de performance à des composants spécifiques du tokeniseur.
• Fragilité au décalage de distribution : Les modèles peinent souvent à généraliser dans des scénarios réels (nouveaux tissus, nouvelles espèces, panels de gènes limités), et la cause racine de cette fragilité n'est pas bien comprise.

2. Méthodologie : Le cadre HEIMDALL

Pour résoudre ces problèmes, les auteurs introduisent HEIMDALL, un cadre unifié et modulaire conçu pour disséquer et redéfinir les tokeniseurs dans les scFMs.

Architecture du Tokeniseur HEIMDALL
Le cadre décompose tout tokeniseur en trois modules fonctionnels distincts, permettant une réimplémentation et une comparaison équitable :

1. $F_G$ (Encodage de l'identité du gène) : Encode l'identité du gène en intégrant des connaissances biologiques préalables (ex: séquences protéiques, co-expression, descriptions textuelles).
2. $F_E$ (Encodage de l'expression) : Encode la valeur d'expression du gène en intégrant des statistiques au niveau de la cellule ou du jeu de données (ex: binning, encodage continu, ou absence d'encodage explicite).
3. $F_C$ (Construction de la cellule) : Intègre les sorties de $F_G$ et $F_E$ pour assembler la représentation finale de la cellule. Ce module est lui-même subdivisé en :
   • ORDER : Définit l'ordre intrinsèque des tokens (ex: tri par expression, ordre chromosomique, aléatoire).
   • SEQUENCE : Sélectionne quels gènes inclure et construit la séquence (ex: troncature, échantillonnage pondéré).
   • REDUCE : Combine les encodages d'identité et d'expression (ex: somme, identité).

Protocole Expérimental

• Réimplémentation : Les auteurs ont réimplémenté les tokeniseurs de cinq scFMs leaders (scGPT, Geneformer, scFoundation, scBERT, UCE) dans ce cadre commun.
• Contrôle des variables : Pour isoler l'effet du tokeniseur, tous les modèles sont entraînés de zéro (sans pré-entraînement préalable) avec une architecture de transformateur minimale et fixe, sur des hyperparamètres identiques.
• Ablations : Le cadre permet de remplacer des modules spécifiques (ex: remplacer l'encodage d'identité de Geneformer par celui de scBERT) pour identifier les composants responsables de la performance.

3. Contributions Clés

1. Cadre de désassemblage modulaire : HEIMDALL fournit une infrastructure réutilisable pour décomposer les stratégies de tokenisation complexes en axes de conception testables.
2. Preuve de l'importance du tokeniseur : L'étude démontre que le choix du tokeniseur est un déterminant critique de la généralisation, souvent plus important que le pré-entraînement ou la taille du modèle dans des scénarios de décalage de distribution.
3. Identification des axes de conception critiques : L'analyse révèle que la robustesse dépend principalement de trois axes : l'identité du gène ($F_G$), l'encodage de l'expression ($F_E$) et l'ordre des tokens ($ORDER$).
4. Tokeniseurs hybrides : Les auteurs montrent qu'en combinant les meilleurs éléments de différents tokeniseurs (ex: l'encodage d'identité de scBERT avec l'ordre de Geneformer), on peut créer des modèles surpassant les stratégies individuelles existantes.

4. Résultats Principaux

Les expériences ont été menées sur quatre benchmarks de transfert difficiles :

• Généralisation inter-tissus (Cross-tissue) :

  • Entraînement sur le côlon/intestin, test sur le cerveau.
  • Résultat : Dans des conditions identiques (train/test matchés), le choix du tokeniseur a peu d'impact. Cependant, sous décalage de distribution, le tokeniseur devient décisif.
  • Découverte : La performance est principalement pilotée par le module ORDER (tri par expression), qui injecte implicitement des informations d'expression, même sans encodage explicite ($F_E$).

• Généralisation inter-espèces (Cross-species) :

  • Entraînement sur l'humain, test sur la souris.
  • Résultat : Les tokeniseurs utilisant des embeddings basés sur la séquence protéique (ESM2 dans UCE) sont les seuls à performer sans mapping d'orthologie.
  • Découverte : Le mapping par orthologie améliore tous les modèles, mais un $F_G$ agnostique à l'espèce (basé sur la séquence) reste crucial pour les espèces non modélisées ou mal annotées.

• Généralisation aux panels de gènes (Spatial Transcriptomics) :

  • Scénarios où les gènes d'entraînement et de test ne se chevauchent que partiellement.
  • Résultat : scBERT-tok (utilisant des embeddings Gene2vec basés sur la co-expression) surpasse significativement les autres.
  • Découverte : Les priors de co-expression dans $F_G$ stabilisent les représentations pour les gènes non vus à l'entraînement. L'encodage de l'expression ($F_E$) et l'ordre ($ORDER$) apportent également des gains significatifs.

• Prédiction de perturbation inverse (Reverse Perturbation) :

  • Prédire la perturbation à partir de l'état cellulaire cible.
  • Résultat : scBERT-tok est le meilleur. L'ablation montre que l'ajout d'encodages d'expression explicites ($F_E$) et d'un ordre basé sur l'expression ($ORDER$) est essentiel pour les modèles qui en manquent par défaut (comme UCE).

5. Signification et Conclusion

L'article établit que la tokenisation est un axe de conception critique et sous-estimé dans les modèles de fondation pour la biologie cellulaire.

• Universalité vs Interface : Bien que les scFMs visent une transférabilité universelle, leur robustesse dépend d'une interface de tokenisation qui n'est pas universelle mais doit être adaptée aux priors biologiques spécifiques (identité, expression, ordre).
• Implications pratiques : Pour les utilisateurs, cela signifie qu'il n'existe pas un "meilleur" tokeniseur global. Le choix doit être guidé par le scénario de déploiement (ex: utiliser un $F_G$ basé sur la séquence pour les nouvelles espèces, ou un $F_G$ basé sur la co-expression pour les panels de gènes restreints).
• Futur : HEIMDALL fournit les bases pour concevoir des modèles plus robustes en rendant explicites et interchangeables les biais inductifs biologiques introduits lors de la tokenisation. Cela ouvre la voie à des "cellules virtuelles" intégrant des données multi-omiques de manière cohérente.

En résumé, HEIMDALL transforme la tokenisation d'un choix heuristique en une composante ingénierie systématique, essentielle pour débloquer le plein potentiel des modèles de fondation en biologie cellulaire.