The Geometry of Projection Heads: Conditioning, Invariance, and Collapse
Abstract
We develop a geometric theory of projection heads in self-supervised learning by interpreting the head as a trainable metric on the backbone representation manifold. Our analysis reveals that head curvature and architectural asymmetry induce negative eigenvalues of the Hessian at collapsed equilibria in networks with smooth activation functions, yielding a destabilization mechanism which explains collapse avoidance in non-contrastive methods. We further show that linear heads perform implicit subspace whitening under induced metric geometry, while nonlinear heads adapt local metrics to satisfy the specific topological constraints of the loss. Finally, we characterize how metric degeneracy governs the information-invariance trade-off in learned representations. Our results apply to both contrastive and non-contrastive objectives including InfoNCE, BYOL, SimSiam, and decorrelation-based methods, demonstrating that the projection head acts as a universal geometric buffer that decouples the semantic backbone from the rigid constraints of the training objective.