Abstract
Understanding the geometric structure of pre-trained language model embeddings matters for interpretability and safety. We investigate whether sentence-level classification signal resides in the Riemannian geometry of contextual token embeddings. This is probed by extracting per-token pullback metrics from a learned encoder's analytical Jacobian and aggregating them with the Fréchet mean on the symmetric positive definite (SPD) manifold, referred to as Riemannian Mean Pooling (RMP).
Across three datasets with non-trivial linguistic structure (CoLA, CREAK, RTE), RMP outperforms Euclidean mean pooling, while on FEVER-Symmetric, a benchmark designed to eliminate annotation-driven lexical artifacts, the method correctly remains at chance. Ablations show that a randomly initialized encoder combined with Fréchet aggregation already surpasses Euclidean pooling on two of the three signal-bearing datasets, localizing the source of the gain to geometric aggregation rather than learned manifold structure; the trained encoder contributes additional signal specifically on CREAK, the most knowledge-heavy of the three signal-bearing datasets.
Blogger's Review: This paper successfully reveals the geometric properties of pre-trained language model embeddings through the introduction of Riemannian Mean Pooling (RMP), demonstrating the potential of geometric aggregation in enhancing model performance. This research offers a new perspective on the interpretability and safety of language models, warranting further exploration.