summaries: 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.

summaries: Practical datasets often come with corruptions, such as missing items or noisy observations, thus needs models enable to recover the corrupted data automatically. This paper presents the tomographic auto-encoder (TVAE), which conducts inference over the data space $x$. Because the prior regularization acts over the data space, TVAE is enforced to generate diverse samples from the corrupted observations. Empirically, the paper demonstrates that TVAE can indeed generate diverse samples and can achieve superior test ELBO compared to the previous baselines.

summaries: This paper proposes a novel approach to handle the recovery of dirty data in fully unsupervised scenarios. The corrupted data considers both missing data and noisy samples. They derive a VAE model with a novel reduced entropy condition inference method that results in richer posteriors. This is a very challenging problem, since the model cannot use clean examples as part of their training procedure.

abstractText: We propose a new probabilistic method for unsupervised recovery of corrupted data. Given a large ensemble of degraded samples, our method recovers accurate posteriors of clean values, allowing the exploration of the manifold of possible reconstructed data and hence characterising the underlying uncertainty. In this setting, direct application of classical variational methods often gives rise to collapsed densities that do not adequately explore the solution space. Instead, we derive our novel reduced entropy condition approximate inference method that results in rich posteriors. We test our model in a data recovery task under the common setting of missing values and noise, demonstrating superior performance to existing variational methods for imputation and de-noising with different real data sets. We further show higher classification accuracy after imputation, proving the advantage of propagating uncertainty to downstream tasks with our model.

authors:
Jonas Adler
Ozan Öktem

title: Deep Bayesian inversion

venue: arXiv preprint arXiv:1811.05910,

authors:
Samuel K. Ainsworth
Nicholas J. Foti
Emily B. Fox

title: Disentangled VAE representations for multi-aspect and missing data

venue: arXiv preprint arXiv:1806.09060,

authors:
D. Anguita
A. Ghio
L. Oneto
X. Parra
J.L. Reyes-Ortiz

title: Human activity recognition on smartphones using a multiclass hardware-friendly support vector machine

venue: In Proceedings of the International Workshop on Ambient Assisted Living,

authors:
Rushil Anirudh
Jayaraman J. Thiagarajan
Bhavya Kailkhura
Timo Bremer

title: An unsupervised approach to solving inverse problems using generative adversarial networks

venue: arXiv preprint arXiv:1805.07281,

authors:
Angel Chang
Angela Dai
Thomas Allen Funkhouser
Maciej Halber
Matthias Niebner
Manolis Savva
Shuran Song
Andy Zeng
Yinda Zhang

title: Matterport3D: Learning from RGB-D data in indoor environments

venue: In 7th IEEE International Conference on 3D Vision,

authors:
Xu Chu
Ihab F. Ilyas
Sanjay Krishnan
Jiannan Wang

title: Data cleaning: Overview and emerging challenges

venue: In Proceedings of the 2016 International Conference on Management of Data,

authors:
Chris Cremer
Xuechen Li
David Duvenaud

title: Inference suboptimality in variational autoencoders

venue: In Proc. 35th Inter. Conference on Machine Learning,

authors:
Antonia Creswell
Anil Anthony Bharath

title: Denoising adversarial autoencoders

venue: IEEE Transactions on neural networks and learning systems,

authors:
Adrian V. Dalca
John Guttag
Mert R Sabuncu

title: Unsupervised data imputation via variational inference of deep subspaces

title: The MNIST database of handwritten digit images for machine learning research [best of the web

venue: IEEE Signal Processing Magazine,

authors:
Emily Denton
Sam Gross
Rob Fergus

title: Semi-supervised learning with context-conditional generative adversarial networks

venue: arXiv preprint arXiv:1611.06430,

authors:
Piotr Dollár
C. Lawrence Zitnick

title: Structured forests for fast edge detection

venue: In Proceedings of the IEEE international conference on computer vision,

authors:
Marta Garnelo
Dan Rosenbaum
Christopher Maddison
Tiago Ramalho
David Saxton
Murray Shanahan
Yee Whye Teh
Danilo Jimenez Rezende
SM Ali Eslami

title: Conditional neural processes

authors:
Lucas L. Geyer
U. Joseph Schoepf
Felix G Meinel
John W. Nance Jr.
Gorka Bastarrika
Jonathon A. Leipsic
Narinder S. Paul
Marco Rengo
Andrea Laghi
Carlo N. De Cecco

title: State of the art: iterative CT reconstruction techniques

title: Towards unsupervised single-channel blind source separation using adversarial pair unmix-and-remix

venue: IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP),

authors:
Daniel Im Jiwoong Im
Sungjin Ahn
Roland Memisevic
Yoshua Bengio

title: Denoising criterion for variational auto-encoding framework

venue: In Thirty-First AAAI Conference on Artificial Intelligence,

authors:
Phillip Isola
Jun-Yan Zhu
Tinghui Zhou
Alexei A. Efros

title: Image-to-image translation with conditional adversarial networks

authors:
Hirokazu Kameoka
Li Li
Shota Inoue
Shoji Makino

title: Semi-blind source separation with multichannel variational autoencoder

venue: arXiv preprint arXiv:1808.00892,

authors:
Durk P. Kingma
Shakir Mohamed
Danilo Jimenez Rezende
Max Welling

title: Semi-supervised learning with deep generative models

authors:
Alexander Krull
Tim-Oliver Buchholz
Florian Jug

title: Noise2void-learning denoising from single noisy images

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Alexander Krull
Tomas Vicar
Florian Jug

title: Probabilistic Noise2Void: unsupervised contentaware denoising

venue: arXiv preprint arXiv:1906.00651,

authors:
Sang Kyu Kwak
Jong Hae Kim

title: Statistical data preparation: management of missing values and outliers

venue: Korean journal of anesthesiology,

authors:
Jaakko Lehtinen
Jacob Munkberg
Jon Hasselgren
Samuli Laine
Tero Karras
Miika Aittala
Timo Aila

title: Noise2noise: Learning image restoration without clean data

venue: arXiv preprint arXiv:1803.04189,

authors:
Steven Cheng-Xian Li
Bo Jiang
Benjamin Marlin

title: Misgan: Learning from incomplete data with generative adversarial networks

venue: arXiv preprint arXiv:1902.09599,

authors:
Yonghong Luo
Xiangrui Cai
Ying Zhang
Jun Xu

title: Multivariate time series imputation with generative adversarial networks

authors:
Chao Ma
Sebastian Tschiatschek
Konstantina Palla
José Miguel Hernández-Lobato
Sebastian Nowozin
Cheng Zhang

title: Eddi: Efficient dynamic discovery of high-value information with partial vae

venue: arXiv preprint arXiv:1809.11142,

authors:
Pierre-Alexandre Mattei
Jes Frellsen

title: Leveraging the exact likelihood of deep latent variable models

title: MIWAE: deep generative modelling and imputation of incomplete data sets

venue: In International Conference on Machine Learning,

authors:
Mehdi Mirza
Simon Osindero

title: Conditional generative adversarial nets

venue: arXiv preprint arXiv:1411.1784,

title: A study on the importance of and time spent on different modeling steps

venue: ACM SIGKDD Explorations Newsletter,

authors:
Alfredo Nazabal
Pablo M. Olmos
Zoubin Ghahramani
Isabel Valera

title: Handling incomplete heterogeneous data using VAEs

venue: arXiv preprint arXiv:1807.03653,

authors:
Anh Nguyen
Jeff Clune
Yoshua Bengio
Alexey Dosovitskiy
Jason Yosinski

title: Plug & play generative networks: Conditional iterative generation of images in latent space

venue: In CVPR,

authors:
Mary Phuong
Max Welling
Nate Kushman
Ryota Tomioka
Sebastian Nowozin

title: The mutual autoencoder: Controlling information in latent code representations

venue: In ICLR,

authors:
Daniel Scharstein
Richard Szeliski

title: High-accuracy stereo depth maps using structured light

venue: IEEE Computer Society Conference on Computer Vision and Pattern Recognition,

authors:
Chao Shang
Aaron Palmer
Jiangwen Sun
Ko-Shin Chen
Jin Lu
Jinbo Bi

title: VIGAN: Missing view imputation with generative adversarial networks

venue: IEEE International Conference on Big Data (Big Data),

authors:
N. Silberman
R. Fergus

title: Indoor scene segmentation using a structured light sensor

venue: In Proceedings of the International Conference on Computer Vision - Workshop on 3D Representation and Recognition,

authors:
Nathan Silberman
Derek Hoiem
Pushmeet Kohli
Rob Fergus

title: Indoor segmentation and support inference from RGB-D images

venue: In European conference on computer vision,

authors:
Kihyuk Sohn
Honglak Lee
Xinchen Yan

title: Learning structured output representation using deep conditional generative models

authors:
Michalis K. Titsias
Francisco Ruiz

title: Unbiased implicit variational inference

venue: In The 22nd International Conference on Artificial Intelligence and Statistics,

authors:
Francesco Tonolini
Jack Radford
Alex Turpin
Daniele Faccio
Roderick Murray-Smith

title: Variational inference for computational imaging inverse problems

venue: Journal of Machine Learning Research,

authors:
Ian R. White
Patrick Royston
Angela M Wood

title: Multiple imputation using chained equations: issues and guidance for practice

venue: Statistics in medicine,

authors:
H. Xiao
K. Rasul
R. Vollgraf

title: Fashion-MNIST: a novel image dataset for benchmarking machine learning algorithms

authors:
Mingzhang Yin
Mingyuan Zhou

title: Semi-implicit variational inference

authors:
Jinsung Yoon
James Jordon
Mihaela Van Der Schaar

title: Gain: Missing data imputation using generative adversarial nets

venue: arXiv preprint arXiv:1806.02920,

title: Non-linear programming via penalty functions

venue: Management science,

authors:
Lina Zhou
Shimei Pan
Jianwu Wang
Athanasios V Vasilakos

title: Machine learning on big data: Opportunities and challenges

heading: 1 INTRODUCTION

text:
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.

heading: 2 METHOD

text:
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.

heading: 2.1 VAES AND THE POSTERIOR COLLAPSE PROBLEM

text:
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.

heading: 2.2 SEPARATING POSTERIOR AND PRIOR: THE TOMOGRAPHIC AUTO-ENCODER

text:
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.

heading: 3 RELATED WORK

heading: 3.1 SUPERVISED BAYESIAN RECONSTRUCTION

