English

Riemannian Geometry for Pre-trained Language Model Embeddings

Computation and Language 2026-07-08 v1 Artificial Intelligence

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}
}