より良いモデル、より速い学習：単一細胞基盤モデルのためのシグモイド・アテンション

要旨: 安定した生物学的基盤モデルを学習するには、注意（attention）機構を再考する必要があることを示します。具体的には、softmax 注意のドロップイン置換として sigmoid 注意を用いると、a) より良い学習済み表現が得られることがわかります。6種類の多様な単一細胞データセットにおいて、sigmoidは細胞タイプ分離で25%高く、細胞タイプの凝集（cohesion）指標も良好で、さらに検証損失が低いです。b) 学習が速く、sigmoid注意を用いたモデルはsoftmaxの対応物より最大10%速く学習できます。c) softmax注意に内在する不安定性の要因を除去することで、より安定した学習が可能になります。sigmoid注意は、softmaxとは対照的に、導関数がグローバルに有界（\leq 0.25）であること、そしてsoftmaxの密な結合に対して対角（diagonal）ヤコビアン構造を持つことを示します。これらは学習の不安定性を緩和するのに役立ちます。8Kトークン系列で、勾配クリッピングなしに学習した1.6億パラメータの双方向注意モデルに対するストレステストでは、softmaxは壊滅的に発散し、勾配が4桁のオーダーで爆発する一方で、sigmoidは安定したままです。最後に、効率的なGPUカーネルであるTritonSigmoidを実装し、オープンソースとして公開します。これはH100 GPUで515 TFLOPSを達成し、FlashAttention-2とFlashSigmoidの両方を上回ります。さらに、生物学的系列に不可欠なネイティブなパディング対応を備えています。これらの結果により、sigmoid注意は理論的に裏付けられているだけでなく、生物学的基盤モデルに対して経験的にも優れていることが示されました。コードは https://github.com/MSDLLCpapers/triton-sigmoid で利用可能です