text:
The reconstruction of posterior densities from incomplete measurements has been recently investigated in supervised situations, where examples of clean data are available. In particular, conditional generative models were demonstrated with high dimensional data (Parmar et al., 2018). These methods work by exploiting an LVM to generate diverse realisations of targets conditioned on associated observations (Isola et al., 2017; Nguyen et al., 2017). Both conditional generative adversarial networks (CGANs) (Mirza & Osindero, 2014; Isola et al., 2017) and conditional VAEs (CVAEs) (Sohn et al., 2015; Nguyen et al., 2017) have been studied in this context. In both cases, the samples generated by conditioning on an observation can be interpreted as samples from the corresponding conditional posterior densities.
These approaches proved successful in a range of recovery tasks: reconstruction of images with missing groups of pixels (Nguyen et al., 2017), super-resolution (Parmar et al., 2018), medical computed tomography reconstructions (Adler & Öktem, 2018) and semi supervised situations, where examples of clean data and conditions are available in different amounts (Kingma et al., 2014; Denton et al., 2016; Tonolini et al., 2020). Other works reconstruct manifolds of solutions from observations only, but can be considered supervised, as they exploit pre-trained generators (Anirudh et al., 2018).
These works make the important observation that when learning to recover data from corrupted or partial observations, there is not a single right solution, but many differently likely ones. We aim to extend this ability to completely unsupervised scenarios, where no clean data examples are available.

heading: 3.2 UNSUPERVISED BAYESIAN RECONSTRUCTION

text:
Reconstructing posteriors in the unsupervised case is largely still an open problem. However, several tasks that fall within this definition have been recently approached with Bayesian machine learning methods. Arguably the most investigated is de-noising. Several works solve this problem by exploiting the natural tendency of neural networks to regularise outputs (Lehtinen et al., 2018; Krull et al., 2019a;b). Other methods build LVMs that explicitly model the noise process in their decoder, retrieving clean samples upon encoding and generation (Im et al., 2017; Creswell & Bharath, 2018).
A second notable example is that of missing value imputation. Corrupted data corresponds to samples with missing entries. Recent works have explored the use of LVMs to perform imputation, both with GANs (Li et al., 2019; Yoon et al., 2018; Luo et al., 2018) and VAEs (Nazabal et al., 2018; Mattei & Frellsen, 2019; Ma et al., 2018). In the former, the discriminator of the GAN is trained to distinguish real values from imputed ones, such that the generator is induced to synthesise realistic imputations. In the latter, the encoder of a VAE maps incomplete samples to a latent space, to then generate complete samples. Successful unsupervised Bayesian missing value imputation has also been demonstrated with neural processes, where a global latent representation is learned to generate input-output models used to impute in each example (Garnelo et al., 2018).
Finally, Bayesian LVM methods have been used on other unsupervised tasks that can be cast as special cases of data recovery problems. Amongst these, we find Multi-view generation (Shang et al., 2017; Ainsworth et al., 2018), where the target clean data includes all views for each samples, but the observed data only presents subsets. Blind source separation can also be cast as a recovery problem and has been approached with GANs and VAEs (Kameoka et al., 2018; Hoshen, 2019).
These models proved to be successful at reconstructing data in their specific domain. However, in our work, we show how exploiting a standard VAE inference structure, similarly to several of the aforementioned methods, often leads to posteriors of clean data that collapse on single estimates, sacrificing the probabilistic capability of LVMs.

heading: 3.3 POSTERIOR COLLAPSE IN VARIATIONAL INFERENCE

text:
The posterior collapse problem we approach with the reduced entropy condition method presented in this paper has some analogy with the latent posterior collapse encountered when using implicit distributions in variational inference to obtain flexible recognition models. The main issue in training these models successfully without collapse is the computation of density rations between the latent prior and the implicit variational posterior. This problem is analogous to the difficulty in estimating the LVM entropy in our method. Yin & Zhou (2018) proposed to use a further lower bound on the ELBO and add a term encouraging diversity to avoid collapse. This term is obtained by drawing K Gaussian components from the LVM posterior and computing the KL divergence of an individual component with the mixture distribution, i.e. the sum of the drawn Gaussians. Titsias & Ruiz (2019) build on this work by deriving an unbiased estimator for the ELBO gradient, instead of using a surrogate lower bound.
These methods rely on estimating the LVM posterior through sampling and aggregating K explicit distribution components. This was demonstrated to work well for posteriors in artificial latent space by drawing only a few components. However, in our data recovery setting, we need to capture the posteriors in clean data space, rather than a latent space. Posteriors capturing the uncertainty in natural data are expected to be much more complex and higher dimensional, leading to the number K of drawn Gaussians needed to approximate the true LVM with these methods to become rather large, making optimisation inefficient or even intractable in extreme cases. The reduced entropy condition method we derive in this paper avoids the posterior collapse without having to estimate the LVM posterior through sampling and is therefore specially suited for the data recovery setting, where we are required to capture posteriors in clean data space.

heading: 4 EXPERIMENTS

heading: 4.1 POSTERIOR RECOVERY

text:
We corrupt the MNIST dataset (Deng, 2012) by introducing missing values and additive Gaussian noise on the observed entry. We then train both a missing value imputation VAE (MVAE), analogous to those presented in (Nazabal et al., 2018) and (Dalca et al., 2019), and our TAE model with the corrupted data sets. The VAE and TAE are constructed such that the structure of their posteriors q(x|y), i.e. the functions mapping corrupted data to distributions of clean data at test time, are exactly the same. In this way, we can ensure that differences in performance are due to the variational inference method employed and not the choice of posterior model. The resulting variational posteriors are used to perform data recovery from the corrupted samples. Fig. 3(a) shows examples of mean reconstruction and posterior draws. See analogous experiments for grouped missings in suppl. D.2.
We evaluate the accuracy of mean reconstruction at different ratios of observed entries by measuring the peak signal to noise ratio (PSNR) between the ground truth data and mean recoveries (Figure 3(b)). To evaluate probabilistic performance we approximately measure the likelihood assigned by the recovered posteriors to the ground truth data through a reconstruction ELBO, by training a new inference function with the clean ground truths, but leaving the posterior fixed, as is common for evaluating ELBOs in unsupervised settings (Cremer et al., 2018; Mattei & Frellsen, 2018; 2019). A detailed description of this approach is given in supplementary C.3. We also carry out analogous experiments testing de-noising and missing value imputation separately. THese results are reported in supplementary D.4 and D.5. Results are shown in figure 3(c). We further evaluate our TAE with Fashion-MNIST – 28 × 28 grey-scale images of clothing (Xiao et al., 2017), and the UCI HAR dataset, which consists of filtered accelerometer signals from mobile phones worn by different people
during common activities (Anguita et al., 2012). As before, we test the recovery of these data sets from a version affected by missing values and additive noise. In addition to the MVAE baseline, we compared against the recently proposed missing values importance weighted auto encoder (MIWAE) (Mattei & Frellsen, 2019), which optimises an importance weighted ELBO in place of the standard one. For each model and settings we compute the ELBO assigned to the ground truth data. Results are shown in Table 2. Experimental details in Sec. C of suppl. mat.

heading: 4.2 DOWNSTREAM TASKS

text:
To investigate the advantage of capturing complex uncertainties with our TAE model, we are interested in testing performance in downstream tasks. We test classification performance on subsets of the MNIST and Fashion-MNIST data sets, after recovery with our TAE. With both sets, we consider situations in which 10.000 examples are available, but corrupted with missing entries and noise. 1, 000 of these are labelled with one of 10 possible classes and we wish to classify the remaining 9, 000. To do so, we first train the TAE model on the full set, then use the recovered posteriors to generate multiple possible cleaned data for the labelled sub-set and use them to train a classifier.
To perform classification on the 9, 000 remaining examples, we generate multiple possible cleaned data with the variational posteriors. Then, for each posterior sample, we perform classification and histogram the results. Examples are shown in figure 4. To evaluate the performance, we take the class with the largest histogram as the inferred one. We repeat this experiment for different ratios of
missing values and several repetitions, varying the subsets of labelled and unlabelled data to be used. Classification accuracy results are shown in figure 5.

heading: 4.3 MISSING VALUES IN THE NYU DEPTH MAPS

text: As a final practical application, we use a convolutional version of our TAE to perform structured missing value imputation on depth maps of indoors rooms collected with a Kinect depth sensor. Missing entries are very common in depth maps recorded with such structured light sensors (Scharstein & Szeliski, 2003). We use raw depth data from the NYU rooms dataset, commonly used to test various computer vision systems (Silberman & Fergus, 2011; Silberman et al., 2012; Dollár & Zitnick, 2013; Chang et al., 2018). A large portion of the set is available only as raw data, which presents missing entries. These are especially concentrated around objects’ edges and reflecting surfaces, breaking the common assumption of missing at random, making this task particularly challenging. We train our TAE with a subset of this raw data set to perform imputation. Examples of results are shown in figure 11. Additional examples are shown in supplementary section D.6.

heading: 5 CONCLUSION

text: We presented tomographic auto-encoders; a variational inference method for recovering posterior distributions of clean data from a corrupted data set alone. We derive the reduced entropy condition method; a novel inference strategy that results in rich distributions of clean data given corrupted observations, thereby capturing the uncertainty of the task, while standard variational methods often collapse on single answers. In our experiments, we demonstrate this capability and show the advantage of capturing uncertainty with the TAE in downstream tasks, outperforming the state-of-the-art VAE based recovery methods.

heading: ACKNOWLEDGEMENTS

text: F.T. and R.M-S. acknowledge funding support from Amazon and EPSRC grants EP/M01326X/1, EP/T00097X/1 (QuantIC – UK Quantum Technology Hub in Quantum Enhanced Imaging). R.MS. acknowledges funding support from EP/R018634/1 (Closed-Loop Data Science for Complex, Computationally- and Data-Intensive Analytics).

heading: Tomographic Auto-Encoder - Supplementary material

heading: A PROOFS AND DERIVATIONS

heading: A.1 DERIVATION OF VAE ELBO FOR DATA RECOVERY

