Riemannian Geometry for Pre-trained Language Model Embeddings
Abstract
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\'echet 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\'echet 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.
Cite
@article{arxiv.2607.07047,
title = {Riemannian Geometry for Pre-trained Language Model Embeddings},
author = {Szczepan Konior and Alexandre Quemy and Przemysław Klocek and Grégoire Cattan and Bartłomiej Sobieski},
journal= {arXiv preprint arXiv:2607.07047},
year = {2026}
}