Riemannian Geometry for Pre-trained Language Model Embeddings

背景

预训练语言模型（PLM）的嵌入表示（embeddings）是理解模型行为与决策的基础。尽管这些嵌入在高维欧几里得空间中通常被直接用于分类或相似度计算，但越来越多的研究表明，语言的内在结构可能并非欧几里得式的。例如，层级关系、语义相似性等往往更自然地嵌入在非欧几里得几何中，尤其是黎曼几何（Riemannian Geometry）框架下。然而，对于句子级别的分类信号是否存在于上下文 token 嵌入的黎曼几何结构中，此前缺乏系统性探究。该论文正是针对这一缺口，提出了一种基于黎曼几何的池化方法，并验证其有效性。

核心内容

论文标题为“Riemannian Geometry for Pre-trained Language Model Embeddings”，作者来自 arXiv cs.CL 领域，提交于 2026 年 7 月 8 日。该研究旨在回答一个核心问题：句子级别的分类信号是否存在于预训练语言模型上下文 token 嵌入的黎曼几何结构中。

为了探测这一假设，作者提出了一种名为 Riemannian Mean Pooling (RMP) 的方法。其具体步骤为：首先，从一个已学习的编码器（encoder）中提取每个 token 的解析雅可比矩阵（analytical Jacobian），进而得到每个 token 的拉回度量（pullback metric）；然后，将这些度量视为对称正定（SPD）流形上的点，并在该流形上通过 Fréchet 均值（Fréchet mean）进行聚合，得到句子级别的表示。这一过程将传统的欧几里得池化（如平均池化）替换为黎曼几何下的聚合。

实验在三个具有非平凡语言结构的数据集上进行：CoLA（语言可接受性判断）、CREAK（常识推理与知识）和 RTE（文本蕴含识别）。结果显示，RMP 在这三个数据集上的表现均优于欧几里得平均池化（Euclidean mean pooling）。此外，论文还使用了 FEVER-Symmetric 数据集——该基准经过专门构造，去除了由标注驱动的词汇假象（lexical artifacts）。在 FEVER-Symmetric 上，RMP 的表现正确地位于随机水平（chance level），说明该方法不会受虚假特征的干扰，具有良好的鲁棒性。

为了进一步剖析性能提升的来源，作者进行了消融实验（ablations）：他们将一个随机初始化的编码器（而非预训练编码器）与 Fréchet 聚合结合，结果发现，在三个含信号的数据集中，有两个数据集（CoLA 和 RTE）上该组合就已经优于欧几里得池化。这表明，性能增益的主要来源是几何聚合本身（即 Fréchet 均值）而非预训练模型学习的流形结构。而在第三个数据集 CREAK（知识密集型）上，经过训练的编码器（含有预训练知识）贡献了额外的信号，表明在需要丰富常识知识的任务中，预训练编码器的黎曼结构进一步增强了分类能力。

关键要点

• 提出 Riemannian Mean Pooling (RMP)，利用每个 token 的拉回度量（来自编码器解析雅可比矩阵）在 SPD 流形上通过 Fréchet 均值聚合，替代欧几里得平均池化。
• 在 CoLA、CREAK、RTE 三个非平凡数据集上，RMP 均优于欧几里得平均池化。
• 在去除了词汇假象的 FEVER-Symmetric 基准上，RMP 表现恰为随机水平，表明其不会过拟合虚假特征。
• 消融实验揭示：仅使用随机初始化的编码器 + Fréchet 聚合即可在两个信号数据集上超越欧几里得池化，说明几何聚合本身是性能提升的主要原因。
• 在知识密集型任务 CREAK 上，预训练编码器贡献了额外信号，表明训练好的模型能够提供更丰富的黎曼几何结构。
• 该方法为理解 PLM 嵌入的几何结构提供了新视角，并可能增强模型的可解释性与安全性。

意义与影响

该研究首次系统性地将句子级分类信号与上下文 token 嵌入的黎曼几何联系起来，并提出了一个简单有效的池化方法 RMP。其意义体现在以下几个方面：

• 理论层面：验证了 PLM 嵌入中确实存在与分类任务相关的黎曼几何结构，且这种结构可以通过拉回度量和 Fréchet 均值被有效利用。这为理解语言模型的内在表示空间提供了几何学视角。
• 方法层面：RMP 不需要对模型架构进行大幅度修改，仅替换池化层即可提升性能，具有轻量级和通用性。同时，在无虚假特征的数据集上保持随机水平，说明该方法对数据集中的偏差具有天然鲁棒性。
• 可解释性与安全：通过几何聚合，模型可能更关注语义结构而非表面统计关联，这有助于提升模型在安全关键应用中的可靠性。后续工作可进一步探索如何利用该几何框架进行对抗性检测或知识注入。
• 未来方向：论文的消融结果暗示，即使不使用预训练编码器，仅靠几何聚合也能带来收益，这提示我们可探索更高效的黎曼池化策略，或将其推广到其他类型的序列建模任务（如文本生成、问答等）。此外，结合黎曼优化和几何学习的全端到端训练也值得关注。