Riemannian Geometry for Pre-trained Language Model Embeddings

This paper investigates whether sentence-level classification signal is embedded within the Riemannian geometry of contextual token representations. Traditional NLP models typically treat embeddings as points in flat Euclidean space, ignoring potential non-Euclidean structures. The authors propose a method to extract local geometric information from pre-trained language models and test whether this information improves performance on downstream classification tasks.

The researchers introduce Riemannian Mean Pooling (RMP). First, they train an encoder using Intrinsic Green’s Learning (IGL) to map token embeddings into a latent space. They then compute the Jacobian of this encoder to derive a pullback metric, which describes how the encoder locally stretches or compresses the input space. These per-token metrics are aggregated into a single sentence-level descriptor using the Fréchet mean on the symmetric positive definite (SPD) manifold. Finally, the resulting descriptors are projected into a tangent space for classification using standard linear models.

RMP consistently outperforms Euclidean mean pooling and standard [CLS] token aggregation on datasets like CoLA, CREAK, and RTE. To validate the method, the authors used FEVER-Symmetric as a negative control, where the method correctly performed at chance levels, confirming it does not rely on spurious lexical artifacts. Ablation studies revealed that even a randomly initialized encoder combined with Fréchet aggregation outperforms Euclidean pooling, suggesting that the geometric aggregation process is the primary source of the performance gain. The trained encoder provides additional benefits specifically on knowledge-heavy tasks like CREAK.

This work demonstrates that the internal geometry of language model embeddings contains valuable, underutilized information for downstream tasks. By moving beyond flat Euclidean assumptions, researchers can better interpret model representations and potentially improve performance on tasks that require capturing complex, non-linear relationships in language data.