Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models

1. Énoncé du problème

L'article aborde deux défis critiques dans l'entraînement des modèles de fondation biologiques (spécifiquement pour le séquençage de l'ARN de cellules uniques, ou scRNA-seq) :

• Instabilité de l'entraînement du Softmax : Les modèles Transformer standards utilisent une attention softmax, qui normalise les poids d'attention sur les tokens pour qu'ils somment à 1. Cela crée un mécanisme compétitif « à somme nulle » où l'augmentation de l'attention sur un gène force une diminution sur les autres. Dans les contextes biologiques, les gènes sont souvent co-régulés par plusieurs facteurs indépendants, rendant cette hypothèse compétitive biologiquement irréaliste. De plus, le softmax est sujet à une instabilité numérique (explosion des gradients) à mesure que les longueurs de séquence augmentent, particulièrement lorsque l'entropie de l'attention s'effondre sur quelques tokens.
• Inefficacité face à la variabilité des données biologiques : Les séquences biologiques (cellules) présentent une variabilité extrême en longueur (allant de centaines à milliers de gènes exprimés). Contrairement au texte, où les séquences peuvent être regroupées en longueurs fixes, les échantillons biologiques doivent être traités individuellement avec un remplissage (padding) jusqu'à la longueur maximale d'un lot. Les noyaux d'attention efficaces existants (comme FlashAttention-2 ou FlashSigmoid) manquent soit de support natif pour le remplissage, soit échouent à gérer efficacement la nature « irrégulière » des données biologiques, entraînant un gaspillage computationnel significatif.

2. Méthodologie

Les auteurs proposent de remplacer le softmax par une Attention Sigmoid et de développer un noyau GPU spécialisé et haute performance pour la supporter.

A. Mécanisme d'Attention Sigmoid

Au lieu de normaliser les scores d'attention sur la ligne (softmax), l'attention sigmoid applique une fonction sigmoid élément par élément ($\sigma(x) = \frac{1}{1+e^{-x}}$) aux scores de produit scalaire :
$$\text{SigmoidAttn}(X) = \sigma\left(\frac{QK^\top}{\sqrt{d}} + b\right)V$$

• Découplage : Chaque poids d'attention est calculé indépendamment. Une requête peut porter une attention forte sur plusieurs clés simultanément, modélisant mieux la co-régulation.
• Terme de biais : Un biais $b = -\log(n)$ est ajouté pour approximer le comportement de mise à l'échelle du softmax tout en maintenant l'indépendance élément par élément.

B. Analyse théorique de la stabilité

L'article fournit une justification théorique de la stabilité du sigmoid :

• Dérivées bornées : La dérivée de la fonction sigmoid est globalement bornée par $0,25$($\sigma'(x) = \sigma(x)(1-\sigma(x))$). En revanche, la constante de Lipschitz du softmax croît exponentiellement avec l'amplitude des scores.
• Jacobienne diagonale : La jacobienne de l'attention sigmoid est diagonale, ce qui signifie que les gradients pour différents tokens sont découplés. Le softmax possède une jacobienne dense, couplant tous les gradients de tokens et conduisant à une instabilité lorsque les scores deviennent grands.

C. Implémentation du noyau TritonSigmoid

Pour rendre l'attention sigmoid pratique pour les données biologiques, les auteurs ont développé TritonSigmoid, un noyau GPU personnalisé utilisant le langage Triton. Les innovations clés incluent :

• Support natif du remplissage : Il implémente un calcul par blocs dispersés (block-sparse), sautant entièrement les blocs remplis à la fois du côté des requêtes et des clés pour éviter des FLOPs gaspillés.
• Opérations fusionnées : Il fusionne le calcul de l'attention en un seul noyau, utilisant une approximation matérielle optimisée tanh pour la fonction sigmoid.
• Décomposition de la passe arrière : Il divise la passe arrière en deux noyaux (un pour les gradients de $Q$, un pour les gradients de $K/V$) pour éliminer les opérations atomiques et recalculer les activations à la volée pour une efficacité mémoire.
• Optimisation matérielle : Conçu pour les architectures NVIDIA H100/Blackwell, supportant la précision BF16.

3. Contributions clés

1. Noyau TritonSigmoid : Une implémentation haute performance atteignant 515 TFLOPS sur les GPU NVIDIA H100. Il surpasse FlashAttention-2 (361 TFLOPS) et FlashSigmoid (440 TFLOPS) et est le premier à supporter un remplissage efficace pour les séquences biologiques de longueur variable.
2. Gains de performance empiriques : Il démontre que l'attention sigmoid produit des représentations apprises supérieures au softmax sur six ensembles de données de cellules uniques divers.
3. Stabilité de l'entraînement : Il prouve que l'attention sigmoid prévient la divergence catastrophique de l'entraînement dans des scénarios à contexte long (8K tokens) sans clipping de gradient, un mode de défaillance observé dans les modèles softmax.

4. Résultats

A. Efficacité computationnelle

• Débit : TritonSigmoid atteint jusqu'à un accélération de 7,15x dans la passe avant et 2,44x dans la passe arrière par rapport à l'attention PyTorch standard.
• Efficacité du remplissage : Avec 25 % de remplissage, TritonSigmoid est 29 % plus rapide dans la passe avant que FlashAttention-2. La chute des TFLOPS due au remplissage est minimale (environ 9,3 %), prouvant que le noyau saute efficacement les régions masquées.
• Entraînement de bout en bout : Les modèles entraînés avec une attention sigmoid sont 3 à 10 % plus rapides à entraîner que leurs équivalents softmax, l'écart s'élargissant à mesure que la longueur de contexte augmente (par exemple, 7,5 % plus rapide à un contexte de 8K pour un modèle de 1,4 milliard de paramètres).

B. Qualité du modèle (Apprentissage de représentations)

Évalué sur six ensembles de données de cellules uniques retenus (Cerveau, Sang, Côlon, Poumon, Cœur) :

• Perte de validation : Les modèles sigmoid ont constamment atteint une perte de validation plus faible sur tous les ensembles de données et longueurs de contexte (2K et 4K).
• Séparation des types cellulaires : Les modèles sigmoid ont montré un Décalage Moyen Maximum (MMD) de 25 % plus élevé entre les types cellulaires, indiquant une meilleure séparation dans l'espace d'incrustation (embedding).
• Conservation biologique : Le sigmoid a atteint une cohésion des types cellulaires supérieure (score de Silhouette) sur les 6 ensembles de données et de meilleurs scores agrégés de conservation biologique sur 4/6 ensembles de données.
• Visualisation : Les visualisations UMAP ont confirmé que les incrustations sigmoid regroupent les types cellulaires plus étroitement et distinctement que le softmax.

C. Test de stress de stabilité

• Configuration : Modèles de 160 millions de paramètres entraînés sur des séquences de 8K tokens sans clipping de gradient.
• Résultat :
  • Softmax : Divergence catastrophique autour de l'étape 55k. La norme globale du gradient a explosé de quatre ordres de grandeur ($100 \to 1,6 \times 10^6$), et les scores d'attention ont croisé exponentiellement.
  • Sigmoid : Entraînement stable sur 80k étapes avec des gradients bornés (plage 10–100) et des scores d'attention contrôlés, confirmant les garanties de stabilité théoriques.

5. Signification

Cet article établit l'attention sigmoid comme une alternative théoriquement fondée et empiriquement supérieure au softmax pour les modèles de fondation biologiques.

• Pertinence biologique : Elle s'aligne mieux sur la biologie de la régulation des gènes (co-activation) en supprimant la contrainte compétitive artificielle du softmax.
• Déploiement pratique : En résolvant le problème d'efficacité des séquences de longueur variable via TritonSigmoid, elle permet l'entraînement de modèles biologiques à grande échelle et à contexte long qui étaient auparavant instables ou computationnellement prohibitifs.
• Open Source : Les auteurs publient l'implémentation, offrant un remplacement direct du softmax qui offre un entraînement plus rapide, une meilleure stabilité et des représentations de plus haute qualité pour la communauté de la biologie des cellules uniques.

Disponibilité du code : https://github.com/MSDLLCpapers/triton-sigmoid