[논문 리뷰] BasisVAE: Translation-invariant feature-level clustering with Variational Autoencoders

Kaspar Märtens, Christopher Yau|arXiv (Cornell University)|2020. 03. 06.

Natural Language Processing Techniques참고 문헌 20인용 수 8

한 줄 요약

BasisVAE는 비선형 차원 감소와 특징 수준의 군집화를 동시에 수행할 수 있도록 Variational Autoencoders (VAEs)와 디코더 내 베이지안 혼합 사전확률을 조합한 공동 확률적 프레임워크를 제안한다. 학습된 기저 함수의 선형 조합으로 특징을 모델링하고, 군집별 스케일링 및 이동 파라미터를 사용함으로써, 고차원 표본 데이터에 대해 해석 가능하고 이동 불변성 있는 군집화를 가능하게 한다. 단일세포 유전자 발현 데이터에서 표준 군집화 및 VAE보다 뛰어난 성능을 보이며 생물학적으로 의미 있는 유전자 발현 패턴을 드러내었다.

ABSTRACT

Variational Autoencoders (VAEs) provide a flexible and scalable framework for non-linear dimensionality reduction. However, in application domains such as genomics where data sets are typically tabular and high-dimensional, a black-box approach to dimensionality reduction does not provide sufficient insights. Common data analysis workflows additionally use clustering techniques to identify groups of similar features. This usually leads to a two-stage process, however, it would be desirable to construct a joint modelling framework for simultaneous dimensionality reduction and clustering of features. In this paper, we propose to achieve this through the BasisVAE: a combination of the VAE and a probabilistic clustering prior, which lets us learn a one-hot basis function representation as part of the decoder network. Furthermore, for scenarios where not all features are aligned, we develop an extension to handle translation-invariant basis functions. We show how a collapsed variational inference scheme leads to scalable and efficient inference for BasisVAE, demonstrated on various toy examples as well as on single-cell gene expression data.

연구 동기 및 목표

• 고차원 표본 데이터(예: 유전체학)에서 표준 VAE가 제공하는 군집 수준의 통찰력이 부족한 문제를 해결하기 위해.
• 차원 감소와 특징 군집화를 하나의 확률적 모델로 통합하여, 부호화된 두 단계 워크플로우를 피하기 위해.
• 유사한 형태이지만 다른 오프셋을 가진 특징들을 군집화하기 위해, 이동 파라미터(δ)를 명시적으로 모델링함으로써 이동 불변성 군집화를 가능하게 하기 위해.
• 특히 단일세포 옴믹스 데이터에서 공통 기능 역학을 가진 특징 그룹을 식별하기 위한 확장 가능하고 해석 가능한 프레임워크를 제공하기 위해.
• 추론된 파라미터(λ 및 δ)를 통해 연구자들이 특징 간 형태 유사성뿐 아니라 상대적인 시간 또는 단계 이동(Phase shift)까지도 탐지할 수 있도록 하기 위해.

제안 방법

• 기본적으로 특징을 군집에 할당하는 범주형 잠재 변수를 사용한 계층적 베이지안 군집 사전확률을 적용한 수정된 VAE 디코더를 도입한다.
• 각 특징의 반응을 기저 함수의 선형 조합으로 모델링하고, 군집별 스케일링(λ) 및 이동(δ) 파라미터를 사용한다.
• 과도하게 군집 수 K를 설정하더라도 효율적이고 희박한 추론을 가능하게 하기 위해 축약된 변분 추론 기법을 적용한다.
• 비음수이고 0이 많은 단일세포 유전자 발현 데이터를 모델링하기 위해 디코더에서 0-중심 음이항 분포 우도를 사용한다.
• 군집별 이동 파라미터(δ)를 허용함으로써, 유사한 형태이지만 다른 단계 오프셋을 가진 특징들을 군집화할 수 있도록 이동 불변성을 명시적으로 모델링한다.
• 합성 유전자 발현 및 단일세포 정자생성 데이터를 포함한 실제 데이터에 모델을 적용하여 잠재적 구조와 특징 그룹화를 동시에 탐지함을 보여준다.

실험 결과

연구 질문

• RQ1 VAE 프레임워크 내에서 차원 감소와 특징 군집화를 공동으로 모델링할 수 있는가? 이를 통해 고차원 표본 데이터에서의 해석 가능성은 향상되는가?
• RQ2 딥 생성 모델에서 이동 불변성을 명시적으로 모델링하면, 유사한 기능적 형태이지만 다른 단계 이동을 가진 특징들을 군집화할 수 있는가?
• RQ3 VAE 디코더에 확률적 군집 사전확률을 통합하면, 후처리 군집화 방법에 비해 더 정확하고 희박한 군집 할당이 가능한가?
• RQ4 BasisVAE는 허위 시간(pseudotime)이나 군집 구조에 대한 사전 지식 없이도 단일세포 데이터에서 알려진 생물학적 유전자 발현 패턴을 어느 정도 회복할 수 있는가?
• RQ5 추론된 스케일링(λ) 및 이동(δ) 파라미터는 유전자 발현의 상대적인 시간 또는 역학적 특성에 대해 생물학적으로 의미 있는 통찰을 제공하는가?

주요 결과

• BasisVAE는 합성 유전자 발현 데이터에서 다섯 개의 명확한 군집을 성공적으로 식별했으며, 이동 불변성 버전은 동일한 형태이지만 이동된 버전의 특징들을 정확히 같은 군집에 할당했다.
• 단일세포 정자생성 데이터에서 이동 불변성 BasisVAE는 Tex101, Ly6k, Sdc4, Tubb5, Ccnd2 유전자들을 동일한 군집에 할당하여, 서로 다른 오프셋을 가짐에도 불구하고 공통된 발현 역학을 드러냈다.
• 모델이 추론한 δ 파라미터를 통해, 같은 군집 내에서 한 유전자의 발현이 다른 유전자보다 지연됨을 설명할 수 있었다.
• 스케일 불변성 버전의 BasisVAE는 동일한 형태이지만 다른 오프셋을 가진 유전자들을 별개의 군집에 할당하여, 정확한 생물학적 군집화를 위해 이동 불변성의 중요성을 강조했다.
• 축약된 변분 추론 기법을 통해 과도하게 군집 수 K=50를 설정하더라도 효율적이고 희박한 추론이 가능했으며, 의미 있고 안정된 군집 할당이 이루어졌다.
• 재정렬된 특징과 잠재 변수를 통해, 합성 및 실제 단일세포 데이터에서 k-means 및 상관계수 기반 유사도 메트릭보다 BasisVAE가 더 정확한 기반이 되는 구조를 드러내었다.