HEIMDALL: Disentangling tokenizer design for robust transfer in single-cell foundation models

이 논문은 단일 세포 기반 모델 (scFMs) 의 성능이 학습 데이터와 일치하는 경우보다 분포 변화가 있는 환경에서 토큰화 설계에 크게 의존함을 보여주며, HEIMDALL 이라는 통합 프레임워크를 통해 유전자 식별, 발현 인코딩, 순서 등 소수의 핵심 설계 축을 최적화함으로써 모델의 강건한 전이 능력을 향상시킬 수 있음을 제시합니다.

핵심

🍳 비유: 세포를 요리하는 인공지능

생물학자들은 수백만 개의 세포 데이터를 가지고 인공지능 (AI) 을 훈련시켜, 세포가 어떤 종류인지 분류하거나 질병을 예측하려고 합니다. 하지만 이 AI 들이 실제로 잘 작동하지 않을 때가 많습니다.

이 연구의 핵심은 **"AI 가 요리를 잘 하려면, 재료를 어떻게 다듬고 (Tokenization) 접시에 담느냐가 가장 중요하다"**는 것을 발견했다는 점입니다.

1. 문제점: "무엇을 먹일지"에 대한 혼란

세포는 수천 개의 유전자로 이루어져 있습니다. 이 유전자들을 AI 가 읽을 수 있는 언어로 바꾸는 과정을 **'토크나이저 (Tokenizer)'**라고 합니다.
지금까지 과학자들은 각자 다른 방식으로 재료를 다듬었습니다.

• A 연구소는 유전자를 알파벳 순서로 나열했습니다.
• B 연구소는 유전자가 얼마나 활발한지 (발현량) 에 따라 순서를 바꿨습니다.
• C 연구소는 유전자의 이름만 보고 의미를 부여했습니다.

이렇게 방식이 제각각이라서, **"어떤 AI 가 더 좋은가?"**를 비교할 때, AI 의 두뇌 (모델 구조) 가 좋은지, 아니면 재료 준비 (토크나이저) 가 잘 된 건지 알 수 없었습니다. 마치 "요리 실력이 좋은지, 아니면 손질한 재료가 좋은지"를 구분할 수 없는 상황과 같습니다.

2. 해결책: HEIMDALL (하이멜) 이라는 '모듈형 주방'

저자들은 HEIMDALL이라는 새로운 프레임워크를 만들었습니다. 이는 마치 모듈형 주방처럼, 토크나이저를 세 가지 주요 단계로 쪼개서 분석할 수 있게 해줍니다.

1. 재료 식별 (Gene Identity): "이 유전자가 누구인가?" (예: 유전자 이름, 단백질 구조, DNA 서열 등)
2. 재료 상태 표현 (Expression): "이 유전자가 얼마나 활발한가?" (예: 발현량을 숫자로 변환하는 방식)
3. 접시 구성 (Cell Construction): "이 재료들을 어떤 순서로 접시에 담을 것인가?" (예: 유전자 순서, 중요도에 따른 정렬)

HEIMDALL 을 사용하면 연구자들은 이 세 가지 요소를 자유롭게 섞고 맞출 수 있습니다. "A 모델의 재료 식별법 + B 모델의 순서 정렬법"처럼 조합해 볼 수 있는 거죠.

3. 주요 발견: "평소엔 비슷하지만, 이질적인 상황에선 천차만별"

이 연구를 통해 놀라운 사실이 밝혀졌습니다.

• 같은 환경 (훈련 데이터와 테스트 데이터가 같을 때):
  어떤 방식을 쓰든 AI 의 성능은 비슷했습니다. 마치 "한국인만 있는 식당에서 한국 음식을 팔 때, 국을 어떤 그릇에 담든 맛은 비슷하다"는 것과 같습니다.

• 이질적인 환경 (새로운 조직, 다른 종, 새로운 유전자 패널):
  하지만 새로운 환경으로 넘어가면 (예: 인간 세포로 훈련된 AI 가 쥐 세포를 분석하거나, 새로운 장기 데이터를 분석할 때), 재료 준비 방식 (토크나이저) 이 성패를 가릅니다.

  • 어떤 방식은 완전히 망하고, 어떤 방식은 훌륭하게 적응했습니다.
  • 특히 **유전자의 순서를 어떻게 정렬하느냐 (발현량 순 vs 무작위 vs 염색체 순)**와 유전자의 의미를 어떻게 부여하느냐가 가장 중요했습니다.

4. 결론: "완벽한 한 가지 방식은 없다"

이 논문은 "세상에서 가장 좋은 토크나이저 하나를 찾아라"라고 말하지 않습니다. 대신, **"상황에 따라 가장 적합한 재료 준비법을 선택하라"**고 조언합니다.

