Gradient Origin Networks

1 INTRODUCTION . Observable data in nature has some parameters which are known , such as local coordinates , but also some unknown parameters such as how the data is related to other examples . Generative models , which learn a distribution over observables , are central to our understanding of patterns in nature and allow for efficient query of new unseen examples . Recently , deep generative models have received interest due to their ability to capture a broad set of features when modelling data distributions . As such , they offer direct applications such as synthesising high fidelity images ( Karras et al. , 2020 ) , super-resolution ( Dai et al. , 2019 ) , speech synthesis ( Li et al. , 2019 ) , and drug discovery ( Segler et al. , 2018 ) , as well as benefits for downstream tasks like semi-supervised learning ( Chen et al. , 2020 ) . A number of methods have been proposed such as Variational Autoencoders ( VAEs , Figure 1a ) , which learn to encode the data to a latent space that follows a normal distribution permitting sampling ( Kingma & Welling , 2014 ) . Generative Adversarial Networks ( GANs ) have two competing networks , one which generates data and another which discriminates from implausible results ( Goodfellow et al. , 2014 ) . Variational approaches that approximate the posterior using gradient descent ( Lipton & Tripathi , 2017 ) and short run MCMC ( Nijkamp et al. , 2020 ) respectively have been proposed , but to obtain a latent vector for a sample , they require iterative gradient updates . Autoregressive Models ( Van Den Oord et al. , 2016 ) decompose the data distribution as the product of conditional distributions and Normalizing Flows ( Rezende & Mohamed , 2015 ) chain together invertible functions ; both methods allow exact likelihood inference . Energy-Based Models ( EBMs ) map data points to ∗Authors contributed equally . energy values proportional to likelihood thereby permitting sampling through the use of Monte Carlo Markov Chains ( Du & Mordatch , 2019 ) . In general to support encoding , these approaches require separate encoding networks , are limited to invertible functions , or require multiple sampling steps . Implicit representation learning ( Park et al. , 2019 ; Tancik et al. , 2020 ) , where a network is trained on data parameterised continuously rather than in discrete grid form , has seen a surge of interest due to the small number of parameters , speed of convergence , and ability to model fine details . In particular , sinusoidal representation networks ( SIRENs ) ( Sitzmann et al. , 2020b ) achieve impressive results , modelling many signals with high precision , thanks to their use of periodic activations paired with carefully initialised MLPs . So far , however , these models have been limited to modelling single data samples , or use an additional hypernetwork or meta learning ( Sitzmann et al. , 2020a ) to estimate the weights of a simple implicit model , adding significant complexity . This paper proposes Gradient Origin Networks ( GONs ) , a new type of generative model ( Figure 1b ) that do not require encoders or hypernetworks . This is achieved by initialising latent points at the origin , then using the gradient of the log-likelihood of the data with respect to these points as the latent space . At inference , latent vectors can be obtained in a single step without requiring iteration . GONs are shown to have similar characteristics to convolutional autoencoders and variational autoencoders using approximately half the parameters , and can be applied to implicit representation networks ( such as SIRENs ) allowing a space of implicit functions to be learned with a simpler overall architecture . 2 PRELIMINARIES . We first introduce some background context that will be used to derive our proposed approach . 2.1 EMPIRICAL BAYES . The concept of empirical Bayes ( Robbins , 1956 ; Saremi & Hyvarinen , 2019 ) , for a random variable z ∼ pz and particular observation z0 ∼ pz0 , provides an estimator of z expressed purely in terms of p ( z0 ) that minimises the expected squared error . This estimator can be written as a conditional mean : ẑ ( z0 ) = ∫ zp ( z|z0 ) dz = ∫ z p ( z , z0 ) p ( z0 ) dz . ( 1 ) Of particular relevance is the case where z0 is a noisy observation of z with covariance Σ . In this case p ( z0 ) can be represented by marginalising out z : p ( z0 ) = ∫ 1 ( 2π ) d/2|det ( Σ ) |1/2 exp ( − ( z0 − z ) TΣ−1 ( z0 − z ) /2 ) p ( z ) dz . ( 2 ) Differentiating this with respect to z0 and multiplying both sides by Σ gives : Σ∇z0p ( z0 ) = ∫ ( z− z0 ) p ( z , z0 ) dz = ∫ zp ( z , z0 ) dz− z0p ( z0 ) . ( 3 ) After dividing through by p ( z0 ) and combining with Equation 1 we obtain a closed form estimator of z ( Miyasawa , 1961 ) written in terms of the score function∇ log p ( z0 ) ( Hyvärinen , 2005 ) : ẑ ( z0 ) = z0 + Σ∇z0 log p ( z0 ) . ( 4 ) This optimal procedure is achieved in what can be interpreted as a single gradient descent step , with no knowledge of the prior p ( z ) . By rearranging Equation 4 , a definition of ∇ log p ( z0 ) can be derived ; this can be used to train models that approximate the score function ( Song & Ermon , 2019 ) . 2.2 VARIATIONAL AUTOENCODERS . Variational Autoencoders ( VAEs ; Kingma & Welling 2014 ) are a probabilistic take on standard autoencoders that permit sampling . A latent-based generative model pθ ( x|z ) is defined with a normally distributed prior over the latent variables , pθ ( z ) = N ( z ; 0 , Id ) . pθ ( x|z ) is typically parameterised as a Bernoulli , Gaussian , multinomial distribution , or mixture of logits . In this case , the true posterior pθ ( z|x ) is intractable , so a secondary encoding network qφ ( z|x ) is used to approximate the true posterior ; the pair of networks thus resembles a traditional autoencoder . This allows VAEs to approximate pθ ( x ) by maximising the evidence lower bound ( ELBO ) , defined as : log pθ ( x ) ≥ LVAE = −DKL ( N ( qφ ( z|x ) ) ||N ( 0 , Id ) ) + Eqφ ( z|x ) [ log pθ ( x|z ) ] . ( 5 ) To optimise this lower bound with respect to θ and φ , gradients must be backpropagated through the stochastic process of generating samples from z′ ∼ qφ ( z|x ) . This is permitted by reparameterising z using the differentiable function z′ = µ ( z ) + σ ( z ) , where ∼ N ( 0 , Id ) and µ ( z ) and σ ( z ) 2 are the mean and variance respectively of a multivariate Gaussian distribution with diagonal covariance . 3 METHOD . Consider some dataset x ∼ pd of continuous or discrete signals x ∈ Rm , it is typical to assume that the data can be represented by low dimensional latent variables z ∈ Rk , which can be used by a generative neural network to reconstruct the data . These variables are often estimated through the use of a secondary encoding network that is trained concurrently with the generative network . An encoding network adds additional complexity ( and parameters ) to the model , it can be difficult to balance capacities of the two networks , and for complex hierarchical generative models designing a suitable architecture can be difficult . This has led some to instead approximate latent variables by performing gradient descent on the generative network ( Bojanowski et al. , 2018 ; Nijkamp et al. , 2020 ) . While this addresses the aforementioned problems , it significantly increases the run time of the inference process , introduces additional hyperparameters to tune , and convergence is not guaranteed . 3.1 GRADIENT ORIGIN NETWORKS . We propose a generative model that consists only of a decoding network , using empirical Bayes to approximate the posterior in a single step . That is , for some data point x and latent variable z ∼ pz , we wish to find an approximation of p ( z|x ) . Given some noisy observation z0 = z +N ( 0 , Id ) of z then empirical Bayes can be applied to approximate z . Specifically , since we wish to approximate z conditioned on x , we instead calculate ẑx , the least squares estimate of p ( z|x ) ( proof in Appendix A ) : ẑx ( z0 ) = z0 +∇z0 log p ( z0|x ) . ( 6 ) Using Bayes ’ rule , log p ( z0|x ) can be written as log p ( z0|x ) = log p ( x|z0 ) + log p ( z0 ) − log p ( x ) . Since log p ( x ) is a normalising constant that does not affect the gradient , we can rewrite Equation 6 in terms only of the decoding network and p ( z0 ) : ẑx ( z0 ) = z0 +∇z0 ( log p ( x|z0 ) + log p ( z0 ) ) . ( 7 ) It still remains , however , how to construct a noisy estimate of z0 with no knowledge of z . If we assume z follows a known distribution , then it is possible to develop reasonable estimates . For instance , if we assume p ( z ) = N ( z ; 0 , Id ) then we could sample from p ( z0 ) = N ( z0 ; 0 , 2Id ) however this could be far from the true distribution of p ( z0|z ) = N ( z0 ; z , Id ) . Instead we propose initialising z0 at the origin since this is the distribution ’ s mean . Initialising at a constant position decreases the input variation and thus simplifies the optimisation procedure . Naturally , how p ( x|z ) is modelled affects ẑx . While mean-field models result in ẑx that are linear functions of x , conditional autoregressive models , for instance , result in non-linear ẑx ; multiple gradient steps also induce non-linearity , however , we show that a single step works well on high dimensional data suggesting that linear encoders , which normally do not scale to high dimensional data are effective in this case . 3.2 AUTOENCODING WITH GONS . Before exploring GONs as generative models , we discuss the case where the prior p ( z ) is unknown ; such a model is referred to as an autoencoder . As such , the distribution p ( z0|z ) is also unknown thus it is again unclear how we can construct a noisy estimate of z . By training a model end-to-end where z0 is chosen as the origin , however , a prior is implicitly learned over z such that it is reachable from z0 . Although p ( z ) is unknown , we do not wish to impose a prior on z0 ; the term which enforces this is in Equation 7 is log p ( z0 ) , so we can safely ignore this term and simply maximise the likelihood of the data given z0 . Our estimator of z can therefore be defined simply as ẑx ( z0 ) = z0 +∇z0 log p ( x|z0 ) , which can otherwise be interpreted as a single gradient descent step on the conditional log-likelihood of the data . From this estimate , the data can be reconstructed by passing ẑx through the decoder to parameterise p ( x|ẑx ) . This procedure can be viewed more explicitly when using a neural network F : Rk → Rm to output the mean of p ( x|ẑx ) parameterised by a normal distribution ; in this case the loss function is defined in terms of mean squared error loss LMSE : Gx = LMSE ( x , F ( −∇z0LMSE ( x , F ( z0 ) ) ) ) . ( 8 ) The gradient computation thereby plays a similar role to an encoder , while F can be viewed as a decoder , with the outer loss term determining the overall reconstruction quality . Using a single network to perform both roles has the advantage of simplifying the overall architecture , avoiding the need to balance networks , and avoiding bottlenecks ; this is demonstrated in Figure 1b which provides a visualisation of the GON process .

Summary: This paper introduces a "new" inference method for autoencoder-type models, where the encoder is taken as a gradient of the decoder with respect to a zero-initialized latent variable. The method is evaluated for both a deterministic autoencoder and a VAE on toy image data (cifar10 being the most complex of them) and applied to convolutional decoder and to SIREN-type implicit representation networks. This is, for all intents and purposes, a single step iterative inference setup. In its VAE variant it is extremely similar to old-school iterative inference, albeit with a single gradient step.