🤖 machine learning

MapFormer: Self-Supervised Learning of Cognitive Maps with Input-Dependent Positional Embeddings

本文介绍了 MapFormer，这是一种新颖的基于 Transformer 的架构，它利用源自李代数生成元的输入依赖型位置嵌入来学习无监督认知地图，从而在多种认知任务中相比现有 AI 模型实现了更优越的分布外泛化能力和可扩展性。

原作者： Victor Rambaud, Salvador Mascarenhas, Yair Lakretz

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Victor Rambaud, Salvador Mascarenhas, Yair Lakretz

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你正在一座新城市中行走。你并非死记硬背所见到的每一栋建筑，而是构建一张心理地图。你理解“向左转”总会改变你相对于街道的位置，无论你在巴黎还是东京。这种将“你在何处”（结构）与“你看到什么”（内容）区分开来的能力，正是科学家所称的“认知地图”。

当前的 AI 系统，例如你可能正在与之对话的大型语言模型，在记忆模式方面表现出色。然而，它们往往非常“脆弱”。如果你要求它们解决一个与训练数据略有不同的问题，它们往往会失败。它们难以理解情境背后的规则，而是依赖基于表面相似性的猜测。

本文介绍了一种名为MapFormers的新型 AI 架构。可以将 MapFormers 想象为赋予 AI 一个“心理 GPS"，使其学习道路规则，而不仅仅是记忆风景。

核心理念：将“何处”与“何物”分离

作者认为，要变得真正智能，AI 需要解开两件事：

内容：它看到的具体事物（例如，一个红苹果，一辆蓝车）。
结构：事物如何移动或相互关联的规则（例如，“向右移动”总是让你向右移动一步）。

标准 AI 模型将这两者混合在一起。MapFormers 的设计旨在将它们保持分离。它们利用涉及李群（一种描述平滑旋转和运动的复杂方式）的数学技巧来实现这一点。AI 不是硬编码规则，而是学习根据输入生成“移动矩阵”。

两种类型的 MapFormers

本文展示了该系统的两个版本，作者将其比作人类的记忆类型：

MapEM（情景记忆）：这就像一本日记。它保存一份独立的、专门的清单，记录“我在哪里”和“我看到了什么”。它非常擅长回忆特定的过去事件（例如，准确记得上周二早餐吃了什么），但处理速度稍慢。
MapWM（工作记忆）：这就像大脑的活跃草稿纸。它即时混合“何处”与“何物”。它更快、更高效，类似于 RoPE（一种流行的当前 AI 技术）的工作原理，但有一个关键升级：它实际上能够学习在地图中动态移动。

它们是如何学习的（训练场）

为了证明这些模型有效，研究人员在三个特定的“电子游戏”风格挑战中测试了它们，这些挑战中的规则是隐藏的：

“忽略噪音”游戏：AI 必须复制一个物品列表，但忽略中间的空白。标准 AI 会被空白搞糊涂；MapFormers 学会了“门控”（屏蔽）噪音，只专注于要复制的物品。
“蒙眼导航员”游戏：AI 收到一系列“移动”和“看见”指令，但未被告知哪个是哪个。它必须弄清楚“上”会让你向上移动，而“看见一棵树”只是更新你的视野。一旦它学会了地图，即使序列比它见过的任何序列都要长，它也能准确预测如果回到曾经访问过的地点会看到什么。
“嵌套括号”游戏：这测试了处理深层逻辑层的能力（如 ((()))）。当序列变长时，标准 AI 难以计算有多少个等待闭合的左括号。MapFormers 将打开括号视为“向前移动”，将关闭括号视为“向后移动”，从而使它们能够完美地跟踪堆栈，即使序列深度远超其训练范围。

结果：为何这很重要

本文声称，MapFormers 实现了近乎完美的泛化能力。

