The Geometry of Projection Heads: Conditioning, Invariance, and Collapse

We develop a geometric theory of projection heads in self-supervised learning by modeling the head as a trainable Riemannian metric on the backbone representation manifold. We show that linear heads perform implicit subspace whitening, while nonlinear heads adapt local metrics to satisfy the specific topological constraints of the loss, with head depth empirically dictating this capacity. Analyzing dimensional collapse, we prove that smooth nonlinear heads natively induce negative eigenvalues in the Hessian at collapsed equilibria, making them unstable. We empirically validate this by continuously tracking the optimization geometry during training, which reveals that smooth activations like Swish can generate explicit negative curvature to escape collapse, whereas linear and ReLU heads under continuous-time gradient flow cannot, relying instead on discrete-time optimization dynamics and BatchNorm. Finally, we geometrically characterize how metric degeneracy governs the information-invariance trade-off, explaining why the head must be discarded. Evaluated across contrastive and decorrelation-based objectives on foundation models, our results demonstrate that the projection head acts as a universal geometric buffer, decoupling the semantic backbone from the rigid, destructive constraints of the pretraining objective.

Artificial intelligence models are often trained by comparing altered versions of the same image, like a rotated photo or a color-shifted picture, and learning to recognize that they represent the same underlying object. However, if an AI is forced to perfectly ignore all these alterations (for example, becoming completely blind to color), it loses valuable information it might need for future tasks, such as classifying animal breeds by their coat. To solve this, modern AI systems use a "projection head," an extra, temporary block of computational layers attached to the end of the model during training. Our research explains the exact mechanics of how and why this architectural trick works. We demonstrate that this extra block acts as a sacrificial shock absorber. It takes on the extreme, destructive pressure of the training rules by heavily distorting the data to perfectly ignore the image alterations. This allows the main AI model to remain safely shielded, retaining a rich, detailed understanding of the world.