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

• URL: http://arxiv.org/abs/2604.27124v1
• Date: Wed, 29 Apr 2026 19:23:30 GMT
• Title: Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models
• Abstract summary: Training stable biological foundation models requires rethinking attention mechanisms.<n>We find that using sigmoid attention as a drop in replacement for softmax attention a) produces better learned representations.<n>We also implement and open-source TritonSigmoid, an efficient GPU kernel.
• Score: 3.784026537291876
• License: http://creativecommons.org/licenses/by/4.0/
• 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

Related papers

• Rethinking Language Model Scaling under Transferable Hypersphere Optimization [67.38433364607897]
  We introduce HyperP, the first framework for transferring optimal learning rates across model width, depth, training tokens, and expert-of-Experts (MoE)<n>A single base rate tuned data transfers across all compute budgets under HyperP, yielding $1.58times over a strong Muon baseline at $6times1021$ FLOPs.<n>We also propose SqrtGate, an MoE gating mechanism derived from the hypersphere constraint that preserves output RMS across MoE granularities.
  arXiv  Detail & Related papers  (2026-03-30T17:51:47Z)
• Exact Discrete Stochastic Simulation with Deep-Learning-Scale Gradient Optimization [0.0]
  Exact simulation of continuous-time Markov chains (CTMCs) is essential when discreteness and noise drive system behavior, but the hard categorical event selection in Gillespie-type algorithms blocks gradient-based learning.<n>We eliminate this constraint by decoupling forward simulation from backward differentiation, with hard categorical sampling generating exact trajectories and gradients propagating through a continuous massively-parallel Gumbel-Softmax straight-through surrogate.<n>Our results enable high-dimensional parameter inference and inverse design across systems biology, chemical kinetics, physics, and related CTMC-governed domains.
  arXiv  Detail & Related papers  (2026-02-23T12:29:43Z)
• MSign: An Optimizer Preventing Training Instability in Large Language Models via Stable Rank Restoration [48.446476072756276]
  Training instability remains a critical challenge in large language model pretraining.<n>We study training failures in a 5M NanoGPT model scaled via $$P.<n>We propose MSign, a new norm that periodically applies matrix sign operations to restore stable rank.
  arXiv  Detail & Related papers  (2026-02-02T07:18:45Z)
• LINA: Linear Autoregressive Image Generative Models with Continuous Tokens [56.80443965097921]
  Autoregressive models with continuous tokens form a promising paradigm for visual generation, especially for text-to-image (T2I) synthesis.<n>We study how to design compute-efficient linear attention within this framework.<n>We present LINA, a simple and compute-efficient T2I model built entirely on linear attention, capable of generating high-fidelity 1024x1024 images from user instructions.
  arXiv  Detail & Related papers  (2026-01-30T06:44:33Z)
• Beyond Force Metrics: Pre-Training MLFFs for Stable MD Simulations [5.913538953257869]
  Machine-learning force fields (MLFFs) have emerged as a promising solution for speeding up ab initio molecular dynamics (MD) simulations.<n>In this work, we employ GemNet-T, a graph neural network model, as an MLFF and investigate two training strategies.<n>We find that lower force errors do not necessarily guarantee stable MD simulations.
  arXiv  Detail & Related papers  (2025-06-17T00:58:56Z)
• Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free [81.65559031466452]
  We conduct experiments to investigate gating-augmented softmax attention variants.<n>We find that a simple modification-applying a head-specific sigmoid gate after the Scaled Dot-Product Attention (SDPA)-consistently improves performance.
  arXiv  Detail & Related papers  (2025-05-10T17:15:49Z)
• Sigmoid Self-Attention has Lower Sample Complexity than Softmax Self-Attention: A Mixture-of-Experts Perspective [69.72942835553228]
  This paper theoretically demonstrates that sigmoid self-attention is more sample-efficient than its softmax counterpart.<n>We represent the self-attention matrix as a mixture of experts and show that experts'' in sigmoid self-attention require significantly less data to achieve the same approximation error as those in softmax self-attention.
  arXiv  Detail & Related papers  (2025-02-01T02:36:14Z)
• Theory, Analysis, and Best Practices for Sigmoid Self-Attention [16.73166377436999]
  We revisit sigmoid attention and conduct an in-depth theoretical and empirical analysis.<n>We prove that transformers with sigmoid attention are universal function approximators.<n>We introduce FLASHSIGMOID, a hardware-aware and memory-efficient implementation of sigmoid attention.
  arXiv  Detail & Related papers  (2024-09-06T17:53:26Z)
• A Multi-Grained Symmetric Differential Equation Model for Learning Protein-Ligand Binding Dynamics [73.35846234413611]
  In drug discovery, molecular dynamics (MD) simulation provides a powerful tool for predicting binding affinities, estimating transport properties, and exploring pocket sites. We propose NeuralMD, the first machine learning (ML) surrogate that can facilitate numerical MD and provide accurate simulations in protein-ligand binding dynamics. We demonstrate the efficiency and effectiveness of NeuralMD, achieving over 1K$times$ speedup compared to standard numerical MD simulations.
  arXiv  Detail & Related papers  (2024-01-26T09:35:17Z)
• Branched Latent Neural Maps [0.0]
  Branched Latent Neural Maps (BLNMs) learn finite dimensional input-output maps encoding complex physical processes. BLNMs show excellent generalization properties with small training datasets and short training times on a single processor. In the online phase, the BLNM allows for 5000x faster real-time simulations of cardiac electrophysiology on a single core standard computer.
  arXiv  Detail & Related papers  (2023-08-04T04:04:58Z)
• A Stable, Fast, and Fully Automatic Learning Algorithm for Predictive Coding Networks [65.34977803841007]
  Predictive coding networks are neuroscience-inspired models with roots in both Bayesian statistics and neuroscience. We show how by simply changing the temporal scheduling of the update rule for the synaptic weights leads to an algorithm that is much more efficient and stable than the original one.
  arXiv  Detail & Related papers  (2022-11-16T00:11:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.