"OOD"测试：在 AI 术语中，“分布外”（OOD）意味着在模型从未见过的场景上测试模型。当序列变长或环境改变时，标准模型惨败，而 MapFormers 则成功了。它们不仅仅是记忆；它们学习了问题的几何结构。
现实世界测试：研究人员还将此应用于一段真实的互联网文本（OpenWebText）。虽然在逻辑游戏上的提升幅度较小，但 MapFormers 的表现仍略优于标准模型，这表明这些原则可能扩展到现实世界的语言任务中。

“秘密武器”：数学作为指南针

本文解释说，MapFormers 使用一种特定的数学方法，将动作（如“向右移动”）视为旋转。

想象你在旋转一个拨盘。如果你转动 90 度，再转动 90 度，最终会到达 180 度。
MapFormers 学习到“右”是一种特定的旋转，而“左”是相反的旋转。
因为它们使用这种数学框架，所以可以通过简单地将角度相加来计算长旅程（路径积分）的结果，而不必为每一步进行复杂、缓慢的计算。这使得它们能够并行处理信息（非常快），同时仍能理解序列。

总结

简而言之，本文提出，为了使 AI 更加稳健和适应性强，我们应该停止将其视为巨大的模式匹配机器，转而赋予它一张认知地图。通过教导 AI 将移动规则与其所见事物分离开来，MapFormers 能够以当前 AI 系统所缺乏的灵活性，导航新的、未见过的情况。它们不仅仅是猜测；它们理解地图。

技术摘要：MapFormers

问题陈述

当前的基于 Transformer 的人工智能系统，尽管在分布内任务上表现出令人印象深刻的性能，但在面对分布外（OOD）场景时却表现出显著的脆弱性。它们往往无法泛化到人类和动物能够轻松处理的新颖情境中。作者将这一差距归因于缺乏结构性归纳偏置，而生物大脑正是利用这种偏置来构建认知地图：即能够解耦环境抽象不变结构与具体可变内容的内部模型。

现有的认知地图建模尝试，例如在循环神经网络（RNN）中使用动作依赖矩阵的 Tolman-Eichenbaum 机器（TEM），存在两个关键局限性：

缺乏可扩展性：它们需要顺序矩阵乘法来更新位置，从而阻碍了并行处理。
预定义动作：它们通常要求动作集合（例如基本方向）被硬编码，而不是从数据中学习。

此外，标准的位置编码（如 RoPE）是静态的，或者依赖于不支持未知环境中路径积分所需动态动作组合的启发式方法。

方法论

本文介绍了MapFormers，这是一类旨在通过无监督的下一词元预测来学习认知地图并执行路径积分的 Transformer 架构。其核心创新在于利用输入依赖矩阵来更新位置编码，这些矩阵源自李群理论。

理论框架

李群形式化：作者将动作建模为李群元素（特别是像 $SO(2)$ 这样的紧致群，用于旋转）。模型不是为每个可能的动作学习单独的矩阵，而是学习李代数生成元。
通过指数映射进行路径积分：动作被表示为李代数中的无穷小生成元。路径积分（组合一系列动作）通过在代数中求和这些生成元（线性累积）来实现，然后应用单个指数映射以获得李群中的有限变换。
- 对于交换映射（例如二维网格导航），这允许通过累积求和随后进行单个矩阵指数运算，以并行方式计算路径积分。
- 这种方法将归纳偏置从动作数量减少到了代数的结构复杂性。

架构变体

本文提出了 MapFormers 的两种变体，将绝对和相对位置编码与不同的记忆模型统一起来：

MapEM（情景记忆）：
- 使用绝对位置嵌入来表示认知地图。
- 将内容（观测值）和结构（位置）解耦到独立的神经群体中。
- 在内容与位置的合取上计算注意力（ $Q_G = \text{vec}(Q^T \cdot P)$ ），有效地充当注意力掩码，其中位置编码独立于内容更新。
- 适用于需要稳健回忆特定状态的任务。
MapWM（工作记忆）：
- 使用相对位置编码（类似于 RoPE 但是动态的）。
- 通过输入依赖的旋转矩阵（ $R_{\theta PI}$ ）将神经活动（键和查询）旋转到不同的子空间，从而隐式地表示结构。
- 旋转角度 $\theta$ 源自输入词元的低秩投影，这使得模型能够学习哪些词元充当动作（更新位置），哪些是观测值（保持位置不变）。

