🧬 biology

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

본 논문은 단일 세포 기반 모델에서 소프트맥스를 시그모이드 어텐션으로 대체하면 이론적으로 경계된 미분값을 통해 표현의 질이 향상되고 학습 속도가 빨라지며 안정성이 강화됨을 보여주며, 이는 고도로 최적화된 Triton 커널 구현으로 뒷받침됩니다.

원저자: Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

게시일 2026-05-01

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Vijay Sadashivaiah, Georgios Dasoulas, Judith Mueller, Soumya Ghosh

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ⚕️ 이것은 동료 심사를 거치지 않은 프리프린트의 AI 생성 설명입니다. 의학적 조언이 아닙니다. 이 내용을 바탕으로 건강 관련 결정을 내리지 마세요. 전체 면책 조항 읽기

이 논문은 간단한 언어와 창의적인 비유를 사용하여 설명합니다.

큰 그림: 생명의 '언어'를 읽는 더 나은 방법

인간 세포의 '언어'를 컴퓨터가 이해하도록 가르치려 한다고 상상해 보세요. 이 언어에서 각 유전자는 단어이고, 전체 세포는 문장입니다. 과학자들은 에세이를 쓰거나 당신과 대화하는 AI 와 동일한 종류의 강력한 도구인 트랜스포머 (Transformer) 를 사용하여 이러한 세포 문장을 읽어 왔습니다.

그러나 이러한 컴퓨터가 단어 간의 관계를 이해하는 데 사용하는 표준 도구는 소프트맥스 어텐션 (Softmax Attention) 입니다. 이 논문의 저자들은 생물학의 맥락에서 이 표준 도구는 살아있는 세포의 messy 한 현실에 맞지 않는 엄격하고 경직된 규칙집을 사용하려는 것과 같다고 주장합니다. 그들은 이를 시그모이드 어텐션 (Sigmoid Attention) 이라는 새로운 도구로 교체할 것을 제안하며, 이는 더 빠르고 안정적이며 실제로 AI 가 더 잘 학습하도록 돕는다고 말합니다.

다음과 같이 설명합니다:

1. 문제: '제로섬 (Zero-Sum)' 게임

옛 방식 (소프트맥스):
교실의 학생들에게 나눠줄 '어텐션 토큰 (예: 스티커 100 개)'이라는 제한된 양을 가진 교사로 상상해 보세요. 만약 학생 A 에게 스티커 50 개를 준다면, 나머지 학생들로부터 그 스티커를 가져와야만 합니다. 이것이 소프트맥스가 작동하는 방식입니다. 이는 AI 에게 "유전자 X 에 집중할 것이므로 유전자 Y 는 무시해야 한다"고 결정하도록 강요합니다.

생물학에 왜 나쁜가:
실제 세포에서 유전자는 종종 함께 작동합니다. 단일 유전자가 동시에 여러 다른 조절 인자에 의해 제어될 수 있습니다. 이는 한 학생이 동시에 세 명의 다른 교사로부터 도움을 필요로 하는 것과 같습니다. 소프트맥스의 '제로섬' 규칙은 AI 가 교사 한 명만 선택하도록 강요하는데, 이는 실제 생물학이 작동하는 방식과 맞지 않습니다.

새로운 방식 (시그모이드):
시그모이드 어텐션은 자격이 있는 모든 학생에게 스티커를 주는 것과 같으며, 제한이 없습니다. 학생 A, B, C 가 모두 도움이 필요하다면 교사는 모두에게 스티커를 줄 수 있습니다. 이는 AI 가 "이 유전자는 중요하고, 저 유전자도 중요하다"고 말하면서 한 가지를 무시하고 다른 하나에 집중할 필요가 없도록 합니다.

2. 안정성 문제: '줄타기꾼'

옛 방식 (소프트맥스):
소프트맥스는 바람에 매우 민감한 줄타기꾼과 같습니다. 숫자가 너무 커지면 (긴 유전자 서열을 다룰 때 발생함) 수학이 엉망이 될 수 있습니다. '바람'이 줄타기꾼을 줄에서 떨어뜨려 훈련이 붕괴됩니다. 논문에서 그들은 안전망 (기울기 클리핑) 을 제거하고 매우 긴 시퀀스를 사용하여 이를 테스트했습니다. 소프트맥스 모델은 절벽에서 떨어졌으며, 수학이 10,000 배나 폭발했습니다.

