より優れたモデル、より高速な学習：単一細胞基盤モデルのためのシグモイド注意機構

要旨

安定した生物学的基盤モデルの訓練には、アテンション機構の再考が必要である：我々は、ソフトマックスアテンションの代替としてシグモイドアテンションを使用することが、以下の利点をもたらすことを発見した。a) 優れた学習表現の生成：6つの多様な単一細胞データセットにおいて、シグモイドは25%高い細胞タイプ分離、優れた細胞タイプ凝集性指標、およびより低い検証損失を達成する。b) 高速な訓練：シグモイドアテンションを用いたモデルは、ソフトマックス版と比較して最大10%高速に訓練される。c) 訓練の安定性向上：ソフトマックスアテンションに内在する不安定性の原因を排除する。我々は、シグモイドアテンションがソフトマックスとは異なり大域的に有界な微分（≤ 0.25）を持ち、ソフトマックスの密な結合性とは対照的な対角ヤコビアン構造を持つことを立証し、これらが共同して訓練の不安定性を軽減することを示す。8Kトークン系列で勾配クリッピングなしで訓練された1億6000万パラメータ双方向アテンションモデルに対するストレステストでは、ソフトマックスは勾配が4桁も爆発的に増大する壊滅的な発散を起こすのに対し、シグモイドは安定したままである。最後に、H100 GPUで515 TFLOPSを達成する効率的なGPUカーネルであるTritonSigmoidを実装しオープンソース化した。これはFlashAttention-2およびFlashSigmoidを上回り、生物学的系列に不可欠なネイティブのパディングサポートを備えている。我々の結果は、シグモイドアテンションが生物学的基盤モデルにおいて理論的根拠があり、経験的にも優れていることを立証する。コードはhttps://github.com/MSDLLCpapers/triton-sigmoid で利用可能である。

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

より優れたモデル、より高速な学習：単一細胞基盤モデルのためのシグモイド注意機構

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

要旨

Support