Guiding Neural Collapse: Optimising Towards the Nearest Simplex Equiangular Tight Frame

TL;DR

This paper leverages the Neural Collapse phenomenon, specifically the convergence to a Simplex ETF, to develop a Riemannian optimisation method that accelerates training convergence and improves stability in neural networks.

Contribution

It introduces a novel Riemannian optimisation approach based on nearest simplex ETF geometry to guide neural network training.

Findings

Accelerates convergence in neural network training

Enhances training stability across architectures

Effective on both synthetic and real-world datasets

Abstract

Neural Collapse (NC) is a recently observed phenomenon in neural networks that characterises the solution space of the final classifier layer when trained until zero training loss. Specifically, NC suggests that the final classifier layer converges to a Simplex Equiangular Tight Frame (ETF), which maximally separates the weights corresponding to each class. By duality, the penultimate layer feature means also converge to the same simplex ETF. Since this simple symmetric structure is optimal, our idea is to utilise this property to improve convergence speed. Specifically, we introduce the notion of nearest simplex ETF geometry for the penultimate layer features at any given training iteration, by formulating it as a Riemannian optimisation. Then, at each iteration, the classifier weights are implicitly set to the nearest simplex ETF by solving this inner-optimisation, which is encapsulated within a declarative node to allow backpropagation. Our experiments on synthetic and real-world architectures for classification tasks demonstrate that our approach accelerates convergence and enhances training stability.