새로운 방식 (시그모이드):
시그모이드는 넓은 평평한 다리를 걷는 사람과 같습니다. 내장된 난간이 있습니다. 숫자가 아무리 커져도 수학은 안전하고 예측 가능한 범위 내에 머뭅니다. 동일한 '안전망 없음' 테스트에서 시그모이드 모델은 완벽하게 걷기를 계속하며 결코 붕괴되지 않았습니다.

3. 속도 문제: '불규칙한 퍼즐'

도전 과제:
생물학적 데이터는 messy 합니다. 한 세포는 500 개의 유전자 (짧은 문장) 를 가질 수 있는 반면, 다른 세포는 10,000 개의 유전자 (긴 소설) 를 가질 수 있습니다. 이를 함께 처리하기 위해 컴퓨터는 보통 짧은 것들을 같은 길이가 되도록 '빈 공간 (0)'으로 채워 넣어야 합니다. 이는 짧은 시와 긴 소설을 포장.peanuts 로 채워 넣는 같은 상자에 넣으려는 것과 같습니다.

해결책:
저자들은 TritonSigmoid라는 새로운 초고속 엔진을 구축했습니다.

비유: 보통 빈 집까지 포함해 동네의 모든 집으로 운전해야 하는 배송 트럭을 상상해 보세요. TritonSigmoid 는 어떤 집이 비어 있는지 정확히 알고 완전히 건너뛰는 똑똑한 트럭입니다.
결과: 이 새로운 엔진은 놀라울 정도로 빠릅니다. 최신 슈퍼컴퓨터 칩 (NVIDIA H100) 에서 초당 515 TFLOPS(초당 조단위 계산) 속도로 실행됩니다. 이는 현재 최고의 소프트맥스 엔진보다 빠르고 이전 시그모이드 시도들보다도 더 빠릅니다. 이는 이러한 생물학적 모델의 훈련을 최대 **10%**까지 빠르게 만듭니다.

4. 결과: 더 나은 두뇌, 더 빠른 훈련

이 팀은 어떤 AI 모델이 더 잘 학습하는지 보기 위해 네 가지 다른 AI 모델을 훈련시켰습니다. 그들은 1 억 3,100 만 개의 세포로 구성된 방대한 데이터셋을 사용했습니다.

더 나은 학습: 시그모이드를 사용한 모델은 단순히 더 빠르게 학습한 것이 아니라, 더 잘 학습했습니다. 그들은 심장 세포와 뇌 세포를 구별하는 등 서로 다른 세포 유형을 구별하는 능력이 소프트맥스 모델보다 25% 더 뛰어났습니다.
낮은 오류: 시그모이드 모델은 유전자 패턴을 예측할 때 더 적은 실수를 했습니다.
붕괴 없음: 스트레스 테스트에서 소프트맥스 모델은 붕괴되어 다시 시작해야 했고, 시간과 돈을 낭비했습니다. 시그모이드 모델은 단일 오류 없이 작업을 완료했습니다.

요약

이 논문은 단일 세포 생물학의 복잡하고 messy 하며 가변적인 데이터를 이해하는 데 있어 표준 '소프트맥스' 어텐션 메커니즘은 너무 경직되고 불안정하다고 주장합니다. 시그모이드 어텐션으로 전환함으로써 그들은 다음과 같은 시스템을 만들었습니다:

여러 연결 허용 (유전자는 한 번에 여러 요인의 영향을 받을 수 있음).
수학이 격렬해져도 안정성 유지 (더 이상 붕괴되지 않음).
빈 데이터 (패딩) 를 똑똑히 건너뛰어 더 빠르게 실행.

그들은 이 새로운 더 빠른 엔진 (TritonSigmoid) 을 무료로 제공하여 다른 과학자들이 더 나은 생물학적 모델을 구축할 수 있도록 했습니다.

"단일 세포 기초 모델용 시그모이드 어텐션: 더 나은 모델, 더 빠른 학습" 논문의 상세 기술 요약입니다.

1. 문제 제기

이 논문은 생물학적 기초 모델 (특히 단일 세포 RNA 시퀀싱, scRNA-seq) 학습에서 두 가지 중요한 과제를 다룹니다.

