🤖 AI Summary
This work investigates the role of projection heads in self-supervised learning, focusing on their impact on representation geometry, dimensional collapse, and the information–invariance trade-off. Modeling the projection head as a trainable Riemannian metric on the representation manifold induced by the backbone network, we offer the first differential-geometric characterization of its mechanism: nonlinear heads avoid representational collapse by inducing negative curvature, whereas linear or ReLU-based heads rely on BatchNorm and discrete optimization dynamics. Through Hessian spectral analysis, continuous-time gradient flows, and geometric tracking of optimization trajectories, we establish a link between metric degeneracy and information loss, and demonstrate that smooth activations such as Swish explicitly generate negative curvature. Our findings further reveal that projection heads act as universal geometric buffers and should be removed prior to downstream tasks to enhance performance.
📝 Abstract
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.