Résumé technique : Lois d'échelle pour les Transformers de reconstruction masquée sur la transcriptomique en cellule unique

Énoncé du problème

Bien que les lois d'échelle neuronales — les relations de type loi de puissance entre la perte, la taille du modèle et le volume de données — soient bien documentées pour les Transformers de langage et de vision, leur existence et leur comportement dans le domaine de la génomique en cellule unique restent largement inexplorés. La communauté de la cellule unique a développé de grands modèles de fondation (par exemple, scVI, scGPT, Geneformer), mais ceux-ci sont typiquement évalués sur des tâches en aval (annotation, correction de lot) plutôt que sur la relation fondamentale entre la perte de pré-entraînement et l'échelle du modèle. Sans lois d'échelle établies, les praticiens ne peuvent pas déterminer si les goulots d'étrangle de performance sont dus à une capacité de modèle insuffisante, au volume de données ou au bruit biologique/technique intrinsèque. De plus, le plancher de perte asymptotique des lois d'échelle offre un potentiel de limite supérieure empirique sur l'erreur réductible, ce qui peut être lié à l'entropie intrinsèque des données d'expression génique.

Méthodologie

Les auteurs présentent la première étude systématique du comportement d'échelle des Transformers de reconstruction masquée entraînés sur des données de séquençage d'ARN en cellule unique (scRNA-seq).

• Source de données : Les données sont issues du CELLxGENE Census. L'expérience principale utilise un sous-ensemble standardisé de 200 000 cellules et 512 gènes hautement variables (HVG). Des analyses supplémentaires utilisent des sous-ensembles à vocabulaire fixe avec des nombres de cellules variables ($D \in \{25k, 50k, 100k, 200k\}$).
• Prétraitement : Le pipeline implique une normalisation de la taille de la bibliothèque (somme cible de $10^4$), une transformation log(1+x), et une division stratifiée en ensembles d'entraînement (90 %), de validation (5 %) et de test (5 %).
• Architecture : Un encodeur Transformer invariant par permutation est employé, omettant les encodages positionnels car l'ordre des gènes n'a pas de signification biologique. Les entrées consistent en des plongements d'identité de gènes, des valeurs d'expression observées (ou un jeton de masque appris), et un masque binaire. Le modèle utilise une normalisation de couche pré-établie (Pre-LN) et une tête de prédiction linéaire pour reconstruire les valeurs de gènes masquées.
• Objectif d'entraînement : Le modèle est entraîné selon un objectif de reconstruction masquée, minimisant l'erreur quadratique moyenne (MSE) entre les valeurs d'expression prédites et réelles aux positions masquées (taux de masquage de 15 %).
• Conception expérimentale :
  • Balayage de paramètres : Six tailles de modèles (du XXS au L) allant de 533 à $1,0 \times 10^8$ paramètres sont entraînées sur le jeu de données standardisé $V=512, D=200k$. Chaque taille est exécutée sur trois graines aléatoires pendant 60 000 étapes.
  • Balayage de la taille des données : Une analyse de taille de données fixe (taille de modèle $M$) faisant varier la taille du jeu de données ($D$) tout en maintenant le nombre total d'étapes d'entraînement constant.
  • Diagnostics : Vérifications de l'hétérogénéité des résidus inter-gènes et calibrages directs de la log-vraisemblance négative (NLL) gaussienne à petite échelle.
• Formulation de la loi d'échelle : La relation entre la MSE de validation ($L$) et le nombre de paramètres ($P$) est modélisée par $L(P) = aP^{-\alpha} + c$, où $c$ représente le plancher de perte irréductible. Les ajustements sont effectués via une recherche par grille sur $c$ et une régression des moindres carrés ordinaires sur la forme linéarisée.

Contributions clés

1. Première étude systématique : Ce travail constitue la première investigation systématique des lois d'échelle neuronales pour le pré-entraînement par reconstruction masquée sur les données scRNA-seq.
2. Échelle démontrée : Il démontre qu'une mise à l'échelle de type loi de puissance émerge dans un balayage standardisé $V=512, D=200k$, l'évaluation sur l'ensemble de test suivant la même tendance qualitative que la validation.
3. Analyse complète : L'étude inclut des analyses supplémentaires de $V$ fixe, comprenant un balayage de la taille des données, un diagnostic de l'hétérogénéité des résidus inter-gènes, et un calibrage de la NLL gaussienne directe à petite échelle.
4. Estimation préliminaire de l'entropie : Elle fournit une estimation préliminaire du plancher d'incertitude asymptotique, interprétée en unités d'information sous des hypothèses spécifiques.
5. Feuille de route pour la rigueur : Elle identifie les ingrédients expérimentaux nécessaires (balayages appariés $V/D$, objectifs basés sur la vraisemblance, comptabilité de calcul) requis pour passer des courbes perte-vs-paramètres aux estimations rigoureuses de l'entropie transcriptomique.

