Neural Collapse versus Low-rank Bias: Is Deep Neural Collapse Really Optimal?

TL;DR

This paper investigates the limitations of neural collapse in deep neural networks, revealing that beyond two layers or classes, a low-rank bias prevents neural collapse from being optimal, supported by theoretical analysis and experiments.

Contribution

It extends the analysis of neural collapse to non-linear, multi-layer models and uncovers a low-rank bias that challenges the optimality of neural collapse in these settings.

Findings

Neural collapse is not optimal beyond two layers or classes.

Low-rank bias of regularization schemes influences solutions.

Experimental evidence shows low-rank structures emerge in trained models.

Abstract

Deep neural networks (DNNs) exhibit a surprising structure in their final layer known as neural collapse (NC), and a growing body of works has currently investigated the propagation of neural collapse to earlier layers of DNNs -- a phenomenon called deep neural collapse (DNC). However, existing theoretical results are restricted to special cases: linear models, only two layers or binary classification. In contrast, we focus on non-linear models of arbitrary depth in multi-class classification and reveal a surprising qualitative shift. As soon as we go beyond two layers or two classes, DNC stops being optimal for the deep unconstrained features model (DUFM) -- the standard theoretical framework for the analysis of collapse. The main culprit is a low-rank bias of multi-layer regularization schemes: this bias leads to optimal solutions of even lower rank than the neural collapse. We support our theoretical findings with experiments on both DUFM and real data, which show the emergence of the low-rank structure in the solution found by gradient descent.