[论文解读] Learning Conditional Variational Autoencoders with Missing Covariates

Siddharth Ramchandran, Gleb Tikhonov+3 more|arXiv (Cornell University)|Mar 2, 2022

Machine Learning in Healthcare被引用 4

一句话总结

本文提出了一种新颖的方法，通过引入一种摊销变分推断框架，在辅助协变量存在缺失值的情况下训练条件变分自编码器（CVAEs），该框架在最大化证据下界（ELBO）的同时，对缺失协变量的不确定性进行边缘化。该方法能够在非时间序列、时间序列和纵向数据集上实现CVAEs和GP先验VAEs的有效学习，在缺失协变量数据下实现了生成和插补任务的最先进性能。

ABSTRACT

Conditional variational autoencoders (CVAEs) are versatile deep generative models that extend the standard VAE framework by conditioning the generative model with auxiliary covariates. The original CVAE model assumes that the data samples are independent, whereas more recent conditional VAE models, such as the Gaussian process (GP) prior VAEs, can account for complex correlation structures across all data samples. While several methods have been proposed to learn standard VAEs from partially observed datasets, these methods fall short for conditional VAEs. In this work, we propose a method to learn conditional VAEs from datasets in which auxiliary covariates can contain missing values as well. The proposed method augments the conditional VAEs with a prior distribution for the missing covariates and estimates their posterior using amortised variational inference. At training time, our method marginalises the uncertainty associated with the missing covariates while simultaneously maximising the evidence lower bound. We develop computationally efficient methods to learn CVAEs and GP prior VAEs that are compatible with mini-batching. Our experiments on simulated datasets as well as on a clinical trial study show that the proposed method outperforms previous methods in learning conditional VAEs from non-temporal, temporal, and longitudinal datasets.

研究动机与目标

• 为解决现有方法在辅助协变量存在缺失值时无法处理的条件VAE学习空白问题。
• 开发一种概率框架，将缺失协变量和观测数据联合建模于CVAE结构中。
• 通过摊销推断实现CVAEs和GP先验VAEs在缺失协变量情况下的高效、小批量兼容训练。
• 在多种数据类型（包括纵向和临床数据集）上，展示在表征学习和数据插补方面优于以往方法的性能。

提出的方法

• 该方法引入了对潜在变量和缺失协变量的联合变分后验分布，支持在存在缺失数据的情况下进行端到端训练。
• 采用摊销变分推断来近似缺失协变量的后验分布，避免对每个样本进行优化。
• 通过蒙特卡洛近似对缺失协变量中的不确定性进行边缘化，同时最大化证据下界（ELBO）。
• 对于GP先验VAEs，该方法推导出一种计算高效的、支持小批量的高斯过程先验相关KL散度项的近似方法。
• 该方法将缺失协变量估计整合到编码器-解码器架构中，实现表征和插补的联合优化。
• 推导出一种新型无偏KL散度项估计器，通过近似拆分患者数据，实现小批量训练的稳定性。

实验结果

研究问题

• RQ1 当辅助协变量包含缺失值时，能否在不假设其完全可观测的前提下，有效训练条件VAEs？
• RQ2 在训练过程中，如何正确地对缺失协变量的不确定性进行边缘化，以提升模型泛化能力？
• RQ3 所提出的方法是否在从部分观测数据中学习（含缺失协变量）方面优于现有的基于VAE的方法？
• RQ4 该方法是否能通过小批量处理高效扩展至大规模数据集，特别是GP先验VAEs？
• RQ5 该模型在非时间序列、时间序列和纵向数据上生成真实样本和插补缺失值方面的表现如何？

主要发现

• 所提出的方法在具有缺失辅助协变量的数据集上，实现了CVAEs和GP先验VAEs的最先进性能，在生成和插补任务中均优于先前方法。
• 在模拟的时间序列数字数据集上，尽管存在协变量缺失，模型仍成功学习了潜在生成过程，实现了旋转、平移和对比度调整数字的准确重构与生成。
• 在一项临床试验研究中，输出数据的观测测量值仅占25.01%，且在辅助协变量中人为引入了5%至40%的缺失值，模型在所有缺失水平下均表现出稳健性能。
• 通过推导KL散度项的无偏估计器，该方法实现了小批量训练的稳定性，从而在大规模数据上实现高效优化。
• 与将缺失值替换为零或使用启发式插补的基线VAE相比，该模型在纵向和复杂相关结构数据中表现显著更优。
• 消融研究证实，对缺失协变量的不确定性进行边缘化，相比朴素插补策略，能带来更好的泛化能力和更准确的后验估计。