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

Vijay Sadashivaiah; Georgios Dasoulas; Judith Mueller; Soumya Ghosh

Source-linked AI summary

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

Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

arXiv:2604.27124v1Apr 29, 2026cs.LGq-bio.QM

TL;DR

Biological foundation models must handle heterogeneous, long cell sequences while avoiding attention-related training instability. This paper evaluates sigmoid attention and develops TritonSigmoid, a padding-aware GPU kernel, finding better representations, faster training, and greater stability than softmax.

Problem
Heterogeneous single-cell sequences require padding and long contexts, while existing sigmoid implementations lack the practical support needed for biological foundation models.
Method
The paper replaces softmax with element-wise sigmoid attention, analyzes its bounded derivatives and diagonal Jacobian, and implements TritonSigmoid with native padding support.
Results
Across six datasets, sigmoid attention yields lower validation loss, 25% higher cell-type separation, stronger cohesion, faster training, and stable stress-test training where softmax diverges.
Takeaways & Limitations
Sigmoid attention is presented as a practical alternative to softmax for biological foundation models, combining representation quality, efficiency, and training stability.
Takeaways & Limitations
A 2,048-token context window truncates genetic context for 43% of cells, motivating longer-context biological models.

Abstract

from arXiv · show

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

1 Introduction

Single-cell foundation models face biological sequence variability and regulatory relationships that challenge conventional softmax attention. This paper investigates sigmoid attention, an efficient and stable alternative, and reports better representations, faster training, and improved stability.

Single-cell models represent each cell’s expressed genes as a sequence, using self-attention to capture co-expression and regulatory relationships.
43% of cells would lose genetic context under a 2,048-token window, while cells must be processed individually with padding because they cannot be split or concatenated.
TritonSigmoid provides a padding-aware GPU kernel that achieves 515 TFLOPS on H100 GPUs, exceeding FlashAttention-2 and FlashSigmoid.
Across six held-out datasets, sigmoid attention achieves lower validation loss, 25% higher cell-type separation, and superior cell-type cohesion.
In 160M-parameter stress tests without gradient clipping on 8K-token sequences, softmax diverges with four-order-of-magnitude gradient explosions while sigmoid remains stable.
The paper releases sigmoid attention as an open-source practical replacement for softmax in biological foundation models.

2 Related Work

Prior biological foundation-model implementations largely target softmax and existing sigmoid kernels lack the padding support or hardware compatibility required by heterogeneous single-cell sequences. The paper positions its work as a systematic biological evaluation and implementation response to these gaps.

Single-cell foundation models learn representations from millions of cells represented as gene-expression-token sequences and trained with masked language modeling.
Biological inputs are heterogeneous and sparse, with heavy-tailed expression, dropout, and library-size variation that can exacerbate softmax instability.
Sigmoid attention offers element-wise bounded derivatives and a diagonal Jacobian, avoiding softmax’s cross-token gradient coupling.
Sigmoid attention had not been systematically studied in biological foundation models, whose sequence lengths range from hundreds to thousands of gene tokens.
Existing sigmoid implementations either lack padding and modern-GPU compatibility or deliver low H100 throughput, leaving no efficient solution combining sigmoid attention with padding support.
TritonSigmoid addresses these limitations using block-sparse computation, fused operations, and backward-pass recomputation within PyTorch.

3 Background on Sigmoid Attention

Sigmoid attention replaces softmax’s across-token normalization with independent element-wise weighting. This decouples token interactions and gives sigmoid a bounded, diagonal derivative structure associated with greater training stability.

Softmax attention applies row-wise normalization, so each query’s attention weights lie on the probability simplex.
Sigmoid attention applies an element-wise sigmoid to attention scores instead of softmax normalization.
Each sigmoid attention weight is computed independently, and the paper sets the bias to b = −log(n) to approximate softmax normalization.
Softmax weights competitively depend on every score in a row, whereas sigmoid weights depend only on each query-key dot product and can attend strongly to multiple keys.
Sigmoid produces a diagonal Jacobian with a maximum derivative of 1/4, independent of input magnitude, unlike softmax’s dense coupling and score-sensitive local Lipschitz behavior.

