Deep Neural Collapse Is Provably Optimal for the Deep Unconstrained Features Model

Peter Súkeník, Marco Mondelli, Christoph H Lampert

Keywords: neural collapse, unconstrained features model, deep learning

TL;DR: We prove the optimality of deep neural collapse for the deep unconstrained features model - an extension of the unconstrained features model to arbitrarily many layers - for binary classification, and we provide experiments to support our theory.

Abstract: Neural collapse (NC) refers to the surprising structure of the last layer of deep neural networks in the terminal phase of gradient descent training. Recently, an increasing amount of experimental evidence has pointed to the propagation of NC to earlier layers of neural networks. However, while the NC in the last layer is well studied theoretically, much less is known about its multi-layered counterpart - deep neural collapse (DNC). In particular, existing work focuses either on linear layers or only on the last two layers at the price of an extra assumption. Our work fills this gap by generalizing the established analytical framework for NC - the unconstrained features model - to multiple non-linear layers. Our key technical contribution is to show that, in a deep unconstrained features model, the unique global optimum for binary classification exhibits all the properties typical of DNC. This explains the existing experimental evidence of DNC. We also empirically show that (i) by optimizing deep unconstrained features models via gradient descent, the resulting solution agrees well with our theory, and (ii) trained networks recover the unconstrained features suitable for the occurrence of DNC, thus supporting the validity of this modeling principle.

Submission Number: 7518