Neighborhood Reconstructing Autoencoders
Year: 2021

Abstract:
Vanilla autoencoders often produce manifolds that overfit to noisy training data, or have the wrong local connectivity and geometry. Autoencoder regularization techniques, e.g., the denoising autoencoder, have had some success in reducing overfitting, whereas recent graph-based methods that exploit local connectivity information provided by neighborhood graphs have had some success in mitigating local connectivity errors. Neither of these two approaches satisfactorily reduce both overfitting and connectivity errors; moreover, graph-based methods typically involve considerable preprocessing and tuning. To simultaneously address the two issues of overfitting and local connectivity, we propose a new graph-based autoencoder, the Neighborhood Reconstructing Autoencoder (NRAE). Unlike existing graph-based methods that attempt to encode the training data to some prescribed latent space distribution – one consequence being that only the encoder is the object of the regularization – NRAE merges local connectivity information contained in the neighborhood graphs with local quadratic approximations of the decoder function to formulate a new neighborhood reconstruction loss. Compared to existing graph-based methods, our new loss function is simple and easy to implement, and the resulting algorithm is scalable and computationally efficient; the only required preprocessing step is the construction of the neighborhood graph. Extensive experiments with standard datasets demonstrate that, compared to existing methods, NRAE improves both overfitting and local connectivity in the learned manifold, in some cases by significant margins. Code for NRAE is available at https://github.com/Gabe-YHLee/NRAE-public.

heading: 1 Introduction
text: 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.

heading: 2 Neighborhood Reconstructing Autoencoder
text: 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.

heading: 2.1 Mathematical Description
text: 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.

heading: 2.2 Algorithmic Details
text: 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}.

heading: 3 Related Work: Regularization of Autoencoders
text: 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.

heading: 4 Experiments
text: In this section, through extensive experiments with both synthetic and real-world image data, we compare NRAE with a range of existing regularization methods: the Variational Autoencoder (VAE) [12], Wasserstein Autoencoder (WAE) [19], Denoising Autoencoder (DAE) [21], Contractive Autoencoder (CAE) [17], Geometry Regularized Autoencoder (GRAE) [10], and Structure-Preserving Variational Autoencoder (Sp-VAE) [7]. For comparison with Sp-VAE, we augment the regularization term introduced in [7] to a vanilla autoencoder rather than the variational autoencoder – we refer to this autoencoder as the Structure-Preserving Autoencoder (SPAE) – since a direct comparison with the variational autoencoder is already made, and our intent is to examine the effects of the structure-preserving regularization term.
For NRAE, we use both the local quadratic and linear approximations, respectively denoted NRAE-Q and NRAE-L. Our focus is on comparing the manifold smoothing property, geometry preserving property, and generalization ability of NRAE against baseline methods. We refer the reader to the Supplementary Material for a description of the network architectures used in the experiments, together with implementation details including the hyperparameter tuning strategy.

heading: 4.1 Manifold Smoothing Property
text: Smoothness of learned manifolds. We first consider a one-dimensional manifold embedded in R2, {x, sin(x)|x ∈ [−π, π]}. Given 50 randomly sampled data points, we add isotropic Gaussian noise with a standard deviation of 0.2. We train the NRAE and baseline AEs with one-dimensional latent spaces. Figure 4 and Table 1 show that NRAE-L and NRAE-Q successfully re-produce smooth manifolds compared to other baseline methods.
Table 1: The mean-squared reconstruction errors of 10,000 clean test data. The averages and standard deviations are computed over ten times run with different noises, and multiplied by 100 (the highest and lowest scores are ignored). The best and second-best results are colored red and blue, respectively.
AE VAE WAE DAE CAE GRAE SPAE NRAE-L NRAE-Q
1.90 ± 0.23 1.45 ± 0.39 2.02 ± 0.76 1.09 ± 0.21 1.33 ± 0.30 1.56 ± 0.43 1.28 ± 0.30 0.29 ± 0.06 0.30 ± 0.07
Second, we rotate 100 MNIST images of the digit eight 100 times by 3.6 degrees, obtaining a set of 10,000 training data. We train the NRAE and baseline AEs with two-dimensional latent spaces, normalize each latent space, and generate rotating images by sampling regular grids of the latent spaces. NRAE-L and NRAE-Q generate smoothly varying images compared to other baselines (Figure 5). The DAE does not generate smoothly varying samples because the noise statistics (Gaussian noise) used in training is different from the statistical noise in the rotating image data.
Finally, to confirm the denoising effect of NRAE on common image data, we train the NRAE and baseline AEs on MNIST and CIFAR10 data corrupted with various levels of Gaussian noise (we use sufficiently large epochs for convergence). We then compare denoising performance measured by the Peak Signal-to-Noise Ratio (PSNR). As shown in Table 2, NRAE-L and NRAE-Q outperform other baselines including DAE which performed the third best (the noise statistics used in training is the same as the added noise statistics).
Curvature of learned manifolds. We train the NRAE and baseline AEs on MNIST images with two-dimensional latent spaces, normalize each latent space, then visualize the scalar curvature field (i.e., twice the Gaussian curvature [9]) (Figure 6). As shown in Figure 6, NRAE-L and NRAE-Q both learn flatter manifolds compared to our baseline AEs.