4 An Efficient Sigmoid Attention Kernel

TritonSigmoid is a padding-aware Triton kernel designed for jagged biological sequences. It combines specialized computation with strong kernel and end-to-end training performance across context lengths and model sizes.

TritonSigmoid builds on FlashAttention-2 and FlashSigmoid strategies while adding native padding support and performance optimizations.
The kernel skips fully padded query and key blocks, reducing computation on masked regions.
Fused attention operations avoid materializing the attention matrix, using a hardware-optimized tanh approximation for sigmoid.
Backward-pass decomposition removes atomic operations, while activation recomputation improves memory efficiency and numerical stability.
At 16K context, TritonSigmoid reaches 515.6 TFLOPS forward and 373.5 TFLOPS backward, outperforming FlashSigmoid, FlashAttention2, and PyTorch attention.
With 25% padding, TritonSigmoid is 29% faster than FlashAttention2 forward and 13% faster backward, while FlashSigmoid cannot handle padded sequences.
End-to-end training is consistently faster than softmax, reaching a 9% speedup for the 160M model at 2K context and 7.5% for the 1.4B model at 8K context.
The advantage grows with context length because attention contributes a larger share of computation at longer sequences, although end-to-end gains are smaller than kernel-level improvements.

5 Sigmoid attention produces better single-cell foundation models

The study evaluates sigmoid attention in single-cell foundation models across representation quality, computational efficiency, and training stability. Sigmoid consistently improves validation and biological-structure metrics, trains faster, and remains stable under stress conditions where softmax diverges.

Experimental setup: The study trains 160M-parameter sigmoid and softmax models on CellxGene at 2K and 4K contexts, evaluating six held-out datasets with complementary representation metrics.The evaluation covers validation loss, SCIB biological-conservation metrics, UMAP visualizations, and MMD-based cell-type separation.
Computational efficiency: Sigmoid reduces projected training compute relative to softmax, with the 1.4B-model advantage increasing from 2.1% at 2K to 7.5% at 8K context.At 4K, measured speedups include 5.1% for 400M, 3.0% for 600M, and 4.0% for 1.4B models.
Validation performance: Sigmoid achieves lower validation loss than softmax across all six datasets and both context lengths, while 4K context outperforms 2K for both mechanisms.All four models converge smoothly, with sigmoid at 4K achieving the lowest training loss.
Biological conservation: Sigmoid achieves the best cell-type cohesion on all six datasets and favors aggregate biological conservation on four of six datasets.For Leiden clustering metrics, sigmoid performs better on four datasets.
Biological conservation: 25% higher mean MMD indicates stronger cell-type separation for sigmoid across all 28 Heart OFT pairwise comparisons.MMD uses 1,000 bootstrap resamples, and higher MMD indicates more distinguishable cell-type representations.
Training stability: Under 8K-token training without gradient clipping, softmax diverges while sigmoid completes 80,000 steps with monotonic loss decrease and bounded gradients.Softmax gradient norms rise from approximately 100 to 1.6 × 10^6, whereas sigmoid remains between 10 and 100.

6 Conclusion

The paper establishes sigmoid attention as a practical alternative to softmax for biological foundation models, combining efficient long-context implementation with improved representations and training stability.

Sigmoid attention learns better representations than softmax across six single-cell datasets, including lower validation loss, 25% higher cell-type separation, and stronger cell-type cohesion.
TritonSigmoid addresses long contexts and extreme sequence-length variability with a padding-aware Triton kernel achieving 515 TFLOPS on H100 GPUs.
Together, the results position sigmoid attention as efficient for jagged long-context sequences, stronger in representation quality, and more stable to train.

A.1 Theoretical Foundations

This section frames the theoretical analysis of sigmoid attention's training stability by comparing Lipschitz constants, Jacobian structure, and spectral norm bounds with softmax attention.

The theoretical analysis characterizes Lipschitz constants, Jacobian structure, and spectral norm bounds for softmax and sigmoid attention.

A.1.1 Lipschitz Structure and Gradient Coupling

