GREmLN: A Cellular Graph Structure Aware Transcriptomics Foundation Model

GREmLN 은 무질서한 유전자 특성을 그래프 신호 처리를 통해 생물학적 인과관계를 반영하는 어텐션 메커니즘에 직접 통합함으로써, 단일 세포 전사체 데이터에서 복잡한 장기적 조절 의존성을 포착하고 세포 유형 주석 및 역 교반 예측 등 다양한 작업에서 최첨단 성능을 달성하는 새로운 기초 모델입니다.

핵심

1. 문제: "무질서한 단어장"과 "지시 없는 지도"

기존의 인공지능 (특히 언어 모델) 은 글을 읽을 때 순서가 중요하다고 배웠습니다. "사과"라는 단어가 "먹는다"라는 단어 앞에 오면 의미가 통하죠. 하지만 세포 안의 유전자 데이터는 다릅니다.

• 비유: 세포의 유전자 데이터를 한 권의 책이라고 상상해 보세요. 그런데 이 책은 단어들이 책장에 아무렇게나 흩어져 있는 상태입니다. "A 유전자, B 유전자, C 유전자"가 순서대로 있는 게 아니라, 그냥 무작위로 쌓여 있습니다.
• 기존 모델의 한계: 기존의 AI 는 이 무질서한 단어들을 읽을 때 "순서"를 찾으려 애쓰다가 헷갈려 합니다. 유전자 A 와 유전자 B 가 서로 멀리 떨어져 있어도, 실제로는 아주 친한 친구 (상호작용) 일 수 있는데, AI 는 그 관계를 놓치기 쉽습니다.

2. 해결책: "유전자 지도"를 활용한 GREmLN

연구팀은 이 문제를 해결하기 위해 GREmLN이라는 새로운 모델을 만들었습니다. 이 모델의 핵심 아이디어는 "유전자들 사이의 관계 지도 (그래프)"를 AI 가 읽는 방식에 직접 넣었다는 점입니다.

• 비유:
  • 기존 AI: 무작위로 흩어진 단어들을 보고 "이게 무슨 뜻일까?"라고 추측만 합니다.
  • GREmLN: 이 모델은 세포 안에 **유전자들 사이의 친구 관계도 (지도)**가 그려져 있다고 가정합니다. "유전자 A 는 유전자 B 와 아주 친하고, 유전자 C 는 멀리 있지만 간접적으로 연결되어 있어"라는 정보를 AI 가 처음부터 알고 있는 것입니다.

이 모델은 **확산 (Diffusion)**이라는 개념을 사용합니다.

• 비유: 한 유전자에 물방울을 떨어뜨리면, 그 물방울이 친구 관계도를 타고 다른 유전자들에게 퍼져나갑니다. 이 과정을 통해 AI 는 "이 유전자는 저 유전자와 얼마나 깊은 관계가 있는지"를 자연스럽게 이해하게 됩니다.

3. 이 모델이 얼마나 뛰어난가? (실험 결과)

연구팀은 GREmLN 이 다른 유명한 모델들 (scGPT, Geneformer 등) 보다 훨씬 잘한다고 증명했습니다.

1. 세포의 정체성 파악 (Cell Type Annotation):

   • 상황: 낯선 세포를 보고 "이게 무슨 세포야?"라고 물어보는 테스트입니다.
   • 결과: GREmLN 은 마치 숙련된 의사가 눈만 보고도 병명을 정확히 진단하듯, 새로운 세포의 종류를 다른 모델들보다 훨씬 정확하게 맞췄습니다. 특히 훈련할 때 보지 못했던 새로운 종류의 세포 (예: 암세포 속 면역세포) 에 대해서도 잘 알아맞혔습니다.

2. 유전자 지도 이해 (Graph Structure Understanding):

   • 상황: 유전자들 사이의 연결 고리 (누가 누구를 조절하는지) 를 복원하는 테스트입니다.
   • 결과: 지도의 일부가 지워졌을 때, GREmLN 은 나머지 조각을 보고 잃어버린 부분을 완벽하게 재구성했습니다. 이는 모델이 유전자 간의 복잡한 관계를 깊이 이해하고 있다는 뜻입니다.

3. 약물 반응 예측 (Reverse Perturbation):

   • 상황: "이 세포의 유전자가 변했다면, 어떤 약을 줬을까?"를 추측하는 테스트입니다.
   • 결과: GREmLN 은 약간의 변화만으로도 어떤 약이 작용했는지를 다른 모델들보다 훨씬 잘 찾아냈습니다. 이는 신약 개발이나 맞춤형 치료에 큰 도움이 될 수 있습니다.

4. 왜 이 연구가 중요한가?

• 효율성: GREmLN 은 다른 거대 모델들보다 매우 가볍습니다. (파라미터 수가 1/3~1/10 수준). 복잡한 지도를 활용했기 때문에, 무작정 모델 크기를 키우지 않아도 똑똑해질 수 있었습니다.
• 해석 가능성: 이 모델은 단순히 "정답"만 알려주는 게 아니라, 어떤 유전자들이 서로 연결되어 있는지를 보여주므로, 생물학자들이 새로운 발견을 하는 데 도움을 줍니다.

