BioM-JEPA: joint-embedding prediction of graph-connected gene blocks in single cells

기술 요약: BioM-JEPA

문제 정의
단일 세포 RNA 시퀀싱(scRNA-seq)은 생물학적 프로그램에 대해 희소하고 깊이 의존적인 관측치를 제공한다. 기존의 자기 지도 학습 모델들은 개별 유전자를 재구성(토큰 예측)하거나 발현 값을 예측하는 방식으로 학습한다. 저자들은 이러한 접근 방식에 한계가 있다고 주장한다. 단일 유전자 측정값은 고립되었을 때 희소하고 노이즈가 많으며 생물학적으로 모호하기 때문이다. 또한, 유전자 수준의 예측을 성공적으로 최적화한다고 해서 반드시 견고하고 고차원적인 세포 표현(representation)이 나타나는 것은 아니다. 모델은 수렴하더라도 표현(embedding)이 저차원이거나 검출된 유전자 수(sequencing depth)와 같은 기술적 아티팩트에 강하게 결합된 상태로 남을 수 있다. 핵심적인 설계 질문은 다음과 같다: 재사용 가능하고 생물학적으로 의미 있는 세포 표현을 학습하기 위해 모델은 무엇을 예측해야 하는가?

방법론: BioM-JEPA
저자들은 단일 세포를 위해 설계된 BioM-JEPA(Joint-Embedding Predictive Architecture)를 제안한다. 마스킹된 이미지 영역이나 비디오의 시공간 블록을 예측하는 기존의 JEPA와 달리, BioM-JEPA는 **그래프 연결 유전자 블록(graph-connected gene blocks)**의 집계된 표현을 예측한다.

• 예측 단위 (그래프 연결 블록): 개별 유전자를 마스킹하는 대신, 모델은 이진 유전자 그래프(binary gene graph)로부터 샘플링된 연결된 유전자 집합을 타겟 블록으로 정의한다. 이 그래프는 높은 신뢰도의 단백질 연관성(STRING v12)과 미라벨링된 사전 학습 코퍼스에서 도출된 전사체 전체의 공발현 증거를 결합하여 구축된다.
• 학생-교사(Student-Teacher) 구조:
  • 학생(Student): 세포 내 보완적인 유전자 집합(컨텍스트)을 관찰하고, 숨겨진 타겟 블록의 집계된 표현을 예측한다.
  • 교사(Teacher): 전체 관측된 유전자 세트를 수신하고 타겟 표현을 제공한다. 결정적으로, 교사의 타겟은 개별 유전자 값을 재구성하는 것이 아니라, 관측된 유전자들로부터 계산된 집계된 상태(aggregate state)이다.
  • 학습 목표: 모델은 학생의 예측된 블록 표현과 교사의 집계된 타당 타겟 표현 사이의 거리를 최소화한다. 여기에는 여러 임베딩 차원 간의 변동성을 보장하고 중복된 공분산을 줄이기 위한 정규화 항이 포함된다. 개별 유전자 재구성 손실은 사용되지 않는다.
• 선형 어텐션(Linear Attention): 밀집 소프트맥스 어텐션(dense softmax attention)의 이차적(quadratic) 계산 비용 없이 방대한 수의 관측된 유전자를 처리하기 위해, 모델은 학생, 교사 및 예측 네트워크 전반에 걸쳐 선형 어텐션을 채택한다. 이를 통해 모델은 유전자 수에 따라 선형적으로 증가하는 계산량으로 전체 관측된 유전자 시퀀스를 처리할 수 있다.
• 데이터 및 규모: 모델은 2,210만 개의 세포 컬렉션(5억 개 세포 규모의 scBaseCount 리소스 중 일부)에서 사전 학습되었다. 최종 표현은 1,024만 번의 세포 제시(presentation) 후에 분석되었으며, 이는 Geneformer, scGPT, scFoundation, CellFM과 같은 유사 모델들에 비해 현저히 적은 학습 노출량이다.

주요 기여

