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

• URL: http://arxiv.org/abs/2605.17180v1
• Date: Sat, 16 May 2026 22:32:05 GMT
• Title: The Geometry of Projection Heads: Conditioning, Invariance, and Collapse
• Abstract summary: We develop a geometric theory of projection heads in self-supervised learning.<n>We show that linear heads perform implicit subspace whitening, while nonlinear heads adapt local metrics to satisfy the specific topological constraints of the loss.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Variance Reduction for Stochastic Gradient Generalized Non-reversible Langevin Monte Carlo Algorithms [7.336637417508747]
  We study the leading-order fluctuation of gradient-Maruyama estimators for non-reversible Langevin dynamics.<n>We prove that the empirical average over a horizon of order the inverse squared stepsize satisfies a central limit theorem in the vanishing-stepsize regime.
  arXiv  Detail & Related papers  (2026-06-27T08:25:48Z)
• Generalization in Nonlinear Least Squares via Learned Feature Geometry [15.966075413898336]
  We study the generalization of ridge-regularized nonlinear least-squares models via on-average algorithmic stability.<n>We deriving error bounds for local minimizers in terms of a data-dependent effective dimension that reflects the geometry of the gradient model.
  arXiv  Detail & Related papers  (2026-06-07T19:33:36Z)
• Simulated Adoption: Decoupling Magnitude and Direction in LLM In-Context Conflict Resolution [3.0242762196828448]
  Large Language Models (LLMs) frequently prioritize conflicting in-context information over pre-existing parametric memory.<n>We show that models do not "unlearn" or suppress the magnitude of internal truths but rather employ a mechanism of geometric displacement.
  arXiv  Detail & Related papers  (2026-02-04T06:13:11Z)
• Diagonalizing the Softmax: Hadamard Initialization for Tractable Cross-Entropy Dynamics [29.85277126753054]
  Cross-entropy (CE) loss dominates deep learning, yet existing theory often relies on simplifications.<n>We provide an in-depth characterization of a canonical network with standard neural-basis vectors.
  arXiv  Detail & Related papers  (2025-12-03T17:45:09Z)
• Generalization Below the Edge of Stability: The Role of Data Geometry [60.147710896851045]
  We show how data geometry controls generalization in ReLU networks trained below the edge of stability.<n>For data distributions supported on a mixture of low-dimensional balls, we derive generalization bounds that provably adapt to the intrinsic dimension.<n>Our results consolidate disparate empirical findings that have appeared in the literature.
  arXiv  Detail & Related papers  (2025-10-20T21:40:36Z)
• On the Dynamics Under the Unhinged Loss and Beyond [104.49565602940699]
  We introduce the unhinged loss, a concise loss function, that offers more mathematical opportunities to analyze closed-form dynamics. The unhinged loss allows for considering more practical techniques, such as time-vary learning rates and feature normalization.
  arXiv  Detail & Related papers  (2023-12-13T02:11:07Z)
• Linear Regression on Manifold Structured Data: the Impact of Extrinsic Geometry on Solutions [4.8234611688915665]
  We study linear regression applied to data structured on a manifold. We analyze the impact of the manifold's curvatures on the uniqueness of the regression solution.
  arXiv  Detail & Related papers  (2023-07-05T17:51:26Z)
• Last-Iterate Convergence of Adaptive Riemannian Gradient Descent for Equilibrium Computation [52.73824786627612]
  This paper establishes new convergence results for textitgeodesic strongly monotone games.<n>Our key result shows that RGD attains last-iterate linear convergence in a textitgeometry-agnostic fashion.<n>Overall, this paper presents the first geometry-agnostic last-iterate convergence analysis for games beyond the Euclidean settings.
  arXiv  Detail & Related papers  (2023-06-29T01:20:44Z)
• Learning Linear Causal Representations from Interventions under General Nonlinear Mixing [52.66151568785088]
  We prove strong identifiability results given unknown single-node interventions without access to the intervention targets. This is the first instance of causal identifiability from non-paired interventions for deep neural network embeddings.
  arXiv  Detail & Related papers  (2023-06-04T02:32:12Z)
• Linear Adversarial Concept Erasure [108.37226654006153]
  We formulate the problem of identifying and erasing a linear subspace that corresponds to a given concept.<n>We show that the method is highly expressive, effectively mitigating bias in deep nonlinear classifiers while maintaining tractability and interpretability.
  arXiv  Detail & Related papers  (2022-01-28T13:00:17Z)
• GELATO: Geometrically Enriched Latent Model for Offline Reinforcement Learning [54.291331971813364]
  offline reinforcement learning approaches can be divided into proximal and uncertainty-aware methods. In this work, we demonstrate the benefit of combining the two in a latent variational model. Our proposed metrics measure both the quality of out of distribution samples as well as the discrepancy of examples in the data.
  arXiv  Detail & Related papers  (2021-02-22T19:42:40Z)
• Semiparametric Nonlinear Bipartite Graph Representation Learning with Provable Guarantees [106.91654068632882]
  We consider the bipartite graph and formalize its representation learning problem as a statistical estimation problem of parameters in a semiparametric exponential family distribution. We show that the proposed objective is strongly convex in a neighborhood around the ground truth, so that a gradient descent-based method achieves linear convergence rate. Our estimator is robust to any model misspecification within the exponential family, which is validated in extensive experiments.
  arXiv  Detail & Related papers  (2020-03-02T16:40:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.