更好的模型、更快训练：单细胞基础模型中的 Sigmoid 注意力
机器学习
2026-05-01 v1 定量方法
摘要
训练稳定的生物学基础模型需要重新思考注意力机制：我们发现将sigmoid注意力作为softmax注意力的替代方案可获得更好的学习表示：在六个多样化的单细胞数据集上，sigmoid实现了25%更高的细胞类型分离率，更好的细胞类型内聚指标，以及更低的验证损失；更快的训练，采用sigmoid注意力的模型在训练速度上可快达10%；更稳定的训练通过消除softmax注意力中固有的不稳定性源实现。我们确立了sigmoid注意力具有全局有界导数（≤0.25），与softmax不同，其雅可比矩阵呈对角结构，这与softmax稠密耦合不同，共同帮助缓解训练不稳定性。在16000万参数的双向注意力模型上进行的压力测试中，未使用梯度裁剪在8K token序列上训练，softmax在梯度爆炸四个数量级后会 catastrophic 发散，而sigmoid保持稳定。最后，我们实现并开源了TritonSigmoid，一种高效GPU内核，在H100 GPU上实现515 TFLOPS的性能，超越FlashAttention-2和FlashSigmoid，具有原生填充支持，这对于生物学序列至关重要。我们的结果确立了sigmoid注意力在生物学基础模型中既在理论上有依据，又在实践中优于softmax。代码已公开：https://github.com/MSDLLCpapers/triton-sigmoid
引用
@article{arxiv.2604.27124,
title = {Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models},
author = {Vijay Sadashivaiah and Georgios Dasoulas and Judith Mueller and Soumya Ghosh},
journal= {arXiv preprint arXiv:2604.27124},
year = {2026}
}