Neighborhood Reconstructing Autoencoders

1 Introduction . Autoencoders are widely used to identify , and to generate samples from , the underlying lowdimensional manifold structure of a given data distribution [ 14 , 1 ] . It has been widely observed that vanilla autoencoders quite often produce manifolds that ( i ) are highly sensitive to noisy training data ( see Figure 1 ( a ) ) , or ( ii ) have the wrong local connectivity and geometry ( see Figure 1 ( b ) ) , significantly impairing their performance . Regularization techniques have had some success in mitigating the former , e.g. , the Denoising Autoencoder [ 21 ] , which uses deliberately corrupted inputs to train the autoencoder , typically learns manifolds that are robust to noise , but not always with the correct local geometry . Recently , autoencoder regularization methods that use neighborhood graphs have had some success in addressing the incorrect connectivity issue [ 16 , 18 , 10 , 7 ] . Notwithstanding the additional computational overhead of constructing local neighborhood graphs , the local geometric information obtained from these graphs can significantly reduce errors in the local geometry of the learned manifold , and make learning more well-behaved and robust . The underlying premise behind these methods is that since the local geometry and topology of the data is captured in the latent space distribution , which is determined entirely by the encoder , regularizing only the encoder should be sufficient ; little if any consideration needs to be given to the decoder . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . The flaw with this premise is that although the encoder learns the correct latent space representation of the manifold data , the decoder is still susceptible to overfitting of the type shown in Figure 1 ( a ) . These existing methods moreover rely on computation-intensive preprocessing steps like manifold learning [ 10 ] , linear coefficients computation [ 7 ] , or computing topological features using persistent homology at each training iteration [ 16 , 2 ] , each of whose computational requirements can grow significantly with problem dimension and scale . The main contribution of this paper is a new graph-based autoencoder training method that addresses both the overfitting and connectivity issues illustrated in Figure 1 ( a ) - ( b ) . Like current methods , our method also employs local graphs that capture the local geometry of the data distribution . The key idea behind our method , which we call the Neighborhood Reconstructing Autoencoder ( NRAE ) , is to employ a local quadratic ( and in some cases linear ) approximation of the decoder function to formulate a new neighborhood reconstruction loss in lieu of the point reconstruction loss typically used for autoencoder training . This idea leads to learning the correct geometry and reducing noise sensitivity , significantly improving the robustness of autoencoder training . To make things more explicit , let gφ : Rn → Rm be the encoder ( parametrized by φ ) and fθ : Rm → Rn be the decoder ( parametrized by θ ) . Whereas vanilla autoencoders are trained to minimize the sum of the point reconstruction errors ∑ i ‖xi − fθ ( gφ ( xi ) ) ‖2 , NRAE minimizes a reconstruction error of the form ∑ i ∑ x∈N ( xi ) ∥∥∥x− f̃θ ( gφ ( x ) ; gφ ( xi ) ) ∥∥∥2 , ( 1 ) where N ( xi ) is the set of neighborhood points of xi ( including xi ) and f̃θ ( · ; gφ ( xi ) ) is a local quadratic ( or linear ) approximation of fθ about gφ ( xi ) . The vanilla autoencoder is obtained by setting N ( xi ) = { xi } for all i . The key idea here is to locally approximate the decoder only , and to exploit the local geometric information extracted from the decoded manifold represented by the image of fθ . Like other neighborhood graph-based methods , NRAE also learns the correct local geometry of the decoded manifold . At the same time , the local quadratic ( or linear ) approximation of fθ considerably reduces any overfitting to noisy training data or sensitivity to outliers , while maintaining computational efficiency – rather than the entire Jacobian or Hessian of fθ , only the more easily computed Jacobianvector and Hessian-vector products are needed for the approximation . Compared to existing graph-based autoencoder regularization methods , NRAE is easy to implement , computationally efficient , and scalable , requiring only a single prior construction of the graph without additional pre-processing steps . Experiments with both synthetic and image data ( MNIST , FashionMNIST , KMNIST , Omniglot , SVHN , CIFAR10 , CIFAR100 , CelebA ) confirm that overall our method better learns the correct geometry of manifolds , showing improved generalization performance vis-á-vis existing graph-based and other autoencoder regularization methods . 2 Neighborhood Reconstructing Autoencoder . In this section , we first provide a high-level mathematical description of the Neighborhood Reconstructing Autoencoder ( NRAE ) , followed by algorithmic details and a discussion of the NRAE ’ s properties and behavior . Throughout we consider a deterministic autoencoder with an encoder function gφ : Rn → Rm and decoder function fθ : Rm → Rn ( m ≤ n ) , with their composition denoted by Fθ , φ : = fθ ◦ gφ . We use the notation D : = { xi ∈ Rn } Mi=1 to denote the set of observed data points . 2.1 Mathematical Description . In what follows we use the notation N ( x ) to denote the set of neighborhood points of x , with x included in N ( x ) . We begin with the following definition : Definition 1 Let F̃θ , φ ( · ; x ) : = f̃θ ( gφ ( · ) ; gφ ( x ) ) , where f̃θ ( · ; z ) is a local quadratic ( or in some cases linear ) approximation of fθ at z = ( z1 , z2 , ... zm ) : f̃θ ( z ′ ; z ) : = fθ ( z ) + m∑ i=1 ∂fθ ∂zi ( z ) dzi + m∑ i , j=1 1 2 ∂2fθ ∂zi∂zj ( z ) dzidzj , ( 2 ) where dz = z′ − z. F̃θ , φ ( N ( x ) ; x ) is said to be a neighborhood reconstruction of N ( x ) . If instead of a quadratic approximation we use the linear approximation of fθ , the image of F̃θ , φ ( · ; x ) is the tangent space of the decoded manifold at Fθ , φ ( x ) , and the neighborhood reconstruction of N ( x ) is a subset of the tangent space ; the neighborhood reconstruction in this case contains first-order local geometric information about the decoded manifold . The key idea behind Definition 1 is that we locally approximate the decoder , and not the encoder , to extract and exploit local geometric information on the decoded manifold , which is captured in the image of F̃θ , φ ( · ; x ) ( i.e. , the local approximation of the decoded manifold ) . Figure 2 illustrates an example where the autoencoder reconstructs the points almost perfectly , but the neighborhood reconstruction ofN ( x ) , whose elements lie in the tangent space ( here we use the linear approximation of fθ ) is considerably different from N ( x ) . Given that the neighborhood reconstruction of N ( x ) reflects the local geometry of the decoded manifold , minimizing a loss function that measures the difference between N ( x ) and its image F̃θ , φ ( N ( x ) ; x ) is one means of training an autoencoder to preserve the local geometry of the original data distribution . With that goal in mind , we formulate a neighborhood reconstruction loss L as follows : L ( θ , φ ; D ) = 1 |D| ∑ x∈D 1 |N ( x ) | ∑ x′∈N ( x ) K ( x′ , x ) · ‖x′ − F̃θ , φ ( x′ ; x ) ‖2 , ( 3 ) whereK ( x′ , x ) is a positive symmetric kernel function that determines the weight for each x′ ∈ N ( x ) . Figure 3 illustrates how the neighborhood reconstruction loss can differentiate among the quality of the learned manifolds whose point reconstruction losses are all the same ( close to zero ) : Case 3 has the smallest neighborhood reconstruction loss compared to Case 1 ( wrong local geometry ) and Case 2 ( overfitting ) . NRAE converges to the vanilla AE – that is , the neighborhood reconstruction loss recovers the point reconstruction loss – if one of the following conditions is met : ( i ) N ( x ) = { x } , ( ii ) K ( x′ , x ) = δ ( x′ , x ) , or ( iii ) fθ is linear . It is reasonable to ask whether there are any advantages to using an approximation for both the encoder and decoder , i.e. , to use a local quadratic ( or linear ) approximation for the composition map Fθ , φ rather than just the decoder . As verified below in Section 4.3 , using a quadratic approximation for both the encoder and decoder results in minimal to no performance improvement ( the results for this Extended NRAE ( E-NRAE ) case are nearly identical to those obtained for NRAE ) , but the computational requirements increase substantially . As intuition suggests , applying a quadratic approximation for the decoder is sufficient in evaluating the neighborhood reconstruction loss . 2.2 Algorithmic Details . Graph construction . The problem of inferring the geometric structure of a data distribution is typically posed as a graph construction problem [ 4 ] . We use one of the simplest graph construction methods , the k-NN graph with the Euclidean distance metric . The robustness of our algorithm with respect to the choice of k is tested in the Supplementary Material . Kernel design . We choose the following simple kernel K ( x′ , x ) = λ+ ( 1− λ ) δ ( x′ , x ) , ( 4 ) where 0 ≤ λ < 1 and δ ( x′ , x ) = 1 if x′ = x and zero otherwise . This assigns the weight 1 for the center x ∈ N ( x ) and the weight λ for the remaining neighborhood points . Batch sampling . To estimate the gradient of the proposed loss function , we use batch sampling for both summations over D and N ( x ) . Given a batch B ⊂ D , we again sample a batch Bx from N ( x ) . We empirically find that forcing each batch Bx to include x improves convergence . In this paper , we set Bx = { x , xn } where xn is uniformly sampled from N ( x ) − { x } . 3 Related Work : Regularization of Autoencoders . In this section , we review some standard autoencoder regularization techniques and their relation to NRAE : i ) regularizing latent space distributions , ii ) Jacobian regularization , and iii ) regularization using neighborhood graphs . Regularizing latent space distributions . One popular regularization strategy is to enforce the latent space distribution ( an approximate posterior ) to be close to some user-specified prior distribution ; some examples include the Variational Autoencoder ( VAE ) [ 12 ] , Adversarial Autoencoder ( AAE ) [ 15 ] , and Wasserstein Autoencoder ( WAE ) [ 19 ] . The Gaussian distribution is a popular choice for the prior , but this choice often leads to over-regularization . Several works improve the VAE by learning more complex priors , or an optimal prior in terms of maximizing the training objective function of the VAE [ 8 , 6 , 20 , 13 ] . These techniques are developed to make latent space distributions close to some easy-to-sample prior distribution , but are not designed to resolve the two main issues ( local geometry and overfitting ) addressed in this paper . Jacobian regularization . The Contractive Autoencoder ( CAE ) attempts to enhance robustness of representation by penalizing the Jacobian norm of the encoder function [ 17 ] . However , like other encoder-only regularization methods , it often learns a manifold that overfits the data . More recently [ 5 ] regularizes the Jacobian of the decoder function to learn a flat manifold . The trained decoder function may produce a smooth manifold but with the wrong local geometry . The main difference between these methods and NRAE is the use of a neighborhood graph containing local connectivity information . Regularization using neighborhood graphs . Several recent works have developed graph-based autoencoder regularization methods that , at least implicitly , try to learn manifolds with the right local geometry : Generalized Autoencoder ( GAE ) [ 22 ] , Topological Autoencoder ( TopoAE ) [ 16 ] , Witness Autoencoder ( W-AE ) [ 18 ] , Geometry Regularized Autoencoder ( GRAE ) [ 10 ] , and StructurePreserving Variational Autoencoder ( Sp-VAE ) [ 7 ] . These works focus almost entirely on the encoder when regularizing the latent space distribution ( with the exception of GAE [ 22 ] ) ; using the regularization term that does not depend on the decoder lead to noise sensitivity and overfitting issues as described in detail above . These methods are also computationally intensive , especially for large scale high-dimensional problems , as a result of expensive pre-processing [ 10 , 7 ] or computation in training [ 16 , 18 ] . In contrast , by focusing more on the geometry of the decoded manifold , NRAE is able to learn smooth manifolds with the correct local geometry , simpler to implement , and more scalable since it requires only a single prior construction of the graph without other additional pre-processing steps . Indeed , with the exception of Sp-VAE [ 7 ] , all other approaches test their algorithms with two-dimensional latent spaces , while we test NRAE for latent spaces with up to 128 dimensions . For general regularization methods that focus on latent space distributions , we surmise that using our neighborhood reconstruction loss in lieu of the point reconstruction loss may lead to several performance improvements ; this is verified in Section 4.3 , where we combine our neighborhood reconstruction loss with some existing encoder regularization methods and compare their performance against the original methods .

Summary: Rather than reconstructing a given input $x$, the first- or second- order* Taylor approximations (centered at $x$) of the decoder are optimized to accurately reconstruct some neighborhood $\mathcal{N}(x)$ of $x$. More formally, the loss for a single $x, x'$ pair, where $x' \in \mathcal{N}(x)$ is written as $K(x, x') \cdot \left \Vert x - \tilde{f}(g(x')) \right \Vert^2$ where $g$ is the encoder and $\tilde{f}(x')$ is a Taylor approximation of the decoder $f$, centered at $x$. The authors evaluate the method across a wide array of empirical experiments and show that it produces superior autoencoding performance by many measures.  *There are two variants of the algorithm -- one that uses a first-order Taylor approximation, and one that uses a second-order Taylor approximation.