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

2026-04-29

AI Summary

This work addresses the limitations of traditional softmax attention in single-cell foundation model training—namely, restricted representational capacity, slow convergence, and numerical instability—by proposing sigmoid-based attention as a replacement. The sigmoid formulation offers globally bounded gradients and a diagonal Jacobian structure, which intrinsically enhance training stability and representation quality. The authors introduce a scalable bidirectional attention architecture supporting long sequences and native padding, along with TritonSigmoid, an efficient open-source GPU kernel. Experiments across six single-cell datasets demonstrate a 25% improvement in cell-type separability, lower validation loss, and 10% faster training, with stable performance maintained at a scale of 160 million parameters. On H100 GPUs, TritonSigmoid achieves a peak throughput of 515 TFLOPS.

Abstract

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

Problem

• single-cell foundation models
• attention mechanism
• training stability
• softmax attention
• representation learning

Innovation

• sigmoid attention
• foundation models
• single-cell genomics
• training stability
• efficient GPU kernel