소프트맥스의 학습 불안정성: 표준 트랜스포머 모델은 토큰 간 어텐션 가중치의 합이 1 이 되도록 정규화하는 소프트맥스 어텐션을 사용합니다. 이는 한 유전자에 대한 어텐션을 높이면 다른 유전자에 대한 어텐션이 감소해야 하는 "제로섬" 경쟁 메커니즘을 만듭니다. 생물학적 맥락에서 유전자는 종종 여러 독립적인 요인에 의해 공동 조절되므로, 이러한 경쟁적 가정은 생물학적으로 비현실적입니다. 또한, 어텐션 엔트로피가 소수의 토큰으로 붕괴될 때 시퀀스 길이가 증가함에 따라 소프트맥스는 수치적 불안정성 (기울기 폭발) 에 취약합니다.
생물학적 데이터 변동성에 대한 비효율성: 생물학적 시퀀스 (세포) 는 발현된 유전자의 수 (수백 개에서 수천 개) 에 따라 길이가 극단적으로 다양합니다. 텍스트와 달리 시퀀스를 고정된 길이로 패킹할 수 없으므로, 생물학적 샘플은 배치 내 최대 길이에 맞춰 패딩하여 개별적으로 처리해야 합니다. 기존 효율적 어텐션 커널 (FlashAttention-2 또는 FlashSigmoid 등) 은 네이티브 패딩 지원을 제공하지 않거나 "불규칙한" 생물학적 데이터를 효율적으로 처리하지 못해 상당한 계산 낭비가 발생합니다.

2. 방법론

저자들은 소프트맥스를 시그모이드 어텐션으로 대체하고 이를 지원하는 전용 고성능 GPU 커널을 개발했습니다.

A. 시그모이드 어텐션 메커니즘

소프트맥스처럼 행을 따라 어텐션 점수를 정규화하는 대신, 시그모이드 어텐션은 dot-product 점수에 요소별 시그모이드 함수 ( $\sigma(x) = \frac{1}{1+e^{-x}}$ ) 를 적용합니다:
$\text{SigmoidAttn}(X) = \sigma\left(\frac{QK^\top}{\sqrt{d}} + b\right)V$

분리 (Decoupling): 각 어텐션 가중치는 독립적으로 계산됩니다. 쿼리는 여러 키에 동시에 강하게 어텐션할 수 있어 공동 조절을 더 잘 모델링합니다.
편향 항 (Bias Term): 요소별 독립성을 유지하면서 소프트맥스의 스케일링 동작을 근사하기 위해 편향 $b = -\log(n)$ 이 추가됩니다.

B. 이론적 안정성 분석

논문은 시그모이드의 안정성에 대한 이론적 근거를 제시합니다.

