Melhores Modelos, Treinamento Mais Rápido: Atenção Sigmóide para Modelos de Base de Célula Única

Resumo

O treinamento de modelos de fundação biológicos estáveis exige uma reformulação dos mecanismos de atenção: descobrimos que o uso da atenção sigmoide como substituta direta da atenção softmax a) produz representações aprendidas superiores: em seis conjuntos de dados de célula única diversos, a sigmoide alcança uma separação de tipos celulares 25% maior, melhores métricas de coesão de tipos celulares e menor perda de validação; b) treinamento mais rápido: modelos com atenção sigmoide treinam até 10% mais rápido que suas contrapartes com softmax; e c) treinamento mais estável, eliminando fontes inerentes de instabilidade na atenção softmax. Demonstramos que a atenção sigmoide possui derivadas globalmente limitadas (≤ 0,25), em oposição à softmax, e uma estrutura diagonal do Jacobiano, contrastando com o acoplamento denso da softmax, o que em conjunto ajuda a aliviar as instabilidades de treinamento. Em testes de estresse com modelos de atenção bidirecionais de 160 milhões de parâmetros, treinados sem corte de gradiente em sequências de 8 mil tokens, a softmax diverge catastrophicamente, com gradientes explodindo em quatro ordens de magnitude, enquanto a sigmoide permanece estável. Por fim, implementamos e disponibilizamos como código aberto o TritonSigmoid, um kernel de GPU eficiente que atinge 515 TFLOPS em GPUs H100, superando tanto o FlashAttention-2 quanto o FlashSigmoid, com suporte nativo a preenchimento (padding), essencial para sequências biológicas. Nossos resultados estabelecem a atenção sigmoide como teoricamente fundamentada e empiricamente superior para modelos de fundação biológicos. O código está disponível em 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