High-dimensional Asymptotics of Denoising Autoencoders

Summary

This paper addresses the problem of denoising data from a Gaussian mixture using a two-layer non-linear autoencoder with tied weights and a skip connection. The authors consider the high-dimensional limit where the number of training samples and the input dimension jointly tend to infinity while the number of hidden units remains bounded. They provide closed-form expressions for the denoising mean-squared test error. Building on this result, they quantitatively characterize the advantage of the considered architecture over the autoencoder without the skip connection that relates closely to principal component analysis. They further show that their results accurately capture the learning curves on a range of real data sets.

Main contributions

The present work considers the problem of denoising data sampled from a Gaussian mixture by learning a two-layer DAE with a skip connection and tied weights via empirical risk minimization. Throughout the manuscript, the high-dimensional limit is considered where the number of training samples n and the dimension d are large (n, d → ∞) while remaining comparable, i.e. α ≡ n/d = Θ(1). The main contributions are:

• Leveraging the replica method, sharp, closed-form formulae for the mean squared denoising test error (MSE) for DAEs are provided, as a function of the sample complexity α and the problem parameters. A sharp characterization for other learning metrics including the weights norms, skip connection strength, and cosine similarity between the weights and the cluster means is also provided. These formulae encompass as a corollary the case of RAEs. These formulae also describe quantitatively rather well the denoising MSE for real data sets, including MNIST and FashionMNIST.

• It is found that PCA denoising (namely denoising by projecting the noisy data along the principal component of the training samples) is widely sub-optimal compared to the DAE, leading to a MSE superior by a difference of Θ(d), thereby establishing that DAEs do not simply learn to perform PCA.

• Building on the formulae, the role of each component of the DAE architecture (skip connection and the bottleneck network) in its overall performance is quantified. The two components are found to have complementary effects in the denoising process –namely preserving the data nuances and removing the noise– and the training of the DAE is discussed as resulting from a tradeoff between these effects.

Setting

The problem of denoising data x ∈ Rd corrupted by Gaussian white noise of variance ∆ is considered, with the noisy data point defined as x̃ = √(1 − ∆)x + √∆ξ, where ξ ∼ N (0, Id) is the additive noise. The clean data x is assumed to be drawn from a Gaussian mixture distribution P with K clusters, x ∼ Σ k=1 K ρ k N (µ k, Σ k). The k−th cluster is centered around µ k ∈ Rd, has covariance Σ k, and relative weight ρ k.

The DAE model is a two-layer DAE with tied weights and a trainable skip-connection, defined as f b,w(x̃) = b × x̃ + (w>/√d) σ(w x̃/√d). The DAE is parametrized by the scalar skip connection strength b ∈ R and the weights w ∈ R p×d, with p the width of the DAE hidden layer. The normalization of the weight w by √d ensures for high dimensional settings d >> 1 that the argument of the non-linearity σ(·) stays Θ(1). The case with p << d is focused on. Two other simple architectures are also considered: the bottleneck network component u v(x̃) = (v>/√d) σ(v x̃/√d), and the rescaling component r c(x̃) = c × x̃. Note that the full DAE f b,w = r b + u w.

To train the DAE, a training set D = x̃ µ, x µ µ=1 n is assumed available, with n clean samples x µ drawn i.i.d from P and the corresponding noisy samples x̃ µ = x µ + ξ µ. The DAE is trained by minimizing the empirical risk R̂(b, w) = Σ µ=1 n x µ − f b,w(x̃ µ)2 + g(w), where g: R p×d → R+ is an arbitrary convex regularizing function. The minimizers are denoted b̂, ŵ and the corresponding trained DAE fˆ ≡ f b̂,ŵ. The components (3) are also trained independently via empirical risk minimization.