text:
We aim to maximise the log likelihood of the observed corrupted data y
log p(y) = log ∫ x ∫ z
p(z)p(x|z)dz︸ ︷︷ ︸ p(x) p(y|x)dx. (7)
We can introduce a variational distribution in both clean data space and latent space q(x, z|y) and define a lower bound as
log p(y) ≥ ∫ x ∫ z q(x, z|y) log p(z)p(x|z)dz q(x, z|y) dzdx
+ ∫ x ∫ z q(x, z|y) log p(y|x)dzdx. (8)
To obtain the VAE ELBO used in data recovery settings, the choice of the variational posterior is q(x, z|y) = q(z|y)p(x|z). The ELBO can then be simplified to give
log p(y) ≥ ∫ x ∫ z q(z|y)p(x|z) log p(z)p(x|z)dz q(z|y)p(x|z) dzdx
+ ∫ x ∫ z q(z|y)p(x|z) log p(y|x)dzdx
= ∫ x
p(x|z)dx︸ ︷︷ ︸ =1
∫ z q(z|y) log p(z)dz q(z|y) dz
+ ∫ x ∫ z q(z|y)p(x|z)dz log p(y|x)dx
(9)
For situations in which the observations’ likelihood ∫ x p(x|z)p(y|x)dx has a closed form, such as additive noise and missing entries, we can define a tighter bound to the likelihood by moving the integral in x in the second term inside the logarithm
log p(y) ≥ ∫ z q(z|y) log p(z)dz q(z|y) dz
+ ∫ z q(z|y) log [∫ x p(y|x)p(x|z)dx ] dz =−KL(q(z|y)||p(z)) + Eq(z|y) log p(y|z)dx.
(10)
Because p(x|z) simplifies in the KL term, this ELBO avoids variational inference in the space of clean data x.

heading: A.2 DERIVATION OF TAE ELBO

text:
In our TAE model we defined separate LVMs for prior and posterior. To distinguish between the posterior latent variable and the prior latent variable, we name the former z and the latter zp. The likelihood we aim to maximise is
log p(y) = log ∫ x ∫ zp
p(zp)p(x|zp)dzp︸ ︷︷ ︸ p(x) p(y|x)dx. (11)
Similarly to the VAE ELBO case, we define a variational posterior q(x, zp|y) to find a lower bound log p(y) ≥ ∫ x ∫ zp q(x, zp|y) log p(zp)p(x|zp)dzp q(x, zp|y) dzpdx
+ ∫ x ∫ zp q(x, zp|y) log p(y|x)dzpdx. (12)
However, in our model we do not make the assumption that the variational posterior has the special form described in section A.1 and instead set it to have the form q(x, zp|y) = q(x|y)q(zp|x), separating posterior inference from observations y to clean data x and inference of prior latent variables zp. The resulting lower bound is
log p(y) ≥ ∫ x ∫ zp q(x|y)q(zp|x) log p(zp)p(x|zp) q(x|y)q(zp|x) dzpdx+ ∫ x ∫ zp q(x|y)q(zp|x) log p(y|x)dzpdx
= ∫ x q(x|y) ∫ zp q(zp|x) log p(zp)p(x|zp) q(zp|x)
dzp︸ ︷︷ ︸ ≥log p(x) dx+
∫ x ∫ zp
q(zp|x)︸ ︷︷ ︸ =1 dzpq(x|y) log p(y|x)dx
− ∫ x ∫ zp
q(zp|x)dzp︸ ︷︷ ︸ =1 q(x|y) log q(x|y)dx
=Eq(x|y) [ Eq(zp|x) log p(x|zp)−KL(q(zp|x)||p(zp)) ] + Eq(x|y) log p(y|x) +H(q(x|y)).
(13)
A.3 PROOF OF THEOREM 1
q(z|x, y) q(z|y) = Bδ(z − g(x, y)) =⇒ q(z|x, y) q(z|y) q(z′|x, y) q(z′|y) = 0, ∀x, z 6= z′
=⇒ q(x|z, y) q(x|y) q(x|z′, y) q(x|y) = 0, ∀x, z 6= z′
=⇒ q(x|z, y)q(x|z′, y) = 0, ∀x, z 6= z′
=⇒ q(x|z′, y) = 0, ∀x ∼ q(x|z, y), z 6= z′
(14)
Using the result of equation 14, we can derive the form of the entropy H(q(x|y)) for this special case as the following:
H(q(x|y)) =− ∫ x [∫ z q(z|y)q(x|z, y)dz ] · log [ ∫ z′ q(z′|y)q(x|z′, y)dz′ ] dx
=− ∫ x ∫ z q(z|y)q(x|z, y) · log [ ∫ z′=z q(z′|y)q(x|z′, y)dz′
+ ∫ z′ 6=z
q(z′|y)q(x|z′, y)dz′︸ ︷︷ ︸ eq.14 =⇒ =0
] dzdx
= − ∫ z ∫ x q(z|y)q(x|z, y) log [q(z|y)q(x|z, y)] dxdz
= − ∫ z ∫ x q(z|y)q(x|z, y) log q(z|y)dxdz − ∫ z ∫ x q(z|y)q(x|z, y) log q(x|z, y)dxdz
= − ∫ z q(z|y) log q(z|y)dz − ∫ z q(z|y) ∫ x q(x|z, y) log q(x|z, y)dxdz = H(q(z|y)) + Eq(z|y)H(q(x|z, y)).
(15)

heading: A.4 PROOF OF THE EQUIVALENCE BETWEEN CONDITIONS

text:
proof of necessary condition:
Eq(x,z|y) log q(z|x, y) q(z|y) = ∫ z q(z|y) ∫ x q(x|z, y) log q(z|x, y) q(z|y) dxdz
= ∫ x q(x|y) ∫ z q(z|x, y) log q(z|x, y) q(z|y) dzdx
= ∫ x q(x|y) ∫ z q(z|x, y) log q(z|x, y)dzdx
− ∫ x q(x|y) ∫ z q(z|x, y) log q(z|y)dzdx
= ∫ x q(x|y) ∫ z q(z|x, y) log q(z|x, y)dzdx
− ∫ x
q(x|z, y)dx︸ ︷︷ ︸ =1
∫ z q(z|y) log q(z|y)dz
=Eq(x|y) ∫ z
q(z|x, y) log q(z|x, y)dz︸ ︷︷ ︸ −H(q(z|x,y))
− ∫ z
q(z|y) log q(z|y)dz︸ ︷︷ ︸ −H(q(z|y)) .
(16)
If the above expression tends to infinity, either H(q(z|x, y))→ −∞ or H(q(z|y))→∞, meaning that either q(z|x, y)→ a Delta function, or q(z|y)→ uniform. Either condition implies q(z|x,y)q(z|y) = Bδ(z − g(x, y)). proof of sufficient condition:
Eq(x,z|y) log q(z|x, y) q(z|y) = ∫ x q(x|y) ∫ z q(z|x, y) log q(z|x, y) q(z|y) dzdx
= ∫ x q(x|y) ∫ z q(z|y)q(z|x, y) q(z|y) log q(z|x, y) q(z|y) dzdx.
(17)
Now we set q(z|x,y)q(z|y) = Bδ(z − g(x, y)):
∫ x q(x|y) ∫ z q(z|y)Bδ(z − g(x, y)) logBδ(z − g(x, y))dzdx
= ∫ x
q(x|y)q(g(x, y)|y) logB δ(g(x, y)− g(x, y))︸ ︷︷ ︸ →∞,∀x dx. (18)
Therefore, q(z|x,y)q(z|y) = Bδ(z − g(x, y)) is a sufficient condition for Eq(x,z|y) log q(z|x,y) q(z|y) →∞.

heading: A.5 PROOF OF EQUATION 6

text:
Eq(x,z|y) log q(z|x, y) q(z|y) = ∫ z ∫ x q(x, z|y) log q(z|x, y)dzdx
− ∫ z ∫ x q(x, z|y) log q(z|y)dzdx
= ∫ x q(x|y) ∫ z q(z|x, y) log q(z|x, y)dzdx
− ∫ z ∫ x q(x, z|y) log q(z|y)dzdx
≥ ∫ x q(x|y) ∫ z q(z|x, y) log r(z|x, y)dzdx
− ∫ z ∫ x q(x, z|y) log q(z|y)dzdx =Eq(x,z|y) log r(z|x, y) q(z|y) ,
(19)
Where the inequality derives from the positivity of the KL divergence KL(q(z|x, y)||r(z|x, y)).

heading: B ALGORITHM

heading: B.1 DETAILS OF TRAINING

text:
As detailed in section 3.2, to train our variational posterior q(x|y), we maximise through gradient ascent the TAE ELBO with the reduced entropy penalty function
argmax Eq(x|y) log p(y|x) + Eq(x|y) [ Eq(zp|x) log p(x|zp)−KL(q(zp|x)||p(zp))︸ ︷︷ ︸
Prior ELBO, ≥p(x)
]
+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 ∣∣∣∣ . (20)
All expectations in the above expression are computed and optimised by sampling the corresponding conditional distributions using the re-parametrisation trick characteristic of VAEs.
Because the prior LVM p(x) = ∫ p(zp)p(x|zp)dzp is training entirely with samples from the posterior LVM, which is also training, the model can easily obtain high values for the prior ELBO by generating collapsed samples x with the posterior and get stuck in an unfavourable local minimum. TO avoid this, we employ a warm up strategy. We define a positive parameter γ that multiplies the expectation of the prior ELBO and the entropy H(x|z, y):
argmax Eq(x|y) log p(y|x) + γEq(x|y) [ Eq(zp|x) log p(x|zp)−KL(q(zp|x)||p(zp))︸ ︷︷ ︸
Prior ELBO, ≥p(x)
]
+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 ∣∣∣∣ . (21)
The value of γ is initially set to zero. After a set number of iterations it is linearly increased to reach one and kept constant for the remaining training iterations.

heading: B.2 COMPLETE OBJECTIVE FUNCTION

text:
Observation Parameters: In the general case, the corruption process p(y|x), mapping clean data x to degraded samples y, is controlled by parameters that differ from sample to sample. We can distinguish these into observed parameters α and unobserved parameters β. For example, in the case of missing values and noise, the indexes of missing entries in each sample are often observed parameters, while the noise level is an unobserved parameter. The complete form of the corruption likelihood for a clean sample xi is then p(y|xi, αi, βi).
Objective Function: With the parameters conditionals described in subsection 3.2.4 and explicitly showing the parameters to be optimised, the objective function we maximise is the following
argmax θ,φ Eqφ(x,β|y,α) log p(y|x, α, β) +γEqφ(x|y,α) [ Eqφ3 (zp|x) log pθ(x|zp)−KL(qφ3(zp|x)||p(zp)) ] +H(qφ1(z|y, α)) + γEqφ1 (z|y,α)H(qφ2(x|z, y, α))
−λ ∣∣∣∣Eqφ(z,x|y,α) log rφ4(z|x)qφ1(z|y, α) − C ∣∣∣∣ , (22)
qφ(x, β|y, α) = ∫ z qφ1(z|y, α)qφ2(x|z, y, α)qφ5(β|z, y, α)dz, qφ(x|y, α) =∫
z qφ1(z|y, α)qφ2(x|z, y, α)dz, qφ(z, x|y, α) = qφ1(z|y, α)qφ2(x|z, y, α), φ = {φ1:5} are
the parameters of the inference models and θ are the parameter of the prior model.

