Лучшие модели, быстрее обучение: Сигмоидная функция внимания для базовых моделей одноклеточных данных

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

April 29, 2026

Авторы: Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

cs.AI

Аннотация

Обучение стабильных базовых биологических моделей требует переосмысления механизмов внимания: мы обнаружили, что использование сигмоидного внимания в качестве прямой замены софтмакс-внимания а) дает лучшие обученные представления: на шести разнородных наборах данных по единичным клеткам сигмоидное внимание достигает на 25% лучшего разделения типов клеток, улучшает метрики когезии типов клеток и снижает потери на валидации; б) ускоряет обучение — модели с сигмоидным вниманием обучаются до 10% быстрее, чем их аналоги с софтмакс-вниманием; и в) повышает стабильность обучения за счет устранения присущих софтмакс-вниманию источников нестабильности. Мы установили, что сигмоидное внимание имеет глобально ограниченные производные (≤ 0,25) в отличие от софтмакса, а также диагональную структуру матрицы Якоби в противовес плотной связности софтмакса, что в совокупности способствует снижению нестабильности обучения. В стресс-тестах на двунаправленных моделях внимания с 160 миллионами параметров, обученных без обрезки градиентов на последовательностях длиной 8 тысяч токенов, софтмакс-внимание катастрофически расходится с взрывом градиентов на четыре порядка величины, в то время как сигмоидное внимание остается стабильным. Наконец, мы реализовали и открыли исходный код TritonSigmoid — эффективного GPU-ядра, которое достигает 515 TFLOPS на GPU H100, превосходя как 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