The analysis attributes softmax-attention instability to score magnitude and probability-mass concentration, both tied to its normalization structure.

Softmax-attention instability arises from interacting score-magnitude and probability-mass-concentration mechanisms.The score term depends on query-key dot-product norms, while concentration is governed by the softmax-normalization Jacobian.
The Frobenius norm of the attention matrix measures the entropy of the attention-weight distribution.
When attention concentrates entirely on one token, the attention matrix has minimum entropy with Frobenius norm equal to 1.
Softmax's local Lipschitz constant scales exponentially with the maximum score magnitude, contributing to sensitivity during training.
These observations indicate that softmax instability is fundamentally tied to its normalization structure.

A.1.2 Jacobian Structure and Decoupling

The Jacobian comparison shows that softmax couples token interactions through off-diagonal derivatives, whereas sigmoid uses an element-wise, diagonal structure with a global derivative bound.

Softmax derivatives contain off-diagonal terms, making the Jacobian dense and allowing its spectral norm to grow when probability mass concentrates.
Sigmoid attention operates element-wise, so its partial derivative is nonzero only for the matching input-output index.
Sigmoid attention is globally 1/4-Lipschitz with respect to the score matrix because its diagonal Jacobian's spectral norm equals its largest diagonal entry.
The bounded derivative and diagonal Jacobian give sigmoid attention a uniform spectral-norm bound independent of attention sharpness and score magnitude.This structure avoids the exponential sensitivity and cross-token gradient coupling associated with softmax normalization.

A.1.3 Spectral Norm Bound for Sigmoid Attention

Sigmoid attention has a Jacobian spectral norm bounded independently of attention sharpness, score magnitude, and sequence length. Its element-wise nonlinearity and bounded derivative avoid softmax’s exponential sensitivity and cross-token gradient coupling.

The Jacobian spectral norm is bounded by a constant depending only on projection norms and sigmoid bias scale.The bound is independent of sequence length and score magnitude.
The sigmoid Jacobian is diagonal because sigmoid acts element-wise on attention scores.This contrasts with the cross-token coupling induced by softmax normalization.
Globally bounded sigmoid derivatives and diagonal Jacobian structure yield a uniformly bounded, well-conditioned attention operator.The theorem states that the bound remains independent of attention sharpness and score magnitude.
Sigmoid attention avoids the exponential sensitivity and cross-token gradient coupling inherent to softmax normalization.

A.2 Experimental Hyperparameters

The experiments compare 160M-parameter Transformer models using softmax or sigmoid attention across standard training, validation, embedding, and kernel-performance evaluations. The protocol spans six held-out single-cell datasets, multiple context lengths, masking-based validation, UMAP visualization, MMD separation, and GPU benchmarking.

Model architecture and training: All experiments use 160M-parameter Transformer encoder models trained on the CellxGene dataset.The performance comparison varies two context lengths and two attention mechanisms under standard training conditions.
Model architecture and training: The stress-test variants share hyperparameters and differ only in attention mechanism, with gradient clipping removed.
Evaluation datasets: Six held-out CellxGene datasets cover diverse tissues, developmental stages, and disease contexts for independent generalization assessment.
Evaluation metrics: Validation loss uses 15 independent random masking trials per cell, with lower loss indicating better masked-language-modeling performance.The best loss is the lowest value for each dataset.
Embedding analysis: UMAP compares learned embedding structure after PCA retains 95% of variance, using n_neighbors = 15 and min_dist = 0.5.The Heart OFT visualization contrasts softmax and sigmoid embeddings at 4K context.
Embedding analysis: MMD compares within-cell-type similarity against cross-cell-type similarity, with higher values indicating greater separation.The Heart OFT analysis evaluates 28 pairs across 8 cell types and reports a 25.0% mean sigmoid improvement across all pairs.
Kernel benchmarking: The TritonSigmoid implementation is benchmarked using theoretical FLOPs and achieved TFLOPS, including padding-aware forward and backward calculations.The benchmark accounts for valid non-padded tokens and assigns the backward pass 2.5× the forward FLOPs.