heading: B.3 PSEUDO-CODE

heading: B.4 POSTERIOR COLLAPSE IN VARIATIONAL INFERENCE

heading: C EXPERIMENTAL DETAILS

heading: C.1 MODELS’ ARCHITECTURES

text:
In all experiments we carry out comparing our TAE with competitive methods, we make the independence assumption q(x|z, y) = q(x|z), consequentially making r(z|x, y) = r(z|x). In this way, the reconstruction posterior LVMs q(x|y) we compare between TAE, MVAE and MIWAE all present identical structure and differences in performance are a result of the model constructed to train them alone. However, we note that, unlike the two competing method, the TAE is not formally limited to this choice and can infer conditionals q(x|z, y) in the general case. We hereafter detail the architecture used for all quantitative experiments of section 4.1 and 4.2.
MVAE and MIWAE models: The MVAE model is built as an LVM having a unit Gaussian prior in the latent space p(z) = N (z;0,1) and isotropic Gaussian clean data likelihood p(x|z) = N (x;µx, σ2x), where the moments µx and σ2x are outputs of a neural network having as input the latent variables z. Because we only observe corrupted data y, rather than clean data x, the recognition model q(z|y) is conditioned on observed corrupted data y and also has the form of an isotropic Gaussian q(z|y) = N (z;µz, σ2z), where the moments µz and σ2z are outputs of a
Algorithm 1 Training the TAE Model
Inputs: Corrupted observations Y = {y1:N}; Observed Parameters A = {α1:N} initial model parameters, {θ(0), φ(0)}; user-defined posterior latent dimensionality, J ; user-defined prior latent dimensionality, Jp; user-defined condition strength λ; user-defined condition parameter C; userdefined latent prior p(zp); user-defined initial warm-up coefficient γ0; user-defined final warm-up coefficient γf ; warm-up start Nw0; warm-up end Nwf ; user-defined number of iterations, Niter.
1: γ(k=0) ← γ0 2: for the k’th iteration in [0 : Niter − 1] 3: for the i’th observation 4: zi ∼ qφ(k)1 (z|yi, αi) 5: xi ∼ qφ(k)2 (x|zi, yi, αi) 6: βi ∼ qφ(k)5 (β|zi, yi, αi) 7: zp,i ∼ qφ(k)3 (zp|xi) 8: E(k)i ← log p(yi|xi, βi) 9: P(k)i ← log pθ(k)(xi|zp,i)
10: K(k)i ← DKL(qφ(k)3 (zp|xi)||p(zp)) 11: Hz(k)i ← H(qφ(k)1 (z|yi, αi)) 12: Hx(k)i ← H(qφ(k)2 (x|zi, yi, αi)) 13: R(k)i ← log rφ(k)4 (zi|xi, yi, αi) 14: Q(k)i ← log qφ(k)1 (zi|yi, αi) 15: end 16: F(k) = ∑ i ( E(k)i + γ(k) [ P(k)i −K (k) i + Hx (k) i
] 17: +Hz(k)i − λ
∣∣∣R(k)i −Q(k)i − C∣∣∣ ) 18: θ(k+1), φ(k+1) ← argmax(F(k))
19: if k > Nw0 and k < Nwf 20: γ(k+1) ← γ(k) + (γf − γ0)/(Nwf −Nw0) 21: else 22: γ(k+1) ← γ(k) 23: end 24: end
neural network having as input the corrupted observations y. The corrupt data likelihood p(y|z) is obtained by simply selecting the likelihood p(x|z) over the observed indexes, i.e. for missing values corruption the integral p(y|z) = ∫ p(x|z)p(y|x)dx simply masks out the unobserved entries. The model is then trained by maximising the ELBO of equation 1. The MIWAE is built with the same structure, but instead of optimising the MVAE ELBO of equation 1, an importance weighted lower bound is maximised, as described in (Mattei & Frellsen, 2019). The precise architectures used for the neural networks are described for the different experiments throughout the rest of this section. One important point to notice is that, in each experiments, the structures of p(x|z) and q(z|y) are chosen such that the resulting reconstruction model after training, i.e. the model taking as input a test corrupt observation y and generating clean samples x, is identical for the TAE and the two tested competing models. That is ∫ q(z|y)p(x|z)dz for the MVAE and MIWAE have identical structure to∫
q(z|y)q(x|z)dz for the TAE. In this way, performance differences can be attributed solely to the difference in inference strategy and not reconstruction model’s capacity.
Posteriors structure: The posterior parametric components are qφ1(z|y, α) and qφ2(x|z) (pφ2(x|z) in the case of the MVAE and MIWAE). qφ1(z|y, α) consists in a fully connected two layers neural network with leaky ReLu non-linearities, taking as input concatenated corrupted observations y and a binary mask that labels the missing entries α and returning as output a vector of latent means and a vector of latent log variances. The two intermediate deterministic layers have 400 hidden units, while the latent space z is 20-dimensional.
qφ2(x|z), and pφ2(x|z) in the case of the MVAE and MIWAE, are similarly constructed, consisting in a fully connected two layers neural network with leaky ReLu non-linearities, taking as input latent variables z and returning a vector of means and a vector of log variances of clean samples x. The two intermediate deterministic layers have 400 hidden units.
TAE Prior LVM Structure: The TAE prior encoder qφ3(zp|x) has the same general structure as the posterior encoder, with two fully connected layers and leaky ReLu non-linearities, taking as input generated clean data x and returning as outputs a vector of latent means and a vector of latent log variances for the prior latent variable zp. As this model has less capacity than the posterior LVM, the two deterministic hidden layers have 50 hidden units each and the latent variables zp are 5-dimensional.
pθ(x|zp) is similarly constructed, consisting in a fully connected two layers neural network with leaky ReLu non-linearities, taking as input latent variables zp and returning a vector of means and a vector of log variances of clean samples x. The two intermediate deterministic layers have 50 hidden units.
Approximate Latent Posterior Structure: The approximate latent posterior r(z|x) has the same structure as the posterior encoder, consisting in a fully connected two layers neural network with leaky ReLu non-linearities, taking as input generated clean data x and returning as outputs a vector of latent means and a vector of latent log variances. The two intermediate deterministic layers have 400 hidden units.
Convolutional TAE Structure: For the imputation of NYU missing data we use convolutional conditionals in our model, instead of fully connected ones. In this version. we do not make the independence assumption, using q(x|z, y) and r(z|x, y). q(z|y, α) takes concatenated y and α and passes them through 4 recurrent convolutional layers with filters of size 3× 3 and 5 channels, each time down-sampling by two. the last layer is mapped to means and standard deviation of latent images z, which are 1/32 of the original size in each axis and have 10 channels, through two convolutional filter banks with strieds 1 × 1. q(x|z, y, α) is built to mirror this structure, with the addition of accepting inputs from y and α. Three recurrent transpose convolutional layers with 3× 3 filters, 5 channels and 2× 2 upsampling each map z to a deterministic layer with 1/2× 1/2 of the original images size. concatenated y and α are mapped to the same size with a single convolutional layer, downsampling it by 1/2× 1/2 and 5 channels. The two are concatenated and the resulting layer is finally upsampled to inferred clean image x by a last convolution with a filter bank. All non-linearities are Elu.
The prior networks are built in a similar way, but with shallower structures to give less capacity. q(zp|x) passes x through 2 convolutional layers, each with down-sampling of 4× 4 and 5 channels. as before, mens and standard deviations of latent images zp are generated from this last layers with 2 × 2 down-sampling and, in this case, 5 channels. The prior generator p(x|zp) is built to exactly mirror this structure. r(z|x, y, α) has the same structure as q(z|y, α), with the only difference being that it accepts as input concatenated x, y and α.

heading: C.2 EXPERIMENTAL CONDITIONS

text:
Posterior Recovery: All posterior recovery experiments, with each of the three data sets tested, are performed on samples that have been re-scaled from 0 to 1. In all cases, the sets are injected with additive Gaussian noise having standard deviation 0.1. Subsequently, random binary masks are generated to block out some entries, resulting in missing values. The proportion of missing entries in the masks was set as described in the main body in each case.
Experiments were repeated with re-generated binary masks 5 times. The means and error bars shown in figure 4 and the uncertainty reported in table 1 were computed from these. The MIWAE was trained with 20 weights per sample. After training, all posteriors q(x|y) have identical structure and
are tested in the same way, by training an inference network on the test set to compute the ELBO values.
Classification Experiments: The TAE models for the MNIST and Fashion-MNIST experiments were trained in the conditions described above. In each case, a random subset of 10, 000 samples is taken from the corrupted set and the TAE and MVAE models are trained with it. A random subset of 1, 000 of these is selected and ground-truth lables for these samples are made available.
A classifier consisting in a single fully connected layer with leaky ReLu non-linearity is trained to perform classification on this subset. For each stochastic training iteration of this classifier, we generate samples associated with the corrupted observations and provide the associated labels. After the classifier is trained, we test classification performance on the remaining 9, 000 examples, by running the train classifier 400 times per sample, each time generating clean data from a corrupted observation with the TAE and the MVAE. The histograms shown in figure 5 are built by aggregating the resulting classification.
The above procedure is repeated 15 times. The resulting means and standard deviations of the tested classification performance are shown in figure 6.
Training Conditions: Hyper-parameters of optimisation for the models were cross validated with the MNIST data set at a proportion of missing entries of 0.9. Hyper-parameters common to all models were determined by obtaining best performance with the MVAE model. Hyper- parameters specific to the TAE model were obtained by fixing the common parameters and cross validating these. The resulting optimal hyper parameters were then used in all other experiments of section 4.1 and 4.2, including those with different data sets. Common parameters are as follows: 500, 000 iterations with the ADAM optimiser in Tensorflow, an initial training of 2−4 and batch size of 20. The hyper-parameters specific to the TAE are instead: γ initially set to 0.01 and then linearly increased to 1 between 50, 000 and 100, 000 iterations, λ = 2 and C = 10. All experiments were performed using a TitanX GPU.
NYU Rooms Experiments: For these experiments, we take a subset of 3612 depth maps from the NYU raw data set. We slightly crop these in one dimension to be 480×608 images. The convolutional TAE and MVAE to obtain the results of figure 7, were trained for 100, 000 iterations using the ADAM optimiser in Tensorflow, with a batch size of 20 images and an initial training rate of 2× 10−2. For the warm up, we initially set γ = 0.01 and linearly increase it to 1 between 10, 000 and 20, 000 iterations. For these experiments, λ = 2 and C = 15.

