Tomographic Auto-Encoder: Unsupervised Bayesian Recovery of Corrupted Data

1 INTRODUCTION . Data sets are rarely clean and ready to use when first collected . More often than not , they need to undergo some form of pre-processing before analysis , involving expert human supervision and manual adjustments ( Zhou et al. , 2017 ; Chu et al. , 2016 ) . Filling missing entries , correcting noisy samples , filtering collection artefacts and other similar tasks are some of the most costly and time consuming stages in the data modeling process and pose an enormous obstacle to machine learning at scale ( Munson , 2012 ) . Traditional data cleaning methods rely on some degree of supervision in the form of a clean dataset or some knowledge collected from domain experts . However , the exponential increase of the data collection and storage rates in recent years , makes any supervised algorithm impractical in the context of modern applications that consume millions or billions of datapoints . In this paper , we introduce a novel variational framework to perform automated data cleaning and recovery without any example of clean data or prior signal assumptions . The Tomographic auto-encoder ( TAE ) , is named in analogy with standard tomography . Tomographic techniques for signal recovery aim at reconstructing a target signal , such as a 3D image , by algorithmically combining different incomplete measurements , such as 2D images from different view points , subsets of image pixels or other projections ( Geyer et al. , 2015 ) . The TAE extends this concept to the reconstruction of data manifolds ; our target signal is a clean data set , where corrupted data is interpreted as incomplete measurements . Our aim is to combine these to reconstruct the clean data . More specifically , we are interested in performing Bayesian recovery , where we do not simply transform degraded samples into clean ones , but recover probabilistic functions , with which we can generate diverse clean signals and capture uncertainty . Uncertainty is considerably important when cleaning data . If we are over-confident about specific solutions , errors are easily ignored and passed on to downstream tasks . For instance , in the example of figure 1 ( a ) , some corrupted observations are consistent with multiple digits . If we were to impute a single possibility for each sample , the true underlying solution may be ignored early on in the modeling pipeline and the digit will be consistently mis-classified . If we are instead able to recover accurate probability densities , we can remain adequately uncertain in any subsequent processing task . Several variational auto-encoder ( VAE ) models have been proposed for applications that can be considered special cases of this problem ( Im et al. , 2017 ; Nazabal et al. , 2018 ; Ainsworth et al. , 2018 ) and , in principle , they are capable of performing Bayesian reconstruction . However , we show that surrogating variational inference ( VI ) in a latent space with VAEs results in collapsed distributions that do not explore the different possibilities of clean samples , but only return single estimates . The TAE performs approximate VI in the space of recovered data instead , through our reduced entropy condition method . The resulting posteriors adequately explore the manifold of possible clean samples for each corrupted observation and , therefore , adequately capture the uncertainty of the task . In our experiments we focus on data recovery from noisy samples and missing entries . This is one of the most common data corruption settings being encountered in a wide range of domains with different types of data ( White et al. , 2011 ; Kwak & Kim , 2017 ) . By testing our approach in this prevalent scenario , we can closely compare with recently proposed VAE approaches ( Nazabal et al. , 2018 ; Dalca et al. , 2019 ; Mattei & Frellsen , 2019 ) . We show how the existing VAE models exhibit the posterior collapse problem while the TAE produces rich posteriors that capture the underlying uncertainty . We further test TAEs on classification subsequent to imputation , demonstrating superior performance to existing methods in these downstream tasks . Finally , we use a TAE to perform automated missing values imputation on raw depth maps from the NYU rooms data set . 2 METHOD . In order to frame the problem and understand the issues with standard variational methods in this context , we view the task from a signal reconstruction prospective . The final scope of a Bayesian data recovery method is to build and train a parametric probability density function ( PDF ) q ( x|y ) , which takes as inputs corrupted samples y and generates different possible corresponding clean data x ∼ q ( x|y ) through sampling . There are two aspects we need to design : i ) the structure of this conditional PDF and ii ) the way it will be trained to perform the recovery task . Regarding the former , as natural data often lies on highly non-linear manifolds , we need the conditional PDF to capture complicated modalities , e.g . the distribution of plausible images consistent with one of the corrupted observations in figure 1 ( a ) . A suitable recovery PDF q ( x|y ) needs to be able to capture such complexity . A natural choice to achieve high capacity and tractability is to construct q ( x|y ) as a conditional latent variable model ( LVM ) . Conditional LVM neural networks have achieved efficient and expressive variational inference in many recovery settings , capturing complex solution spaces in high dimensional problems , such as image reconstruction ( Nguyen et al. , 2017 ; Mirza & Osindero , 2014 ; Adler & Öktem , 2018 ) . The conditional LVM consists of a first conditional distribution q ( z|y ) mapping input corrupted data y to latent variables z , and a second inference q ( x|z , y ) mapping latent variables to output clean data x . The resulting PDF is q ( x|y ) = ∫ q ( z|y ) q ( x|z , y ) dz , where both q ( z|y ) and q ( x|z , y ) are simple distributions , such as isotropic Gaussians , whose moments are inferred by neural networks taking the respective conditional arguments as inputs . Figure 2 ( a ) shows a graphical model for the conditional LVM . While the choice of structure is fairly straightforward , the main difficulty lies in training the recovery LVM in the absence of clean ground truths x . In the supervised case , several established methods exist ; the observed distributions of clean data x conditioned on paired observations y can be matched by parametric ones through a VAE or GAN training strategy ( Sohn et al. , 2015 ; Adler & Öktem , 2018 ; Tonolini et al. , 2020 ) . However , we are instead interested in the unsupervised situation , where we only have corrupted data Y = { y1 : N } and a functional form for the corrupted data likelihood p ( y|x ) , e.g. , missing values and additive noise . Training a conditional LVM to fit posteriors without any ground truth examples x is rather challenging , as we do not have data to encode from , in the case of VAE architectures , or adversarially compare with , in the case of GAN models . 2.1 VAES AND THE POSTERIOR COLLAPSE PROBLEM . Variational auto-encoders ( VAEs ) have been proposed for several problems within this definition of unsupervised reconstruction ( Dalca et al. , 2019 ; Im et al. , 2017 ; Ainsworth et al. , 2018 ) . These methods lead to good single estimates of the underlying targets . However , they easily over-fit their posteriors resulting in collapsed PDFs q ( x|y ) . Put differently , they are often unable to explore different possible solutions to the recovery problem and return single estimates instead . Figure 1 ( b-ii ) shows this pathology in a two dimensional experiment . The reason for this can be explained considering what the reconstruction LVM q ( x|y ) is and how it is trained when directly employing a VAE in the unsupervised recovery scenario . The VAE encodes latent vectors z from corrupted observations y with an encoder q ( z|y ) and reconstructs clean data x with a decoder p ( x|z ) . These two functions constitute the reconstruction LVM q ( x|y ) =∫ q ( z|y ) p ( x|z ) dz . As we do not have clean ground truths x , data likelihood is maximised by mapping reconstructed clean samples x back to corrupted samples y with a corruption process likelihood p ( y|x ) , e.g . zeroing out missing entries , to maximise reconstruction of the observations y. Concurrently , regularisation in the latent space is induced with a user defined prior p ( z ) ( e.g . a unit Gaussian ) . The resulting lower bound to be maximised during training can be expressed as follows : LV AE = Eq ( z|y ) log p ( y|z ) −KL ( q ( z|y ) ||p ( z ) ) , ( 1 ) where the observations likelihood is p ( y|z ) = ∫ p ( x|z ) p ( y|x ) dx and in some cases , such as for missing values and additive noise , it is analytical . A derivation is given in supplementary A.1 . Viewing the VAE training from a signal reconstruction prospective , where our reconstruction model is q ( x|y ) = ∫ q ( z|y ) p ( x|z ) dz , we can see that we are not introducing any prior directly on the hidden targets x , but only in the LVM latent space z . While regularising only in z may be computationally desirable , if the decoder p ( x|z ) is of sufficient capacity , the model can learn to collapse regularised distributions in z to single estimates in x , failing to capture uncertainty . In fact , this is induced by the objective function of equation 1 ; the model finds broad distributions in the latent space q ( z|y ) , which minimise the KL divergence with p ( z ) , but the generator p ( x|z ) can learn to collapse them back to single maximum likelihood solutions in x , maximising Eq ( z|y ) log ∫ p ( x|z ) p ( y|x ) dx . This effect may be counteracted by reducing the capacity of p ( x|z ) or the dimensionality of z , but doing so also reduces the capacity of the reconstruction model q ( x|y ) , resulting in an undesirable coupling between regularisation and posterior capacity . 2.2 SEPARATING POSTERIOR AND PRIOR : THE TOMOGRAPHIC AUTO-ENCODER . The premise of our model to address the aforementioned problem is simple : Introduce a prior p ( x ) in the hidden clean signal space . In particular , we propose to use an empirical prior , having itself the form of an LVM p ( x ) = ∫ p ( zp ) p ( x|zp ) dzp . In this way , we perform approximate variational inference in clean data space x , instead of surrogating it to the reconstruction function ’ s latent space z . By doing so , we can control the capacity of the prior p ( x ) to induce regularisation independently of the capacity of our reconstruction model q ( x|y ) = ∫ q ( z|y ) q ( x|z , y ) dz . For this framework , We can formulate the following ELBO : LTAE = Eq ( x|y ) log p ( y|x ) + Eq ( x|y ) [ Eq ( zp|x ) log p ( x|zp ) −KL ( q ( zp|x ) ||p ( zp ) ) ] +H ( q ( x|y ) ) . The above ELBO is derived in detail in supplementary section A.2 . The main technical challenge and focus of this paper is how to compute and maximise the self entropy of the approximate posterior H ( q ( x|y ) ) , as this conditional distribution is an LVM of the form q ( x|y ) = ∫ q ( z|y ) q ( x|z , y ) dz . Reduced Entropy Condition : Direct computation of the entropy of an LVM model q ( x|y ) =∫ z q ( z|y ) q ( x|z , y ) dz is intractable in the general case . Titsias & Ruiz ( 2019 ) proposed an approximate inference method to compute the gradient of the LVM ’ s entropy for variational inference in latent spaces . However , this involves multiple samples to be drawn and evaluated with the LVM , which is expected to scale in complexity as the dimensionality and capacity of the target distribution increase . In our case , we aim to approximately compute and optimise the entropy H ( q ( x|y ) ) for a distribution capturing natural data , which can be high-dimensional and lie on complicated manifolds . In order to maintain efficiency in the entropy estimation , we propose a new strategy ; we identify a class of LVM posteriors for which the entropy reduces to a tractable form and then approximately constrain the posterior to such a class in our optimisation . Our main result is summarized in the following theorem : Theorem 1 If q ( z|x , y ) q ( z|y ) = Bδ ( z − g ( x , y ) ) , where δ ( · ) is the Dirac Delta function , B is a real positive parameter and g ( x , y ) is a deterministic function , then H ( q ( x|y ) ) = H ( q ( z|y ) ) + Eq ( z|y ) H ( q ( x|z , y ) ) . We detail the proof in supplementary A.3 . Theorem 1 states that if the posterior over latent variables q ( z|x , y ) is infinitely more localised than the latent conditional q ( z|y ) , then the LVM entropy H ( q ( x|y ) ) has the tractable form given above . This condition imposes the LVM posterior to present non-overlapping conditionals q ( x|z , y ) for different latent variables z , but does not impose any explicit restriction to the capacity of the model . We can also formulate the condition as follows : Eq ( x , z|y ) log q ( z|x , y ) q ( z|y ) = C , C →∞ . ( 2 ) The proof is provided in supplementary section A.4 . To train our posterior q ( x|y ) , we aim to maximise the ELBO LTAE with the reduced entropy , while enforcing the condition of equation 2 : argmax Eq ( x|y ) log p ( y|x ) + Eq ( x|y ) [ Eq ( zp|x ) log p ( x|zp ) −KL ( q ( zp|x ) ||p ( zp ) ) ] +H ( q ( z|y ) ) + Eq ( z|y ) H ( q ( x|z , y ) ) , s.t . Eq ( x , z|y ) log q ( z|x , y ) q ( z|y ) = C , C →∞ . ( 3 ) While the ELBO is now amenable to stochastic optimization , the constraint is intractable since C →∞ and the posterior q ( z|x , y ) is intractable . Relaxed Constraint : To render the constraint tractable , we firstly relax C to be a positive hyperparameter . The higher the value of C , the more localised q ( z|x , y ) is imposed to be compared to q ( z|y ) and the closest the reduced entropy is to the true one . To address the intractability of the posterior q ( z|x , y ) , we employ a variational approximation with a parametric function r ( z|x , y ) . In fact , for any valid probability density r ( z|x , y ) , we can prove that Eq ( x , z|y ) log q ( z|x , y ) q ( z|y ) ≥ Eq ( x , z|y ) log r ( z|x , y ) q ( z|y ) . ( 4 ) The proof is given in supplementary section A.5 . The above bound implicates the following : Eq ( x , z|y ) log r ( z|x , y ) q ( z|y ) = C ⇒ Eq ( x , z|y ) log q ( z|x , y ) q ( z|y ) ≥ C. This means that imposing the condition with a parametric distribution r ( z|x , y ) , which is trained along with the rest of the model , ensures deviation from the set condition only by excess . As the exact condition is met only at Eq ( x , z|y ) log q ( z|x ) q ( z|y ) →∞ , we can never relax the constraint more than already set by the finite value of C. The TAE Objective Function : Having defined a tractable ELBO and a tractable condition , we need to perform the constrained optimisation argmax Eq ( x|y ) log p ( y|x ) + Eq ( x|y ) [ Eq ( zp|x ) log p ( x|zp ) −KL ( q ( zp|x ) ||p ( zp ) ) ] +H ( q ( z|y ) ) + Eq ( z|y ) H ( q ( x|z , y ) ) , s.t . Eq ( x , z|y ) log r ( z|x , y ) q ( z|y ) = C. ( 5 ) We use the commonly adopted penalty function method ( Zangwill , 1967 ; Phuong et al. , 2018 ) and relax equation 5 to an unconstrained optimisation with the use of a positive hyper-parameter λ : argmax Eq ( x|y ) log p ( y|x ) + Eq ( x|y ) [ Eq ( zp|x ) log p ( x|zp ) −KL ( q ( zp|x ) ||p ( zp ) ) ] +H ( q ( z|y ) ) + Eq ( z|y ) H ( q ( x|z , y ) ) − λ ∣∣∣∣Eq ( z , x|y ) log r ( z|x , y ) q ( z|y ) − C ∣∣∣∣ . ( 6 ) To train the model , we perform the maximisation of equation 6 using the ADAM optimiser . Once the model is trained , we can generate diverse reconstructions from a corrupt observation yi by sampling from the posterior q ( x|yi ) . Details of our optimisation are reported in supplementary B.1 . We describe how we handle parameters of the corruption process p ( y|x ) in supplementary B.2 .

Summary: This paper proposes a Tomographic auto-encoder (TAE) for unsupervised recovery of corrupted data. More specifically, TAE takes a Bayesian approach to recover the posterior distribution of a clean image conditioned on an observed corrupted image and thus effectively modeling uncertainty in data recovery. The paper argues that a naive application of VAE is not effective due to the latent variable collapse, and proposes an alternative model where hierarchical latent variable models are used for both prior and variational posterior. Some tricks are introduced to facilitate the stochastic gradient variational inference.