Supervising the Decoder of Variational Autoencoders to Improve Scientific Utility

TL;DR

This paper introduces SOS-VAE, a second order supervision framework for variational autoencoders that improves the reliability of generative parameters and predictive performance in scientific modeling, addressing bias issues in encoder representations.

Contribution

The paper proposes SOS-VAE, a novel second order supervision method that enhances the interpretability and predictive accuracy of VAEs for scientific data analysis.

Findings

SOS-VAE maintains reliable generative parameters while improving prediction.

The method effectively handles missing data in multi-experiment datasets.

Demonstrated success on synthetic and electrophysiological data.

Abstract

Probabilistic generative models are attractive for scientific modeling because their inferred parameters can be used to generate hypotheses and design experiments. This requires that the learned model provide an accurate representation of the input data and yield a latent space that effectively predicts outcomes relevant to the scientific question. Supervised Variational Autoencoders (SVAEs) have previously been used for this purpose, where a carefully designed decoder can be used as an interpretable generative model while the supervised objective ensures a predictive latent representation. Unfortunately, the supervised objective forces the encoder to learn a biased approximation to the generative posterior distribution, which renders the generative parameters unreliable when used in scientific models. This issue has remained undetected as reconstruction losses commonly used to evaluate model performance do not detect bias in the encoder. We address this previously-unreported issue by developing a second order supervision framework (SOS-VAE) that influences the decoder to induce a predictive latent representation. This ensures that the associated encoder maintains a reliable generative interpretation. We extend this technique to allow the user to trade-off some bias in the generative parameters for improved predictive performance, acting as an intermediate option between SVAEs and our new SOS-VAE. We also use this methodology to address missing data issues that often arise when combining recordings from multiple scientific experiments. We demonstrate the effectiveness of these developments using synthetic data and electrophysiological recordings with an emphasis on how our learned representations can be used to design scientific experiments.