heading: 4.2 Geometry Preserving Property.
text: Swiss roll. Consider a one-dimensional Swiss roll r(θ) = 0.1+ 0.9 · θ/(2π) for θ ∈ [0, 2π]. Given 30 randomly sampled data points, we add isotropic Gaussian noise with a standard deviation of 0.01. We then train the NRAE and baseline AEs with one-dimensional latent spaces. Figure 7 shows the reconstruction results for 1000 test data points. Only GRAE, SPAE, and NRAE-Q successfully learn manifolds with the right local geometry. However, unlike NRAE-Q, the regularization terms in GRAE and SPAE have the effect of increasing reconstruction errors (observe the red rectangles in Figure 7). NRAE-L fails to reconstruct points around the point xr circled in red: the nearest neighbor to xr in the training set, marked by the blue circle, lies close to the tangent space of the decoded manifold at xr, i.e., the neighborhood reconstruction loss is small.
Rotated/Shifted MNIST. First, as shown in the first row of the two figures in Figure 8, we generate two sets of 20 training data: the rotated MNIST images of the digit 3 and shifted MNIST images of the digit 7. We train NRAE and baseline AEs with one-dimensional latent spaces, encode the training data to the latent spaces, sort the encoded values in ascending order, and decode them to generate the images (Figure 8). Only SPAE, NRAE-L, and NRAE-Q successfully learn manifolds of the rotated/shifted images with the correct geometry.
Second, we rotate an MNIST image of the digit six 300 times by 1.2 degrees and obtain 300 training data; 40 of these are visualized in the first row of Figure 9. We further generate 1000 test data in a similar manner. We then train the AE, DAE, CAE, and NRAE with one-dimensional circular latent spaces by adding one layer to the end of the encoder corresponding to z → z/‖z‖, and visualize which data are encoded to which point in the latent spaces via the cyclic color map. Only NRAE-Q is able to learn a manifold with the correct local geometry.

heading: 4.3 Generalization
text: Test data reconstruction. With various large-scale standard benchmark image data, we show that NRAE generalizes better to the test data compared to other baseline methods. We use the test reconstruction Mean Square Error (MSE) for quantitative comparison.
First, we train fully-connected neural networks with the MNIST, FMNIST, and KMNIST data by varying the number of training data from 1000, 2000 . . . , 10, 000, and compare the test reconstruction MSEs. The numbers for the validation and test data are fixed at 10, 000 and 50, 000, respectively. Figure 10 shows the MSEs as a function of the number of training data. The MSEs decrease as the number of training data increases, while those of the NRAE are lower than most of the other baselines regardless of the number of training data.
Second, we train convolutional neural networks and compare the test reconstruction MSEs (Table 3). We conduct two experiments i) with the entire public data denoted by L (large) and ii) with the subset
of data denoted by S (small). The number of selected subsets are around 20 ∼ 30% of the entire data. NRAE-L and NRAE-Q produce lower reconstruction MSEs than most of the other baselines, meaning that the manifolds learned with our methods better generalize to the test data. The CAE experiment on CELEBA data is excluded because of its high computational cost for high-dimensional data. Other measures such as the Frechet-Inception Distance (FID) scores [11] and the Evidence Lower Bound (ELBO) are reported in the Supplementary Material.
Compatibility with other regularization methods. In this section, we show that the NRAE loss function can be used together with other existing latent space distribution regularization methods to improve their generalization performance. We train convolutional neural networks with the MNIST and CIFAR10 data (small). As shown in Table 4, the NRAE loss improves WAE, CAE, and SPAE by significant margins. For VAE, there is only little improvement, primarily because the VAE regularization term is too dominant in training. Interestingly, some cases of WAE, CAE, and SPAE combined with the NRAE loss show even better generalization performance compared to NRAE alone. These results imply that a proper combination of the neighborhood reconstruction loss and existing regularization methods can lead to further improvements in performance.
Extended NRAE. Instead of using a quadratic (or linear) approximation of fθ, we can approximate the composition function Fθ,φ = fθ ◦ gφ to define the neighborhood reconstruction loss. While NRAE only requires the computation of the Jacobian-vector product or Hessian-vectors product of fφ, E-NRAE requires the computation of these quantities for Fθ,φ. We train convolutional neural networks on the MNIST and CIFAR10 data (small) and compare their generalization abilities and computational requirements. As shown in Table 5, the extended versions show comparable generalization performance to the original version, yet take a longer per-epoch runtime (100 batch size and 10,000 training data).