heading: C.3 EVALUATION ELBO

text:
To evaluate the probabilistic performance of our method compared to others, we compute an evaluation ELBO which relies on test ground truths. After each model is trained unsupervisedly, we obtain a posterior of the form q(x|y) = ∫ q(z|y)q(x|z)dz, where for the MVAE and MIWAE, q(x|z) = p(x|z). Given the a test set of paired clean and corrupted samples xt and yt, we construct a new parametric recognition model, which encodes latent distributions from ground-truths qη(z|x). We then optimise the following:
argmax η Eqη(z|xt) log q(xt|z) +KL(qη(z|xt)|q(z|yt)). (23)
The above is a conditional VAE ELBO with conditional prior q(z|y) and is a lower bound to the test likelihood we are interested in q(xt|yt). Note that we optimise over η only, therefore the new recognition model q(z|x) is the only one which is affected by this optimisation and the components of our reconstruction model q(z|y) and q(x|z) remain the same as trained with the unsupervised training set. As a result, this new optimisation only tightens the bound, rather than maximising the likelihood, which we want to evaluate as trained previously.

heading: D ADDITIONAL EXPERIMENTS

text:
D.1 C AND λ CROSS-VALIDATION
C and λ in equation 22 are hyper-parameters of our inference algorithm and need to be user defined. In our experiments, we determine the optimal values by cross-validation, as described in section C.
We report in figure 7 a cross validation study where we measure the TAE ELBO for MNIST with 90% missing values and additive noise.
As shown in figure 7, the performance of TAEs is robust to variations in hyper-parameters C and λ over a broad range of values. They also have an intuitive meaning that helps in their selection. In practice, C controls the final value of localisation and is desirable to be as high as stability of the optimisation allows. λ controls how fast we are imposing the model to approach C.

heading: D.2 STRUCTURED MISSINGS

text:
We test a TAE in a situation analogous to that shown in figure 4 of section 4, but with structured missing values instead of randomly missing ones. For each sample in MNIST, we only make visible a small 10 × 10 pixels window, randomly placed in each example, while the rest of the images remain hidden. In addition, the values in the observed window are subject to additive Gaussian noise, similarly to the missing-at-random case. Reconstructions with the comparative MVAE and our TAE are shown in figure 8.
Similarly to the missing-at-random case, the MVAE collapses on single solutions, giving draws from the posterior that are all very similar to each other. Contrarily, the TAE gives more variation in the possible solutions exploring more appropriately the uncertainty in the solution space. The MVAE ELBO over the clean data for this problem is 428, while the TAE one is 638. The performance improvement provided by the TAE is analogous to that observed with missing-at-random experiments.

heading: D.3 MORE ELBO EVALUATIONS

text:
D.4 IMPUTATION WITHOUT NOISE
We carry out experiments on MNIST analogous to those shown in figure 3, but in the absence of noise, in order to test performance on imputation alone. Each tested ratio of observed entries is
repeated three times, re-generating the patterns of missings each repeat in order to obtain error bars. Results are shown in figure 10.

heading: D.5 DE-NOISING

text:
We carry out experiments on MNIST analogous to those shown in figure 3, but testing fully observed images at different levels of noise. Each tested ratio of observed entries is repeated three times, re-generating the patterns of missings each repeat in order to obtain error bars. Results are shown in figure 10.
D.6 NYU ROOMS RECOVERY EXAMPLES

title: TOMOGRAPHIC AUTO-ENCODER: UNSUPERVISED BAYESIAN RECOVERY OF CORRUPTED DATA


summaries: This paper proposes MAGNA, a multi-hop self-attention mechanism for attention based graph neural networks. The proposed method increases the receptive field at each layer, requiring less layers to achieve a large receptive field. Also, with the proposed method the attention coefficient between two nodes is not just a function of the two nodes but also of their neighbourhood. The proposed MAGNA method is an extension of GAT networks that introduces a diffusion step on the computed attention coefficients, following a similar approach (Diffusion-GCNs) that has been used for GCNs.

summaries: The authors propose a novel attention-based GNN called MAGNA. The main contribution consists in considerably increasing the receptive field by considering a multi-hop neighborhood instead of the standard one hop. The technical challenge consists in obtaining attention scores for all relevant nodes in an efficient way. MAGNA solves this by using a diffusion-based technique combined with a geometric distribution. The authors show that the latter further allows for approximations, and also give interesting theoretical insights (e.g., show a relation to page rank).

summaries: Conventional Graph Neural Networks (GNNs) learn node representations that encode information from multiple hops away by iteratively aggregating information through their immediate neighbors. Self-Attention modules have been adopted to GNNs to selectively aggregate information coming through the immediate neighbors at different propagation stages. However, current self-attention mechanisms are limited to only attend over the nodes' immediate neighbors and not directly over their neighbors that are multiple hops away. Here in this work, the authors intend to address this issue and propose a means to obtain attention scores over indirectly connected neighbors.

abstractText: Self-attention mechanism in graph neural networks (GNNs) led to state-of-theart performance on many graph representation learning task. Currently, at every layer, attention is computed between connected pairs of nodes and depends solely on the representation of the two nodes. However, such attention mechanism does not account for nodes that are not directly connected but provide important network context, which could lead to improved predictive performance. Here we propose Multi-hop Attention Graph Neural Network (MAGNA), a principled way to incorporate multi-hop context information into attention computation, enabling long-range interactions at every layer of the GNN. To compute attention between nodes that are not directly connected, MAGNA diffuses the attention scores across the network, which increases the “receptive field” for every layer of the GNN. Unlike previous approaches, MAGNA uses a diffusion prior on attention values, to efficiently account for all paths between the pair of disconnected nodes. This helps MAGNA capture large-scale structural information in every layer, and learn more informative attention. Experimental results on node classification as well as the knowledge graph completion benchmarks show that MAGNA achieves stateof-the-art results: MAGNA achieves up to 5.7% relative error reduction over the previous state-of-the-art on Cora, Citeseer, and Pubmed. MAGNA also obtains the best performance on a large-scale Open Graph Benchmark dataset. On knowledge graph completion MAGNA advances state-of-the-art on WN18RR and FB15k237 across four different performance metrics.

authors:
Dzmitry Bahdanau
Kyunghyun Cho
Yoshua Bengio

title: Neural machine translation by jointly learning to align and translate

authors:
Ivana Balazevic
Carl Allen
Timothy Hospedales

title: Tucker: Tensor factorization for knowledge graph completion

venue: In EMNLP,

authors:
Trapit Bansal
Da-Cheng Juan
Sujith Ravi
Andrew McCallum

title: A2n: Attending to neighbors for knowledge graph inference

venue: In ACL,

authors:
Peter W Battaglia
Jessica B Hamrick
Victor Bapst
Alvaro Sanchez-Gonzalez
Vinicius Zambaldi
Mateusz Malinowski
Andrea Tacchetti
David Raposo
Adam Santoro
Ryan Faulkner

title: Relational inductive biases, deep learning, and graph networks

venue: arXiv preprint arXiv:1806.01261,

authors:
Antoine Bordes
Nicolas Usunier
Alberto Garcia-Duran
Jason Weston
Oksana Yakhnenko

title: Translating embeddings for modeling multi-relational data

venue: In NeurIPS,

authors:
Ines Chami
Adva Wolf
Da-Cheng Juan
Frederic Sala
Sujith Ravi
Christopher Ré

title: Lowdimensional hyperbolic knowledge graph embeddings

authors:
Michaël Defferrard
Xavier Bresson
Pierre Vandergheynst

title: Convolutional neural networks on graphs with fast localized spectral filtering

authors:
Tim Dettmers
Pasquale Minervini
Pontus Stenetorp
Sebastian Riedel

title: Convolutional 2d knowledge graph embeddings

venue: In AAAI,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova. Bert

title: Pre-training of deep bidirectional transformers for language understanding

venue: In NAACL,

authors:
Hongyang Gao
Shuiwang Ji

title: Graph u-nets

venue: In ICML,

authors:
Hongyang Gao
Zhengyang Wang
Shuiwang Ji

title: Large-scale learnable graph convolutional networks

venue: In KDD,

authors:
Aditya Grover
Jure Leskovec

title: node2vec: Scalable feature learning for networks

authors:
Will Hamilton
Zhitao Ying
Jure Leskovec

title: Inductive representation learning on large graphs

authors:
Elvin Isufi
Fernando Gama
Alejandro Ribeiro

title: Edgenets: Edge varying graph neural networks

venue: arXiv preprint arXiv:2001.07620,

authors:
Thomas N Kipf
Max Welling

title: Semi-supervised classification with graph convolutional networks

authors:
Johannes Klicpera
Aleksandar Bojchevski
Stephan Günnemann

title: Predict then propagate: Graph neural networks meet personalized pagerank

authors:
Johannes Klicpera
Stefan Weißenberger
Stephan Günnemann

title: Diffusion improves graph learning

authors:
Zhenzhong Lan
Mingda Chen
Sebastian Goodman
Kevin Gimpel
Piyush Sharma
Radu Soricut

title: Albert: A lite bert for self-supervised learning of language representations

authors:
Qimai Li
Zhichao Han
Xiao-Ming Wu

title: Deeper insights into graph convolutional networks for semi-supervised learning

authors:
Renjie Liao
Zhizhen Zhao
Raquel Urtasun
Richard S. Zemel

title: Lanczosnet: Multi-scale deep graph convolutional networks

authors:
Yinhan Liu
Myle Ott
Naman Goyal
Jingfei Du
Mandar Joshi
Danqi Chen
Omer Levy
Mike Lewis
Luke Zettlemoyer
Veselin Stoyanov

title: Roberta: A robustly optimized bert pretraining approach

venue: arXiv preprint arXiv:1907.11692,

authors:
Ziqi Liu
Chaochao Chen
Longfei Li
Jun Zhou
Xiaolong Li
Le Song
Yuan Qi

