用改进的sigmoid注意力提升单细胞模型训练速度与稳定性
Better Models, Faster Training: Sigmoid Attention for single-cell Foundation Models
- 用sigmoid替代softmax注意力，避免梯度爆炸问题
- 在6个单细胞数据集上细胞类型区分度提高25%，验证损失更低
- 开源高效GPU核TritonSigmoid，H100实测达515 TFLOPS
训练稳定的生物基础模型需重新思考注意力机制：将sigmoid注意力作为softmax的即插即用替代品，a) 在六个多样化的单细胞数据集上，显著提升细胞类型分离度（25%）、细胞类型凝聚性指标，并降低验证损失；b) 实现更快训练速度，相比softmax模型提速最高达10%；c) 增强训练稳定性，因其全局导数有界（≤0.25）且雅可比矩阵对角结构，而softmax存在密集耦合与固有不稳定性。在8K token序列、无梯度裁剪条件下，1.6亿参数双向注意力模型中，softmax出现灾难性发散，梯度暴涨四个数量级，而sigmoid保持稳定。我们实现并开源TritonSigmoid，一种高效的GPU内核，在H100上达到515 TFLOPS，优于FlashAttention-2和FlashSigmoid，且支持原生填充，对生物序列至关重要。结果表明，sigmoid注意力在理论与实践上均优于传统softmax，适用于生物基础模型。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。