🧬 biology

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

This paper introduces sigmoid attention as a theoretically grounded and empirically superior alternative to softmax for single-cell foundation models, demonstrating that it yields better representations, faster training, and superior stability while being supported by a highly efficient Triton kernel.

Original authors: Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

Published 2026-05-01

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are trying to teach a computer to understand the complex "language" of a single cell. In this language, every gene is a word, and the cell's behavior is a sentence made of thousands of these words. To understand the sentence, the computer uses a mechanism called Attention, which decides which words are important to focus on when trying to predict the next one.

For years, the standard way to do this has been Softmax Attention. The authors of this paper argue that for biological data, this old method is like trying to run a race while carrying a heavy, unstable backpack. They propose a new, lighter, and more stable method called Sigmoid Attention.

Here is a breakdown of their findings using simple analogies:

1. The Problem with the Old Way (Softmax)

Think of Softmax Attention like a strict teacher with a limited supply of "attention points" (like 100 points total). If the teacher gives 50 points to the word "Gene A," they must take 50 points away from all the other words.

The Biological Issue: In a cell, genes often work together in teams. A single gene might be controlled by three different regulators at the same time. The "strict teacher" model forces a competition where focusing on one gene means ignoring the others. This doesn't match how biology actually works.
The Stability Issue: As the sentences get longer (cells have thousands of genes), this "strict teacher" gets nervous. The math behind it can become unstable, causing the computer's calculations to explode, like a balloon inflating until it pops. This leads to failed training runs and wasted time.

2. The New Solution (Sigmoid Attention)

The authors propose Sigmoid Attention, which is like a teacher who has an unlimited supply of attention points.

How it works: Instead of taking points away from one word to give to another, this teacher can give a high score to Gene A, Gene B, and Gene C all at the same time without any trade-offs.
The Benefit: This fits biology better because it allows the model to recognize that multiple genes can be active simultaneously. It also acts like a "shock absorber" for the math, preventing the calculations from exploding even when the sentences get very long.

3. The Three Big Wins

The paper claims this new method wins in three specific areas:

