[논문 리뷰] Compositional perturbation autoencoder for single-cell response modeling

논문 바로가기

Introduction

최근 single-cell RNA-seq 기술의 발전으로 특정 약물 처리나 유전자 조작과 같은 perturbation이 세포의 유전자 발현에 어떤 영향을 미치는지를 정밀하게 분석할 수 있게 되었다. 그러나 실제 실험 환경에서는 가능한 perturbation의 종류, 용량(dose), 시간(time), 그리고 세포 유형(cell type)과 같은 조건의 조합이 매우 다양하기 때문에 모든 경우를 직접 측정하는 것은 사실상 불가능하다. 따라서 관측되지 않은 새로운 perturbation 조건에서의 gene expression을 예측할 수 있는 계산 모델의 필요성이 커지고 있다.

기존의 mechanistic 모델이나 gene regulatory network 기반 접근법은 생물학적 해석 가능성은 높지만 genome-wide 수준의 고차원 single-cell 데이터에 확장하기 어렵고, 시간 정보가 충분히 주어지지 않는 경우 적용이 제한적이라는 문제가 있다. 한편 differential expression 분석과 같은 전통적인 통계적 방법은 특정 perturbation과 control 조건 사이의 차이를 비교하는 데는 효과적이지만, 여러 perturbation의 조합 효과나 새로운 조건에서의 gene expression을 예측하는 데에는 적합하지 않다. 최근에는 variational autoencoder(VAE) 기반의 딥러닝 모델들이 제안되었지만, 이러한 방법들은 latent space의 해석 가능성이 낮고 perturbation과 covariate의 효과를 명확하게 분리하지 못한다는 한계를 가진다.

이러한 문제를 해결하기 위해 CPA(Compositional Perturbation Autoencoder)는 gene expression이 세포의 고유한 basal state에 perturbation 효과와 covariate 효과가 가법적으로 결합된 결과로 생성된다고 가정하고, 이를 분해 가능한(latent disentanglement) 구조로 학습하는 모델을 제안한다. 이 접근법은 perturbation과 covariate의 효과를 각각 독립적인 embedding으로 표현함으로써 관측되지 않은 perturbation 조합이나 새로운 조건에서도 gene expression을 예측할 수 있도록 하며, 동시에 모델의 해석 가능성까지 확보한다는 점에서 기존 방법들과 차별화된다.

Method

이 방법론의 핵심 가정은 다음과 같다

• $z^{basal}_i$는 관측되지 않은 latent 변수이다. 즉, 각 cell에는 perturbation 이전의 "기본 상태"가 존재한다.
• $z_i^{basal} \perp (d_i,c_i)$로 basal state는 perturbation과 covariate와 독립이다.

$\rightarrow$ gene expression = basal state + perturbation + covariate

사용하는 데이터셋은 다음과 같이 정의된다.

• $x_i \in \mathbb{R}^G$ : 세포 $i$의 $G$개 gene expression
• $d_i = (d_{i,1},...,d_{i,M})$ : 세포 $i$에 적용된 약물 $j$의 용량 ($d_{i,j} \ge 0$, 이때 0은 약물 $j$가 세포 $i$에 적용되지 않았음을 의미한다.)
• $c_i = (c_{i,1},...,c_{i,K})$ : 세포 유형 또는 종 같은 추가적인 이산형 공변량을 포함함. $c_{i,j}$는 K차원의 one-hot vector이다.

