단일세포 파운데이션 모델 Tahoe-x1: 약물 반응 예측으로 신약 개발을 앞당기는 법

핵심 메시지

• Tahoe-x1은 약 2억 5천만 개 세포 데이터로 사전학습한 단일세포 파운데이션 모델로, 유전자 발현값을 마스킹하고 복원하는 BERT 방식의 트랜스포머 구조를 쓴다.
• 세포들을 한데 모아 시퀀싱한 뒤 유전적 정보로 다시 분리하는 기법 덕분에 같은 비용으로 약 50배 많은 데이터를 만들 수 있었고, 이것이 대규모 학습의 출발점이 되었다.
• 30억 파라미터 모델의 학습 비용은 시장 가격 기준 약 4만 5천 달러로, 작은 모델로 수많은 실험을 먼저 돌려 불필요한 구성 요소를 걷어낸 결과다.
• 유전자 필수성 예측, 유전자 기능 분류, 약물 반응 예측 등 여러 하위 과제에서 모델 크기가 커질수록 성능이 함께 올라가는 경향이 확인됐다.
• 완전히 처음 보는 세포주에 대한 예측(제로샷)은 여전히 어렵지만, 약물 5~10종 정도의 소량 데이터만 있어도 성능이 크게 회복된다.

쉽게 이해하기

버클리 BioML 세미나에서 Tahoe Bio의 머신러닝 책임자가 단일세포 파운데이션 모델 Tahoe-x1을 소개했다. 출발점은 데이터였다. 수십에서 수백 개의 세포주에서 얻은 세포를 3차원 구형체로 모아 한꺼번에 시퀀싱한 뒤, 유전적 차이를 이용해 각 세포가 어느 세포주에서 왔는지 되짚는 방식이다. 발표자는 이 방법으로 같은 비용에 약 50배의 처리량을 얻었고, 그 결과 만들어진 데이터셋이 이전까지 공개된 모든 섭동 데이터를 합친 것보다 50배가량 크다고 설명했다.

모델 구조는 의도적으로 단순하다. 세포 하나는 유전자 목록과 각 유전자의 정수 카운트 목록으로 표현되고, 카운트는 구간으로 나눠 백분위로 변환한다. 이미지에서 화소값을 0~255로 고정하는 것과 같은 발상이다. 학습 목표는 발현값 일부를 가린 뒤 맞히는 것이고, 여기에 세포 전체를 대표하는 특수 토큰의 임베딩으로도 가려진 유전자를 예측하게 하는 손실이 더해진다. 이전 모델들이 여러 목적함수를 쌓아 올린 것과 달리, 성능을 실제로 끌어올리는 두 요소만 남겼다.

발표의 상당 부분은 비용을 줄이기 위한 공학적 선택에 할애됐다. 초기에는 가려진 유전자와 보이는 유전자의 어텐션 패턴이 특수해 표준 가속 기법을 쓸 수 없었고, 직접 커널을 작성해 마스크 대신 어텐션 편향을 더하는 방식으로 우회했다. 그런데 이후 실험에서 아예 전체를 밀집 어텐션으로 두고 입력만 가려도 성능이 거의 떨어지지 않는다는 것을 확인해, 최종 모델은 널리 쓰이는 고속 어텐션 구현을 그대로 사용했다. 데이터셋 스트리밍 라이브러리를 도입해 GPU 128장으로 학습해도 한 장으로 학습할 때와 성능 차이가 없게 만든 점도 함께 소개됐다.

성능 검증은 세 갈래로 진행됐다. 첫째는 특정 유전자를 제거했을 때 세포가 죽을 확률을 예측하는 과제로, 같은 원본 데이터를 준 상태에서 여러 임베딩을 비교했을 때 Tahoe-x1 임베딩이 더 나은 점수를 냈다. 둘째는 유전자가 어떤 기능 집합에 속하는지 맞히는 다중 레이블 분류다. 셋째가 실제 목표에 가장 가까운 약물 반응 예측으로, 학습에서 빼둔 세포주와 약물 조합에 대해 유전자 발현 변화를 예측하고 실제 실험값과 비교한다.

가장 큰 과제는 한 번도 본 적 없는 세포주에 대한 예측이다. 같은 세포주에서 다른 약물을 본 적 있는 경우와, 그 세포주 데이터가 전혀 없는 경우 사이에 성능 격차가 컸다. 발표자는 이후 개발한 후속 모델이 이 격차를 크게 줄였다고 밝혔는데, 임베딩 공간 안에서 플로우 매칭으로 섭동된 세포를 직접 생성하는 구조다. 대조군과 처리군 세포를 쌍으로 묶을 필요가 없어 학습 시간이 3~4일에서 두세 시간 수준으로 줄었고, 그만큼 많은 아이디어를 시험할 수 있었다.

주요 인사이트

• 데이터 생산 방식의 개선이 모델 성능보다 먼저 왔다. 같은 예산으로 50배 많은 실험을 돌릴 수 있게 되자 비로소 대규모 사전학습이 가능해졌고, 발표자는 이를 이미지 인식 분야가 대형 데이터셋을 얻었던 순간에 빗댔다.
• 30억 파라미터 모델을 4만 5천 달러에 학습했다는 대목은 규모보다 실험 설계가 중요하다는 점을 보여준다. 7천만 파라미터짜리 작은 모델로 어떤 목적함수와 구성 요소가 실제로 기여하는지 먼저 가려낸 뒤에야 큰 모델을 돌렸다.
• 복잡한 커스텀 최적화를 만들어 놓고도, 더 단순한 방식이 성능을 해치지 않는다는 것을 확인하자 그것을 버렸다. 최종 선택이 항상 더 정교한 쪽은 아니라는 실무적 교훈이다.
• 약물 반응 예측의 상관계수가 0.7 수준이면 신호의 절반가량을 설명하는 것이고, 발표자는 이 정도만 되어도 세포주를 구해 실험을 설계하는 초기 비용 없이 후보를 좁힐 수 있다고 본다. 신약 개발 전체가 빨라지는 것은 아니지만 검증할 가설의 수는 늘어난다.
• 코딩 에이전트에게 데이터와 평가 기준을 주고 밤새 구조를 탐색하게 한 실험이 사람이 만든 최고 성능을 넘어섰다는 대목은 주목할 만하다. 다만 이는 평가 지표가 명확하고 학습이 빠른 좁은 과제에 한정된 결과다.

자주 묻는 질문

관련 AI 소식