Elucidating the Design Space of Generative Models for Single-Cell Perturbation Prediction

Motivation

• Known: Autoregressive next-token prediction은 언어 데이터에서 예측 가능한 scaling law를 보여주었고, scVI 계열의 Gaussian-latent VAE, latent-diffusion 계열(scDiffusion, CFGen, scLDM), 그리고 Geneformer, scGPT, scFoundation, STATE 같은 대규모 perturbation-response foundation model들이 모두 raw count를 continuous vector로 임베딩해 사용해왔다.
• Gap: single-cell RNA-seq count에는 유전자 간 자연스러운 순서가 없어 raw expression에 autoregressive한 left-to-right bias를 직접 적용하면 실패하며, 기존 continuous-latent 및 discrete count-domain 생성모델들은 discrete latent 표현과 유연한 generative prior를 결합하는 조합을 탐구하지 않았다.
• Why: discrete latent가 raw count의 부적절한 순서 편향 문제를 해결하면서도 언어모델처럼 예측 가능한 scaling과 강력한 discrete prior 활용이 가능한 virtual cell model의 새로운 설계축을 제시하기 때문에, in-silico 신약 발굴 등 가상세포 모델링 발전에 중요하다.
• Approach: 각 세포를 FSQ bottleneck을 통해 짧은 discrete code 시퀀스로 압축하는 ExpressionVAE를 학습한 뒤, 이 코드에 대해 perturbation-conditioned discrete prior(AR 또는 masked discrete diffusion)를 학습하는 2단계 파이프라인을 제안한다.

Achievement

Figure 1. ExpressionVAE on the outer envelope of both metric families. Six-axis comparison of our representative model (

1. 분포 지표 SOTA 달성: Replogle와 Parse 1M 데이터셋에서 eVAE가 모든 distributional metric에서 새로운 state-of-the-art를 기록했으며, Fr\'echet distance와 MMD^2가 가장 강력한 continuous-latent baseline 대비 약 3~20배 낮았다.
2. Prior 종류 무관성 입증: autoregressive prior와 masked discrete diffusion prior를 교체해도 성능이 거의 동일하여, 성능 향상이 prior architecture가 아닌 discrete latent 표현 자체에서 비롯됨을 격리하여 보였다.
3. Decoder-head 설계축 규명: decoder-head ablation을 통해 inference-time predictive distribution의 richness라는 단일 설계축을 발견하고, 표준 metric들이 variance-sensitive와 mean-sensitive 두 그룹으로 나뉘어 이 축을 따라 상반되게 움직임을 규명했다.
4. 일반화된 latent 유용성 검증: 염증성 cytokine stress 하 1,732개 perturbation으로 구성된 held-out CRISPRi reversion benchmark에서, frozen eVAE encoder가 UMAP과 differential expression을 능가하고 훨씬 적은 데이터로 10배 큰 foundation model인 scGPT와 perturbation ranking에서 대등한 성능을 보였다.

How

Figure 2. A two-stage pipeline. Stage 1 (the row of boxes): a cross-attention transformer encoder maps raw count vectors

• Stage 1 (ExpressionVAE 학습): cross-attention transformer 인코더가 각 유전자의 log1p-normalized count와 학습된 gene-identity token을 임베딩하여 L개의 learnable latent query token이 G개의 유전자 임베딩에 cross-attend하도록 하고, FSQ bottleneck을 통해 discrete code 시퀀스로 양자화한다.
• 디코더 헤드 다양화: cross-entropy, MSE, hurdle, negative-binomial 등 다양한 decoder output head를 동일 backbone 위에서 비교하는 within-architecture ablation을 수행한다.
• Stage 2 (조건부 prior 학습): 학습된 VAE를 고정하고, 타겟 유전자의 ESM2-3B protein-language-model embedding과 학습된 additive cell-type signal에 조건화된 perturbation-conditioned discrete prior(AR 또는 MDLM)를 code sequence 위에서 학습한다.
• 평가: distribution metric(W2, MMD^2, FD)과 cell-eval perturbation metric(PDS, Sp-LFC, ovl@N)으로 Replogle와 Parse 1M에서 성능을 측정하고, 별도의 CRISPRi reversion benchmark에서 frozen encoder의 perturbation-ranking enrichment AUC를 UMAP, differential expression, scGPT와 비교한다.

Originality

• 단일세포 perturbation modeling에 discrete latent representation(FSQ 기반)을 도입하고, prior family(AR vs masked discrete diffusion)를 통제 변수로 삼아 성능 향상의 원인을 discrete latent 자체로 명확히 귀속시킨 통제 실험 설계가 독창적이다.
• 기존 문헌에서 상충되어 보이던 decoder-head 평가 지표들의 순위를 "predictive distribution richness"라는 단일 측정이론적(measurement-theoretic) 축으로 통합 설명한 점이 새롭다.
• CRISPRi phenotypic reversion benchmark를 활용해 frozen encoder의 latent 표현이 대규모 foundation model과 견줄 만한 일반화 능력을 갖는지 검증한 실용적 접근이 돋보인다.

Limitation & Further Study

• 본문 발췌만으로는 FSQ codebook 크기, latent sequence 길이 L, 학습 규모(scaling law 검증)에 대한 구체적 분석이 충분히 제시되지 않아, 언어모델과 같은 predictable scaling이 실제로 성립하는지는 추가 검증이 필요하다.
• 평가가 Replogle, Parse 1M, 그리고 하나의 CRISPRi reversion benchmark(TeloHAEC)에 국한되어 있어, 다양한 세포 유형·조직·질병 맥락에 대한 일반화 가능성은 추가 벤치마크로 확인해야 한다.
• decoder-head의 variance-sensitive/mean-sensitive 이분법이 다른 metric이나 downstream task(예: 약물 스크리닝 실제 적용)에서도 동일하게 유지되는지에 대한 추가 검증이 필요하다.
• scGPT와 "matches"라는 표현이 정확히 어떤 조건(데이터량, task 난이도)에서 성립하는지 더 세밀한 분석이 요구된다.

Evaluation

총평: discrete latent와 유연한 generative prior의 결합이라는 설계공간을 체계적으로 통제 실험한 견고한 연구로, 단일세포 perturbation modeling에서 discrete latent 표현의 가치를 명확히 입증했으며 후속 virtual cell model 설계에 실질적 지침을 제공한다.