Meilleurs modèles, entraînement plus rapide : l'attention sigmoïde pour les modèles de fondation en biologie cellulaire unique

Résumé

L'entraînement de modèles de fondation biologiques stables nécessite de repenser les mécanismes d'attention : nous constatons que l'utilisation de l'attention sigmoïde en remplacement direct de l'attention softmax a) produit de meilleures représentations apprises : sur six jeux de données unicellulaires diversifiés, la sigmoïde permet d'atteindre une séparation des types cellulaires 25 % supérieure, de meilleures métriques de cohésion des types cellulaires et une perte de validation plus faible, b) un entraînement plus rapide, les modèles avec attention sigmoïde s'entraînent jusqu'à 10 % plus vite que leurs équivalents softmax, et c) un entraînement plus stable en éliminant les sources intrinsèques d'instabilité de l'attention softmax. Nous établissons que l'attention sigmoïde possède des dérivées globalement bornées (≤ 0,25) contrairement à la softmax, et une structure jacobienne diagonale contrastant avec le couplage dense de la softmax, ce qui contribue ensemble à atténuer les instabilités d'entraînement. Dans des tests de stress sur des modèles d'attention bidirectionnelle de 160 millions de paramètres entraînés sans écrétage du gradient sur des séquences de 8 000 tokens, la softmax diverge de manière catastrophique, avec des gradients qui explosent de quatre ordres de grandeur, tandis que la sigmoïde reste stable. Enfin, nous implémentons et ouvrons le code de TritonSigmoid, un noyau GPU efficace qui atteint 515 TFLOPS sur les GPU H100, surpassant à la fois FlashAttention-2 et FlashSigmoid, avec une prise en charge native du padding, essentielle pour les séquences biologiques. Nos résultats établissent l'attention sigmoïde comme étant à la fois théoriquement fondée et empiriquement supérieure pour les modèles de fondation biologiques. Le code est disponible à l'adresse https://github.com/MSDLLCpapers/triton-sigmoid.

English

Training stable biological foundation models requires rethinking attention mechanisms: we find that using sigmoid attention as a drop in replacement for softmax attention a) produces better learned representations: on six diverse single-cell datasets, sigmoid achieves 25% higher cell-type separation, better cell-type cohesion metrics, and lower validation loss, b) faster training, models with sigmoid attention train up to 10% faster than their softmax counterparts, and c) more stable training by eliminating inherent sources of instability in softmax attention. We establish that sigmoid attention has globally bounded derivatives (leq 0.25) as opposed to softmax, and a diagonal Jacobian structure in contrast with softmax's dense coupling, which together help alleviate training instabilities. In stress tests on 160M-parameter bidirectional attention models trained without gradient clipping on 8K-token sequences, softmax diverges catastrophically, with gradients exploding by four orders of magnitude, while sigmoid remains stable. Finally, we implement and open-source TritonSigmoid, an efficient GPU kernel that achieves 515 TFLOPS on H100 GPUs, outperforming both FlashAttention-2 and FlashSigmoid, with native padding support, which is essential for biological sequences. Our results establish sigmoid attention as both theoretically grounded and empirically superior for biological foundation models. Code is available at https://github.com/MSDLLCpapers/triton-sigmoid

Meilleurs modèles, entraînement plus rapide : l'attention sigmoïde pour les modèles de fondation en biologie cellulaire unique

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

Résumé

Support