[논문 리뷰] Using VAEs to Learn Latent Variables: Observations on Applications in cryo-EM

Daniel G. Edelberg, Roy R. Lederman|arXiv (Cornell University)|2023. 03. 13.

Computational Physics and Python Applications인용 수 5

한 줄 요약

이 논문은 냉동전자현미경(cryo-EM) 응용 분야에서 변분 오토인코더(VAE)의 분산 추론(amortized inference) 성질을 조사한다. 특히 연속적인 구조적 이종성(conformational heterogeneity)을 모델링하는 데 초점을 맞춘다. 연구 결과, CryoDRGN에서 VAE 인코더는 기울임 또는 이동된 입자 영상에 대해 일반화가 잘 되지 않으며, 이는 노이즈 특징에 과적합된다는 것을 시사한다. 반면, 명시적 잠재변수 추정은 동등하거나 더 나은 성능을 보이며, 과학적 영상에서 VAE의 일반화 가정에 도전한다.

ABSTRACT

Variational autoencoders (VAEs) are a popular generative model used to approximate distributions. The encoder part of the VAE is used in amortized learning of latent variables, producing a latent representation for data samples. Recently, VAEs have been used to characterize physical and biological systems. In this case study, we qualitatively examine the amortization properties of a VAE used in biological applications. We find that in this application the encoder bears a qualitative resemblance to more traditional explicit representation of latent variables.

연구 동기 및 목표

• 분산 인코더가 VAE에서 훈련 중에 보지 못한 또는 증강된 냉동전자현미경 입자 영상에 잘 일반화되는지 평가하는 것.
• 냉동전자현미경 데이터에서 VAE 기반 분산 추론과 명시적 잠재변수 추정 간의 성능을 비교하는 것.
• CryoDRGN의 인코더가 진정한 구조적 이종성보다는 노이즈나 구조적 잔상에 과적합되는지 조사하는 것.
• 데이터의 불변성(예: 회전, 이동)이 고노이즈 생물학적 영상에서 VAE의 일반화에 어떤 영향을 미치는지 평가하는 것.
• 냉동전자현미경과 같은 과학적 응용에서 VAE가 일반적으로 생각하는 것보다 명시적 잠재변수 모델에 더 가까운 방식으로 작동할 수 있는지 탐색하는 것.

제안 방법

• 공유 인코더 네트워크를 사용하여 냉동전자현미경 입자 영상에 VAE를 훈련시켰으며, 각 입력에 대해 정규 잠재분포를 생성한다.
• 냉동전자현미경 데이터의 알려진 불변성을 활용하여 평면 내 회전과 이동을 통한 데이터 증강을 적용하여 일반화 능력을 테스트한다.
• 증강된 영상의 잠재 표현을 원본 영상의 표현과 비교하여 일관성과 일반화 능력을 평가한다.
• 공유 인코더 파rameter를 회피하기 위해 각 입자 영상마다 별도의 잠재분포를 피팅하여 명시적 변분 추론을 수행한다.
• 증강된 시각화에 대한 잠재공간의 구조와 일관성에 대한 정성적 비교를 통해 성능을 평가한다.
• 기본 비교를 위해 표준 딥러닝 환경에서 MNIST VAE를 사용한 탈거 실험(Ablation experiments)을 부록 B에서 수행하였다.

Figure 1: Experimental setup of a standard VAE. $x_{i}$ is a datapoint used as input, $z_{i}$ is the associated latent space point, $\hat{x}_{i}$ is the reconstructed output of $x_{i}$ . $\mathcal{L}$ represents the calculation of the ELBO between $x$ and $\hat{x}_{i}$

실험 결과

연구 질문

• RQ1 Cryo-EM에서 VAE 인코더는 훈련 중에 보지 못한 회전 또는 이동된 입자 영상에 잘 일반화되는가?
• RQ2 공유 인코더를 통한 분산 추론 성능과 냉동전자현미경에서의 명시적 잠재변수 추정 간의 성능는 어떻게 비교되는가?
• RQ3 VAE 인코더가 진정한 구조적 상태를 포착하기보다는 노이즈나 구조적 잔상에 얼마나 과적합되는가?
• RQ4 Cryo-EM VAE에서의 일반화 실패 원인은 고노이즈 수준과 영상 변동성 때문인가, 아님 아키텍처적 제약 때문인가?
• RQ5 이 맥락에서 명시적 잠재변수 추정이 분산 VAE보다 정성적으로 더 나은 또는 더 강건한 표현을 제공할 수 있는가?

주요 결과

• CryoDRGN의 VAE 인코더는 기울임 또는 이동된 입자 영상에 대해 일반화가 잘 되지 않으며, 자연스러운 데이터 증강에 대한 일반화 능력이 떨어진다는 것을 시사한다.
• 증강된 영상의 인코더 잠재 표현은 원본 영상과 상당히 다름으로써, 노이즈나 영상 고유의 특징에 과적합된다는 것을 시사한다.
• 공유 인코더 파rameter 없이 명시적 잠재변수 추정을 수행한 결과, 동일한 작업에서 분산 VAE보다 정성적으로 유사하거나 더 나은 성능을 보였다.
• CryoDRGN의 아키텍처에 회전 불변성 등의 불변성이 없기 때문에 일반화 능력이 떨어질 수 있으며, 이러한 불변성이 문제를 완전히 해결할 수는 없을 것임을 시사한다.
• VAE 인코더의 행동은 일반적으로 예상되는 것보다 명시적 잠재변수 모델에 더 가까이 있으며, 이는 분산 추론의 이점에 대한 일반적인 가정을 도전한다.
• MNIST VAE와는 달리, 냉동전자현미경의 경우 고노이즈와 영상 변동성이 심한 상황에서 조건이 성공적으로 학습되었음에도 불구하고 VAE의 일반화 능력이 떨어지는 것으로 나타났다.

Figure 2: The CryoDRGN Pipeline for heterogeneous reconstruction. $x_{i}$ is the cryo-EM image datapoint as input. $\omega_{i}$ , $t_{i}$ , $h_{i}$ are the given rotation, translation, and contrast transfer function of the input image. These variables with latent variable $z_{i}$ are used as input t