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

1. Probleemstelling

Het artikel behandelt twee kritieke uitdagingen bij het trainen van biologische foundation-modellen (specifiek voor single-cell RNA-sequencing, of scRNA-seq):

• Trainingsinstabiliteit van Softmax: Standaard Transformer-modellen gebruiken softmax-attention, die de attentiewaarden over tokens normaliseert zodat ze optellen tot 1. Dit creëert een "zero-sum" competitief mechanisme waarbij het verhogen van de attentie op het ene gen een vermindering op andere genen forceert. In biologische contexten worden genen vaak geco-reguleerd door meerdere onafhankelijke factoren, waardoor deze competitieve aanname biologisch onrealistisch is. Bovendien is softmax vatbaar voor numerieke instabiliteit (gradient-explosie) naarmate de sequentielengte toeneemt, vooral wanneer de attentie-entropie instort op een paar tokens.
• Inefficiëntie bij Variabiliteit van Biologische Data: Biologische sequenties (cellen) vertonen extreme variabiliteit in lengte (variërend van honderden tot duizenden tot expressie gebrachte genen). In tegenstelling tot tekst, waar sequenties in vaste lengtes kunnen worden gepakt, moeten biologische monsters individueel worden verwerkt met padding tot de maximale lengte in een batch. Bestaande efficiënte attention-kernels (zoals FlashAttention-2 of FlashSigmoid) missen ofwel native padding-ondersteuning of falen in het efficiënt hanteren van het "gezaagde" karakter van biologische data, wat leidt tot aanzienlijke computerverspilling.

2. Methodologie

De auteurs stellen voor om softmax te vervangen door Sigmoid Attention en ontwikkelen een gespecialiseerde, hoogpresterende GPU-kernel om dit te ondersteunen.

A. Sigmoid Attention Mechanisme

In plaats van attentiescores over de rij te normaliseren (softmax), past sigmoid attention een element-voor-element sigmoid-functie ($\sigma(x) = \frac{1}{1+e^{-x}}$) toe op de dot-product scores:
$$\text{SigmoidAttn}(X) = \sigma\left(\frac{QK^\top}{\sqrt{d}} + b\right)V$$

• Ontkoppeling: Elke attentiewaarde wordt onafhankelijk berekend. Een query kan tegelijkertijd sterk attenderen op meerdere keys, wat co-regulatie beter modelleert.
• Bias Term: Een bias $b = -\log(n)$ wordt toegevoegd om de schaalgedraging van softmax te benaderen terwijl de element-voor-element onafhankelijkheid behouden blijft.

B. Theoretische Stabiliteitsanalyse

Het artikel biedt een theoretische rechtvaardiging voor de stabiliteit van sigmoid:

• Begrensde Afgeleiden: De afgeleide van de sigmoid-functie is globaal begrensd door $0.25$($\sigma'(x) = \sigma(x)(1-\sigma(x))$). Daarentegen groeit de Lipschitz-constante van softmax exponentieel met de grootte van de scores.
• Diagonale Jacobiaan: De Jacobiaan van sigmoid attention is diagonaal, wat betekent dat gradienten voor verschillende tokens ontkoppeld zijn. Softmax heeft een dichte Jacobiaan, die alle token-gradienten koppelt en leidt tot instabiliteit wanneer scores groot worden.

C. TritonSigmoid Kernel Implementatie

Om sigmoid attention praktisch te maken voor biologische data, hebben de auteurs TritonSigmoid ontwikkeld, een aangepaste GPU-kernel geschreven in de Triton-taal. Belangrijke innovaties zijn:

• Native Padding-ondersteuning: Het implementeert block-sparse berekening, waarbij volledig gepadded blokken aan zowel de query- als de key-zijde volledig worden overgeslagen om verspilde FLOPs te voorkomen.
• Gefuseerde Operaties: Het fuseert de attention-berekening in een enkele kernel, waarbij een hardware-geoptimaliseerde tanh-benadering wordt gebruikt voor de sigmoid-functie.
• Decompositie van de Backward Pass: Het splitst de backward pass in twee kernels (één voor $Q$-gradienten, één voor $K/V$-gradienten) om atomaire operaties te elimineren en activaties op de vlucht opnieuw te berekenen voor geheugenefficiëntie.
• Hardware-optimalisatie: Ontworpen voor NVIDIA H100/Blackwell-architecturen, met ondersteuning voor BF16-precisie.

3. Belangrijkste Bijdragen

1. TritonSigmoid Kernel: Een hoogpresterende implementatie die 515 TFLOPS bereikt op NVIDIA H100 GPU's. Het presteert beter dan FlashAttention-2 (361 TFLOPS) en FlashSigmoid (440 TFLOPS) en is de eerste die efficiënte padding ondersteunt voor variabele lengte biologische sequenties.
2. Empirische Prestatiewinst: Aangetoond dat sigmoid attention superieure geleerde representaties oplevert ten opzichte van softmax op zes diverse single-cell datasets.
3. Trainingsstabiliteit: Bewezen dat sigmoid attention catastrofale trainingsdivergentie voorkomt in scenario's met lange context (8K tokens) zonder gradient clipping, een falingsmodus die bij softmax-modellen wordt waargenomen.

4. Resultaten

A. Computerefficiëntie

• Doorvoer: TritonSigmoid bereikt tot 7.15x snelheidswinst in de forward pass en 2.44x in de backward pass in vergelijking met standaard PyTorch attention.
• Padding-efficiëntie: Met 25% padding is TritonSigmoid 29% sneller in de forward pass dan FlashAttention-2. De daling in TFLOPS door padding is minimaal (slechts ~9.3%), wat bewijst dat de kernel effectief gemaskeerde gebieden overslaat.
• End-to-End Training: Modellen getraind met sigmoid attention zijn 3–10% sneller te trainen dan hun softmax-gegenen, waarbij de kloof groter wordt naarmate de contextlengte toeneemt (bijv. 7.5% sneller bij 8K context voor een model van 1.4B parameters).

B. Modelkwaliteit (Representatieleer)

Geëvalueerd op zes terughoudende single-cell datasets (Brain, Blood, Colon, Lung, Heart):

• Validatieverlies: Sigmoid-modellen bereikten consistent lager validatieverlies over alle datasets en contextlengtes (2K en 4K).
• Scheiding van Celtypen: Sigmoid-modellen vertoonden een 25% hogere Maximum Mean Discrepancy (MMD) tussen celtypen, wat wijst op betere scheiding in de embedding-ruimte.
• Biologische Behoud: Sigmoid bereikte superieure cohesie van celtypen (Silhouette-score) op alle 6 datasets en betere geaggregeerde biologische behoudsscores op 4/6 datasets.
• Visualisatie: UMAP-visualisaties bevestigden dat sigmoid-embeddings celtypen strakker en duidelijker clusteren dan softmax.

C. Stabiliteitsstress-test

• Opzet: Modellen van 160M parameters getraind op 8K-token sequenties zonder gradient clipping.
• Uitkomst:
  • Softmax: Divergeerde catastrofaal rond stap 55k. De globale gradientnorm explodeerde met vier ordes van grootte ($100 \to 1.6 \times 10^6$) en attentiescores groeiden exponentieel.
  • Sigmoid: Trainde stabiel gedurende 80k stappen met begrensde gradienten (bereik 10–100) en gecontroleerde attentiescores, wat de theoretische stabiliteitsgaranties bevestigt.

5. Betekenis

Dit artikel vestigt sigmoid attention als een theoretisch onderbouwde en empirisch superieur alternatief voor softmax voor biologische foundation-modellen.

• Biologische Relevantie: Het sluit beter aan bij de biologie van genregulatie (co-activatie) door de kunstmatige competitieve beperking van softmax te verwijderen.
• Praktische Implementatie: Door het efficiëntieprobleem van variabele lengte-sequenties op te lossen via TritonSigmoid, maakt het het trainen van grootschalige, lange-context biologische modellen mogelijk die voorheen instabiel of computereconomisch onhaalbaar waren.
• Open Source: De auteurs publiceren de implementatie, die een directe vervanging biedt voor softmax en snellere training, betere stabiliteit en hogere kwaliteit representaties biedt voor de single-cell biologie community.

Code Beschikbaarheid: https://github.com/MSDLLCpapers/triton-sigmoid