• 새로운 종 (Species) 을 다룰 때: 유전자의 DNA 서열 자체를 이해하는 방식이 가장 안전합니다.
• 새로운 유전자 패널 (Gene Panel) 을 다룰 때: 유전자들이 서로 어떻게 협력하는지 (공발현) 를 아는 방식이 유리합니다.
• 약물 반응 (Reverse Perturbation) 을 예측할 때: 유전자의 활동량 (발현량) 을 정교하게 표현하는 방식이 핵심입니다.

💡 한 줄 요약

"인공지능이 세포를 잘 이해하려면, 단순히 더 많은 데이터를 먹이는 것보다, 그 데이터를 어떤 '레시피'로 준비하느냐가 훨씬 더 중요합니다. HEIMDALL 은 그 레시피를 과학적으로 설계하고 최적화할 수 있게 해주는 도구입니다."

이 연구는 앞으로 더 강력하고 신뢰할 수 있는 생물의학 AI 를 만들기 위해, **데이터를 어떻게 입력하느냐 (토크나이저 설계)**에 집중해야 함을 강조합니다.

기술 요약

1. 문제 정의 (Problem)

단일 세포 RNA 시퀀싱 (scRNA-seq) 데이터를 위한 기초 모델 (Foundation Models, scFMs) 이 등장하면서 생물학적 분석에 혁신을 가져오고 있지만, 실제 성능은 모델이 세포를 어떻게 토큰화 (tokenization) 하느냐에 크게 의존합니다.

• 표준 부재: 텍스트나 이미지와 달리 단일 세포 데이터는 토큰화 표준이 없으며, 기존 모델들은 서로 다른 휴리스틱 (heuristic) 설계와 생물학적 가정을 기반으로 토큰화 방식을 채택하고 있습니다.
• 평가의 어려움: 기존 벤치마크는 모델 크기, 아키텍처, 학습 데이터, 토크나이저 등 여러 요소가 동시에 달라지기 때문에, 성능 차이를 특정 설계 (예: 토크나이저) 로 귀결시키기 어렵습니다.
• 전이 학습의 취약성: 토크나이저 설계가 분포 이동 (distribution shift, 예: 다른 조직, 다른 종, 다른 유전자 패널) 상황에서 모델의 견고한 일반화 능력을 결정짓는 핵심 요소임에도 불구하고, 이를 체계적으로 분석하거나 개선할 수 있는 프레임워크가 부족했습니다.

2. 방법론 (Methodology: HEIMDALL)

저자들은 scFMs 의 토크나이저 설계를 체계적으로 분석하고 재설계하기 위해 HEIMDALL이라는 통합 모듈형 프레임워크를 제안했습니다.

• 모듈화 설계: 기존 토크나이저를 세 가지 기능적 모듈로 분해하여 재구성합니다.
  1. FG (Gene Identity Encoding): 유전자의 정체성을 인코딩 (예: 무작위 임베딩, ESM2, Gene2vec, GenePT 등).
  2. FE (Expression Encoding): 유전자 발현 값을 인코딩 (예: No-op, 연속형, 양분할, 정수 분할 등).
  3. FC (Cell Construction): FG 와 FE 출력을 통합하여 최종 세포 시퀀스를 구성.
     • ORDER: 유전자 토큰의 내재적 순서 정의 (발현량 기반, 염색체 순서, 무작위 등).
     • SEQUENCE: 포함할 유전자 선택 및 시퀀스 구성 (트러네이션, 가중 샘플링 등).
     • REDUCE: 유전자 및 발현 인코딩을 통합하는 방식 (합산, 식별 등).
• 실험 설정:
  • 5 가지 주요 scFM (scGPT, Geneformer, scFoundation, scBERT, UCE) 의 토크나이저를 HEIMDALL 프레임워크 내에서 재구현했습니다.
  • 모델 아키텍처 (Transformer), 하이퍼파라미터, 데이터 크기를 고정하고 토크나이저만 변경하여 공정한 비교를 수행했습니다.
  • 사전 학습 (pretraining) 을 배제하고 처음부터 학습 (from scratch) 하여 토크나이저의 순수한 영향을 측정했습니다.
• 벤치마크:
  • 조직 간 전이 (Cross-tissue): 장 (대장/소장) 데이터로 학습하여 뇌 데이터로 테스트.
  • 종 간 전이 (Cross-species): 인간 데이터로 학습하여 마우스 데이터로 테스트.
  • 유전자 패널 전이 (Gene-panel shift): 훈련과 테스트에서 겹치는 유전자가 적은 공간 전사체 (Spatial Transcriptomics) 데이터.
  • 역적 교란 예측 (Reverse perturbation): 목표 세포 상태를 보고 어떤 교란 (유전자 녹아웃) 이 발생했는지 추론.

