Mejores Modelos, Entrenamiento Más Rápido: Atención Sigmoide para Modelos Fundacionales de Célula Única

Resumen

Entrenar modelos fundacionales biológicos estables requiere replantear los mecanismos de atención: encontramos que usar atención sigmoide como reemplazo directo de la atención softmax a) produce mejores representaciones aprendidas: en seis conjuntos de datos unicelulares diversos, sigmoide logra un 25% más de separación de tipos celulares, mejores métricas de cohesión de tipos celulares y menor pérdida de validación, b) entrenamiento más rápido: los modelos con atención sigmoide se entrenan hasta un 10% más rápido que sus contrapartes con softmax, y c) entrenamiento más estable al eliminar fuentes inherentes de inestabilidad en la atención softmax. Establecemos que la atención sigmoide tiene derivadas globalmente acotadas (≤ 0.25) a diferencia de softmax, y una estructura jacobiana diagonal en contraste con el acoplamiento denso de softmax, lo cual ayuda conjuntamente a aliviar las inestabilidades del entrenamiento. En pruebas de estrés con modelos de atención bidireccional de 160 millones de parámetros entrenados sin recorte de gradientes en secuencias de 8.000 tokens, softmax diverge catastróficamente, con gradientes que explotan en cuatro órdenes de magnitud, mientras que sigmoide permanece estable. Finalmente, implementamos y publicamos como código abierto TritonSigmoid, un núcleo GPU eficiente que alcanza 515 TFLOPS en GPUs H100, superando tanto a FlashAttention-2 como a FlashSigmoid, con soporte nativo para relleno (padding), esencial para secuencias biológicas. Nuestros resultados establecen la atención sigmoide como teóricamente fundamentada y empíricamente superior para modelos fundacionales biológicos. El código está disponible en 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

Mejores Modelos, Entrenamiento Más Rápido: Atención Sigmoide para Modelos Fundacionales de Célula Única

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

Resumen

Support