[论文解读] Variational AutoEncoder For Regression: Application to Brain Aging Analysis

Qingyu Zhao, Ehsan Adeli+3 more|arXiv (Cornell University)|Apr 11, 2019

Generative Adversarial Networks and Image Synthesis参考文献 13被引用 11

一句话总结

该论文提出了一种统一的变分自编码器（VAE）框架用于回归任务，通过显式建模潜在表征在年龄上的条件分布，实现了可分离的、可解释的脑老化模式。通过联合训练VAE与回归器，采用变分推断方法，模型在3D磁共振图像上实现了最先进的年龄预测精度（rMSE为6.9年），并能直观可视化与年龄相关的脑结构变化，例如脑室扩大。

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体素数据的情况下，提升现有方法的年龄预测准确率。
• 通过年龄特异性潜在表征的生成建模，提供直观的发育脑模式可视化。

提出的方法

• 该模型引入了一个条件先验 p(z|c)，显式建模潜在表征 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卷积神经网络处理全体积图像。
• 训练采用变分下界（ELBO），结合重建损失与KL散度项，实现端到端优化。

实验结果

研究问题

• RQ1 能否设计一种统一的VAE框架，联合学习解耦的潜在表征并实现神经影像数据的准确回归？
• RQ2 与标准VAE或独立回归器相比，将潜在先验条件化于回归目标（年龄）是否能提升预测准确率？
• RQ3 解耦的潜在空间能否提供可解释、可可视化的大脑老化模式，例如脑室扩大？
• RQ4 在使用ROI测量值和3D MR体素数据的情况下，该模型的性能与现有最先进方法相比如何？
• RQ5 VAE与回归器的联合训练在多大程度上实现了正则化并提升了泛化能力？

主要发现

• 所提出的模型在3D T1加权MR图像上的年龄预测实现了6.9年的rMSE，优于所有基线方法，包括深度神经网络和传统回归器。
• 在基于ROI的实验中，模型实现了R2为0.666、rMSE为9.0，甚至优于单一前馈神经网络（rMSE为10.3）。
• 该模型表现出更优的解耦能力：潜在空间中一个维度明显与年龄对齐，而传统VAE未能产生如此清晰的解耦。
• 从年龄特异性潜在表征（z = u^T c）生成的合成脑图像清晰显示出随年龄增长的脑室扩大，与已知的神经解剖学老化模式一致。
• 18岁与86岁大脑之间配准的雅可比行列式图证实，主要结构变化为脑室扩张，与生成模型的输出一致。
• 联合训练过程实现了相互正则化，提升了泛化能力与鲁棒性，尤其在基于ROI的设置中表现最为显著。