Reasoning Across Space: Tiny Recursive Models for Spatial Omics

Motivation

• Known: 기존 spatial omics 모델들은 graph 기반 방법(SpaGCN, STAGATE, GraphST, 3d-OT)이나 patch 기반 histology-to-expression 예측 방법(STEM, BLEEP, TRIPLEX, HisToGene)처럼 대부분 slice-specific으로 학습되어 지역적(local) 영역에서 단일 pass로 spatial reasoning을 수행한다. 최근 recurrent forward pass(Geiping et al., 2025; Goyal et al., 2026)나 Tiny Recursive Models(TRM, Jolicoeur-Martineau, 2025)와 같은 재귀적 연산이 파라미터 수 증가 없이 유효 깊이(depth)를 늘려 data-scarce 환경에서 reasoning 능력을 향상시킬 수 있음이 알려져 있다.
• Gap: 기존 spatial omics 방법들은 local processing, 부족한 annotation, 조직 section 간 낮은 일반화 성능이라는 한계를 가지며, 단일 pass 방식으로는 많은 생물학적 신호의 근간이 되는 장거리 의존성(long-range dependency)을 통합할 수 없다. 또한 대부분의 방법이 test sample 자체에 직접 fitting하는 방식(test-time training)에 의존해 slice 및 tissue type 간 일반화가 취약하다는 문제가 있다.
• Why: 재귀적 연산이 spatial omics에 자연스러운 inductive bias를 제공하는지 검증함으로써, annotation이 부족하고 test-time training 없이도 조직 slice 및 tissue type 간 일반화가 가능한 확장 가능하고 data-efficient한 프레임워크의 가능성을 제시하기 때문에 중요하다.
• Approach: TRM을 core architecture로 채택하여 latent reasoning state와 prediction state를 반복적으로 정제(iterative refinement)하는 구조를 SDI와 H2G 두 과제에 적용하고, 엄격한 leave-one-slice-out(LOSO) 설정 하에서 parameter-matched baseline 및 다수의 state-of-the-art per-slice 방법들과 비교 평가하였다.

Achievement

Figure 2. H&E overlay of ground truth and TRM predictions with 1,

1. LOSO SDI 성능 우위: TRM이 12개 DLPFC LOSO fold 전체에서 No Recursion 2-Layer와 Universal Transformer baseline을 모두 능가했으며(mean ARI 0.559 vs 0.461, 0.483), SpaGCN, STAGATE, GraphST 등 다수의 slice-specific unsupervised 방법도 상회했다(단, test slice 학습을 활용하는 3d-OT는 예외).
2. H2G SOTA 달성: 지도학습 방법들 중 gene expression prediction에서 state-of-the-art 성능을 기록했다.
3. Test-time scaling 및 adaptive computation 효과 검증: 재귀 단계(total processing block, latent recursion step) 수를 늘릴수록 ARI가 향상되고 약 3 block에서 saturate하는 양상을 확인했으며, latent recursion step 증가가 피질층 경계(cortical layer boundary) 예측을 선명하게 함을 시각적으로 보였다.
4. 부분적 tissue-type 간 일반화: TRM이 학습한 reasoning latent space가 domain identification에서 다른 tissue type으로 zero-shot 전이가 가능함을 보였다.

How

• Task 정의: (1) Visium 기반 spot-level gene expression을 spatial domain으로 clustering하는 SDI task, (2) H&E histology image로부터 spot별 gene expression profile을 예측하는 H2G task
• Architecture: 2-layer transformer 백본(896K 파라미터)에 NZL=6 latent step, NZH=3 deep step, Nouter=8 outer step을 적용한 TRM을 사용하며, 하나의 forward step에서 (NZL+1)·NZH·2=42회의 forward pass, 최대 rollout은 336회
• Baseline: capacity를 통제한 No Recursion 2-Layer(NZL=NZH=Nouter=1)와, depth 및 weight-sharing을 매칭한 Universal Transformer(2-layer를 21회 반복, effective depth 42)를 구성하여 recursion의 기여를 분리 검증
• Adaptive Computation Time(ACT): 연속적인 회귀 문제인 H2G task를 위해 PCC-halt, RVD-halt, MSE-halt, per-class PCC-halt 등 4가지 halt 전략을 설계하여 halt head가 panel-wide metric 또는 gene별 halt probability를 기준으로 조기 종료를 학습하도록 함
• 평가: 12개 DLPFC slice에 대한 leave-one-slice-out(LOSO) 검증 및 HER2ST 데이터에서 halt-step 분포 분석, human breast cancer로의 zero-shot transfer 실험 수행

Originality

• Spatial omics 분야에 test-time training 없이 재귀적 연산(recursive computation)을 도입한 최초의 시도로, TRM을 global spatial reasoning의 inductive bias로 활용
• 연속적 회귀 문제(gene expression prediction)에 맞춘 새로운 adaptive computation time(ACT) 전략들(PCC-halt, RVD-halt, MSE-halt, per-class PCC-halt)을 제안하여 기존 이산적 ACT(Graves, 2016)를 확장
• Depth-matched, weight-sharing-matched 등 다양한 통제된 baseline 설계를 통해 성능 향상의 원인이 recursion 구조 자체(latent-state와 solution refinement의 구조적 분리)에 있음을 입증하려는 실험 설계

Limitation & Further Study

• Workshop paper 분량의 발췌본으로 실험 세부사항(예: fully depth-matched baseline이 학습에 실패한 이유에 대한 분석)이 충분히 제시되지 않음
• Tissue-type 간 일반화가 "부분적(partially generalizable)"이라고 명시되어 있어 완전한 domain-agnostic 일반화에는 한계가 있음을 저자들 스스로 인정
• 평가가 DLPFC, HER2ST 등 제한된 데이터셋에 국한되어 있어 다양한 조직 유형과 플랫폼(Visium 외 다른 spatial platform)으로의 확장성 검증이 추가로 필요
• TRM의 896K라는 작은 파라미터 규모에서의 결과이므로, 더 큰 스케일이나 다른 spatial omics modality(예: single-cell resolution 데이터)에 대한 추가 검증이 요구됨

Evaluation

총평: 재귀적 연산이라는 비교적 단순한 아이디어를 spatial omics의 local processing 및 일반화 문제에 창의적으로 적용하여 견고한 LOSO 검증 결과를 제시한 흥미로운 workshop 논문으로, 향후 더 넓은 범위의 조직 유형과 platform에 대한 검증이 뒷받침된다면 spatial omics 분야에 실질적인 임팩트를 줄 수 있을 것으로 보인다.