Understanding Generalization from Embedding Dimension and Distributional Convergence

Deep neural networks often generalize well despite heavy over-parameterization, challenging classical parameter-based analyses. We study generalization from a representation-centric perspective and analyze how the geometry of learned embeddings is associated with generalization performance for a fixed trained model. We derive a post-hoc generalization bound that relates the gap between population risk and held-out empirical risk to two factors: (i) the intrinsic dimension of the embedding, which determines the convergence rate of the empirical embedding distribution to its population counterpart in Wasserstein distance, and (ii) the sensitivity of the downstream mapping from embeddings to predictions, characterized by Lipschitz constants. Together, these provide a post-hoc explanation of generalization for trained models. At the final embedding layer, architectural sensitivity disappears and the bound is dominated by embedding dimension, explaining its strong empirical correlation with generalization performance. Experiments across architectures and datasets validate the theory and demonstrate the utility of embedding-based diagnostics.

Recent studies have found that AI models that perform well on new data often learn simpler, lower-dimensional internal representations. However, it has remained unclear why these low-dimensional representations help models generalize. In this work, we study this question by viewing a model’s representations as a distribution: the pattern of internal features produced across many examples. We show that when this distribution has lower intrinsic dimension, a limited sample can capture the overall structure of the data more faithfully. This makes performance estimates more accurate and helps explain why the model behaves reliably on unseen examples. We also show that if later network layers are less sensitive to small representation changes, generalization improves further. This provides a clearer explanation of why compact representations are linked to better AI performance.