[논문 리뷰] Conditional out-of-sample generation for unpaired data using trVAE

Mohammad Lotfollahi, Mohsen Naghipourfar외 2명|arXiv (Cornell University)|2019. 10. 04.

Single-cell and spatial transcriptomics참고 문헌 26인용 수 12

한 줄 요약

이 논문은 복수의 조건 간에 쌍이 맞지 않는 데이터에서 샘플 외부 생성을 향상시키기 위해 농심층에 최대 평균 차이(MMD) 정규화를 적용한 조건부 변동 자동인코더(trVAE)를 제안한다. 다양한 조건 간에 분포 유사성을 강제로 부여함으로써 trVAE는 고차원 샘플 생성에서 유의미하게 높은 정확도와 강건성을 달성한다. 이를 통해 평균에 대한 피어슨 상관계수 0.97, 분산에 대한 피어슨 상관계수 0.87 향상 및 소수 세포 유형과 새로운 간섭 조건에 대한 단세포 유전자 발현 예측 작업에서 뛰어난 성능을 입증하였다.

ABSTRACT

While generative models have shown great success in generating high-dimensional samples conditional on low-dimensional descriptors (learning e.g. stroke thickness in MNIST, hair color in CelebA, or speaker identity in Wavenet), their generation out-of-sample poses fundamental problems. The conditional variational autoencoder (CVAE) as a simple conditional generative model does not explicitly relate conditions during training and, hence, has no incentive of learning a compact joint distribution across conditions. We overcome this limitation by matching their distributions using maximum mean discrepancy (MMD) in the decoder layer that follows the bottleneck. This introduces a strong regularization both for reconstructing samples within the same condition and for transforming samples across conditions, resulting in much improved generalization. We refer to the architecture as \emph{transformer} VAE (trVAE). Benchmarking trVAE on high-dimensional image and tabular data, we demonstrate higher robustness and higher accuracy than existing approaches. In particular, we show qualitatively improved predictions for cellular perturbation response to treatment and disease based on high-dimensional single-cell gene expression data, by tackling previously problematic minority classes and multiple conditions. For generic tasks, we improve Pearson correlations of high-dimensional estimated means and variances with their ground truths from 0.89 to 0.97 and 0.75 to 0.87, respectively.

연구 동기 및 목표

• 학습 데이터에 이러한 쌍이 없는 조합의 도메인과 조건에 조건부로 고차원 샘플을 생성하는 문제를 해결하기 위해.
• 기본적인 CVAE는 조건 간의 명시적 교차 조건 정규화가 부족하여 조건 간의 밀집 공통 분포를 학습하지 못하는 한계를 극복하기 위해.
• 하드코딩된 벡터 산술이나 히스토그램 매칭을 넘어서, 동시에 복수의 조건을 지원하는 데이터 기반의 종단간 프레임워크를 개발하기 위해.
• 실제 생물학적 데이터에서, 학습 데이터에 존재하지 않는 약물 치료에 대한 세포 반응을 예측하는 등 샘플 외부 예측을 향상시키기 위해.

제안 방법

• 조건부 변동 자동인코더(CVAE)의 농심층 이후의 디코더 레이어에 MMD 정규화를 도입하여, 다양한 범주형 조건 간의 잠재 표현에 대해 정규화한다.
• 커널 기반 MMD 손실을 사용하여 모든 조건 간의 잠재 표현 분포를 매칭함으로써, 모델이 공통의 압축된 표현을 학습하도록 유도한다.
• 재구성 손실과 KL 발산을 포함한 변동 하한 경계 목표를 사용하여 모델을 종단간으로 훈련시키며, 잠재 공간에 MMD를 정규화로 추가한다.
• 성별, 머리카락 색상, 질병 상태 등의 조건을 포함하여 이미지 및 단세포 유전자 발현 데이터와 같은 고차원 데이터에 모델을 적용한다.
• 농심층에 다중헤드 어텐션 메커니즘을 도입하여 효과적인 교차 조건 표현 학습을 가능하게 하여 '트랜스포머 VAE'(trVAE) 아키텍처의 기초를 마련한다.
• 예측된 고차원 데이터의 평균/분산과 진짜 값 간의 피어슨 상관계수와 유전자 발현 프로파일에 대한 정성적 분석과 같은 지표를 사용해 성능을 평가한다.

실험 결과

연구 질문

• RQ1 CVAE의 농심층에 MMD 정규화를 적용하면 쌍이 맞지 않는 데이터에서 복수의 조건 간에 샘플 외부 생성을 향상시킬 수 있는가?
• RQ2 조건 간에 분포 유사성을 강제로 적용하면, 특히 소수 또는 새로운 조건에서 고차원 샘플의 생성이 더 강건하고 정확해지는가?
• RQ3 학습 중에 관찰되지 않은 간섭 조건에 대한 세포 반응을 예측할 때 trVAE는 cycle GAN, SAUCIE, 기존 CVAE와 비교해 어떻게 성능을 냈는가?
• RQ4 단세포 RNA-seq 데이터에서 MMD 정규화는 예측된 유전자 발현 평균과 분산의 정밀도를 진짜 값과 비교해 얼마나 향상시키는가?
• RQ5 학습 중에 해당 조건에 노출되지 않은 상태에서 trVAE는 자연 살해세포에서 IFN-β 자극 효과를 예측하는 복잡한 생물학적 과제에 일반화될 수 있는가?

주요 결과

• trVAE는 1,000차원의 예측된 평균 발현과 진짜 평균 발현 간의 피어슨 상관계수를 0.89에서 0.97로 향상시켜 데이터 분포 평균의 유의미한 개선을 보였다.
• 예측된 분산에 대한 상관계수는 0.75에서 0.87로 향상되어 조건 간 데이터 산산 분포 모델링 능력이 향상됨을 보여주었다.
• 단세포 유전자 발현 예측에서, trVAE는 학습 데이터에 포함되지 않은 바이오스피어 조건에서도 자연 살해세포에서 IFN-β 자극 후 ISG15의 상승 조절을 정확히 포착하였다.
• trVAE는 cycle GAN과 MMD 정규화된 오토인코더(SAUCIE)보다 유전자 발현의 평균과 분산을 예측하는 데에서 뛰어난 성능을 보이며, 모든 모델 중에서 가장 높은 R² 값을 기록하였다.
• 각 조건을 학습 중에 제거했을 때도 trVAE는 모든 8개의 세포 유형에 대해 모든 간섭 조건에서 성공적으로 예측하였으며, 복수의 샘플 외부 조합에 대한 강건성과 일반화 능력을 입증하였다.
• 정성적 분석 결과, trVAE는 자극 후 상위 10개의 유전자 발현 변화를 정확히 재현하였으며, 변화 방향성과 크기 모두가 정확하였다.