scCBGM: Single-Cell Editing via Concept Bottlenecks

Motivation

• Known: 기존 연구들은 조건부 분포 모델링을 통해 population-level 처치 효과를 예측하거나(Lotfollahi et al., 2019; Kana et al., 2023), 최근에는 cell-level counterfactual 예측으로 발전했지만(Zhang et al., 2024; Piran et al., 2024), gene program이나 cell type 같은 해석 가능한 concept 단위의 명시적 제어는 부족했다. 또한 기존 interpretability 방법들은 상관관계를 설명할 뿐 실제로 세포 반응을 시뮬레이션하거나 편집하는 데는 사용할 수 없었다.
• Gap: 개별 세포에 대해 생물학적으로 의미 있는 concept(유전자 프로그램, 세포 유형 등)을 이용해 정밀하고 해석 가능한 counterfactual editing을 수행하는 방법이 부재했으며, 특히 single-cell 데이터 특유의 높은 이질성, 기술적 노이즈, 불완전한 concept 주석 하에서 이를 검증할 ground-truth counterfactual 벤치마크도 없었다.
• Why: 개별 세포 수준에서 특정 개입(약물, 사이토카인, 경로 활성화 등)에 대한 반응을 정밀하게 예측하고 편집할 수 있는 능력은 질병 기전 규명, 치료제 설계, 정밀의학에 필수적인 인과적 추론 도구를 제공하기 때문이다.
• Approach: Pearl의 structural causal model (SCM) 틀 안에서 관측 concept C와 잔여 요인 U가 유전자 발현 X를 생성한다고 가정하고, concept bottleneck을 갖춘 encoder-decoder 구조(scCBGM)를 VAE와 flow matching 모델 모두에 적용하여 counterfactual 편집을 수행한다.

Achievement

Figure 4. Counterfactual modeling predicts cellular response to perturbation. Left: UMAP of CD4 T-cell data from the Kan

1. 아키텍처 개선: 차원 제약 없이 disentanglement를 촉진하는 계산 효율적인 cross-covariance penalty와 noisy concept annotation 하에서도 제어 가능한 생성을 지원하는 decoder concept skip connection을 CBGM에 도입했다.
2. 생성 모델 확장: scCBGM을 VAE를 넘어 flow matching 모델까지 확장하여 decoding-only 및 encoding-decoding 방식 모두에서 concept 기반 편집이 가능함을 보였다.
3. 합성 벤치마크 개발: 외생 노이즈와 concept를 분리하여 실제 counterfactual ground-truth에 접근 가능한 합성 데이터 생성 프로세스를 제안, 노이즈·결측 라벨·이질적 집단 등 현실적 조건에서 체계적 평가를 가능하게 했다.
4. 성능 검증: 합성 벤치마크에서 강력한 cell-level 편집 정확도를, 세 개의 실제 데이터셋에서 population-level 지표로 경쟁력 있거나 우수한 성능을 입증했으며, combinatorial 및 zero-shot generalization에서 여러 최신 방법들을 능가했다. 또한 pathway activity와 treatment에 대한 in-silico intervention 사례를 통해 치료 반응의 기전적 가설을 제시했다.

How

Figure 2. Single-cell Concept Bottleneck Generative Model Overview. Top: A concept bottleneck VAE encodes gene expressio

• 데이터 생성 과정을 SCM으로 정식화: U ~ P(U), UC ~ P(UC), C ← fC(UC), X ← fX(C, U), U ⊥ UC 가정
• counterfactual edit을 µx,c′ := EU[fX(c′, U) | X = x]로 정의하고, 모델 fθ(x, c′)가 이를 근사하도록 학습(abduction 기반)
• encoder E(·)가 x를 잠재 변수 z의 posterior q(z|x) = N(µz, Σz)로 매핑
• concept bottleneck module: concept network fc가 z로부터 concept 예측 ĉ = fc(z)를 산출, 별도 네트워크 fu가 설명되지 않는 요인 u = fu(z)를 산출
• decoder D(·)는 h0 = [u, ĉ]에서 시작해 각 layer마다 이전 hidden state와 ĉ를 concatenate하는 concept skip connection 구조(hℓ = Dℓ([hℓ−1, ĉ]))를 사용해 concept 조건을 매 층마다 강화
• β-VAE ELBO 기반 손실 LVAE = 재구성 항 + λcLconcept + βKL(q(z|x)‖p(z))로 학습하며, 결측 concept 주석을 다루기 위한 binary mask 기반 Lconcept 정의
• VAE 구조를 flow matching 모델로도 확장하여 concept 기반 편집을 generation regime에도 적용

Originality

• concept bottleneck model을 단일세포 RNA-seq counterfactual editing이라는 새로운 도메인에 적용하고, 이를 위한 구조적 개선(decoder concept skip connection, cross-covariance penalty)을 제시함
• 기존 CBGM을 VAE에서 flow matching 생성 모델까지 확장하여 프레임워크의 일반성을 입증함
• 외생 노이즈와 조건을 분리해 ground-truth counterfactual을 제공하는 합성 벤치마크를 새로 설계하여, 기존에는 불가능했던 cell-level 정량 평가를 가능하게 함
• 기존 population-level 조건부 생성 모델과 달리, 개별 세포에 대한 causal counterfactual (SCM 기반 abduction) 정의를 명시적으로 도입함

Limitation & Further Study

• 실제 데이터셋에서는 ground-truth counterfactual이 없어 population-level proxy 지표에 의존할 수밖에 없어, 실제 생물학적 타당성 검증에는 한계가 있음
• cross-covariance penalty가 disentanglement를 어느 정도까지 보장하는지에 대한 이론적 분석이 제한적일 수 있음
• concept annotation의 노이즈나 결측이 심한 실제 대규모 데이터셋에서의 강건성에 대한 추가 검증이 필요함
• flow matching 확장이 VAE 대비 어떤 상황에서 실질적 이점을 주는지 더 폭넓은 비교 분석이 필요함
• 후속 연구로는 실제 wet-lab 검증을 통한 in-silico intervention 가설의 실험적 확인, 더 많은 concept 유형(연속형, 계층적 concept)에 대한 확장이 필요함

Evaluation

총평: 단일세포 데이터에서 해석 가능한 개념 기반 counterfactual editing이라는 중요하지만 미해결 문제를 SCM 기반으로 정식화하고, 이를 실현하는 구체적 아키텍처와 평가 벤치마크를 함께 제시한 점에서 완성도 높은 연구이다. 다만 실제 데이터에서의 검증은 여전히 proxy 지표에 의존하므로 생물학적 타당성에 대한 후속 검증이 필요하다.