HEXST: Hexagonal Shifted-Window Transformer for Spatial Transcriptomics Gene Expression Prediction

Motivation

• Known: 기존 연구들은 CNN, graph 기반 모델, Transformer 기반 attention 메커니즘을 활용해 H&E histology 이미지로부터 spatial gene expression을 예측해왔으며, ST-Net, Hist2ST, HisToGene, TCGN, EGNv1/v2, PEKA, NH2ST 등이 대표적이다.
• Gap: 기존 방법들은 Visium 등 널리 쓰이는 spot-array 플랫폼의 hexagonal lattice 샘플링을 무시한 채 Cartesian 혹은 geometry-agnostic locality를 가정하며, point-wise regression objective는 neighborhood aggregation과 결합되어 예측을 local mean으로 수렴시켜 gene-specific spatial heterogeneity를 훼손하는 over-smoothing 문제를 야기한다.
• Why: tumor–stroma boundary, immune niche 등 임상적으로 중요한 생물학적 신호는 gene expression의 절대값뿐 아니라 급격한 공간적 변화(spatial contrast)에 담겨 있으므로, 이를 보존하며 정확히 예측하는 것이 spatially informed transcriptomic 분석과 biomarker 발굴에 필수적이다.
• Approach: 저자들은 hexagonal spot 좌표에 직접 작동하는 shifted-window attention과 hexagonal rotary positional encoding을 결합한 geometry-aligned Transformer를 제안하고, 여기에 deviation-matching objective와 single-cell foundation model 기반 transcriptomic prior alignment를 학습 목표로 추가하였다.

Achievement

Figure 4. Qualitative comparison of spatial gene expression pre-

1. Geometry-aligned Transformer 제안: hexagonal 좌표 기반의 shifted-window attention(HexMSA)을 설계하여 hexagonal lattice의 non-Cartesian 구조에 맞는 local-to-global contextual modeling을 가능케 함.
2. HexRoPE 도입: axial/cube 좌표계에서 정의된 rotary positional encoding을 제안하여 principal lattice direction을 따라 query/key feature를 회전시킴으로써 hexagonal 상에서의 상대적 위치 관계를 효과적으로 인코딩함.
3. Spatial contrast 보존을 위한 deviation-matching objective: point-wise regression을 보완하는 contrast-sensitive 목표를 통해 예측이 local mean으로 collapse되는 over-smoothing 현상을 완화함.
4. Transcriptomic prior 활용: pretrained single-cell foundation model의 transcriptomic embedding을 학습 중 정렬(align)하여 생물학적으로 의미 있는 표현을 주입함.
5. 7개 spatial transcriptomics 데이터셋에서 SOTA 성능 달성: 기존 최신 모델들을 일관되게 능가하며 gene-wise contrast와 spatial heterogeneity를 보존하는 정확하고 robust한 예측을 제공함.

How

Figure 2. Visualization of hexagonal window partitioning. By

• WSI로부터 spot-level image patch를 추출하고, pretrained visual pathology foundation model E를 통해 각 patch를 Din-차원 visual token vi로 인코딩함.
• 다단계(multi-stage) hexagonal shifted-window Transformer 구조를 사용하여 hexagonal spot 좌표 상에서 window partitioning과 shifting을 연속 블록에 걸쳐 적용, local self-attention과 cross-window information exchange를 통해 receptive field를 local에서 tissue-wide로 점진적으로 확장함(HexMSA).
• HexRoPE를 통해 axial/cube hexagonal 좌표에서 정의된 relative positional relationship을 attention에 주입함.
• 학습 시 point-wise regression loss에 더해, 예측과 ground truth 간 spot 간 deviation을 매칭하는 deviation-matching objective를 추가하여 spatial contrast를 보존함.
• 추가로 pretrained single-cell foundation model(예: Geneformer, xTrimoscFoundation 계열)로부터 얻은 transcriptomic embedding과 중간 표현을 정렬하는 transcriptomic feature alignment를 학습 목표에 포함함.
• 7개 ST 데이터셋에 대해 정량적 성능 비교와 정성적 spatial gene expression prediction 시각화, transcriptomics-guided survival 관련 분석을 수행함.

Originality

• Cartesian 가정을 벗어나 hexagonal spot-array 기하 구조 자체를 아키텍처적 inductive bias로 명시적으로 인코딩한 최초의 시도로 보이며, shifted-window Transformer를 hexagonal lattice에 맞춰 재설계함.
• 1D/2D Cartesian 좌표에 국한되었던 RoPE를 hexagonal axial/cube 좌표계로 확장한 HexRoPE는 새로운 positional encoding 방식임.
• point-wise regression의 over-smoothing 문제를 spot 간 deviation-matching이라는 대조적(contrast-sensitive) objective로 직접 겨냥해 해결하려는 접근이 참신함.
• 단일세포 foundation model의 transcriptomic prior를 histology 기반 gene expression 예측 학습에 보조 신호로 결합한 점도 기존 연구와 차별화됨.

Limitation & Further Study

• 논문 발췌에는 hexagonal window partitioning의 계산 복잡도나 실제 inference 속도에 대한 구체적 분석이 부족해, 실제 임상 적용 시 효율성 검증이 추가로 필요함.
• deviation-matching objective와 transcriptomic alignment의 개별 기여도(ablation)에 대한 세부 분석이 발췌 부분에서 충분히 드러나지 않아, 각 구성요소의 상대적 중요도를 명확히 하는 추가 실험이 요구됨.
• 7개 데이터셋이 대부분 Visium 계열 hexagonal 플랫폼에 한정되어 있을 가능성이 있어, 다른 플랫폼(비-hexagonal, 예: MERFISH 등 imaging-based assay)으로의 일반화 가능성에 대한 검증이 필요함.
• pretrained single-cell foundation model에 대한 의존성이 있어, 해당 foundation model의 편향이나 도메인 불일치가 결과에 미치는 영향에 대한 분석이 후속 연구로 요구됨.

Evaluation

총평: hexagonal geometry를 명시적으로 반영한 shifted-window Transformer와 HexRoPE, deviation-matching objective의 결합은 spatial transcriptomics 예측 분야에서 기하학적 정합성과 spatial heterogeneity 보존이라는 실질적 문제를 잘 짚어낸 참신하고 견고한 기여로 평가된다.