Riemannian Mean Pooling: Why Euclidean Space Might Be Oversimplifying Embeddings

Sentence-level classification in NLP typically relies on pooling methods to aggregate token-level embeddings into a single vector. Most of these methods assume a flat Euclidean space, but research into Riemannian geometry suggests that the signals we're looking for might not live there. By investigating the geometry of contextual token embeddings, researchers have identified that moving from Euclidean mean pooling to Riemannian Mean Pooling (RMP) provides a consistent performance boost across several non-trivial linguistic datasets.

Moving Beyond Flat Averages

The core of this approach is the shift from standard averaging to Fréchet mean aggregation on the symmetric positive definite (SPD) manifold. Rather than just summing vectors and dividing by count, RMP extracts per-token pullback metrics from a learned encoder’s Jacobian. This allows the model to account for the underlying curvature of the embedding space rather than treating it as a simple, linear coordinate system.

The results are specific: RMP consistently outperforms Euclidean mean pooling on three datasets characterized by non-trivial linguistic structure: CoLA, CREAK, and RTE. The ablation studies provided in the research are particularly telling because they isolate the source of the improvement. The data indicates that the geometric aggregation itself—the way we combine the information—is the primary driver of the performance gains, rather than just the choice of encoder.

Filtering Noise from Actual Signal

A common pitfall in these types of geometric studies is the risk of the model simply picking up on lexical artifacts or annotation biases. To address this, the researchers used FEVER-Symmetric as a negative control. This dataset is specifically designed to strip away those artifacts, and the RMP method correctly stayed at chance. This is a critical piece of evidence for anyone skeptical of "fancy" geometry; it suggests that RMP is actually capturing linguistic signal rather than just exploiting shortcuts in the data.

Furthermore, the performance gains on the CREAK dataset are noteworthy. CREAK is identified as the most knowledge-heavy of the tested datasets, and the researchers found that a trained encoder provides additional signal specifically for this task. This suggests that while the geometry handles the aggregation, the encoder still needs to do the heavy lifting for high-level knowledge retrieval.

The Practicality of Manifold-Based Inference

From a practitioner’s standpoint, the real story here is about the limitations of our current "flat" assumptions. We use Euclidean mean pooling because it’s computationally cheap and easy to implement, but these results suggest we are essentially smoothing over the very nuances that make complex language difficult to process.

However, we have to be realistic about the overhead. Moving to Riemannian manifolds involves more complex math and potentially higher computational costs during the pooling step. While RMP shows a clear advantage on datasets like CREAK, the leap from a simple average to Fréchet mean aggregation on an SPD manifold isn't free. For production systems where latency is king, the question remains: is the marginal gain on knowledge-heavy tasks significant enough to justify the move away from standard Euclidean operations? The research confirms the signal exists, but it doesn't yet provide the cost-benefit analysis required to swap out standard pipelines at scale.