3. 주요 기여 (Key Contributions)

• HEIMDALL 프레임워크 개발: scFMs 의 토크나이저 설계를 모듈 단위로 해체하고 재조합할 수 있는 최초의 체계적인 인프라를 제공했습니다.
• 설계 원칙 규명: 단일 최적의 토크나이저가 존재하는 것이 아니라, 유전자 정체성 (Gene Identity), 발현 인코딩 (Expression Encoding), 순서 (Ordering) 라는 세 가지 설계 축 (axes) 이 분포 이동 상황에서 일반화 성능을 결정함을 밝혔습니다.
• 하이브리드 토크나이저 제안: 기존 모델들의 가장 강력한 요소들을 결합한 하이브리드 토크나이저가 개별 최상위 모델보다 더 나은 성능을 낼 수 있음을 증명했습니다.

4. 주요 결과 (Key Results)

A. 분포 이동 (Distribution Shift) 상황에서의 토크나이저의 중요성

• 동일 분포 (In-distribution): 훈련과 테스트 데이터가 동일한 경우, 토크나이저 선택에 따른 성능 차이는 미미했습니다.
• 분포 이동 (Distribution Shift): 조직, 종, 유전자 패널이 변경되는 상황에서는 토크나이저 설계가 성능을 결정하는 결정적 요소가 되었습니다.

B. 구체적인 벤치마크 결과

1. 조직 간 전이 (Cross-tissue):
   • Geneformer-tok 이 가장 우수했으나, 그 차이는 주로 발현량 기반 순서 (Expression-based ORDER) 모듈에서 비롯되었습니다.
   • 사전 학습 (Pretraining) 은 토크나이저 설계에 비해 미미한 영향을 미쳤습니다.
2. 종 간 전이 (Cross-species):
   • UCE-tok이 사전 학습 없이도 우수한 성능을 보였는데, 이는 단백질 서열 기반의 유전자 정체성 인코딩 (ESM2, FG) 덕분이었습니다.
   • 다른 모델들도 ESM2 기반 FG 와 오물로그 (orthology) 매핑을 적용하면 성능이 크게 향상되었습니다.
3. 유전자 패널 전이 (Gene-panel shift, 공간 전사체):
   • scBERT-tok이 가장 우수했습니다. 이는 Gene2vec 기반의 FG(공발현 정보 반영) 와 정수 분할 (Integer binning) FE가 겹치지 않는 유전자에 대한 표현력을 안정화시켰기 때문입니다.
   • UCE-tok 은 이 상황에서 가장 성능이 낮았습니다.
4. 역적 교란 예측 (Reverse perturbation):
   • scBERT-tok이 가장 높은 정확도를 보였습니다.
   • UCE-tok 의 경우, 발현 인코딩 (FE) 과 순서 (ORDER) 모듈을 추가하거나 개선하면 성능이 비약적으로 상승했습니다. 이는 UCE 가 기본적으로 발현 정보를 무시하는 설계였기 때문입니다.

C. 설계 축의 발견

• 유전자 정체성 (FG): 종 간 전이와 패널 전이에서 가장 중요한 요소 (ESM2, Gene2vec 등).
• 발현 인코딩 (FE): 교란 예측 및 패널 전이에서 중요한 요소 (연속형, 분할 등).
• 순서 (ORDER): 조직 간 전이에서 중요한 요소 (발현량 정렬).

5. 의의 및 결론 (Significance)

• 설계 패러다임의 전환: scFMs 의 성능은 단순히 모델 크기나 아키텍처가 아니라, 어떻게 생물학적 정보를 모델에 노출시키느냐 (토크나이저) 에 달려 있음을 강조합니다.
• 실무적 가이드라인: 새로운 조직, 종, 혹은 제한된 유전자 패널을 가진 데이터에 scFMs 을 적용할 때, 어떤 토크나이저 설계 (예: ESM2 기반 FG 사용, 발현량 정렬 등) 를 선택해야 할지에 대한 원칙을 제시합니다.
• 향후 연구 방향: "보편적인 전이성 (Universal Transferability)"을 달성하기 위해서는 "보편적이지 않은 토크나이저 인터페이스"가 필요하며, HEIMDALL 은 이를 위한 모듈형 재사용 인프라를 제공합니다.
• 오픈소스: HEIMDALL 은 오픈소스 파이썬 패키지로 공개되어, 연구자들이 새로운 토크나이저를 쉽게 설계하고 평가할 수 있게 했습니다.

요약하자면, 이 논문은 단일 세포 기초 모델의 성능 향상을 위해 토크나이저 설계를 체계적으로 해체하고 최적화해야 함을 증명하며, 분포 이동 상황에 강한 견고한 모델을 구축하기 위한 구체적인 설계 원칙을 제시했습니다.