[논문 리뷰] Variational AutoEncoder For Regression: Application to Brain Aging Analysis

Qingyu Zhao, Ehsan Adeli외 3명|arXiv (Cornell University)|2019. 04. 11.

Generative Adversarial Networks and Image Synthesis참고 문헌 13인용 수 11

한 줄 요약

이 논문은 뇌 노화 패턴의 분리 및 해석 가능성을 보장하기 위해 잠재 표현의 조건부 분포를 연령에 따라 명시적으로 모델링하는 통합 변분 autoencoder (VAE) 프레임워크를 제안한다. 변분 추론을 통해 VAE와 회귀기의 공동 학습을 통해 모델은 3D MR 영상에서 최고 수준의 연령 예측 정확도(6.9년 rMSE)를 달성하였으며, 뇌 구조적 변화, 예를 들어 뇌실 확장과 같은 연령 관련 변화를 직관적으로 시각화할 수 있다.

ABSTRACT

While unsupervised variational autoencoders (VAE) have become a powerful tool in neuroimage analysis, their application to supervised learning is under-explored. We aim to close this gap by proposing a unified probabilistic model for learning the latent space of imaging data and performing supervised regression. Based on recent advances in learning disentangled representations, the novel generative process explicitly models the conditional distribution of latent representations with respect to the regression target variable. Performing a variational inference procedure on this model leads to joint regularization between the VAE and a neural-network regressor. In predicting the age of 245 subjects from their structural Magnetic Resonance (MR) images, our model is more accurate than state-of-the-art methods when applied to either region-of-interest (ROI) measurements or raw 3D volume images. More importantly, unlike simple feed-forward neural-networks, disentanglement of age in latent representations allows for intuitive interpretation of the structural developmental patterns of the human brain.

연구 동기 및 목표

• 비지도 VAE와 지도 학습 회귀 간의 격차를 메우기 위해.
• 영상 데이터에서 잠재 표현을 공동으로 학습하고 회귀를 수행하는 통합 확률 모델을 개발하기 위해.
• 해석 가능한 뇌 노화 분석을 위해 잠재 공간에서 연령을 변화의 요인으로 분리 가능하게 하기 위해.
• ROI 측정치와 원시 3D MR 볼륨을 모두 사용하여 기존 방법보다 높은 연령 예측 정확도를 향상시키기 위해.
• 연령별 잠재 표현의 생성 모델링을 통해 발달적 뇌 패턴을 직관적으로 시각화하기 위해.

제안 방법

• 모델은 잠재 표현 z가 회귀 목표 c(연령)에 따라 달라지도록 명시적으로 모델링하는 조건부 사전 분포 p(z|c)를 도입하여, 연령별로 z를 생성할 수 있도록 한다.
• 생성 과정은 p(x,z,c) = p(x|z) * p(z|c) * p(c)로 정의되며, 여기서 p(x|z)는 신경망 디코더이고, p(z|c)는 u^T c로 매개변수화된 선형 가우시안 생성자이다.
• 공동 모델에 대해 변분 추론을 수행하여 VAE 인코더와 신경망 회귀기의 공동 최적화를 이끌어내며 상호 정규화 효과를 갖는다.
• 잠재 공간은 연령 관련 변동을 분리하도록 학습되어 한 차원을 따라 이동하면 연령별로 특화된 뇌 구조적 모습을 생성할 수 있다.
• 모델은 ROI 특징에 대해 다층 퍼셉트론을, 전체 볼륨 영상에 대해 3D 컨볼루션 신경망을 사용하여 구현된다.
• 학습은 재구성 손실과 KL 발산 항을 포함하는 변분 하한(ELBO)을 사용하며, 엔드 투 엔드 최적화가 이루어진다.

실험 결과

연구 질문

• RQ1 뇌 영상 데이터에서 분리된 잠재 표현을 공동으로 학습하고 정확한 회귀를 수행할 수 있는 통합 VAE 프레임워크를 설계할 수 있는가?
• RQ2 회귀 목표(연령)에 따라 잠재 사전 분포를 조건화하는 것이 기존 VAE나 독립된 회귀기보다 예측 정확도를 향상시키는가?
• RQ3 분리된 잠재 공간이 뇌 노화의 해석 가능하고 시각화 가능한 패턴(예: 뇌실 확장)을 제공할 수 있는가?
• RQ4 ROI 측정치와 3D MR 볼륨을 모두 사용한 연령 예측에서, 이 모델의 성능은 최고 수준의 기존 방법과 비교해 어떻게 되는가?
• RQ5 VAE와 회귀기의 공동 학습이 정규화 효과를 유도하고 일반화 능력을 향상시키는 데 얼마나 기여하는가?

주요 결과

• 제안된 모델은 3D T1-강조 MR 영상에서 6.9년 rMSE의 연령 예측 성능을 달성하여, 모든 기준 모델(딥 네트워크 및 전통적 회귀기 포함)을 초월하였다.
• ROI 기반 실험에서는 R2가 0.666, rMSE가 9.0을 기록하여, 단일 피드포워드 신경망(10.3)조차도 뛰어넘었다.
• 모델은 뛰어난 분리 능력을 보였으며, 잠재 공간의 한 차원이 명확히 연령과 일치하는 반면, 기존 VAE는 이러한 명확한 분리를 이루지 못했다.
• 연령별로 생성된 뇌 영상(예: z = u^T c)은 연령 증가에 따라 뇌실 확장이 뚜렷하게 나타나며, 알려진 뇌해부학적 노화 패턴과 일치하였다.
• 18세와 86세 뇌 영상 간의 정렬에서의 자코비안 행렬식 지도는 주로 뇌실 확장이 구조적 변화의 주요 원인임을 확인하였으며, 이는 생성 모델의 출력과 일치하였다.
• 공동 학습 과정은 상호 정규화 효과를 유도하여 일반화 능력과 강인성을 향상시켰으며, 특히 성능 향상이 두드러진 ROI 기반 설정에서 뚜렷하게 나타났다.