The performance of the DAE is quantified by its reconstruction (denoising) test MSE, defined as mse fˆ ≡ E D E x∼P E ξ∼N (0,Id) x − f b̂,ŵ(√(1 − ∆)x + √∆ξ)2. Another question of interest is how much the DAE manages to learn the structure of the data distribution, as described by the cluster means µ k. This is measured by the cosine similarity matrix θ ∈ R p×K, where θ ik ≡ E D (ŵ i T µ k)/(ŵ iµ k).

The analysis is performed in the high-dimensional limit where the input dimension d and number of training samples n jointly tend to infinity, while their ratio α = n/d stays Θ(1). The hidden layer width p, the noise level ∆, the number of clusters K and the norm of the cluster means µ k are also assumed to remain Θ(1).

Asymptotic formulae for DAEs

The main result is the closed-form asymptotic formulae for the learning metrics mse fˆ and θ for a DAE learnt with the empirical loss, derived using the replica method in its replica-symmetric formulation. The result relies on two assumptions: Assumption II.1 states that the covariances Σ k admit a common set of eigenvectors e i i=1 d, and the eigenvalues and the projection of the cluster means on the eigenvectors are assumed to admit a well-defined joint distribution ν as d → ∞. Assumption II.2 states that g(·) is a `2 regularizer with strength λ, i.e. g(·) = λ/2·2 F.

Result II.3 provides the denoising test MSE mse fˆ as an expression involving the summary statistics q, q k, m k, which are defined as q = lim d→∞ E D (ŵŵ T)/d, q k = lim d→∞ E D (ŵΣ k ŵ T)/d, m k = lim d→∞ E D (ŵµ k)/√d. The learnt skip connection strength b̂ is given by a closed-form expression. The cosine similarity θ admits a compact formula. The summary statistics q, q k, m k can be determined as solutions of a system of equations (13), which involves the solutions of a finite-dimensional optimization problem (14).

Result II.3 encompasses as special cases the asymptotic characterization of the components r̂, û. Corollary II.4 states that the test MSE of r̂ is given by mse r̂ = mse◦, and the learnt value of its single parameter ĉ is given by the same expression as b̂. The test MSE, cosine similarity and summary statistics of the bottleneck network û follow from Result II.3 by setting b̂ = 0. Corollary II.5 states that in the noiseless limit ∆ = 0, the denoising task reduces to a reconstruction task, and Result II.3 also includes RAEs as a special case. The MSE, cosine similarity and summary statistics for an RAE follow from Result II.3 by setting x = 0 in (14), removing the first term in the brackets in the equation of V̂ (13) and taking the limit ∆, q̂, b̂ → 0.

Equations (10) and (12) of Result II.3 thus characterize the statistics of the learnt parameters b̂, ŵ of the trained DAE. These summary statistics are sufficient to fully characterize the learning metrics (5) and (6) via equations (8) and (11). The high-dimensional optimization (4) and the high-dimensional average over the train set D involved in the definition of the metrics are reduced to a simpler system of equations over 4 + 6K variables (13) which can be solved numerically. All the summary statistics involved in (13) are finite-dimensional as d → ∞, and therefore Result II.3 is a fully asymptotic characterization.

Examples of applications

The first example is a synthetic binary Gaussian mixture with K = 2, µ 1 = −µ 2, Σ 1,2 = 0.09 × Id, ρ 1,2 = 1/2, using a DAE with σ = tanh and p = 1. The MSE mse fˆ evaluated from the solutions of the self-consistent equations (13) is plotted and compared to numerical simulations corresponding to training the DAE with the Pytorch implementation of the Adam optimizer, for sample complexity α = 1 and `2 regularization (weight decay) λ = 0.1. The agreement between the theory and simulation is compelling. A particularly striking observation is that due to the non-convexity of the loss (4), there is a priori no guarantee that an Adam-optimized DAE should find a global minimum, as described by the Result II.3, rather than a local minimum. The compelling agreement between theory and simulations suggests that the loss landscape of DAEs trained with the loss (4) for the data model (1) should in some way be benign.

