Emergent Ordinal Geometry in Transformers Trained on Local Comparisons

技术摘要：在基于局部比较训练的 Transformer 中涌现的序数几何结构

问题陈述
本文研究了 Transformer 模型在仅接受隐藏全序关系中局部（相邻）比较训练的情况下，是否会获得与生物系统用于传递推理（Transitive Inference, TI）相同的原始机制。TI 是指从已知的相邻关系（如 $A < B, B < C$）中推断非相邻关系（如 $A < C$）的能力。虽然生物认知通过“模拟心理数轴”来解决这一问题——这由符号距离效应（即远距离比较比相邻比较处理得更快/更准确）所证实——但 Transformer 并没有受到显式的线性表示压力。本研究探讨了这种特定的几何表示及其行为特征是否会在针对经典的 Bryant 和 Trabasso [1971] 范式进行训练时自发产生。

方法论
作者构建了一个包含 $N$ 个实体（$e_0 \prec e_1 \prec \dots \prec e_{N-1}$）的隐藏全序关系的合成任务。

• 训练数据： 模型仅在双向的相邻对上进行训练（$D_{train} = \{(e_i, e_{i+1}), (e_{i+1}, e_i)\}$）。
• 评估数据： 模型在不相交的远距离对集合上进行评估（$D_{OOD} = \{(e_i, e_j) : |i-j| > 1\}$），要求模型仅通过局部数据重建全局顺序。
• 模型架构： 一个单层 Transformer，包含实体嵌入、一个注意力头（在抽象部分）或两个（在结果部分）、一个 4d 宽度的 GELU MLP 以及一个线性解嵌入头（unembedding head）。关系从最终位置的残差中解码。
• 指标：
  1. 几何结构： 对实体嵌入进行主成分分析（PCA），以测量线性度（第一主成分上的方差 $\rho_1$）和秩对齐度（投影位置与真实秩之间的相关性）。
  2. 行为特征： 测量决策置信度（logit 间隔）和几何分离度（$g_{ij}$）相对于秩距离（$|i-j|$）的变化，以检测符号距离效应。

核心贡献

1. 合成任务隔离： 本文引入了一个受控任务，将序数泛化从局部比较中隔离出来，严格模拟了生物学中的 TI 范式。
2. 涌现的 1D 流形： 研究提供了证据，表明 Transformer 会自发地将实体嵌入组织成一个一维流形。该流形的主轴能近乎完美地恢复隐藏的秩顺序。
3. 距离效应的计算类比： 研究表明，即使在准确率达到饱和时，模型也会表现出符号距离效应：决策置信度和几何分离度均随项目间的秩距离（$|i-j|$）单调缩放。

结果
实验在 $N=10, 15$ 和 $20$ 个实体的集合上进行。

• 类 Grokking 动力学： 分布外（OOD）泛化的出现与嵌入的几何重组相一致。这种转变对优化过程（特别是权重衰减）非常敏感，并产生了类似于“grokking”的瞬态动力学。
• 几何恢复： 随着模型的学习，嵌入坍缩到一条线上，其主成分与真实的秩顺序对齐。
• 行为镜像： 在峰值 OOD 准确率下，模型的内部表示和输出置信度严格遵循符号距离效应。与相邻对相比，远距离对不仅被正确分类，而且在嵌入空间中具有更高的置信度和更大的几何分离度。

意义
本文声称将一个存在了 50 年之久的行为规律（TI 中的符号距离效应）建立在学习到的神经表示几何学之上。通过证明 Transformer 在没有显式监督的情况下也能自发发展出“心理数轴”，这项工作为传递推理提供了一个机械论解释，架起了认知科学与现代神经网络可解释性之间的桥梁。它表明，执行抽象关系推理的能力可能是优化全序内局部约束的一个自然副产品。