输入依赖机制

两种架构都采用两阶段投影来确定每个词元的“积分持续时间”（ $\Delta_t$ ）：

内部投影：将输入词元映射到低维动作表示。
外部投影：将该表示投影到每个块的积分持续时间。
最终的旋转角度为 $\theta_t = \omega \odot \Delta_t$ ，其中 $\omega$ 是一个可学习的角速度向量。路径积分计算为这些角度的累积和（ $\theta_{0 \to t} = \text{cumsum}(\theta_t)$ ）。

主要贡献

MapFormers 架构：一种新颖的 Transformer 变体，通过基于李代数生成元的输入依赖位置编码，将结构从内容中解耦，从而学习认知地图。
并行路径积分：通过利用交换李群（例如 $SO(2)^k$ ），模型实现了并行路径积分，克服了先前基于 RNN 的认知地图模型的顺序瓶颈。
记忆模型的统一：这项工作提供了绝对/相对位置编码与情景/工作记忆之间的理论联系，表明两者如何可以从相同的李理论原理中推导出来。
向非交换映射的可扩展性：该框架通过使用多个生成元扩展到非交换问题（例如家谱），尽管这需要顺序路径积分。

实验结果

作者在三个针对不同认知能力的形式化任务上评估了 MapFormers：

选择性复制（门控）：
- 任务：复制序列同时忽略干扰词元。
- 结果：MapFormers（包括 EM 和 WM）实现了完美的 OOD 泛化，而标准基线（RoPE, TAPE）则失败了。这证明了学习动态门控机制的能力。
网格导航（一维和二维）：
- 任务：预测交替动作和随机观测值序列中的观测值，要求模型推断潜在拓扑并路径积分位置。
- 结果：MapFormers 在 OOD 序列（长度和密度变化）上实现了近乎完美的准确率，而基线（包括 CoPE 和 PathAtt）未能泛化。
- 神经分析：模型学会将观测词元投影到零旋转（恒等变换），将动作词元投影到特定的旋转方向，有效地将内容从结构中解耦。
Dyck-2（嵌套层次结构）：
- 任务：预测括号序列的有效延续，需要跟踪栈。
- 结果：MapFormers 仅使用单层即可完美泛化到比训练数据更长、更深的序列。基线需要多层，且在 OOD 数据上仍然失败。
- 机制：模型通过使用正交维度表示不同类型的括号，并使用相反方向表示开/闭括号，将栈直接嵌入到位置编码中。
自然语言扩展：
- 任务：在 OpenWebText 上进行预训练。
- 结果：一个 12 层的 MapWM 模型显示出比 RoPE 基线一致的困惑度改进（$18.79 $对比$ 19.14$）以及更好的长度外推能力。然而，在 BLIMP 句法基准测试中，MapWM 未显示出比 RoPE 显著的增益，表明仅靠交换路径积分可能不足以进行所有句法建模。

意义与主张

本文主张，MapFormers 提供了一种关键的结构性偏置，通过路径积分捕捉潜在结构而非依赖启发式方法，从而实现了系统性的 OOD 泛化。

鲁棒性：模型证明，将抽象关系与内容解耦，可以实现对未见过的序列长度、深度和环境密度的稳健泛化。
可扩展性：在交换映射上使用并行计算使得这些原理能够扩展到大规模、现实世界的领域，自然语言结果即证明了这一点。
生物合理性：该方法弥合了生物认知地图（将“什么”与“哪里”分解）与现代深度学习之间的差距，为构建能够通过理解环境潜在几何结构来推理新颖问题的 AI 系统提供了一条途径。

作者对局限性保持谦逊，指出非交换路径积分无法并行化，并且当前的交换约束可能限制了对语言中某些复杂句法结构的建模。他们建议未来的工作可以探索混合架构或非交换扩展，以进行更复杂的关系推理。