Accurate Imputation and Efﬁcient Data Acquisition with Transformer-based VAEs
Abstract
Predicting missing values in tabular data, with uncertainty, is an essential task by itself as well as for downstream tasks such as personalised data acquisition.
It is not clear whether state-of-the-art deep generative models for these tasks are well equipped to model the complex relationships that may exist between different features, especially when the subset of observed data are treated as a set.
In this work we propose new attention-based models for estimating the joint conditional distribution of randomly missing values in mixed-type tabular data.
The models improve on the state-of-the-art Partial Variational Autoencoder (Ma et al., 2018) on a range of imputation and information acquisition tasks.
Introduction
Real-world data is often incomplete, which indicates a crucial need for extending the applicability of machine learning models beyond complete and homogeneous datasets.
Consider asking a doctor for a disease diagnosis based on a patient’s medical records.
The doctor attempts a diagnosis based on available information, but importantly, they also express their uncertainty regarding missing information that is relevant to the diagnosis.
Introducing such desired behaviour to machine learning models requires the ability to handle incomplete data with many potential patterns of partial observations.
State-of-the-art approaches use the partial variational autoencoder (PVAE) to learn the underlying data distribution from mixed-type tabular data which contains randomly missing values (Ma et al., 2018, 2020; Mattei and Frellsen, 2018).
This enables downstream tasks such as personalised acquisition of relevant information in order to resolve uncertainty in the model’s predictions.
Though successful in handling partially observed data with simple structures, the PVAE’s architecture is less suited to modelling complex interactions between features.
This is due to the use of a simple set-processing network Zaheer et al.
; Qi et al.
(2017) which separately encodes each observed feature and then simply aggregates them, e.g. by taking the mean.
To address this, we improve the PVAE by introducing multi-head attention (Vaswani et al., 2017) to handle the interactions between features.
This is non-trivial due to multiple modeling choices, and we research different modeling choices with experimental analysis.
In particular, in this paper, we introduce attention-based set-processing modules to the PVAE model for probabilistic missing value prediction.
We introduce a batch transformer VAE architecture which generates coherent samples in a single forward pass.
This provides a signiﬁcant speed-up when compared with auto-regressive sampling.
We also present a simple and effective transformer architecture, which we call Rupert, for mixed-type tabular data with missing values, and we present a procedure for personalised data acquisition using Rupert.
We perform experiments on multiple datasets and show that our new models improve on PVAE’s performance, both in terms of raw imputation accuracy, and on a downstream personalised data acquisition task.
Methods
Problem formulation We focus on missing value prediction.
Let x =
[ x 1 , ..., x | I | ] be a set of random variables with joint probability density p ( x ) .
For any set A
⊂ I , we write x A for the subset of variables x i , i ∈
A .
We consider problems where a subset of variables x
O are observed (i.e., their values are given as input to the model) while the remainder x U , U
=
I \
O are unobserved.
For images, x is an image and I is the set of pixel locations.
For tabular data, x is a row in a table and I is the set of columns.
Our models will be used to make inferences about the unobserved values x U given observed values x O .
After training, we use the models to predict randomly masked values in held-out test data and measure their accuracy.
We also test our models on the downstream personalised data acquisition task, in which the models are used to select which variables to observe in order to reduce uncertainty about a given target variable.
We perform the active learning tasks using the same method as in (Ma et al., 2018), with a small modiﬁcation for the one non-VAE model,
Rupert.
2.1 Models
Using pointnet based PVAE (P-PVAE) as baseline, we propose transformer encoder PVAE (TEPVAE) and full transformer PVAE (T-PVAE) as our novel contribution with Rupert as novel reference.
Each of the models that we implement takes as input some partially observed data x O , O ⊂ I , and outputs a probability distribution for each x i .
For continuous variables, the models output a mean and variance of a Gaussian distribution, and for categorical variables the output is passed through a softmax to obtain a discrete probability distribution over class labels.
Figure 1 shows the architectures of the models.
A detailed diagram can be found in the appendix in Figure 3.
Our T-PVAE and TE-PVAE models, described below, follow the design of variational autoencoders (VAE) (Kingma and Welling, 2013).
They assume a latent variable model p ( x | z ) p ( z ) where p ( z ) is the prior of the latent variable z and p ( x | z ) is parametrized by a decoder network.
An encoder network parametrizes q ( z | x ) which approximates p ( z | x ) .
Both networks are trained by maximizing an evidence lower bound (ELBO).
PointNet PVAE (P-PVAE) (Ma et al., 2018)
This baseline PVAE uses PointNet (Qi et al., 2017) as the set encoder which encodes partial observations, each of which may contain values for different numbers of features, to a ﬁxed-size set encoding vector.
PointNet transforms individual input features (points) independently, then aggregates them using a simple sum or max aggregation.
Transformer Encoder PVAE (TE-PVAE)
We improve the encoder of the PVAE ﬁrst.
Improved architectures such as the Set Transformer (Lee et al., 2019) have shown superior performance over
PointNet in expressing complicated permutation-invariant functions.
The main innovation of Set
Transformer is the use of multi-head attention to transform elements of the input set before pooling.
In TE-PVAE, we replace the PointNet with a Set Transformer, which contains several set attention blocks (SABs), followed by pooling by multihead attention (PMA).
The SAB and PMA are described in Lee et al.
(2019), and details of the multi-head attention block are in the appendix in Figure 4.
Transformer PVAE (T-PVAE)
While TE-PVAE has a powerful encoder, it retains the simple multilayer perceptron decoder of P-PVAE, which is a potential bottleneck.
To address this, in TPVAE both the encoder and the decoder are composed of multiple SABs.
The latent dimension (the dimension of z ) is equal to the total number of features (i.e., for tabular data, the number of columns) times the transformer embedding dimension.
The mean and variance of q ( z | x ) are read out using a linear layer from the outputs of the transformer encoder.
Rupert To evaluate the importance of using latent variables to handle partially observed data, we also implemented a simple transformer model to directly estimate p ( x i | x O ) , i ∈ U , which we call Rupert.
The model consists of a stack of SABs between a feature embedder and readout layer borrowed from the PVAE.
Thus, it is similar to BERT (Devlin et al., 2018), but with minimal adaptations so that — like its namesake — it will ingest inputs of various types indiscriminately.
Rupert is not a VAE and does not explicitly represent dependencies between unseen features, but one could sample from p ( x U | x O ) autoregressively, by predicting one feature at a time and adding the already-predicted features to the input of Rupert at each step.
2.2 Implementation
Code is available at https://github.com/microsoft/project-azua.
Training We train all three PVAE models in the same way.
We randomly mask out some of the observed features in each training example and maximise the ELBO:
[Equation]
where A denotes unmasked features whose values are input to the model, B denotes the randomly masked features, and A ∪ B
= O .
This is slightly different from the ELBO used in Ma et al.
, where the likelihood term is just log p ( x A | z ) .
We found that this modiﬁcation, which directly rewards imputation during training, improves imputation performance at test time.
Whereas (Ma et al., 2018) maximises an ELBO for x A , our modiﬁcation means we now maximise an ELBO for x O =
x A ∪ B , subject to a constraint that the inference network must evaluate q ( z | x ) with values of x B masked out, i.e. with some dropout.
We train Rupert by minimizing − (cid:80) i ∈
U log p ( x i | x O ) , similar to masked language model pre-training.
Preprocessing and feature embedding As in Ma et al.
(2018), categorical variables are one-hot encoded and continuous variables are normalised before passing them into any of the models.
We then concatenate each value with a learned feature embedding .
For transformer models, the feature embedding plays a similar role to the positional encoding of Vaswani et al.
Masking In our setting, different numbers of features may be observed in each example within a minibatch, and we handle this using masking.
For TE-PVAE, we use a mask in the multihead attention functions to enforce that the multihead attention does not attend to the masked keys/values.
However, masked values are still passed as queries in the SABs, and it is only after the PMA block that we have outputs that do not depend at all on the masked values.
In T-PVAE and Rupert there is no PMA and so we handle masking differently.
In these models, the feature embedder transforms an observed value x i for feature i to ( x i , e i , 1) (where e i is the learned feature embedding), and an unobserved value to (0 , e i , 0) .
No masking is used the multihead attention functions.
[Table]
Scaling to large numbers of features.
Standard self-attention has complexity O ( N 2 ) in the number of input features N , which becomes prohibitive for some datasets e.g., image datasets where each pixel is a feature.
For such datasets, we replace the SAB with the induced set attention block (ISAB) of Lee et al.
(2019) to reduce the computational and memory complexity of self-attention from quadratic to linear.
Model applications The aforementioned models can be used both for imputing missing data, and for personalised data acquisition similar to (Ma et al., 2018).
To estimate a missing variable x i using a PVAE-based model, we take multiple samples z 1 , ..., z k ∼ q ( z | x O ) , and then report the mean (if continuous) or mode (if categorical) of the equally weighted mixture of p ( x i | z 1 ) , ..., p ( x i | z k ) .
For Rupert, we report the mode of p ( x i | x O ) .
For personalised data acquisition using TE-PVAE and T-PVAE, we can use the same objective and inference method as PVAE Ma et al.
Rupert does not support the same objective as it is not a probabilistic method.
Thus, we introduce predictive variance reduction as an objective for Rupert for the same task.
The algorithm can be found in the appendix.
Evaluation
Datasets We evaluate our models on UCI regression datasets (Dua and Graff, 2017), MNIST (LeCun et al., 1998), as well as data from the online education platform Eedi which consists of students’ answers to hundreds of different multiple-choice maths questions (Wang et al., 2020).
We treat the Eedi data as a table with one row per student and one column per question.
Metrics To evaluate imputation performance, we randomly mask out 30% of variables in each example.
We produce a point estimate for each missing value.
To combine errors for categorical and continuous variables into a single metric, we use the discrete metric for categorical values and report normalised RMSE over all variables.
We evaluate performance on personalised data acquisition using area under the information curve (AUIC) as in Ma et al.
The information curve is the graph of normalised RMSE in the target variable versus the number of variables queried.
Results Table 1 summarises imputation performance of the presented models.
On all datasets,
P-PVAE is outperformed by one of the transformer-based models.
On MNIST, TE-PVAE outperforms
P-PVAE.
Table 2 summarises data acquisition performance for the UCI datasets.
In most cases, when using
EDDI, TE-PVAE gives the best results.
Rupert (with the variance-based algorithm) gives the best results overall, probably because the objective used to select features aligns with the RMSE-based metric more directly.
[Table]
Conclusion
We presented multiple novel deep generative models for missing value prediction by using attention to propagate information between features, and investigated the role of attention in downstream tasks using such models.
Our results show that attention architectures improve performance in both imputation and downstream tasks on multiple datasets.
Our experiments also raise research questions regarding methods for downstream tasks such as Bayesian experimental design.
A Appendix
Require: set I of observable variables, target variable t / ∈
I
Observed variables O ← initially observed variables ⊂
I while O (cid:54)
=
I (or some other stopping criterion) do for i ∈
I \
O do
Sample n values x 1 i , ..., x ni from p ( x i | x O )
[Figure]