Riemannian Geometry for Pre-trained Language Model Embeddings
For researchers in interpretability and safety of pre-trained language models, this work provides a novel geometric perspective on embedding aggregation, though the gains are modest and partially attributable to the aggregation method rather than learned manifold structure.
The paper investigates whether sentence-level classification signal resides in the Riemannian geometry of contextual token embeddings. By extracting per-token pullback metrics and aggregating them via Fréchet mean on the SPD manifold (Riemannian Mean Pooling), they achieve improved performance over Euclidean mean pooling on three datasets (CoLA, CREAK, RTE), while remaining at chance on a debiased benchmark (FEVER-Symmetric).
Understanding the geometric structure of pre-trained language model embeddings matters for interpretability and safety. We ask whether sentence-level classification signal lives in the Riemannian geometry of contextual token embeddings, and probe it 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; we call this procedure 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 constructed to remove annotation-driven lexical artifacts, the method correctly stays at chance. Ablations show that a randomly initialised encoder combined with Fréchet aggregation already beats Euclidean pooling on two of the three signal-bearing datasets, localising the source of the gain to the geometric aggregation rather than to learned manifold structure; the trained encoder contributes additional signal specifically on CREAK, the most knowledge-heavy of the three signal-bearing datasets.