Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Understanding Generalization from Embedding Dimension and Distributional Convergence

Junjie Yu, Zhuoli Ouyang, Haotian Deng, Chen Wei, Wenxiao Ma, Jianyu Zhang, Zihan Deng, Quanying Liu

TL;DR

The paper advances a representation-centric theory of generalization for fixed neural networks, showing that population risk can be bounded by empirical risk plus terms that scale with the intrinsic dimension of intermediate embeddings and the Lipschitz amplification of downstream mappings. A key insight is a dimension-dependent Wasserstein convergence rate, $\mathbb{E}[\mathcal{W}_1(\tilde{P}_k^Z, \hat{\tilde{P}}_{k,n}^Z)] \lesssim C_k \; n^{-1/(d_k+\epsilon)}$, which becomes faster as the embedding distribution concentrates on lower-dimensional manifolds. The bound simplifies at the final layer, explaining why the final-layer embedding dimension robustly correlates with generalization across architectures and modalities. Empirically, the authors validate the theory across synthetic and real networks, including CNNs and large pretrained models, and demonstrate a joint dimension–Lipschitz tradeoff in intermediate layers as a causal driver of generalization. This work provides practical embedding-based diagnostics and a principled framework for analyzing generalization beyond parameter-count or hypothesis-class notions.

Abstract

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 controls predictive performance for a fixed trained model. We show that population risk can be bounded by two factors: (i) the intrinsic dimension of the embedding distribution, which determines the convergence rate of empirical embedding distribution to the population distribution in Wasserstein distance, and (ii) the sensitivity of the downstream mapping from embeddings to predictions, characterized by Lipschitz constants. Together, these yield an embedding-dependent error bound that does not rely on parameter counts or hypothesis class complexity. At the final embedding layer, architectural sensitivity vanishes 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.

Understanding Generalization from Embedding Dimension and Distributional Convergence

TL;DR

The paper advances a representation-centric theory of generalization for fixed neural networks, showing that population risk can be bounded by empirical risk plus terms that scale with the intrinsic dimension of intermediate embeddings and the Lipschitz amplification of downstream mappings. A key insight is a dimension-dependent Wasserstein convergence rate, , which becomes faster as the embedding distribution concentrates on lower-dimensional manifolds. The bound simplifies at the final layer, explaining why the final-layer embedding dimension robustly correlates with generalization across architectures and modalities. Empirically, the authors validate the theory across synthetic and real networks, including CNNs and large pretrained models, and demonstrate a joint dimension–Lipschitz tradeoff in intermediate layers as a causal driver of generalization. This work provides practical embedding-based diagnostics and a principled framework for analyzing generalization beyond parameter-count or hypothesis-class notions.

Abstract

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 controls predictive performance for a fixed trained model. We show that population risk can be bounded by two factors: (i) the intrinsic dimension of the embedding distribution, which determines the convergence rate of empirical embedding distribution to the population distribution in Wasserstein distance, and (ii) the sensitivity of the downstream mapping from embeddings to predictions, characterized by Lipschitz constants. Together, these yield an embedding-dependent error bound that does not rely on parameter counts or hypothesis class complexity. At the final embedding layer, architectural sensitivity vanishes 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.
Paper Structure (78 sections, 12 theorems, 59 equations, 10 figures, 2 tables)

This paper contains 78 sections, 12 theorems, 59 equations, 10 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Works
  3. Classical Generalization Bounds.
  4. Representation-Based Approaches.
  5. Intrinsic Dimension of Representations.
  6. Preliminaries
  7. Measures and Wasserstein distance
  8. Network Decomposition and Embeddings
  9. Lipschitz continuity
  10. Geometric Complexity and Wasserstein Convergence
  11. Wasserstein convergence rates.
  12. Risk and Bayes Predictor
  13. Standing Assumptions
  14. Main Theoretical Results
  15. Dimension-dependent post-training generalization bound
  16. ...and 63 more sections

Key Result

Theorem 3.8

For any $p\in[1,\infty)$ and any $\varepsilon>0$, setting $s = d_p^*(\mu)+\varepsilon$ yields Since $\varepsilon$ may be chosen arbitrarily small, the convergence rate can be made arbitrarily close to $n^{-1/d_p^*(\mu)}$.

Figures (10)

  • Figure 1: Embedding Dimension and Lipschitz Constant of Network Jointly Influence Generalization Error.(A)Generalization error depends jointly on embedding dimension and network's Lipschitz constant. (B) Lower intrinsic dimension accelerates convergence of empirical to population distribution. (C) Smaller Lipschitz constants reduce output sensitivity to perturbations.
  • Figure 2: Scaling of Wasserstein Convergence in Neural Network Embeddings.(A) With fixed sample size, log(Wasserstein distance) increases approximately linearly with embedding dimension. (B) With fixed embedding dimension, log(Wasserstein distance) decreases approximately linearly with log(sample size).
  • Figure 3: Relationship Between Final-Layer Embedding Dimension, Wasserstein Distance and Generalization Error. We evaluate CIFAR-10 (A) and CIFAR-100 (B) and observe a significant correlation between final-layer embedding dimension, Wasserstein distance and generalization error.
  • Figure 4: Relationship between embedding geometry and generalization performance in large models. (A) Vision models. (B) Language models. Across both modalities, embedding dimension and Wasserstein distance are strongly correlated with generalization performance, consistent with the theoretical predictions of our generalization bound.
  • Figure 5: Effect of Network Width on Embedding Dimension and Generalization.(A) Reducing the width of the third layer does not lead to a consistent decrease in generalization error. (B) As the layer width decreases, the embedding dimension gradually decreases, but the network's Lipschitz constant increases, offsetting potential gains.
  • ...and 5 more figures

Theorems & Definitions (29)

  • Definition 3.1: Empirical measure
  • Definition 3.2: Wasserstein distance
  • Definition 3.3: Network decomposition
  • Definition 3.4: Empirical and population embedding distributions
  • Definition 3.5: Lipschitz map
  • Definition 3.6: Covering numbers and measure covering dimension
  • Definition 3.7: Upper Wasserstein dimension
  • Theorem 3.8: Wasserstein convergence governed by intrinsic dimension
  • Definition 3.9: Population and empirical risk
  • Definition 3.10: Bayes predictor
  • ...and 19 more