Neural Collapse by Design: Learning Class Prototypes on the Hypersphere
Abstract
Lay Summary
When we want to classify different objects into different categories, say a table from a helmet, we actually want examples of each category grouped tightly and different categories kept well apart. It is proven that there is one ideal geometric arrangement that does this perfectly, known as Neural Collapse: each category gets an "anchor point," all examples of that category cluster around it, and these anchors sit maximally far from each other, like equally spaced cities on a globe. Despite this elegant theory, the two most popular methods for training such classifiers almost never reach this ideal in practice. We show that these two methods, which look very different on the surface, are actually doing the same thing: both are learning an anchor point for each category, and each fails in its own specific way. We fix both by constraining the geometry to lie on a sphere, where the ideal arrangement is the only option, and we design new training objectives that provably reach it. Models trained this way are more accurate, transfer better to new tasks, handle imbalanced and corrupted data more gracefully, and skip an hours-long step that traditional methods normally require.