揭示自监督学习中投影头的几何本质及其对模型稳定性的影响
The Geometry of Projection Heads: Conditioning, Invariance, and Collapse
- 将投影头建模为可训练的黎曼度量，解析其在特征空间中的几何作用
- 证明平滑非线性头能产生负曲率，有效避免维度坍缩；线性/ReLU头则依赖批归一化
- 解释为何预训练后需丢弃投影头：它在信息与不变性间权衡并缓解目标约束
我们通过将投影头建模为骨干表示流形上的可训练黎曼度量，构建了自监督学习中投影头的几何理论。研究表明，线性头隐式执行子空间去相关，而非线性头则自适应局部度量以满足损失函数的拓扑约束，且头部深度决定该能力。分析维度坍缩时，证明平滑非线性头在坍缩平衡点处天然产生负的海森矩阵特征值，导致不稳定性。实验上持续追踪训练过程的优化几何发现，如Swish等平滑激活可显式生成负曲率以逃离坍缩，而线性或ReLU头在连续时间梯度流下无法实现，必须依赖离散优化动态和批归一化。最后，我们从几何角度刻画了度量退化如何支配信息-不变性权衡，解释为何必须丢弃投影头。在基础模型上对比多种对比和去相关目标的评估表明，投影头充当通用几何缓冲器，解耦语义骨干与预训练目标的刚性破坏性约束。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。