A solvable high-dimensional model where nonlinear autoencoders learn structure invisible to PCA while test loss misaligns with generalization

TL;DR

This work presents a solvable high-dimensional model in which nonlinear autoencoders can recover latent structure invisible to PCA through higher-order correlations, while PCA and linear autoencoders cannot. By combining a spiked cumulant data model with a minimal nonlinear autoencoder, the authors characterize population and empirical risk learning, establish a correlation-exponent framework, and analyze both Bayes-optimal and ERM performance, including AMP-based weak recovery thresholds. A key finding is the misalignment between test loss and representation quality: linear methods can minimize reconstruction loss yet fail to capture the hidden spike, whereas nonlinear autoencoders succeed in latent recovery with higher test loss. The framework offers a tractable testbed for understanding nonlinear representation learning in self-supervised settings and highlights when reconstruction-based validation may mislead, with implications for evaluating downstream usefulness and designing robust objectives.

Abstract

Many real-world datasets contain hidden structure that cannot be detected by simple linear correlations between input features. For example, latent factors may influence the data in a coordinated way, even though their effect is invisible to covariance-based methods such as PCA. In practice, nonlinear neural networks often succeed in extracting such hidden structure in unsupervised and self-supervised learning. However, constructing a minimal high-dimensional model where this advantage can be rigorously analyzed has remained an open theoretical challenge. We introduce a tractable high-dimensional spiked model with two latent factors: one visible to covariance, and one statistically dependent yet uncorrelated, appearing only in higher-order moments. PCA and linear autoencoders fail to recover the latter, while a minimal nonlinear autoencoder provably extracts both. We analyze both the population risk, and empirical risk minimization. Our model also provides a tractable example where self-supervised test loss is poorly aligned with representation quality: nonlinear autoencoders recover latent structure that linear methods miss, even though their reconstruction loss is higher.

Figures (7)

• Figure 1: Cosine similarity of the BO and ERM estimators with, respectively, $\mathbf{u^\star}$ and $\mathbf{v^\star}$, for different choices of the activation of the autoencoder. Here, we consider a distribution $P_{\rm latents}$ with correlation exponent $k^\star=2$, concretely $\lambda \sim \mathcal{N}(0, 1)$ and $\nu$ is $-\sqrt{2}$ if $|\lambda| < \Phi^{-1}(0.75)$ and $+\sqrt{2}$ otherwise, where $\Phi$ is the standard Gaussian cdf. Solid lines are analytical predictions in the high-dimensional limit (see \ref{['sec:ERM']}). Colored dots are numerical simulations of minimization with full-batch Adam DBLP:journals/corr/KingmaB14, with learning rate $\eta=0.1$ over $800$ epochs and no weight decay ($d = 2000$, averaged over $30$ instances, error bars represent one standard deviation). Black dots are runs of AMP for $d=10000$, averaged over $72$ seeds. The vertical dashed line is $\alpha^{\rm lin}_{u}$ given in \ref{['res:bo-weak']}. The theoretical predictions and the simulation match, modulo finite-size effects that are further discussed in Appendix \ref{['app:fixed_alpha_simulations']}. We see that the linear autoencoder does not achieve weak recovery of the $\mathbf{v^\star}$ spike, the one with $\tanh$ non-linearity either, while autoencoders with other depicted non-linearities do.
• Figure 2: Train loss (left panel) and test loss (center panel) of the non-linear autoencoders minus the corresponding loss of the linear autoencoder, $\sigma=\mathrm{id}$, as a function of the sample complexity. The vertical dashed line is $\alpha^{\rm lin}_{u}$. The parameters and $P_{\rm latents}$ are the same as in Fig. \ref{['fig:cosine']}. We see that both train and test losses are smaller for the linear autoencoder. Right panel: Classification error on the downstream task, where we see the superiority of the non-linearities that recovered correlation with the hidden spike.
• Figure 3: Cosine similarity of the ERM estimators with, respectively, $\mathbf{u^\star}$ and $\mathbf{v^\star}$, for different choices of the activation of the autoencoder and dimension $d$. Results are shown at fixed sample ratio $\alpha=n/d=1.7$ and the same setting as Fig. \ref{['fig:cosine']}, showcasing the finite-size effects. The dashed line indicates the standard scaling of the cosine similarity with an independent random direction. Points are numerical simulations, and the stars are the predicted value (Section \ref{['sec:ERM']}), which agree in the high-dimensional limit.
• Figure 4: Train loss (left panel) and test loss (right panel) of the non-linear autoencoder with $\sigma=\tanh$ activation minus the corresponding loss of the linear autoencoder, $\sigma=\mathrm{id}$, as a function of the sample ratio $\alpha$. Colored dots depend on the dimension $d$. Inset: scaling of the training loss at fixed $\alpha=0.48$ with respect to $d$. Solid line corresponds to the theoretical prediction of ERM (Section \ref{['sec:ERM']}).
• Figure 5: Left: Train loss of the non-linear autoencoder with $\sigma=\tanh$ activation minus the corresponding loss of the linear autoencoder, $\sigma=\mathrm{id}$, during training. Center and Right: signed cosine similarity of the autoencoders weights, without the absolute value of the definition, with the $\mathbf{u}^\star$ and $\mathbf{v}^\star$ spikes ($\theta_u$, $\theta_v$). Simulations at fixed sample rate $\alpha =0.48$, $d=6400$ and $30$ instances. Each blue line corresponds to a single instance, different initialization and training dataset, and the black line is the mean of them. Gray fill represents one standard deviation at each epoch.

Theorems & Definitions (15)

• Definition 1.1: Correlation exponent
• Proposition 2.1: Linear autoencoder yields PCA
• Lemma 4.1: Hermite expansion of the population loss
• Theorem 4.2: Gradient flow dynamics on population loss
• proof : Idea of the proof
• proof : Proof of \ref{['prop:linear_pca']}
• Lemma A.1: Empirical covariance eigenvectors do not correlate with $\mathbf{v}^\star$
• proof
• Lemma A.2: Triple Hermite identity
• proof