가상 세포 모델 STATE: 단일세포 교란 반응을 집합 트랜스포머로 예측하는 방법과 평가 결과

핵심 메시지

• 세포는 측정하는 순간 죽기 때문에 같은 세포의 처리 전후를 짝지어 볼 수 없고, 그래서 교란 반응 예측은 개별 세포가 아니라 분포에서 분포로 가는 문제로 다뤄진다.
• STATE는 같은 세포 타입이나 같은 배치로 묶은 세포 '집합'을 통째로 트랜스포머에 넣어, 주석되지 않은 이질성을 데이터 적재 단계에서 걷어내고 교란 효과만 학습하게 한다.
• 집합 크기를 1에서 256까지 키우면 검증 손실이 계속 내려가지만 512에서는 오히려 나빠져, 잡음 상쇄와 정보량 사이에 최적점이 존재한다.
• 이 분야에서 가장 까다로운 기준선은 '교란 평균'과 '문맥 평균'이라는 단순한 방법이며, 기존 딥러닝 모델 대부분이 이를 넘지 못했다.
• 1억 6,700만 개 인간 세포로 사전학습한 별도의 임베딩 모델을 붙이면, 작고 잡음 많은 데이터셋으로의 제로샷 전이 성능이 개선된다.

쉽게 이해하기

버클리 머신러닝 학생 단체가 연 BioML 세미나에서 아크 연구소의 연구자가 교란 반응 예측 모델 STATE를 소개했다. 목표는 분명하다. 배양 접시에서 세포에 약물을 처리하고 결과를 보는 실험을, 신경망의 순전파 한 번으로 대체하는 것이다. 그렇게 되면 종양 조직에서 얻은 세포 집단을 두고 '어떤 처리를 하면 건강한 상태로 되돌아갈까'를 계산으로 물어볼 수 있다.

단백질 구조 예측에서 딥러닝이 통했던 조건이 세포에는 그대로 성립하지 않는다는 것이 이 발표의 출발점이다. 알파폴드는 서열 하나에 구조 하나가 대응하는 개체 수준 데이터를 썼지만, 단일세포 측정은 세포를 죽이기 때문에 처리 전후를 같은 개체로 이을 수 없다. 여기에 세포 주기처럼 데이터에 기록되지 않은 변수와, 실험 배치에서 오는 기술적 잡음이 겹친다. 발표자는 그 결과 이 분야에 아직 뚜렷한 스케일링 법칙이 없고, 데이터가 늘어도 딥러닝이 단순한 방법을 이기지 못하는 상태였다고 정리했다.

STATE는 두 모델로 이뤄진다. 전이 모델은 세포 여러 개를 한 묶음으로 받는 집합 트랜스포머다. 세포마다 위치 개념이 없으니 순서 정보를 쓰지 않고, 대신 한 묶음에 같은 세포 타입만, 혹은 같은 실험 웰의 세포만 담아 모델이 굳이 세포 타입 차이를 배우느라 용량을 쓰지 않게 한다. 세포 타입을 입력 특징으로 넣지 않은 덕분에 학습 때 못 본 세포 타입에도 제로샷으로 예측할 수 있다는 설계상의 이점이 따라온다.

왜 하필 셀프 어텐션인지도 설명된다. 기존 연구들은 여러 세포를 평균 내는 '유사 벌크' 방식으로 잡음을 줄였는데, 셀프 어텐션은 그 평균 내기를 일반화한 연산으로 볼 수 있다는 것이다. 실제 학습된 어텐션 맵에는 모든 위치를 고르게 모으는 집계형 헤드, 대각선 형태로 개별 세포만 보는 헤드, 그리고 잡음 제거처럼 보이는 복잡한 분포형 헤드가 함께 나타난다. 논문에서는 데이터가 무한할 때 STATE가 최적 수송 기반 접근을 포괄한다는 점도 보였다.

평가는 유전자 교란 60만 세포, 신호 전달 교란 1,000만 세포, 화학 약물 교란 1억 세포 세 규모에서 이뤄졌다. 지표는 예측한 반응이 실제 교란과 제대로 짝지어지는지 보는 순위 매기기, 대조군 대비 변화량의 피어슨 상관, 그리고 차등 발현 유전자 판정의 정밀도-재현율 곡선 아래 면적이다. STATE는 세 데이터셋 모두에서 다른 딥러닝 방법을 앞섰고, 데이터 규모가 커질수록 평균 기준선과의 격차가 벌어지는 경향을 보였다. 발표자는 이것이 이 분야가 기다리던 신호라고 말한다.

주요 인사이트

• 임베딩 모델은 서로 다른 컨소시엄이 측정한 유전자 목록이 제각각이라는 문제를 유전자 원핫 인코딩 대신 단백질 언어 모델 임베딩으로 푼다. 유전자별로 단백질 전사체 임베딩을 평균 내 특징으로 쓰고, 마스킹 언어 모델과 비슷하게 입력에 없던 유전자의 발현을 맞히게 학습한다.
• 데이터셋 토큰을 학습 가능한 형태로 넣고 '이 세포가 어느 데이터셋에서 왔는지' 맞히는 보조 과제를 붙이면 학습이 안정된다. 데이터셋 고유의 효과를 그 토큰이 흡수하도록 유도하는 장치다.
• 모델이 두 평균 기준선을 동시에 넘는다는 것은 정의상 세포주 특이적인 반응을 잡아냈다는 뜻이다. 발표자는 실제로 특정 변이를 가진 흑색종 세포주에서, 그 변이에 듣는 것으로 알려진 승인 약물의 세포주 특이 효과를 재현했다고 소개했다.
• 남은 벽은 뚜렷하다. 지금은 변동이 큰 유전자 2,000개 정도만 모델링하는데 데이터셋에는 1만 8,000개 측정값이 있고, 같은 손실과 학습 방식은 그 규모에서 모드 붕괴 같은 수치 문제를 일으킨다. 교란 자체는 아직 원핫으로만 표현된다.

자주 묻는 질문

관련 AI 소식