유계 도함수: 시그모이드 함수의 도함수는 전역적으로 $0.25 $로 유계입니다 ($ \sigma'(x) = \sigma(x)(1-\sigma(x))$). 반면, 소프트맥스의 리프시츠 상수는 점수 크기에 따라 기하급수적으로 증가합니다.
대각 야코비안: 시그모이드 어텐션의 야코비안은 대각 행렬이므로, 서로 다른 토큰에 대한 기울기가 분리됩니다. 소프트맥스는 밀집된 야코비안을 가지며 모든 토큰 기울기를 결합하여 점수가 커질 때 불안정성을 초래합니다.

C. TritonSigmoid 커널 구현

시그모이드 어텐션을 생물학적 데이터에 실용적으로 적용하기 위해 저자들은 Triton 언어를 사용한 맞춤형 GPU 커널인 TritonSigmoid를 개발했습니다. 주요 혁신 사항은 다음과 같습니다.

네이티브 패딩 지원: 쿼리와 키 양쪽에서 완전히 패딩된 블록을 건너뛰어 낭비되는 FLOPs 를 방지하는 블록-희소 계산을 구현했습니다.
퓨전 연산: 어텐션 계산을 단일 커널로 융합하고, 시그모이드 함수를 위해 하드웨어 최적화된 tanh 근사치를 사용했습니다.
역전파 분해: 메모리 효율성을 위해 역전파를 두 개의 커널 ( $Q$ 기울기용, $K/V$ 기울기용) 로 분할하여 원자적 연산을 제거하고 활성화 값을 실시간으로 재계산했습니다.
하드웨어 최적화: NVIDIA H100/Blackwell 아키텍처를 위해 설계되었으며 BF16 정밀도를 지원합니다.

3. 주요 기여

TritonSigmoid 커널: NVIDIA H100 GPU 에서 515 TFLOPS를 달성하는 고성능 구현체입니다. FlashAttention-2(361 TFLOPS) 와 FlashSigmoid(440 TFLOPS) 보다 우수하며, 가변 길이 생물학적 시퀀스에 대한 효율적 패딩을 지원하는 최초의 솔루션입니다.
실증적 성능 향상: 여섯 가지 다양한 단일 세포 데이터셋에서 소프트맥스보다 우수한 학습된 표현을 생성함을 입증했습니다.
학습 안정성: 기울기 클리핑 없이 긴 컨텍스트 (8K 토큰) 시나리오에서 소프트맥스 모델에서 관찰된 치명적인 학습 발산을 시그모이드 어텐션이 방지함을 증명했습니다.

4. 결과

A. 계산 효율성

처리량: TritonSigmoid 는 표준 PyTorch 어텐션 대비 순전파에서 최대 7.15 배, 역전파에서 2.44 배의 속도 향상을 달성했습니다.
패딩 효율성: 25% 패딩 시, TritonSigmoid 는 순전파에서 FlashAttention-2 보다 29% 더 빠릅니다. 패딩으로 인한 TFLOPS 감소는 미미하여 (~9.3%) 커널이 마스킹된 영역을 효과적으로 건너뜀을 입증했습니다.
종단 간 학습: 시그모이드 어텐션으로 학습된 모델은 소프트맥스 대비 3~10% 더 빠르게 학습되며, 컨텍스트 길이가 증가함에 따라 그 격차가 벌어집니다 (예: 14 억 파라미터 모델의 8K 컨텍스트에서 7.5% 빠름).

B. 모델 품질 (표현 학습)

Brain, Blood, Colon, Lung, Heart 등 6 개의 홀드아웃 단일 세포 데이터셋에서 평가:

검증 손실: 모든 데이터셋과 컨텍스트 길이 (2K 및 4K) 에서 시그모이드 모델이 일관되게 더 낮은 검증 손실을 달성했습니다.
세포 유형 분리: 시그모이드 모델은 세포 유형 간 최대 평균 불일치 (MMD) 가 25% 더 높아 임베딩 공간에서 더 나은 분리를 보여주었습니다.
생물학적 보존: 시그모이드는 모든 6 개 데이터셋에서 우수한 세포 유형 응집성 (실루엣 점수) 을 달성했으며, 6 개 중 4 개 데이터셋에서 더 나은 집계 생물학적 보존 점수를 기록했습니다.
시각화: UMAP 시각화는 시그모이드 임베딩이 세포 유형을 소프트맥스보다 더 밀집되고 명확하게 군집화함을 확인했습니다.

C. 안정성 스트레스 테스트

설정: 기울기 클리핑 없이 8K 토큰 시퀀스로 1 억 6 천만 파라미터 모델을 학습.
결과:
- 소프트맥스: 약 55k 스텝에서 치명적으로 발산했습니다. 전역 기울기 노름이 4 개 차수 ( $100 \to 1.6 \times 10^6$ ) 폭발했으며, 어텐션 점수가 기하급수적으로 증가했습니다.
- 시그모이드: 80k 스텝 동안 기울기 (10~100 범위) 와 제어된 어텐션 점수를 유지하며 안정적으로 학습되어 이론적 안정성 보장을 확인했습니다.

5. 의의

이 논문은 시그모이드 어텐션을 생물학적 기초 모델용 소프트맥스의 이론적 근거가 있고 실증적으로 우수한 대안으로 확립합니다.

생물학적 관련성: 소프트맥스의 인위적 경쟁 제약을 제거함으로써 유전자 조절 (공동 활성화) 의 생물학적 특성과 더 잘 부합합니다.
실용적 배포: TritonSigmoid를 통해 가변 길이 시퀀스의 효율성 문제를 해결함으로써, 이전에는 불안정하거나 계산적으로 불가능했던 대규모 장기 컨텍스트 생물학적 모델 학습을 가능하게 합니다.
오픈 소스: 저자들은 구현체를 공개하여 단일 세포 생물학 커뮤니티에 더 빠른 학습, 더 나은 안정성, 더 높은 품질의 표현을 제공하는 소프트맥스의 교체 가능한 대안을 제공합니다.

코드 이용 가능: https://github.com/MSDLLCpapers/triton-sigmoid