ST-JEPA: Joint-Embedding Predictive Architecture for Spatial Transcriptomics

Motivation

• Known: 기존 spatial transcriptomics 분석 방법들은 CellCharter, Banksy, GraphST와 같이 niche identification 등 단일 태스크에 특화되어 단일 스케일에서만 표현을 학습하며, Nicheformer와 Novae 같은 foundation model들도 각각 spatial context를 pretraining에 반영하지 않거나 multi-resolution 계층 구조를 갖추지 못한 한계가 있다.
• Gap: gene, cell, cellular neighborhood라는 여러 생물학적 스케일을 동시에 아우르면서, 이들 간 정보를 latent space에서 관계짓는 self-supervised objective와 서로 다른 gene panel을 가진 플랫폼 간 일반화 메커니즘을 모두 갖춘 tokenization 및 학습 프레임워크가 부재했다.
• Why: 세포 내 분자 프로그램과 조직 내 세포의 공간적 조직이 함께 생물학적 기능을 결정하므로, 이 두 층위를 동시에 모델링하는 multi-scale foundation model은 niche identification, batch integration 등 다양한 spatial biology 하위 분석을 하나의 사전학습 모델로 지원할 수 있어 중요하다.
• Approach: ST-JEPA는 spatial cell graph를 gene-cell-neighborhood 계층을 갖는 구조화된 transformer 시퀀스로 변환하는 multi-scale graph tokenization과, block masking 기반 JEPA objective를 결합하여 context encoder, target encoder(EMA 업데이트), predictor로 구성된 아키텍처에서 masked target의 표현을 latent space에서 예측하도록 학습한다.

Achievement

1. 최고 성능의 niche identification: mouse brain 데이터에서 weighted NMI=0.67로 기존 방법들 대비 가장 우수한 niche identification 성능을 달성했다.
2. 명시적 정렬 목적 없는 batch integration: niche identification 성능이 좋은 방법들 중에서 iLISI 기준 최고의 batch integration 성능을 별도의 integration objective 없이 달성했다.
3. Cross-platform 일반화: technology metatoken을 도입하여 MERFISH와 STARmap처럼 gene panel이 겹치지 않는 서로 다른 기술 플랫폼 간에도 일반화가 가능함을 보였다.
4. 체계적인 ablation을 통한 설계 지침 제공: tokenizer architecture, expression encoding, spatial encoding, sequence length, masking ratio, normalization strategy 등 6가지 설계 축에 대한 systematic ablation을 통해 spatial transcriptomics를 위한 self-supervised learning 설계에 실질적인 가이드를 제시했다.

How

• k=10인 k-nearest neighbor graph를 spatial coordinate 기반으로 구성하고, 각 cell에 대해 index cell과 k개의 neighbor로 이루어진 cellular neighborhood sequence를 추출한다.
• 각 cell 내 gene을 발현량 기준으로 순위화하여 gene vocabulary의 고유 토큰으로 매핑하고(gene-rank tokenizer, Geneformer에서 착안), 각 토큰에 gene identity, sinusoidal positional rank, 거리 기반 sinusoidal segment embedding, MLP 기반 count value embedding(scFoundation에서 착안)을 부여해 L=64 토큰/cell, 총 N=704 토큰의 multi-segment sequence를 구성한다.
• context encoder(transformer), EMA로 업데이트되는 target encoder, 그리고 더 작은 predictor로 구성된 JEPA 아키텍처를 사용하며, neighborhood sequence 전체에 걸쳐 gene 토큰을 block 단위로 무작위 마스킹하여(masking ratio ρ=0.6) context로부터 masked target의 latent 표현을 예측하도록 ℓ1 loss로 학습한다.
• mouse brain의 MERFISH(1,085 genes)와 STARmap(1,015 genes) 데이터를 gene panel을 축소하지 않고 각 기술의 전체 panel로 학습하며, technology metatoken을 통해 서로 다른 기술 간 정보를 통합하고, 추론 시에만 겹치는 gene을 사용해 cross-technology 평가를 수행한다.
• niche identification은 Leiden clustering 기반 NMI/ARI로, batch integration은 iLISI와 MMD로 평가하며, 모든 방법에 동일한 k-NN spatial graph를 적용해 공정하게 비교한다.

Originality

• Spatial transcriptomics에 JEPA 패러다임을 최초로 적용하여, 데이터 공간 재구성이나 augmentation 의존적 contrastive learning의 한계를 피하고 latent space에서의 예측 학습을 도입했다.
• Gene, cell, cellular neighborhood라는 세 가지 생물학적 해상도를 하나의 multi-segment transformer sequence로 동시에 인코딩하는 multi-scale graph tokenization을 새롭게 설계했다.
• Neighborhood sequence 전반에 걸친 gene 토큰 block masking을 통해 분자 정보와 공간 정보를 명시적으로 연결하는 학습 신호를 고안했다.
• gene panel이 겹치지 않는 이종 플랫폼 간 일반화를 가능하게 하는 technology metatoken을 도입해, 명시적 batch correction 없이 cross-platform 통합을 달성했다.

Limitation & Further Study

• 실험이 mouse brain의 두 기술(MERFISH, STARmap) 단일 배치에 국한되어 있어, 더 다양한 조직·종·기술 조합에서의 일반화 성능 검증이 필요하다.
• Nicheformer(110M cells)나 Novae(30M cells) 등 비교 대상 foundation model에 비해 학습 데이터 규모가 제한적일 가능성이 있어, 대규모 사전학습 시의 스케일링 효과가 충분히 검증되지 않았다.
• k=10으로 고정된 k-NN 그래프와 L=64 고정 시퀀스 길이 등 하이퍼파라미터 선택이 다양한 조직 밀도나 세포 유형 분포에 대해 얼마나 강건한지 추가 검증이 필요하다.
• niche identification과 batch integration 외에 세포 유형 주석(cell type annotation), 유전자 발현 예측 등 더 폭넓은 하위 태스크에서의 성능 평가가 향후 과제로 남아 있다.

Evaluation

총평: Spatial transcriptomics를 위한 최초의 multi-scale JEPA 아키텍처로서 tokenization, self-supervised objective, cross-platform generalization을 통합적으로 설계한 참신하고 실용적인 기여이며, workshop paper 특성상 데이터 규모와 검증 범위가 제한적이나 후속 연구를 위한 명확한 방향성을 제시한다.