title: Geniepath: Graph neural networks with adaptive receptive paths

title: Efficient Algorithms for Personalized PageRank

venue: PhD thesis, Stanford University,

authors:
Sitao Luan
Mingde Zhao
Xiao-Wen Chang
Doina Precup

title: Break the ceiling: Stronger multiscale deep graph convolutional networks

authors:
Andrew Y Ng
Michael I Jordan
Yair Weiss

title: On spectral clustering: Analysis and an algorithm

authors:
Kenta Oono
Taiji Suzuki

title: Graph neural networks exponentially lose expressive power for node classification

title: Approximation theory of the mlp model in neural networks

venue: Acta numerica,

authors:
Alec Radford
Jeffrey Wu
Rewon Child
David Luan
Dario Amodei
Ilya Sutskever

title: Language models are unsupervised multitask learners

venue: OpenAI Blog,

authors:
Aliaksei Sandryhaila
José MF Moura

title: Discrete signal processing on graphs: Graph fourier transform

venue: In ICASSP,

authors:
Michael Schlichtkrull
Thomas N Kipf
Peter Bloem
Rianne Van Den Berg
Ivan Titov
Max Welling

title: Modeling relational data with graph convolutional networks

venue: In ESWC,

authors:
Prithviraj Sen
Galileo Namata
Mustafa Bilgic
Lise Getoor
Brian Galligher
Tina Eliassi-Rad

title: Collective classification in network data

venue: AI magazine,

authors:
Chao Shang
Yun Tang
Jing Huang
Jinbo Bi
Xiaodong He
Bowen Zhou

title: End-to-end structureaware convolutional networks for knowledge base completion

authors:
Uday Shankar Shanthamallu
Jayaraman J Thiagarajan
Andreas Spanias

title: A regularized attention mechanism for graph attention networks

authors:
Zhiqing Sun
Zhi-Hong Deng
Jian-Yun Nie
Jian Tang

title: Rotate: Knowledge graph embedding by relational rotation in complex space

authors:
Yun Tang
Jing Huang
Guangtao Wang
Xiaodong He
Bowen Zhou

title: Orthogonal relation transforms with graph context modeling for knowledge graph embedding

authors:
Kiran K Thekumparampil
Chong Wang
Sewoong Oh
Li-Jia Li

title: Attention-based graph neural network for semi-supervised learning

venue: arXiv preprint arXiv:1803.03735,

authors:
Kristina Toutanova
Danqi Chen

title: Observed versus latent features for knowledge base and text inference

venue: In CVSC-WS,

authors:
Théo Trouillon
Johannes Welbl
Sebastian Riedel
Éric Gaussier
Guillaume Bouchard

title: Complex embeddings for simple link prediction

authors:
Ashish Vaswani
Noam Shazeer
Niki Parmar
Jakob Uszkoreit
Llion Jones
Aidan N Gomez
Łukasz Kaiser
Illia Polosukhin

title: Attention is all you need

venue: NeurIPS,

authors:
Petar Veličković
Guillem Cucurull
Arantxa Casanova
Adriana Romero
Pietro Lio
Yoshua Bengio

title: Graph attention networks

authors:
Guangtao Wang
Rex Ying
Jing Huang
Jure Leskovec

title: Improving graph attention networks with large margin-based constraints

venue: In NeurIPS-WS,

authors:
Quan Wang
Pingping Huang
Haifeng Wang
Songtai Dai
Wenbin Jiang
Jing Liu
Yajuan Lyu
Yong Zhu
Hua Wu

title: Coke: Contextualized knowledge graph embedding

venue: arXiv preprint arXiv:1911.02168,

title: Open graph benchmark: Datasets for machine learning on graphs

authors:
Zonghan Wu
Shirui Pan
Fengwen Chen
Guodong Long
Chengqi Zhang
Philip S Yu

title: A comprehensive survey on graph neural networks

venue: IEEE Trans Neural Netw Learn Syst,

authors:
Louis-Pascal A.C. Xhonneux
Meng Qu
Jian Tang

title: Continuous graph neural networks

venue: arXiv preprint arXiv:1901.00596,

authors:
Ruibin Xiong
Yunchang Yang
Di He
Kai Zheng
Shuxin Zheng
Chen Xing
Huishuai Zhang
Yanyan Lan
Liwei Wang
Tie-Yan Liu

title: On layer normalization in the transformer architecture

authors:
Keyulu Xu
Chengtao Li
Yonglong Tian
Tomohiro Sonobe
Ken-ichi Kawarabayashi
Stefanie Jegelka

title: Representation learning on graphs with jumping knowledge

authors:
Bishan Yang
Wen-tau Yih
Xiaodong He
Jianfeng Gao
Li Deng

title: Embedding entities and relations for learning and inference in knowledge bases

authors:
Jiani Zhang
Xingjian Shi
Junyuan Xie
Hao Ma
Irwin King
Dit-Yan Yeung

title: Gaan: Gated attention networks for learning on large and spatiotemporal graphs

venue: arXiv preprint arXiv:1803.07294,

authors:
Shuai Zhang
Yi Tay
Lina Yao
Qi Liu

title: Quaternion knowledge graph embedding

authors:
Chenyi Zhuang
Qiang Ma

title: Dual graph convolutional networks for graph-based semi-supervised classification

venue: In WWW,

heading: 1 INTRODUCTION

text:
The introduction of the self-attention mechanism (Bahdanau et al., 2015; Vaswani et al., 2017), has pushed the state-of-the-art in many domains including graph presentation learning (Radford et al., 2019; Devlin et al., 2019; Liu et al., 2019a; Lan et al., 2019). Graph Attention Network (GAT) (Veličković et al., 2018) and related models (Li et al., 2018; Wang et al., 2019a; Liu et al., 2019b; Oono & Suzuki, 2020) developed attention mechanism for Graph Neural Networks (GNNs), which compute attention scores between nodes connected by an edge, allowing the model to attend to messages of node’s direct neighbors according to their attention scores. However, such attention computation on pairs of nodes connected by edges implies that a node can only attend to its immediate neighbors to compute its (next layer) representation. This implies that receptive field of a single GNN layer is restricted to one-hop network neighborhoods. Although stacking multiple GATs could in principle enlarge the receptive field and learn non-neighboring interactions, such deep GAT architectures suffer from the oversmoothing problem (Wang et al., 2019a; Liu et al., 2019b; Oono & Suzuki, 2020) and do not perform well in practice. Furthermore, edge attention weights in the single GAT layer are based solely on representations of the two nodes at the edge endpoints, and do not depend on their graph neighborhood context. In other words, the one-hop attention mechanism in GATs limits their ability to explore the relationship between the broader graph structure and the attention weights. While previous works (Xu et al., 2018; Klicpera et al., 2019b) have shown advantages in performing multi-hop message-passing in a single layer, these approaches are not graph-attention based. Therefore, incorporating multi-hop neighboring context into the attention computation in graph neural networks had not been explored. Here we present Multi-hop Attention Graph Neural Network (MAGNA), an effective and efficient multi-hop self-attention mechanism for graph structured data. MAGNA uses a novel graph attention diffusion layer (Figure 1), where we first compute attention weights on edges (represented by solid arrows), and then compute self-attention weights (dotted arrows) between disconnected pairs of nodes through an attention diffusion process using the attention weights on the edges.
Our model has two main advantages: 1) MAGNA captures long-range interactions between nodes that are not directly connected but may be multiple hops away. Thus the model enables effective long-range message passing, from important nodes multiple hops away. 2) The attention computation in MAGNA is context-dependent. The attention value in GATs (Veličković et al., 2018) only depends on node representations of the previous layer, and is zero between disconnected pairs of nodes. In contrast, for any pair of nodes within a chosen multi-hop neighborhood, MAGNA computes attention by aggregating the attention scores over all the possible paths (length≥ 1) connecting the two nodes. Theoretically we demonstrate that MAGNA places a Personalized Page Rank (PPR) prior on the attention values. We further apply spectral graph analysis to show that MAGNA has the capability of emphasizing on large-scale graph structure and lowering high-frequency noise in graphs. Specifically, MAGNA enlarges the lower Laplacian eigen-values, which correspond to the large-scale structure in the graph, and suppresses the higher Laplacian eigen-values which correspond to more noisy and fine-grained information in the graph.
We experiment on standard datasets for semisupervised node classification as well as knowledge graph completion. Experiments show that MAGNA achieves state-of-the-art results: MAGNA achieves up to 5.7% relative error reduction over previous state-of-the-art on Cora, Citeseer, and Pubmed. MAGNA also obtains better performance on a large-scale Open Graph Benchmark dataset. On knowledge graph completion, MAGNA advances state-of-the-art on WN18RR and FB15k-237 across four metrics, with the largest gain of 7.1% in the metric of Hit at 1. Furthermore, we show that MAGNA with just 3 layers and 6 hop wide attention per layer significantly out-performs GAT with 18 layers, even though both architectures have the same receptive field. Moreover, our ablation study reveals the synergistic effect of the essential components of MAGNA, including layer normalization and multi-hop diffused attention. We further observe that compared to GAT, the attention values learned by MAGNA have higher diversity, indicating the ability to better pay attention to important nodes.

heading: 2 MULTI-HOP ATTENTION GRAPH NEURAL NETWORK (MAGNA)

text: We first discuss the background and then explain Multi-hop Attention Graph Neural Network’s novel multi-hop attention diffusion module and its overall model architecture.

heading: 2.1 PRELIMINARIES

text:
Let G = (V, E) be a given graph, where V is the set of Nn nodes, E ⊆ V × V is the set of Ne edges connecting M pairs of nodes in V . Each node v ∈ V and each edge e ∈ E are associated with their type mapping functions: φ : V → T and ψ : E → R. Here T and R denote the sets of node types (labels) and edge/relation types. Our framework supports learning on heterogeneous graphs with multiple elements inR. A general Graph Neural Network (GNN) approach learns an embedding that maps nodes and/or edge types into a continuous vector space. Let X ∈ RNn×dn and R ∈ RNr×dr be the node embedding and edge/relation type embedding, where Nn = |V|, Nr = |R|, dn and dr represent the embedding dimension of node and edge/relation types, each row xi =X[i :] represents the embedding of node vi (1 ≤ i ≤ Nn), and rj = R[j :] represents the embedding of relation rj (1 ≤ j ≤ Nr). MAGNA builds on GNNs, while bringing together the benefits of Graph Attention and Diffusion techniques. The core of MAGNA is Multi-hop Attention Diffusion, a principled way to learn attention between any pair of nodes in a scalable way, taking into
account the graph structure and enabling multihop context-dependent attention directly. The key challenge here is how to allow for flexible but scalable context-dependent multi-hop attention, where any node can influence embedding of any other node in a single GNN layer (even if they are far away in the underlying network). Simply learning attention scores over all node pairs is infeasible and would lead to overfitting and poor generalization.