The second example uses real data sets, FashionMNIST (from which only boots and shoes were kept) and MNIST (1s and 7s). For each data set, samples sharing the same label were considered to belong to the same cluster. The mean and covariance thereof were estimated numerically, and combined with Result II.3. The resulting denoising MSE predictions mse fˆ agree very well with numerical simulations of DAEs optimized over the real data sets using the Pytorch implementation of Adam. The observation that the MSEs of real data sets are to such degree of accuracy captured by the equivalent Gaussian mixture strongly hints at the presence of Gaussian universality.

The role and importance of the skip connection

Result II.3 for the full DAE fˆ and Corollary II.4 for its components r̂, û allow to disentangle the contribution of each part, and thus to pinpoint their respective roles in the DAE architecture.

For noise levels ∆ below a certain threshold, the full DAE fˆ yields better MSE than the learnt rescaling r̂. The improvement is more sizeable at intermediate noise levels ∆, and is observed for a growing region of ∆ as the sample complexity α increases. This lower MSE further translates into visible qualitative changes in the result of denoising. The full DAE fˆ yields denoised images with sensibly higher definition and overall contrast, while a simple rescaling r̂ leads to a still largely blurred image. For the isotropic binary mixture, the DAE test error mse fˆ in fact approaches the information-theoretic lowest achievable MSE mse? as the sample complexity α increases. This mse? is given by the application of Tweedie’s formula, that requires perfect knowledge of the cluster means µ k and covariances Σ k – it is, therefore, an oracle denoiser. The DAE MSE approaches the oracle test error mse? as the number of available training samples n grows, and is already sensibly close to the optimal value for α = 8.

Comparing the full DAE fˆ to the bottleneck network component û, it follows from Result II.3 and Corollary II.4 that û leads to a higher MSE than the full DAE fˆ, with the gap being Θ(d). More precisely, a closed-form expression (15) is provided for this gap. The theoretical prediction compares excellently with numerical simulations. Strikingly, PCA denoising yields an MSE almost indistinguishable from û, strongly suggesting that û essentially learns, also in the denoising setting, to project the noisy data x̃ along the principal components of the training set. This echoes the findings of previous works in the case of RAEs that bottleneck networks are limited by the PCA reconstruction performance. Crucially however, it also means that compared to the full DAE fˆ, PCA is sizeably suboptimal, since mse PCA ≈ mse û = mse fˆ + Θ(d). This last observation has an important consequence: in contrast to previously studied RAEs, the full DAE fˆ does not simply learn to perform PCA. In contrast to bottleneck RAE networks, the non-linear DAE hence does not reduce to a linear model after training. The non-linearity is important to improve the denoising MSE. Trained alone, the bottleneck network û only learns to perform PCA; trained jointly with the rescaling component as part of the full DAE f b,w, it learns a richer, non-linear representation.

Result II.3, alongside Corollary II.4 and the discussion in Section III provide a firm theoretical basis for the well-known empirical intuition that skip-connections allow to better propagate information from the input to the output of the DAE, thereby contributing to preserving intrinsic characteristics of the input. This effect is clearly illustrated in Fig. 4, where the resulting denoised image of an MNIST 7 by r̂, fˆ, û and PCA are presented. While the bottleneck network û perfectly eliminates the background noise and produces an image with a very good resolution, it essentially collapses the image to the cluster mean, and yields, like PCA, the average MNIST 7. As a consequence, the denoised image bears little resemblance with the original image – in particular, the horizontal bar of the 7 is lost in the process. Conversely, the rescaling r̂ preserves the nuances of the original image, but the result is still largely blurred and displays overall poor contrast. Finally, the complete DAE manages to preserve the characteristic features of the original data, while enhancing the image resolution by slightly overlaying the average 7 thereupon.