1. 블록 수준 예측: 독립적인 유전자 수준의 타겟에서 벗어나, 전사체 JEPA를 위한 근본적인 예측 단위로서 그래프 연결 유전자 블록을 도입하였다.
2. 아키텍처 효율성: 선형 어텐션이 인위적인 토큰 제한이나 이차적 어텐션 행렬 없이도 광범위한 유전자 컨텍스트를 효율적으로 처리할 수 있음을 입증하였다.
3. 통제된 비교: 공유된 백본(backbone)과 추출 파이프라인을 사용하여 세 가지 특정 베이스라인과 엄격하게 비교하였다:
   • Token-IJEPA: 개별 유전자를 예측함 (표준 JEPA).
   • Random Block: 그래프 연결성 없이 무작위로 샘플링된 유전자 세트를 사용함.
   • Decoder-Only: 그래프 블록으로부터 발현 값을 재구성함.

결과

• 표현 품질: BioM-JEPA는 Token-IJEPA(~19) 및 Decoder-only(<10)에 비해 훨씬 높은 유효 순위(effective rank)(학습 중 약 19에서 약 40으로 상승)를 생성하였다. 이는 모델이 많은 독립적인 방향에 걸쳐 변동성을 유지했음을 나타낸다.
• 기술적 견고성: BioM-JEPA 임베딩은 모든 대조군에 비해 검출된 유전자 깊이(sequencing depth)와의 연관성이 가장 낮게 나타났으며, 이는 모델이 생물학적 신호를 확률적 샘플링 노이즈와 성공적으로 구분했음을 시사한다.
• 다운스트림 성능:
  • 주석(Annotation): hPancreas 및 cortex 데이터셋에 대한 퓨샷(few-shot) 세포 유형 주석 작업에서 가장 높은 평균 Top-5 macro-F1 점수(0.835)를 달 기록했다.
  • 재구성(Reconstruction): 동결된(frozen) 임베딩은 연속적인 발현 구조를 유지하며, 다른 모델보다 고발현 유전자를 더 잘 회복하고 Reactome 경로 상관관관계를 더 잘 보존하였다.
  • 섭동 반응(Perturbation Response): 유전적 섭동(Adamson, Norman, Dixit 데이터셋)에 대한 로그 폴드 변화(log-fold-change) 반응을 예측하는 데 있어 가장 낮은 총 평균 제곱 오차(MSE)를 달성하였고, 상위 반응 유전자를 가장 잘 회복하였다.
• 효율성: hPancreas에 대한 매칭된 1-에포크 미세 조정 실험(배치 크기 8, 학습 가능한 파라미터 약 2,600만 개)에서, BioM-JEPA는 scFoundation에 비해 학습 처리량은 5.75배, 홀드아웃(held-out) 임베딩 추출은 3.76배 더 빨랐다.

의의 및 주장
본 논문은 그래프 연결 유전자 블록이 단일 세포 생물학을 위한 JEPA 스타일의 표현 학습에 효과적인 예측 단위 역할을 한다고 주장한다. 저자들은 다음과 같이 단언한다:

• 유전자 수준의 목표(토큰 예측)의 수렴이 반드시 견고한 세포 표현을 보장하는 것은 아니다. 즉, 예측 단위의 규모가 중요하다.
• 블록 수준의 목표는 연결된 유전자가 공유된 잠재 상태(latent state)를 위한 중복된 증거를 제공하므로, 희소한 전사체 데이터에 매우 적합하며 단일한 희소 측정값에 대한 의존도를 낮춘다.
• 학습된 표현은 세포 유형이나 경로에 대한 명시적인 감독 없이도 생물학적 구조(정형화된 세포 정체성, 경로 관계, 유전적 섭동 조성)를 포착한다.
• 선형 어텐션은 이 접근 방식을 대규모 유전자 컨텍스트에서도 계산적으로 실용적이게 만들며, 다른 대규모 모델들에 비해 상대적으로 압축된 학습 노출을 통해서도 유용한 표현이 나타날 수 있게 한다.

저자들은 이 접근 방식이 유전자를 독립적인 "패치"로 취급하는 대신 분자적 관계(경로, 복합체)에 기반한 예측 단위를 정의하는 직접적인 경로를 제공하며, 미래의 생물학적 표현 모델을 위한 토대를 제공한다고 결론짓는다.