Bessere Modelle, schnellere Trainings: Sigmoid-Attention für Single-Cell-Foundation-Modelle

Zusammenfassung

Die Entwicklung stabiler biologischer Foundation-Modelle erfordert ein Umdenken bei Aufmerksamkeitsmechanismen: Wir stellen fest, dass die Verwendung von Sigmoid-Attention als direkter Ersatz für Softmax-Attention a) bessere gelernte Repräsentationen erzeugt: Auf sechs verschiedenen Einzelzell-Datensätzen erreicht Sigmoid eine 25 % höhere Zelltyp-Trennung, bessere Metriken für Zelltyp-Kohäsion und einen niedrigeren Validierungsverlust, b) schnelleres Training – Modelle mit Sigmoid-Attention trainieren bis zu 10 % schneller als ihre Softmax-Pendants – und c) stabileres Training durch die Beseitigung inhärenter Instabilitätsquellen in der Softmax-Attention. Wir zeigen, dass Sigmoid-Attention global beschränkte Ableitungen (≤ 0,25) besitzt, im Gegensatz zu Softmax, und eine diagonale Jacobi-Matrix-Struktur aufweist, im Kontrast zur dichten Kopplung bei Softmax, was zusammen dazu beiträgt, Trainingsinstabilitäten zu verringern. In Stresstests mit bidirektionalen Aufmerksamkeitsmodellen mit 160M Parametern, die ohne Gradient-Clipping auf Sequenzen mit 8K-Token trainiert wurden, divergiert Softmax katastrophal, wobei die Gradienten um vier Größenordnungen explodieren, während Sigmoid stabil bleibt. Schließlich implementieren und veröffentlichen wir TritonSigmoid, einen effizienten GPU-Kernel, der auf H100-GPUs 515 TFLOPS erreicht und sowohl FlashAttention-2 als auch FlashSigmoid übertrifft, mit nativer Unterstützung für Padding, was für biologische Sequenzen essentiell ist. Unsere Ergebnisse etablieren Sigmoid-Attention als sowohl theoretisch fundiert als auch empirisch überlegen für biologische Foundation-Modelle. Der Code ist verfügbar unter 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

Bessere Modelle, schnellere Trainings: Sigmoid-Attention für Single-Cell-Foundation-Modelle

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

Zusammenfassung

Support