heading: 2.2 MULTI-HOP ATTENTION DIFFUSION

text:
We first introduce attention diffusion to compute the multi-hop attention directly, which operates on the MAGNA’s attention scores at each layer. The input to the attention diffusion operator is a set of triples (vi, rk, vj), where vi, vj are nodes and rk is the edge type. MAGNA first computes the attention scores on all edges. The attention diffusion module then computes the attention values between pairs of nodes that are not directly connected by an edge, based on the edge attention scores, via a diffusion process. The attention diffusion module can then be used as a component in MAGNA architecture, which we will further elaborate in Section 2.3.
Edge Attention Computation. At each layer l, a vector message is computed for each triple (vi, rk, vj). To compute the representation of vj at layer l + 1, all messages from triples incident to vj are aggregated into a single message, which is then used to update vl+1j .
In the first stage, the attention score s of an edge (vi, rk, vj) is computed by the following: s (l) i,k,j = LeakyReLU(v (l) a tanh(W (l) h h (l) i ‖W (l) t h (l) j ‖W (l) r rk)) (1) whereW (l)h ,W (l) t ∈Rd (l)×d(l) ,W (l)r ∈Rd (l)×dr and v(l)a ∈R1×3d (l)
are the trainable weights shared by l-th layer. h(l)i ∈ Rd (l)
represents the embedding of node i at l-th layer, and h(0)i = xi. rk is the trainable relation embedding of the k-th relation type rk (1 ≤ k ≤ Nr), and a‖b denotes concatenation of embedding vectors a and b. For graphs with no relation type, we treat as a degenerate categorical distribution with 1 category 1. Applying Eq. 1 on each edge of the graph G, we obtain an attention score matrix S(l):
S (l) i,j = { s (l) i,k,j , if (vi, rk, vj) appears in G −∞, otherwise
(2)
Subsequently we obtain the attention matrix A(l) by performing row-wised softmax over the score matrix S(l): A(l) = softmax(S(l)). A(l)ij denotes the attention value at layer l when aggregating message from node j to node i. Attention Diffusion for Multi-hop Neighbors. In the second stage, we further enable attention between nodes that are not directly connected in the network. We achieve this via the following attention diffusion procedure. The procedure computes the attention scores of multi-hop neighbors via graph diffusion based on the powers of the 1-hop attention matrixA:
A = ∞∑ i=0 θiA i where ∞∑ i=0 θi = 1 and θi > 0 (3)
where θi is the attention decay factor and θi > θi+1. The powers of attention matrix, Ai, give us the number of relation paths from node h to node t of length up to i, increasing the receptive field of the attention (Figure 1). Importantly, the mechanism allows the attention between two nodes to not only depend on their previous layer representations, but also taking into account of the paths between the
1In this case, we can view that there is only one “pseudo” relation type (category), i.e., Nr = 1
nodes, effectively creating attention shortcuts between nodes that are not directly connected (Figure 1). Attention through each path is also weighted differently, depending on θ and the path length. In our implementation we utilize the geometric distribution θi = α(1 − α)i, where α ∈ (0, 1]. The choice is based on the inductive bias that nodes further away should be weighted less in message aggregation, and nodes with different relation path lengths to the target node are sequentially weighted in an independent manner. In addition, notice that if we define θ0 = α ∈ (0, 1], A0 = I , then Eq. 3 gives the Personalized Page Rank (PPR) procedure on the graph with the attention matrix A and teleport probability α. Hence the diffused attention weights, Aij , can be seen as the influence of node j to node i. We further elaborate the significance of this observation in Section 4.3.
We can also viewAij as the attention value of node j to i since ∑Nn j=1Aij = 1.2 We then define the graph attention diffusion based feature aggregation as AttDiffusion(G,H(l), Θ) = AH(l), (4) where Θ is the set of parameters for computing attention. Thanks to the diffusion process defined in Eq. 3, MAGNA uses the same number of parameters as if we were only computing attention between nodes connected via edges. This ensures runtime efficiency as well as good model generalization. Approximate Computation for Attention Diffusion. For large graphs computing the exact attention diffusion matrix A using Eq. 3 may be prohibitively expensive, due to computing the powers of the attention matrix (Klicpera et al., 2019a). To resolve this bottleneck, we proceed as follows: Let H(l) be the input entity embedding of the l-th layer (H(0) = X) and θi = α(1 − α)i. Since MAGNA only requires aggregation via AH(l), we can approximate AH(l) by defining a sequence Z(K) which converges to the true value of AH(l) (Proposition 1) as K →∞: Z(0) =H(l), Z(k+1) = (1− α)AZ(k) + αZ(0), where 0 ≤ k < K (5) Proposition 1. limK→∞ Z(K) = AH(l)
In the Appendix we give the proof which relies on the expansion of Eq. 5. Using the above approximation, the complexity of attention computation with diffusion is still O(|E|), with a constant factor corresponding to the number of hops K. In practice, we find that choosing the values of K such that 3 ≤ K ≤ 10 results in good model performance. Many realworld graphs exhibit small-world property, in which case even a smaller value of K is sufficient. For graphs with larger diameter, we choose larger K, and lower the value of α.

heading: 2.3 DIRECT MULTI-HOP ATTENTION BASED GNN ARCHITECTURE

text:
Figure 2 provides an architecture overview of the MAGNA Block that can be stacked multiple times. Multi-head Graph Attention Diffusion Layer. Multi-head attention (Vaswani et al., 2017; Veličković et al., 2018) is used to allow the model to jointly attend to information from different representation sub-spaces at different viewpoints. In Eq. 6, the attention diffusion for each head i is computed separately with Eq. 4, and aggregated:
Ĥ(l) = MultiHead(G, H̃(l)) = (∥∥∥M
i=1 headi
) Wo, where
headi = AttDiffusion(G, H̃(l), Θi), H̃(l) = LayerNorm(H(l)), (6)
where ‖ denotes concatenation andΘi are the parameters in Eq. 1 for the i-th head (1 ≤ i ≤M ), and Wo represents a parameter matrix. Since we calculate the attention diffusion in a recursive way in Eq. 5, we add layer normalization which helpful to stabilize the recurrent computation procedure (Ba et al., 2016). Deep Aggregation. Moreover our MAGNA block contains a fully connected feed-forward sublayer, which consists of a two-layer feed-forward network. We also add the layer normalization and residual connection in both sub-layers, allowing for a more expressive aggregation step for each block (Xiong et al., 2020):
Ĥ(l+1) = Ĥ(l) +H(l)
H(l+1) =W (l) 2 ReLU
( W
(l) 1 LayerNorm(Ĥ
(l+1)) ) + Ĥ(l+1)
(7)
MAGNA generalizes GAT. MAGNA extends GAT via the diffusion process. The feature aggregation in GAT is H(l+1) = σ(AH(l)W (l)), where σ represents the activation function. We can
2Obtained by the attention definition A(l) = softmax(S(l)) and Eq. 3.
divide GAT layer into two components as follows: H(l+1) = σ︸︷︷︸
(2) (AH(2)W (l)︸ ︷︷ ︸ (1) ). (8)
In component (1), MAGNA removes the restriction of attending to direct neighbors, without requiring additional parameters as A is induced from A. For component (2) MAGNA uses layer normalization and deep aggregation which achieves significant gains according to ablation studies in Table 1. Compared to the “shallow” activation function elu in GAT, we can view deep aggregation (i.e., two-layer MLP) as a learnable deep activation function as two layer MLP can approximate many different functions (Pinkus, 1999).

heading: 3 ANALYSIS OF GRAPH ATTENTION DIFFUSION

text: In this section, we investigate the benefits of MAGNA from the viewpoint of discrete signal processing on graphs (Sandryhaila & Moura, 2013). Our first result demonstrates that MAGNA can better capture large-scale structural information. Our second result explores the relation between MAGNA and Personalized PageRank (PPR).

heading: 3.1 SPECTRAL PROPERTIES OF GRAPH ATTENTION DIFFUSION

text:
We view the attention matrix A of GAT, and A of MAGNA as weighted adjacency matrices, and apply Graph Fourier transform and spectral analysis (details in Appendix) to show the effect of MAGNA as a graph low-pass filter, being able to more effectively capture large-scale structure in graphs. By Eq. 3, the sum of each row of either A or A is 1. Hence the normalized graph Laplacians are L̂sym = I − A and Lsym = I − A for A and A respectively. We can get the following proposition:
Proposition 2. Let λ̂gi and λ g i be the i-th eigeinvalues of L̂sym and Lsym. λ̂gi λgi = 1− α 1−(1−α)(1−λgi ) λgi = 1 α 1−α + λ g i . (9) Refer to Appendix for the proof. We additionally have λgi ∈ [0, 2] (proved by (Ng et al., 2002)). Eq. 9 shows that when λgi is small such that α 1−α + λ g i < 1, then λ̂ g i > λ g i , whereas for large λ g i , λ̂gi < λ g i . This relation indicates that the use ofA increases smaller eigenvalues and decreases larger eigenvalues3. See Section 4.3 for its empirical evidence. The low-pass effect increases with smaller α. The eigenvalues of the low-frequency signals describe the large-scale structure in the graph (Ng et al., 2002) and have been shown to be crucial in graph tasks (Klicpera et al., 2019b). As λgi ∈ [0, 2] (Ng et al., 2002) and α1−α > 0, the reciprocal format in Eq. 9 will amplify the ratio of lower eigenvalues to the sum of all eigenvalues. In contrast, high eigenvalues corresponding to noise are suppressed.

heading: 3.2 PERSONALIZED PAGERANK MEETS GRAPH ATTENTION DIFFUSION

text:
We can also view the attention matrixA as a random walk matrix on graph G since ∑Nn j=1Ai,j = 1 and Ai,j > 0. If we perform Personalized PageRank (PPR).with parameter α ∈ (0, 1] on G with transition matrixA, the fully Personalized PageRank (Lofgren, 2015) is defined as: Appr = α(I − (1− α)A)−1 (10) Using the power series expansion for the matrix inverse, we obtain
Appr = α ∞∑ i=0 (1− α)iAi = ∞∑ i=0 α(1− α)iAi (11)
Comparing to the diffusion Equation 3 with θi = α(1− α)i, we have the following proposition. Proposition 3. Graph attention diffusion defines a personalized page rank with parameter α ∈ (0, 1] on G with transition matrixA, i.e., A = Appr.
The parameter α in MAGNA is equivalent to the teleport probability of PPR. PPR provides a good relevance score between nodes in a weighted graph (the weights from the attention matrix A). In
3The eigenvalues of A and A correspond to the same eigenvectors, as shown in Proposition 2 in Appendix.
summary, MAGNA places a PPR prior over node pairwise attention scores: the diffused attention between node i and j depends on the attention scores on the edges of all paths between i and j.

heading: 4 EXPERIMENTS

text: We evaluate MAGNA on two classical tasks4. (1) On node classification we achieve an average of 5.7% relative error reduction; (2) On knowledge graph completion we achieve 7.1% relative improvement in the Hit at 1 metric.5 We compare with numbers reported by baseline papers when available.

heading: 4.1 TASK 1: NODE CLASSIFICATION

text:
Datasets. We employ four benchmark datasets for node classification: (1) standard citation network benchmarks Cora, Citeseer and Pubmed (Sen et al., 2008; Kipf & Welling, 2016); and (2) a benchmark dataset ogbn-arxiv on 170k nodes and 1.2m edges from the Open Graph Benchmark (Weihua Hu, 2020). We follow the standard data splits for all datasets. Further information about these datasets is summarized in the Appendix. Baselines. We compare against a comprehensive suite of state-of-the-art GNN methods including: GCNs (Kipf & Welling, 2016), Chebyshev filter based GCNs (Defferrard et al., 2016), DualGCN (Zhuang & Ma, 2018), JKNet (Xu et al., 2018), LGCN (Gao et al., 2018), DiffusionGCN (Klicpera et al., 2019b), APPNP (Klicpera et al., 2019a), Graph U-Nets (g-U-Nets) (Gao & Ji, 2019), and GAT (Veličković et al., 2018). Experimental Setup. For datasets Cora, Citeseer and Pubmed, we use 6 MAGNA blocks with hidden dimension 512 and 8 attention heads. For the large-scale ogbn-arxiv dataset, we use 2 MAGNA blocks with hidden dimension 128 and 8 attention heads. Refer to Appendix for detailed description of all hyper-parameters and evaluation settings.
Results. We report node classification accuracies on the benchmarks. Results are summarized in Tables 1 and 2. MAGNA improves over all methods and achieves the new state-of-the-art on all datasets.
4All datasets used are public, and the code will be released at the time of publication. 5Please see the definitions of these two tasks in Appendix.
Ablation study. We report (Table 1) the model performance after removing each component of MAGNA (layer normalization, attention diffusion and deep aggregation feed forward layers) from every layer of MAGNA. Note that the model is equivalent to GAT without these three components. We observe that both diffusion and layer normalization play a crucial role in improving the node classification performance for all datasets. While layer normalization alone does not benefit GNNs, its use in conjunction with the attention diffusion module significantly boosts MAGNA’s performance. Since MAGNA computes many attention values, layer normalization is crucial in ensuring training stability (Ba et al., 2016). Meanwhile, we also remove both layer normalization and deep aggregation feed forward layer, and only keep the attention diffusion layer (see the next-to-last row of Table 1). Comparing to GAT, attention diffusion allows multi-hop attention in each layer still benefits the performance of node classification.

heading: 4.2 TASK 2: KNOWLEDGE GRAPH COMPLETION

text: Datasets. We evaluate MAGNA on standard benchmark knowledge graphs: WN18RR (Dettmers et al., 2018) and FB15K-237 (Toutanova & Chen, 2015). See the statistics of these KGs in Appendix. Baselines. We compare MAGNA with state-of-the-art baselines, including (1) translational distance based models: TransE (Bordes et al., 2013) and its latest extension RotatE (Sun et al., 2019), OTE (Tang et al., 2020) and ROTH (Chami et al., 2020); (2) semantic matching based models: ComplEx (Trouillon et al., 2016), QuatE (Zhang et al., 2019), CoKE (Wang et al., 2019b), ConvE (Dettmers et al., 2018), DistMult (Yang et al., 2015) and TuckER (Balazevic et al., 2019); (3) GNN-based models: R-GCN (Schlichtkrull et al., 2018), SACN (Shang et al., 2019) and A2N (Bansal et al., 2019). Training procedure. We use the standard training procedure used in previous KG embedding models (Balazevic et al., 2019; Dettmers et al., 2018) (Appendix for details). We follow an encoderdecoder framework: The encoder applies the proposed MAGNA model to compute the entity embeddings. The decoder then makes link prediction given the embeddings, and existing decoders in prior models can be applied. To show the power of MAGNA, we employ the DistMult decoder (Yang et al., 2015), a simple decoder without extra parameters. Evaluation. We use the standard split for the benchmarks, and the standard testing procedure of predicting tail (head) entity given the head (tail) entity and relation type. We exactly follow the evaluation used by all previous works, namely the Mean Reciprocal Rank (MRR), Mean Rank (MR), and hit rate at K (H@K). See Appendix for a detailed description of this standard setup. Results. MAGNA achieves new state-of-the-art in knowledge graph completion on all four metrics (Table 3). MAGNA compares favourably to both the most recent shallow embedding methods (QuatE), and deep embedding methods (SACN). Note that with the same decoder (DistMult), MAGNA using its own embeddings achieves drastic improvements over using the corresponding DistMult embeddings.

heading: 4.3 MAGNA MODEL ANALYSIS

text:
Here we present (1) the spectral analysis results, (2) effect of the hyper-parameters on MAGNA performance, and (3) attention distribution analysis to show the strengths of MAGNA.
Spectral Analysis: Why MAGNA works for node classification? We compute the eigenvalues of the graph Laplacian of the attention matrix A, λ̂gi , and compare to that of the diffused matrixA, λ g i . Figure 3 (a) shows the ratio λ̂gi /λ g i on the Cora dataset. Low eigenvalues corresponding to largescale structure in the graph are amplified (up to a factor of 8), while high eigenvalues corresponding to eigenvectors with noisy information are suppressed (Klicpera et al., 2019b). MAGNA Model Depth. Here we conduct experiments by varying the number of GCN, GAT and our MAGNA layers to be 3, 6, 12, 18 and 24 for node classification on Cora. Results in Figure 3 (b) show that both deep GCN and deep GAT (even with residual connection) suffer from degrading performance, due to the over-smoothing problem (Li et al., 2018; Wang et al., 2019a). In contrast, the MAGNA model achieves consistent best results even with 18 layers, making deep MAGNA model robust and expressive. Notice that GAT with 18 layers cannot out-perform MAGNA with 3 layers and K=6 hops, although they have the same receptive field.
Effect of K and α. Figures 3 (c) and (d) report the effect of hop number K and teleport probability α on model performance. We observe significant increase in performance when considering multihop neighbors information (K > 1). However, increasing the hop number K has a diminishing returns, forK ≥ 6. Moreover, we find that the optimalK is correlated with the largest node average shortest path distance (e.g., 5.27 for Cora). This provides a guideline for choosing the best K. We also observe that the accuracy drops significantly for larger α > 0.25. This is because small α increases the low-pass effect (Figure 3 (a)). However, α being too small results in the model only focusing on large-scale graph structure and ignores too much high-frequency information. Attention Distribution. Last we also analyze the learned attention scores of GAT and MAGNA.
We first define a discrepancy metric over the attention matrix A for node vi as ∆i = ‖A[i,:]−Ui‖ degree(vi)
(Shanthamallu et al., 2020), where Ui is the uniform distribution score for the node vi. ∆i gives a measure of how much the learnt attention deviates from an uninformative uniform distribution. Large ∆i indicates more meaningful attention scores. Fig. 4 shows the distribution of the discrepancy metric of the attention matrix of the 1st head w.r.t. the first layer of MAGNA and GAT. Observe that attention scores learned in MAGNA have much larger discrepancy. This shows that MAGNA is more powerful than GAT in distinguishing important and non-important nodes and assigns attention scores accordingly.

heading: 5 RELATED WORK

text:
Our proposed MAGNA belongs to the family of Graph Neural Network (GNN) models (Battaglia et al., 2018; Wu et al., 2020; Kipf & Welling, 2016; Hamilton et al., 2017), while taking advantage of graph attention and diffusion techniques. Graph Attention Neural Networks (GATs) generalize attention operation to graph data. GATs allow for assigning different importance to nodes of the same neighborhood at the feature aggregation step (Veličković et al., 2018). Based on such framework, different attention-based GNNs have been proposed, including GaAN (Zhang et al., 2018), AGNN (Thekumparampil et al., 2018), GeniePath (Liu et al., 2019b). However, these models only consider direct neighbors for each layer of feature aggregation, and suffer from over-smoothing when they go deep (Wang et al., 2019a).
Diffusion based Graph Neural Network. Recently Graph Diffusion Convolution (GDC) (Klicpera et al., 2019b;a) proposes to aggregate information from a larger (multi-hop) neighborhood at each layer, by sparsifying a generalized form of graph diffusion. This idea was also explored in (Liao et al., 2019; Luan et al., 2019; Xhonneux et al., 2019; Klicpera et al., 2019a) for multi-scale deep Graph Convolutional Networks. However, these methods do not incorporate attention mechanisms which proves to have a significant gain in model performance, and do not make use of edge embeddings (e.g., Knowledge graph) (Klicpera et al., 2019b). Our approach defines a novel multi-hop context-dependent self-attention GNN which resolves the over-smoothing issue of GAT architectures (Wang et al., 2019a). EdgeNets (Isufi et al., 2020) also extends attention mechanism for multi-hop information aggregation, but it needs more parameters to compute the attention scores of multi-hop neighbors. In contrast, our method infers the attention scores of multi-hop neighbors based on the one-hop neighbor attention scores via graph diffusion, and thus not only more parameter efficiency but show better spectral property.

heading: 6 CONCLUSION

text: We proposed Multi-hop Attention Graph Neural Network (MAGNA), which brings together benefits of graph attention and diffusion techniques in a single layer through attention diffusion, layer normalization and deep aggregation. MAGNA enables context-dependent attention between any pair of nodes in the graph in a single layer, enhances large-scale structural information, and learns more informative attention distribution. MAGNA improves over all state-of-the-art methods on the standard tasks of node classification and knowledge graph completion.