heading: 5 Conclusion
text: This paper has proposed a new graph-based autoencoder, the Neighborhood Reconstructing Autoencoder (NRAE), that is capable of learning accurate manifolds that are robust to noisy training data and have the correct local connectivity and geometry, while being easy to implement, scalable, and computationally efficient. Neighborhood graphs that capture the local geometry of the data distribution are combined with local quadratic (or linear) approximations of the decoder function to formulate a new neighborhood reconstruction loss, which turns out to be a generalization of the original point reconstruction loss. Through extensive experiments with both synthetic and standard image datasets, we have demonstrated the manifold smoothing property, geometry preserving property, and the generalization performance advantages – in some cases by significant margins – of our method.
Further, we have empirically verified that (i) a proper combination of the neighborhood reconstruction loss and existing regularization terms that focus on the latent space distributions can lead to further improvements in performance and (ii) in neighborhood reconstruction loss, approximating the decoder only is sufficient for learning the correct manifold while being computationally more efficient than approximating both the encoder and decoder.
Our algorithm can be further enhanced in a number of different ways. First, the current implementation of NRAE uses the k-NN graph construction with Euclidean distance metric and a simple kernel function that outputs binary values. These choices are made for simplicity, and clearly sub-optimal. Other combinations of the graph construction method (e.g., the persistent homology [24]), distance metric (e.g., using deep metric learning), and kernel function (e.g., the Gaussian kernel) are worth exploring. Second, NRAE can be extended to a stochastic model in a number of ways, e.g., using the probabilistic autoencoder [3] that uses the latent variable model, or the energy-based approach as done in [23].
Acknowledgments and Disclosure of Funding
This work was supported by the NAVER LABS AMBIDEX Project, NRF-2016R1A5A1938472, SNUIAMD, SNU BK21+ Program in Mechanical Engineering, and the SNU Institute for Engineering Research.

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

This paper proposes a new graph-based autoencoder dubbed ‘Neighborhood Reconstructing Autoencoder’ (NRAE) for learning smoother manifolds. The goal is to (1) learn an AE that is robust to data noise (2) learn a manifold with the correct geometry. The authors claim the focus of regularization to induce a smooth representation should, unlike the previous work, be on the decoder and not the encoder. The proposed model minimizes an objective that encourages neighbouring xs (based on a predefined graph that is constructed using some kernel K(x,x’)) to also be close by in latent space by having a small ‘neighbourhood reconstruction loss’. The authors finally perform an extensive set of experiments on a variety of datasets which demonstrates that the proposed model is able to learn a more smooth manifold compared to other AE baselines, as well as having a better generalization error. 

The paper goes into detail about two major obstacles that frequent Autoencoders, namely overfitting to training data, and learning the correct local geometry of the data. They then claim that all existing methods struggle with one or the other, or they are so computationally taxing that they aren't scalable for most problems. The paper's main contribution is that they provide a new training method/loss function that uses a local 2nd order approximation of the decoder function. The paper claims that the decoder is more important than the encoder in correctly learning the local geometry of the data. They then go into the algorithm's details and then explore the related work of regularization of autoencoders. Finally, they run experiments to demonstrate the algorithm's properties as well as show that NRAE is compatible with other methods (VAE, WAE, CAE, and SPAE) and often gets better results combining methods.

An auto-encoding (AE) loss that reconstructs data according to a local approximation of the autoencoding function. In particular, for each $x_i$ in the dataset, the model reconstructs data around a neighborhood of $x_i$ using a low-order Taylor expansion of the encoder-decoder composition around $x_i$. The loss to each point in the neighborhood is assigned a weight according to a simple kernel function. The author demonstrates empirical advantage compared to many other AE losses in extensive evaluations.