The optimization of the DAE is therefore described by a tradeoff between two competing effects – namely the preservation of the input nuances by the skip connection, and the enhancement of the resolution/noise removal by the bottleneck network. This allows to discuss the curious non-monotonicity of the cosine similarity θ as a function of the noise level ∆. While it may at first seem curious that the DAE seemingly does not manage to learn the data structure better for low ∆ than for intermediate ∆ (where the cosine similarity θ is observed to be higher), this is actually due to the afore-discussed tradeoff. Indeed, for small ∆, the data is still substantially clean, and there is therefore no incentive to enhance the contrast by using the cluster means –which are consequently not learnt. This phase is thus characterized by a large skip connection strength b̂, and small cosine similarity θ and weight norm ŵ F. Conversely, at high noise levels ∆, the nuances of the data are already lost because of the noise. Hence the DAE does not rely on the skip connection component (whence the small values of b̂), and the only way to produce reasonably denoised data is to collapse to the cluster mean using the network component (whence a large ŵ F).

Conclusion

The paper considers the problem of denoising a high-dimensional Gaussian mixture, by training a DAE via empirical risk minimization, in the limit where the number of training samples and the dimension are proportionally large. A sharp asymptotic characterization of a number of summary statistics of the trained DAE weight, average MSE, and cosine similarity with the cluster means is provided. These results contain as a corollary the case of RAEs. Building on these findings, the role of the skip connection and the bottleneck network in the DAE architecture is isolated and the tradeoff between those two components in terms of preservation of the data nuances and noise removal is characterized – thereby providing some theoretical insight into a longstanding practical intuition in machine learning. The real data experiments hint at the presence of Gaussian universality. The DAE with skip connection provides an autoencoder model which does not just simply learn the principal components of the training set, and therefore affords a genuinely non-linear network model where richer learning settings can be investigated.

Based on the scientific paper, here are the specific improvements I can make to AI systems, along with what the improved systems can do:

1. Improved Denoising Autoencoder (DAE) Architecture Design

• Improvement: Implement a two-layer DAE with a trainable skip connection and tied weights, as described in the paper. The skip connection is not just a residual path but a learned scalar multiplier that adapts to the noise level.

• What the improved AI can do: Achieve significantly lower mean-squared-error (MSE) on denoising tasks compared to standard DAEs without skip connections. The system will automatically balance between preserving fine-grained input details (via the skip connection) and removing noise (via the bottleneck network). For example, on MNIST, this architecture preserves the horizontal bar of a "7" while removing background noise, which a standard DAE or PCA would blur away.

2. Optimal Noise-Level Adaptation

• Improvement: Use the closed-form formula for the skip connection strength (Equation 10) to dynamically adjust the skip connection weight during inference, based on the estimated noise level (∆) of the input.

• What the improved AI can do: Automatically switch between a detail-preserving mode (high skip connection weight, low noise) and a noise-removal mode (low skip connection weight, high noise). This eliminates the need for manual tuning or retraining when the noise level changes, making the system robust to varying signal-to-noise ratios in real-world applications like medical imaging or low-light photography.

3. Sample-Complexity-Aware Training

• Improvement: Use the asymptotic formulae (Result II.3) to predict the test MSE as a function of the number of training samples (α = n/d). This allows the system to determine the minimum number of samples needed to achieve a target denoising performance before training begins.

• What the improved AI can do: Optimize data collection and training budgets. For instance, if the target is to get within 10% of the oracle Tweedie denoiser's performance, the system can predict that α = 8 is sufficient (as shown in Fig. 3) and avoid wasting resources on collecting more data. Conversely, it can flag when more data is critical.

4. Non-Linear Feature Learning Beyond PCA

• Improvement: Replace linear autoencoder bottlenecks (which are proven to be equivalent to PCA) with the non-linear activation (e.g., tanh) and the specific architecture from the paper. The paper proves that this non-linear DAE does not collapse to PCA, unlike standard RAEs.

• What the improved AI can do: Learn genuinely non-linear representations that capture cluster structure (e.g., the difference between two Gaussian mixture clusters) that PCA cannot. This leads to a Θ(d) improvement in MSE over PCA-based denoising (as shown in Equation 15 and Fig. 3), meaning the system can recover signals that linear methods fundamentally miss.

