Understanding Deep Representation Learning via Layerwise Feature Compression and Discrimination

Peng Wang^1,*, Xiao Li^1,*, Can Yaras¹, Zhihui Zhu², Laura Balzano¹, Wei Hu¹, and Qing Qu¹

¹University of Michigan · ²Ohio State University · ^*Equal contribution

TL;DR. Each layer of a trained classification deep network compresses within-class features and discriminates between-class features at precise, predictable rates.

What happens inside a deep network, layer by layer?

Deep networks are believed to perform hierarchical feature learning: each layer builds on the last to extract increasingly meaningful structure from raw data. But despite efforts on empirical study, it remains unclear exactly how this happens: what quantitative pattern, if any, governs the transformation of features from shallow to deep layers.

Empirical studies have observed two consistent trends. First, intermediate layers the intrinsic dimension of features as depth increases. Second, deep networks tend to produce features that are within-class compressed and between-class discriminative — a pattern closely related to the phenomenon observed at the final layer.

Numerical rank and training accuracy plotted against layer index for an MLP and a hybrid network, showing rank rising then falling while accuracy saturates.

Feature rank rises then falls, while accuracy saturates early. Initial layers make features linearly separable; later layers compress them further.

However, no theoretical framework has explained why this happens, or at what rate. We take a first step toward closing that gap.

Why deep linear networks?

Deep linear networks (DLNs) are a kind of network with no nonlinear activation but a cascade of linear layers. But we find that linear layers placed deep in a network behave remarkably like their nonlinear counterparts: once early layers have made features linearly separable, the deeper layers' job is simply to compress and discriminate, a role that linear layers perform just as well.

UMAP visualization of features at layers 1, 2, 4, and 6 for a nonlinear MLP and a hybrid network, showing increasingly separated and compact class clusters.

Linear layers replicate the deep-layer behavior of nonlinear networks. A hybrid network (nonlinear layers followed by linear layers) shows the same progressive class separation as a fully nonlinear MLP.

This lets us study a simple, tractable model — and the simplicity is not a weakness here, since DLNs are over-parameterized: stacking linear layers does not reduce to a single linear operator in any meaningful training sense, and depth in DLNs still meaningfully improves generalization, exactly as in their nonlinear cousins.

Two metrics: compression and discrimination

To make "feature evolution" precise, we define two metrics at each layer $$l$$ . Let $\Sigma_W^l$ and $\Sigma_B^l$ denote the within-class and between-class covariance of features at that layer. We define:

\[C_l = \frac{\mathrm{Tr}(\Sigma_W^l)}{\mathrm{Tr}(\Sigma_B^l)}, \qquad D_l = 1 - \max_{k \neq k'} \frac{\langle \mu_k^l, \mu_{k'}^l \rangle}{\|\mu_k^l\| \|\mu_{k'}^l\|}.\]

$$C_l$$ measures within-class compression (smaller is more compressed), and $$D_l$$ measures between-class discrimination (larger is more separated, computed from the smallest angle between any two class means).

Our main theoretical result, proved under on the data and trained weights, is strikingly simple:

The compression metric decays at a geometric rate, while the discrimination metric increases at a linear rate, with respect to the number of layers.

More precisely, under these assumptions, the ratio of within-class compression between consecutive layers satisfies

\[\frac{C_{l+1}}{C_l} = O\!\left(\frac{\varepsilon^2}{n^{1/L}}\right),\]

while the between-class discrimination metric satisfies

\[D_l \geq 1 - O\!\left(\frac{\theta + \delta}{L}\right) \cdot (\text{linear in } l).\]

Here $$n$$ is the number of samples per class, $$L$$ is the network depth, and $\varepsilon,\theta,\delta$ are the small constants from the assumptions above, controlling the data's near-orthogonality and the weights' balancedness and low-rankness. In words: with more samples or smaller residual constants, compression happens faster per layer; discrimination grows steadily and linearly regardless.

Compression and discrimination metrics plotted against layer index for a linear network and a nonlinear network, both showing geometric decay and linear increase respectively.

Theory matches practice in both linear and nonlinear networks. Compression decays geometrically (log-scale plot, left); discrimination increases linearly (right) — for both a DLN and a trained MLP.

This pattern holds well beyond the assumptions

We validate the theorem with synthetic data matching our assumptions exactly, then show the same trend holds approximately under default PyTorch initialization, on real datasets (CIFAR-10, FashionMNIST), and even in nonlinear MLPs of varying depth, and beyond images, on a text classification benchmark.

Within-class compression metric decaying approximately geometrically across layers for networks of varying depths (L=5,7,9) trained on FashionMNIST and CIFAR-10 with default initialization.

Progressive compression persists with real data and default initialization, across network depths. Our theoretical assumptions are a simplification — the pattern survives without them.

Why this matters

Beyond characterizing a previously undescribed pattern, our result has concrete implications:

Neural collapse, without the unconstrained features model. Most prior neural collapse analyses treat last-layer features as free optimization variables, ignoring the network's hierarchical structure. Our theorem shows that sufficiently deep linear networks naturally exhibit neural collapse as a consequence of layerwise dynamics without the unconstrained features model assumption.
Why projection heads help transfer learning. Contrastive learning methods discard a few extra MLP layers (the "projection head") after pretraining, keeping only the features before them. Our theory explains why: features closer to the final layer are more collapsed and thus less transferable, so backing off by a few layers preserves more usable structure.

Diagram of projection head usage during pretraining and transfer, alongside a plot showing transfer accuracy improving and feature compression decreasing as more projection head layers are added.

More projection head layers mean less feature collapse and better transfer accuracy. This empirically confirms the mechanism predicted by our theory.

In short

Even though deep linear networks lack the expressive power of their nonlinear counterparts, they are powerful enough to reveal — and prove — a precise, quantitative law governing how deep networks transform data into discriminative representations, layer by layer.

Citation

If you find our work useful, please consider citing:

@article{wang2025understanding,
  title={Understanding deep representation learning via layerwise feature compression and discrimination},
  author={Wang, Peng and Li, Xiao and Yaras, Can and Zhu, Zhihui and Balzano, Laura and Hu, Wei and Qu, Qing},
  journal={Journal of Machine Learning Research},
  volume={26},
  number={220},
  pages={1--61},
  year={2025}
}