更好的模型，更快的训练：单细胞基础模型的 Sigmoid Attention

训练稳定的生物基础模型需要重新思考注意力机制：我们发现使用 sigmoid 注意力来替代 softmax 注意力 a) 产生更好的学习表示：在六个不同的单细胞数据集上，sigmoid 实现了 25% 更高的细胞类型分离、更好的细胞类型凝聚力指标和更低的验证损失，b) 更快的训练，具有 sigmoid 注意力的模型比其 softmax 模型训练速度快了 10%，c) 通过消除固有源实现更稳定的训练Softmax 注意力的不稳定性。我们确定 sigmoid 注意力具有与 softmax 不同的全局有界导数 ($\leq 0.25$)，以及与 softmax 的密集耦合相反的对角雅可比结构，这些都有助于减轻训练的不稳定性。在对 8K 词元序列上未进行梯度裁剪训练的 160M 参数双向注意力模型进行压力测试时，softmax 出现了灾难性的发散，梯度爆炸了四个数量级，而 sigmoid 保持稳定。最后，我们实现并开源了 TritonSigmoid，这是一个高效的 GPU 内核，在 H100 GPU 上实现了 515 TFLOPS，性能优于 FlashAttention-2 和 FlashSigmoid，并具有本机填充支持，这对于生物序列至关重要。我们的结果确立了 sigmoid 注意力对于生物基础模型的理论依据和经验优越性。代码可在 https://github.com/MSDLLCpapers/triton-sigmoid 获取