CPA의 목표는 데이터셋 삼중항 $(x_i,d_i,c_i)$와 목표 perturbation $d'$이 주어졌을 때, 세포 $i$가 $d_i$ 대신 $d'$로 perturbation 되었다면 공변량 $c_i$를 가진 유전자 발현 $x'_i$를 예측하는 $x'_i = M((x_i,d_i,c_i),d')$을 학습하는 것이다.

간단하게 설명하면 다음과 같다. 우리는 새로운 perturbation을 적용했을 때의 gene expression이 궁금하다.

그래서 gene expression = basal state + preturbation + covariate이므로 basal state 값만 정확히 알면 perturbation과 covariate를 더해서 gene expression을 알 수 있다.

먼저 basal state를 만들기 위해서 인코더와 discriminator를 기반으로 인코더는 perturbation, covariate 정보를 숨기고 discriminator는 basal state에 맞는 perturbation, covariate를 예측하게 되면서 인코더는 basal state만 예측하도록 점차 학습된다.

인코더가 예측한 basal state에 정답 perturbation, covariate 임베딩을 더해서 $z$를 만들고 디코더를 통해 gene expression 값을 예측한다. 그리고 예측값과 실제 gene expression 값 차이가 최소화되도록 한다.

CPA 모델의 학습 과정은 다음 세 단계로 이루어진다

1. gene expression $x_i$를 perturbation 및 covariate $(d_i,c_i)$에 대한 정보를 포함하지 않는 basal state $z_i^{basal}$로 인코딩한다. (이때 discriminator를 기반으로 encoder가 basal state만 학습하도록 한다)
   $\hat{z}_i^{basal} = \hat{f}^{enc}(x_i)$
2. 이후 추정 basal state를 사용하여 추정된 perturbation 후 $\hat{z}_i$를 계산한다
   $\hat{z}_i = \hat{z}^{basal}_i + \hat{V}^{perturbation} \cdot (\hat{f}_1(d_{i,1}),...,\hat{f}_M(d_{i,M})) + \sum_{j=1,...,K}\hat{V}^{cov_j}\cdot c_{i,j}$

Example - perturbation, covariate

$V^{perturbation} = \begin{bmatrix} 1 & 2 & 0 & -1 \\ 0 & 1 & 3 & 2 \\ 2 & 0 & 1 & 1 \end{bmatrix}$

이때 각 column은 서로다른 drug를 의미함. 그래서 $d_i = (0.5, 0., 1.2, 0)$인 경우 dose response 함수를 적용하여 각 perturbation마다 scaling function을 적용함. $f_1(0.5) = 0.8, f_2(0)=0,...$ 최종적으로는 $f(d_i) = \begin{bmatrix}0.8 \\ 0 \\ 1.5 \\ 0\end{bmatrix}$ 이 되고 $V^{perturbation}$과 $f(d_i)$를 행렬 곱하면 $z_i^{perturbation} = \begin{bmatrix}0.8 \\ 4.5 \\ 3.1\end{bmatrix}$이 된다.

covariate도 동일한 방법으로 적용된다.

$V^{cov} = \begin{bmatrix} 1 & 0 & 2 \\ 0 & 3 & 1 \\ 4 & 2 & 0 \end{bmatrix}$

각 column 은 cell type을 나타낸다. covariate는 perturbation과 다르게 $c_i = \begin{bmatrix}0 \\ 1 \\ 0\end{bmatrix}$는 one hot vector로 어떤 cell type을 사용하는지를 결정해준다. $z^{cov}_i = V^{cov} \cdot c_i = \begin{bmatrix} 0 \\ 3\\ 2\end{bmatrix}$가 된다.

3. 마지막 단계에서는 decoder $\hat{f}^{dec}$는 $\hat{z}_i$를 입력으로 받아 $x'_i$의 추정 평균$(\hat{f}^{dec}_{\mu}(\hat{z}_i))$과 분산$(\hat{f}^{dec}_{\sigma^2}(\hat{z}_i))$을 반환한다.

reconstruction loss

decoder가 gene expression을 잘 복원하도록 유도

Adversarial loss

함수 $\hat{f}^{adv}_d, \hat{f}^{adv}_{c_{i,j}}$는 $\hat{z}^{basal}_i$가 주어졌을 때 $(d_i,c_i)$를 예측하는 discriminator이다.

final

위 두 손실 함수를 바탕으로 다음 두 최적화 단계를 번갈아 수행한다

1. $(x_i,d_i,c_i) \sim \mathcal{D}$를 샘플링하고, 모든 $j = 1,...,K$에 대해 $\hat{f}^{adv}_d$와 $\hat{f}^{adv}_{c_{i,j}}$의 파라미터를 업데이트하여 $\ell^d_i + \sum_j \ell^c_{i,j}$를 최소화한다.
2. $(x_i,d_i,c_i) \sim \mathcal{D}$를 샘플링하고, encoder, decoder, perturbation 임베딩, covariate 임베딩, dose-response function의 파라미터를 업데이트하여 $\ell_i - \lambda \cdot (\ell_i^d + \sum_j \ell_{i,j}^c)$를 최소화한다.

---

Contribution

• perturbation 효과를 분해 가능한 compositional latent space로 모델링
• unseen perturbation 조합에 대한 gene expression 예측 가능
• continous dose / time perturbation modeling 지원
• covaraite 효과를 disentangle하여 조건 변화에 대한 일반화 가능
• perturbation 및 covariate embedding을 통한 해석 가능성 확보

Limitation

• training distribution과 크게 다른 조건에서는 예측 성능 저하 가능
  CPA는 unseen perturbation 조합을 예측할 수 있지만, 학습 데이터와 지나치게 다른 조건이 주어질 경우 성능이 저하될 수 있음
• uncertainty estimation이 heuristic 수준에 머무름
  논문에서는 예측 결과의 불확실성을 추정하기 위해 distance-based heuristic 방법을 사용하지만, 이는 엄밀한 Bayesian uncertainty estimation이 아니며, 보다 정교한 uncertainty modeling이 향후 개선 과제로 제시된다.
• gene expression 데이터만 사용
  CPA는 입력으로 gene expression만 사용함. transcriptomic 수준을 넘어서는 생물학적 메커니즘은 직접적으로 반영하지 못함