요약

GREmLN은 세포 속 유전자 데이터를 읽을 때, 단순히 단어 순서를 따지는 게 아니라 "유전자들 사이의 친구 관계도 (지도)"를 참고하여 읽는 똑똑한 AI입니다.

이 모델은 마치 어둠 속에서 지도를 들고 길을 찾는 사람처럼, 복잡한 세포의 작동 원리를 더 빠르고 정확하게 이해하게 해줍니다. 이는 향후 암 치료, 신약 개발, 그리고 우리 몸의 복잡한 생물학적 현상을 이해하는 데 혁신적인 도구가 될 것입니다.

기술 요약

1. 문제 정의 (Problem)

단일 세포 RNA 시퀀싱 (scRNA-seq) 데이터의 폭발적인 증가로 인해 세포의 특성과 행동을 포착하는 '기초 모델 (Foundation Model)' 개발의 기회가 열렸습니다. 그러나 기존 Transformer 기반 언어 모델은 다음과 같은 한계를 가지고 있습니다.

• 순서 없는 데이터의 특성: scRNA-seq 데이터는 유전자 발현 값의 집합이며, 유전자 토큰들 사이에 언어 모델이 가정하는 것과 같은 선형적 순서나 절대/상대적 위치 정보가 존재하지 않습니다.
• 기존 모델의 접근법: 기존 모델 (scGPT, Geneformer 등) 은 유전자를 이산적 토큰으로 간주하고 임의의 순서로 시퀀스를 구성하거나, 단순한 주의 편향 (attention bias) 을 추가합니다. 이는 유전자 간의 복잡한 비국소적 (non-local) 의존성과 인과 관계를 효과적으로 인코딩하지 못합니다.
• 분자 상호작용의 부재: 유전자 발현은 유전자 조절 네트워크 (GRN) 나 단백질 - 단백질 상호작용 (PPI) 과 같은 분자 상호작용 그래프를 통해 설명될 수 있는데, 기존 모델들은 이러한 생물학적 그래프 구조를 충분히 활용하지 못합니다.

2. 방법론 (Methodology)

저자들은 GREmLN (Gene Regulatory Embedding-based Large Neural model) 을 제안합니다. 이는 그래프 신호 처리 (Graph Signal Processing) 를 활용하여 유전자 토큰의 그래프 구조를 어텐션 메커니즘 내부에 직접 통합한 기초 모델입니다.

핵심 아키텍처 및 기술

1. 토큰화 및 초기 임베딩 (Tokenization & Initial Embeddings):

   • 입력은 유전자 식별자 (Gene ID) 와 유전자 발현 순위 (Gene Rank) 로 구성됩니다.
   • 발현 값은 이진화 (binning) 되어 연속적인 값을 이산적인 순위 토큰으로 변환하며, <MASK> 토큰을 사용하여 마스킹 언어 모델링을 수행합니다.
   • <CLS> 토큰은 모든 유전자를 주시하지만 그래프 연산에는 참여하지 않습니다.

2. 그래프 확산 커널 어텐션 (Graph Diffusion Kernel Attention, GDKA):

   • 핵심 아이디어: 쿼리 (Query) 임베딩을 토큰 그래프 (GRN 또는 PPI) 에서 유도된 확산 커널 (Diffusion Kernel) 그램 행렬로 변환하여 그래프 구조 정보를 주입합니다.
   • 수학적 과정:
     • 그래프 라플라시안 행렬 $L$ 의 고유값 분해를 수행합니다.
     • 스펙트럼 필터 $\kappa$ (예: $e^{-\beta \Lambda}$) 를 적용하여 노이즈를 필터링하고 다중 홉 (multi-hop) 의존성을 포착하는 커널 행렬 $\Phi_L$ 을 생성합니다.
     • 쿼리 벡터 $Q$ 를 $\Phi_L(Q)$ 로 변환하여 어텐션 점수를 계산합니다. 이는 그래프에서 정의된 저주파 (장거리) 감쇠 패턴을 따르도록 어텐션을 유도합니다.
     • 키 (Key) 와 값 (Value) 은 변환하지 않아 원래 토큰 정보를 보존합니다.

3. 대규모 그래프를 위한 근사화 (Approximation):

   • 대규모 생물학적 네트워크에서 행렬 지수 연산과 스펙트럼 분해의 계산 비용을 줄이기 위해 체비셰프 다항식 (Chebyshev Polynomials) 을 사용하여 커널 행렬을 근사합니다. 이를 통해 확장성 (Scalability) 을 확보합니다.