Résultats

• Échelle des paramètres : Dans le balayage standardisé, la MSE de validation diminue régulièrement du modèle XXS au modèle M. Le modèle L ($\sim 10^8$ paramètres) ne montre aucune amélioration par rapport au modèle M ($\sim 1,9 \times 10^7$ paramètres), indiquant que la courbe de perte a commencé à s'aplatir.
• Exposants d'échelle : L'exposant de la loi de puissance ajusté est $\alpha \approx 0,266$ (validation) et $\alpha \approx 0,268$ (test), avec un $R^2$ d'environ 0,86. C'est sensiblement plus grand que les exposants rapportés pour les modèles de langage autorégressifs ($\alpha \approx 0,076$), mais qualitativement cohérent avec l'existence d'une mise à l'échelle de type loi de puissance.
• Plancher irréductible : Le plancher de perte irréductible $c$ est estimé à $\approx 1,44$ (MSE de validation) et $\approx 1,46$ (MSE de test).
• Effets de la taille des données : Un balayage de la taille des données à $V$ fixe (variant $D$ à étapes fixes) produit un ajustement de validation presque plat ($\beta \approx -0,002$), suggérant qu'à nombre d'étapes d'entraînement fixe, de plus grands ensembles de données sont vus pour moins d'époques effectives, empêissant l'émergence d'une loi $L(D)$ propre dans cette configuration spécifique.
• Hétérogénéité des résidus : Les erreurs résiduelles sont fortement dépendantes des gènes. La corrélation entre l'expression moyenne et la MSE par gène est de 0,892, et le coefficient de variation entre les gènes est de 0,730, indiquant que l'hypothèse gaussienne homoscédastique est imparfaite.
• Estimation de l'entropie : Sous une approximation gaussienne homoscédastique, le plancher asymptotique correspond à environ 2,30 bits d'incertitude irréductible par position de gène masqué. Cette estimation est dérivée à la fois du plancher de la MSE et d'un plancher de la NLL gaussienne converti, bien que les auteurs notent que ces deux éléments ne sont pas totalement indépendants en raison de la dérivation a posteriori.

Signification et affirmations

Le papier prétend établir un cadre quantitatif pour raisonner sur le comportement d'échelle des modèles de fondation en cellule unique. Sa principale importance réside dans :

• Orientation pratique : Il suggère que pour des jeux de données d'environ 200 000 cellules, des modèles de quelques dizaines de millions de paramètres sont suffisants pour capturer la structure apprenable, avec des rendements décroissants au-delà de $\sim 20\text{M}$ paramètres dans la configuration actuelle.
• Clarté méthodologique : Il souligne que le volume de données, le budget d'optimisation et l'interprétation de la perte interagissent de manière complexe ; une analyse propre de la loi des données nécessite des budgets d'optimisation appariés entre les tailles de jeux de données, ce qui n'a pas été réalisé dans le balayage à étapes fixes.
• Aperçu biologique : Le plancher de perte irréductible ($\sim 2,3$ bits) est présenté non pas comme une constante biologique universelle, mais comme une quantité spécifique à la tâche reflétant les effets combinés du bruit biologique, du bruit technique et de la perte d'information du pipeline de prétraitement spécifique (transformation log, sélection des HVG).
• Direction future : Le travail trace la voie selon laquelle des estimations rigoureuses de l'entropie transcriptomique nécessiteront des travaux futurs impliquant des matrices expérimentales $V/D/C$ appariées, des objectifs natifs basés sur la vraisemblance (prédiction de $\mu$ et $\sigma^2$), et une comptabilité de calcul complète.

Les auteurs maintiennent un ton modeste, soulignant que le chiffre de 2,3 bits est un benchmark préliminaire sous des hypothèses spécifiques et que l'étude actuelle ne permet pas encore de distinguer définitivement les planchers limités par les données de ceux limités par les modèles.