A. Better Understanding (Better Representations):
When they tested the new method on six different biological datasets, the models learned to separate different cell types much better. Imagine trying to sort a pile of mixed-up Lego bricks. The old method (Softmax) sorted them okay, but the new method (Sigmoid) sorted them so perfectly that the different colors (cell types) were clearly distinct from one another. They found 25% better separation between cell types.
B. Faster Training:
Because the math is simpler (it doesn't have to do the complex "take from Peter to pay Paul" calculation), the computer works faster. The authors found that models using Sigmoid Attention trained up to 10% faster than those using the old method.
C. Unshakeable Stability:
This is the most dramatic result. The authors put the models through a "stress test" by making them learn on very long sequences without any safety nets (like gradient clipping).
- The Result: The old method (Softmax) crashed and burned, with its internal numbers exploding by a factor of 10,000. The new method (Sigmoid) kept calm, learned smoothly, and never crashed. It's like driving a car off a cliff: the old car fell apart; the new car had a parachute and landed safely.

4. The "Engine" Upgrade (TritonSigmoid)

You might think, "If the math is different, isn't it slower to run?" The authors built a custom, high-speed engine called TritonSigmoid to make sure this new method runs as fast as possible on modern graphics cards (GPUs).

The Challenge: Biological data is "jagged." Some cells have 500 genes, others have 10,000. To process them in a batch, you have to pad the short ones with empty space (like filling a short box with bubble wrap).
The Innovation: Most existing fast engines choke on this "bubble wrap" (padding). The authors' new engine is smart enough to skip the empty space entirely. It runs at 515 TFLOPS (a measure of speed), beating the current top performers.

Summary

The paper concludes that for building "Foundation Models" (super-smart AI) for single-cell biology, Sigmoid Attention is superior to the traditional Softmax. It learns better, trains faster, and doesn't crash when the data gets messy or long. The authors have even released their code so others can use this new, more stable "engine" for their own biological discoveries.

1. Problem Statement

Biological foundation models, particularly those processing single-cell RNA sequencing (scRNA-seq) data, face unique challenges that make standard Transformer architectures (relying on Softmax attention) suboptimal:

Biological Reality vs. Competitive Attention: Softmax attention forces a "zero-sum" competition where increasing attention on one gene token necessitates decreasing it on others. However, gene regulation is often modular and overlapping; a single gene can be co-regulated by multiple transcription factors simultaneously. Softmax cannot model this non-competitive, multi-faceted regulation effectively.
Training Instability: As sequence lengths increase (scRNA-seq profiles often contain thousands of genes), Softmax attention suffers from numerical instability. High attention scores lead to entropy collapse, causing gradient explosions and catastrophic training divergence, especially in long-context scenarios without gradient clipping.
Computational Inefficiency with Variable Lengths: Biological sequences exhibit extreme variability in length (from hundreds to thousands of genes per cell). Unlike text, these sequences cannot be easily packed; they require padding. Existing efficient attention kernels (like FlashAttention-2 or FlashSigmoid) either lack native padding support or fail to handle the "jagged" nature of biological data efficiently, leading to significant computational waste.

2. Methodology

The authors propose replacing Softmax attention with Sigmoid Attention and developing a specialized, high-performance GPU kernel to support it.

A. Sigmoid Attention Mechanism

Instead of normalizing attention scores across tokens via Softmax, Sigmoid attention applies an element-wise sigmoid function ( $\sigma$ ) to the query-key dot products:
$\text{SigmoidAttn}(X) = \sigma\left(\frac{QK^\top}{\sqrt{d}} + b\right)V$

Decoupling: Unlike Softmax, Sigmoid weights are computed independently. A query can attend strongly to multiple keys simultaneously, better modeling co-regulatory gene networks.
Stability Guarantees: Theoretically, the derivative of the sigmoid function is globally bounded ( $\le 0.25$ ), and its Jacobian is diagonal. This eliminates the cross-token gradient coupling and exponential sensitivity to score magnitudes inherent in Softmax, preventing gradient explosions.

B. TritonSigmoid Kernel

To make Sigmoid attention practical for biological data, the authors implemented TritonSigmoid, a custom GPU kernel using the Triton programming language. Key innovations include:

Native Padding Support: The kernel utilizes block-sparse computation to identify and skip fully padded blocks entirely, avoiding wasted FLOPs on masked positions. This is critical for handling the variable-length nature of scRNA-seq data.
Optimized Backward Pass: The backward pass is decomposed into two kernels to eliminate atomic operations, using direct gradient accumulation for better numerical stability and performance.
Hardware Optimization: It leverages hardware-optimized tanh primitives (approximating sigmoid) and transposed memory reads to maximize throughput on NVIDIA H100 GPUs.

3. Key Contributions

Efficient Implementation (TritonSigmoid): A high-performance kernel achieving 515 TFLOPS on H100 GPUs (outperforming FlashAttention-2 at 361 TFLOPS and FlashSigmoid at 440 TFLOPS). It is the first implementation to combine Sigmoid attention with efficient, native padding support.
Empirical Superiority: Demonstrated that Sigmoid attention learns better representations than Softmax across six diverse single-cell datasets, achieving lower validation loss and superior biological structure preservation.
Training Stability: Proved that Sigmoid attention prevents catastrophic training failure in stress tests (long sequences, no gradient clipping) where Softmax diverges.

4. Results

A. Computational Performance

Throughput: TritonSigmoid achieves up to 7.15x speedup in forward passes and 2.44x in backward passes compared to standard PyTorch attention.
Padding Efficiency: With 25% padding, TritonSigmoid is 29% faster in the forward pass than FlashAttention-2. The TFLOPS drop due to padding is minimal (~9.3%), proving the effectiveness of the block-sparse skipping mechanism.
End-to-End Training: Models trained with Sigmoid attention are 4–10% faster to train than Softmax counterparts, with the gap widening as context length increases (e.g., saving 645 GPU hours for a 1.4B parameter model at 8K context).

B. Model Quality (Representation Learning)

Evaluated on six held-out datasets (Brain, Blood, Colon, Lung, Heart):

Validation Loss: Sigmoid models consistently achieved lower validation loss across all datasets and context lengths (2K and 4K).
Cell-Type Separation: Sigmoid achieved a 25% higher Maximum Mean Discrepancy (MMD) between cell types, indicating better distinguishability in the embedding space.
Biological Conservation: Sigmoid models showed superior cell-type cohesion (Silhouette scores) and clustering agreement (Leiden NMI/ARI) compared to Softmax.

C. Training Stability (Stress Test)

In a stress test using 160M parameter models on 8K-token sequences without gradient clipping:

Softmax: Diverged catastrophically around step 55k. The global gradient norm exploded by four orders of magnitude ( $100 \to 1.6 \times 10^6$ ), and attention scores grew exponentially ( $2 \times 10^8$ ).
Sigmoid: Remained stable throughout 80k steps. Gradients remained bounded ( $10–100$ ), and attention scores stayed controlled, confirming theoretical stability predictions.

5. Significance

This paper establishes Sigmoid Attention as a theoretically grounded and empirically superior alternative to Softmax for biological foundation models.

Theoretical Impact: It resolves the inherent instability of Softmax in long-context, high-dimensional biological data by leveraging the bounded derivatives and diagonal Jacobian of Sigmoid.
Practical Impact: The release of TritonSigmoid removes the computational barrier to using Sigmoid attention. By solving the "padding problem" for jagged biological sequences, it enables efficient training of large-scale models on real-world scRNA-seq data.
Domain Advancement: The results suggest that the "competitive" nature of Softmax is ill-suited for gene regulation, and that "non-competitive" attention mechanisms better capture the complex, overlapping regulatory networks defining cellular identity.

The authors have open-sourced their implementation at https://github.com/MSDLLCpapers/triton-sigmoid, facilitating immediate adoption in the bioinformatics and AI communities.