5. Predictive Performance Estimation for Real-World Data

• Improvement: Use the paper's finding that Gaussian mixture models with estimated means and covariances accurately predict the DAE's performance on real datasets (MNIST, FashionMNIST). Implement a pre-training diagnostic that fits a Gaussian mixture to the data and uses the closed-form equations to predict the expected MSE.

• What the improved AI can do: Before deploying a denoising system on a new dataset, it can estimate the achievable denoising quality without running expensive training. This is useful for feasibility studies, benchmarking, and selecting between different architectures. The system can also identify when a dataset is too complex for a shallow DAE (when the Gaussian mixture prediction deviates significantly from simulations).

6. Oracle-Baseline Comparison and Benchmarking

• Improvement: Implement the Tweedie's formula-based oracle denoiser (Appendix B) as a built-in benchmark. This provides an information-theoretic lower bound on the achievable MSE for any denoising algorithm.

• What the improved AI can do: Automatically report the gap to optimality for any trained DAE. This allows users to understand how much room for improvement exists. For example, the system can tell a user: Your DAE is within 5% of the theoretical optimum for this noise level, so further architectural changes will yield diminishing returns.

7. Regularization Strength Optimization

• Improvement: Use the paper's finding that the `2 regularization strength (λ) has minimal impact on the qualitative behavior (Fig. 8) to simplify hyperparameter tuning. The system can default to a moderate λ (e.g., 0.1) and only adjust it if the Gaussian mixture prediction deviates from simulations.

• What the improved AI can do: Reduce the hyperparameter search space, saving computational time. The system can also use the theoretical curves to verify that the chosen λ is not causing underfitting or overfitting, by comparing the predicted weight norms and cosine similarities with the actual trained values.

8. Adaptive Architecture Selection

• Improvement: Use Corollary II.4 to compare the theoretical performance of the full DAE, the bottleneck-only network, and the rescaling-only network. The system can then choose the simplest architecture that meets the performance target.

• What the improved AI can do: Automatically select between a simple rescaling (1 parameter), a bottleneck network, or the full DAE based on the noise level and available data. For very low noise, a simple rescaling might suffice; for high noise, the full DAE is necessary. This prevents over-engineering and reduces inference cost.

9. Robustness to Non-Gaussian Data

• Improvement: Use the paper's empirical finding that the Gaussian mixture model accurately predicts performance on MNIST and FashionMNIST to build a Gaussian universality check. The system can test if a new dataset's denoising behavior is well-approximated by its second-order statistics (mean and covariance).

• What the improved AI can do: Flag when a dataset is likely to benefit from a deeper or more complex architecture (when the Gaussian prediction fails). Conversely, it can confidently apply the shallow DAE to datasets that exhibit Gaussian-like second-order statistics, even if they are not strictly Gaussian mixtures.

10. Training Dynamics Validation

• Improvement: Use the paper's observation that Adam-optimized DAEs converge to the global minimum predicted by the replica method (Fig. 1) to add a validation step during training. The system can compare the current weight norms and cosine similarities against the theoretical predictions.

• What the improved AI can do: Detect if training is stuck in a poor local minimum (which would be indicated by a significant deviation from the theoretical values). The system can then automatically restart training with different initialization or adjust the learning rate, saving time and improving final performance.

We address the problem of denoising data from a Gaussian mixture using a two-layer non-linear autoencoder with tied weights and a skip connection. We consider the high-dimensional limit where the number of training samples and the input dimension jointly tend to infinity while the number of hidden units remains bounded. We provide closed-form expressions for the denoising mean-squared test error. Building on this result, we quantitatively characterize the advantage of the considered architecture over the autoencoder without the skip connection that relates closely to principal component analysis. We further show that our results accurately capture the learning curves on a range of real data sets.