4. 학습 목표 (Training Objectives):

   • 그래프 조건부 마스킹 언어 모델링 (Masked Language Modeling) 을 수행합니다. 특정 세포 유형의 정규화된 라플라시안 $L_c$ 와 마스킹되지 않은 유전자 발현을 기반으로 마스킹된 유전자의 발현 값을 예측합니다.

3. 주요 기여 (Key Contributions)

• 생물학적 유도를 통한 어텐션 메커니즘: 유전자 발현 데이터의 순서 없는 특성을 해결하기 위해, 유전자 간의 생물학적 상호작용 (GRN/PPI) 을 그래프 구조로 인코딩하고 이를 어텐션 메커니즘에 통합한 최초의 기초 모델 중 하나입니다.
• 해석 가능성 및 효율성: 그래프 구조를 인덕티브 바이어스 (Inductive Bias) 로 활용하여 파라미터 효율성을 높이고 (기존 모델 대비 1/3 이하의 파라미터), 학습 수렴을 가속화합니다.
• 장거리 의존성 포착: 확산 커널을 통해 유전자 간의 장거리 조절 의존성을 효과적으로 학습하며, 그래프 구조를 이해하는 능력을 입증했습니다.
• 범용성: GRN 과 PPI 네트워크 등 다양한 분자 상호작용 그래프를 입력으로 사용할 수 있는 유연한 프레임워크를 제공합니다.

4. 실험 결과 (Results)

GREmLN 은 scGPT, Geneformer, scFoundation 등 최신 기초 모델들과 비교 평가되었습니다.

• 세포 유형 주석 (Cell Type Annotation):

  • 인간 면역 세포 데이터셋에서 정밀도 (Precision), 재현율 (Recall), F1 점수에서 모든 베이스라인 모델을 능가했습니다.
  • Zero-shot 성능: 학습 중 보지 못한 비면역 세포 (Non-immune cells) 에 대해서도 우수한 Zero-shot 분류 성능을 보여주었습니다.
  • 파라미터 효율성: 1030 만 개의 파라미터로 베이스라인 모델 (scFoundation 등) 보다 훨씬 적은 파라미터로 더 높은 성능을 달성했습니다.

• 그래프 구조 이해 (Graph Structure Understanding):

  • 학습된 임베딩을 사용하여 보이지 않는 유전자 조절 네트워크 (GRN) 의 간선 (edge) 을 예측하는 작업에서 가장 높은 AUROC 및 Average Precision을 기록했습니다.
  • 특히 암 침윤성 골수 세포 (Cancer Infiltrating Myeloid) 와 같이 병리적 조건에서 그래프 구조가 재배선된 경우에도 우수한 일반화 성능을 보였습니다.

• 역변형 예측 (Reverse Perturbation Prediction):

  • Perturb-Seq 데이터를 기반으로 유전자 교란 (CRISPR) 을 예측하는 미세 조정 (Fine-tuning) 작업에서 SOTA 성능을 달성했습니다.
  • PPI 기반 사전 학습은 고정된 임베딩에서 강점을 보였고, GRN 기반 사전 학습은 미세 조정 시 더 우수한 성능을 보여주어, 모델이 세포 유형별 조절 회로를 학습했음을 시사합니다.

• Ablation Study:

  • 그래프 구조를 제거한 일반 Transformer (Vanilla) 와 비교 시, 모든 태스크에서 성능이 크게 저하되었습니다. 이는 그래프 구조가 단순한 정보가 아닌 예측에 필수적인 인덕티브 바이어스임을 증명합니다.

• 확장성 (Scaling Behavior):

  • 모델 레이어 수를 증가시킬수록 (1 층 -> 6 층) 역변형 예측 성능이 단조 증가하여, GREmLN 이 모델 크기 증가에 따라 체계적으로 성능이 향상됨을 확인했습니다.

5. 의의 및 결론 (Significance)

이 논문은 단일 세포 전사체 데이터 분석에 있어 그래프 구조 기반의 기초 모델의 중요성을 입증했습니다.

• 생물학적 통찰력 통합: 단순한 통계적 패턴 학습을 넘어, 생물학적 지식 (GRN/PPI) 을 모델 아키텍처에 직접 통합함으로써 더 해석 가능하고 생물학적으로 타당한 표현을 학습합니다.
• 데이터 효율성: 대규모 파라미터를 요구하는 기존 접근법과 달리, 그래프 구조를 활용한 인덕티브 바이어스를 통해 적은 파라미터로도 높은 성능을 달성하여 계산 자원을 효율적으로 사용합니다.
• 미래 방향: 이 모델은 암의 진행 메커니즘, 약물 반응 예측, 그리고 최적의 치료 표적 발굴을 위한 강력한 도구로 활용될 수 있으며, 특히 복잡한 생물학적 시스템에서의 장거리 조절 관계를 이해하는 데 필수적인 프레임워크를 제공합니다.

결론적으로 GREmLN 은 단일 세포 데이터의 무질서한 특성을 생물학적 그래프 구조로 해결하고, 차세대 전사체 기초 모델의 새로운 표준을 제시한 연구입니다.