用于从图像序列端到端认知地图学习的可微克隆结构因果图

# 可微克隆结构因果图：从图像序列进行端到端认知地图学习
来源：https://arxiv.org/html/2607.12382
可微克隆结构因果图：从图像序列进行端到端认知地图学习

Arash Nikzad¹, Sasan Sarbishegi², Ali Dasmeh³, Muhammad Asif⁴, Parsa Gharavi¹, Erik Husom⁵, Sagar Sen⁵, Andrew B. Lehr⁶,⁷,¹¹, Olivier Penacchio⁸,⁹, Ana Clemente⁴, Tristan M. Stöber⁷,¹⁰,¹¹,*

¹歌德大学法兰克福，德国法兰克福；²独立研究员，伊朗德黑兰；³马克斯·普朗克人类发展研究所，德国柏林；⁴认知神经心理学系，马克斯·普朗克经验美学研究所，德国法兰克福；⁵SINTEF，挪威奥斯陆；⁶神经与感觉生理学系，哥廷根大学医学中心，德国哥廷根；⁷计算机科学研究所及校园数据科学研究所，哥廷根大学，德国；⁸人工智能与神经科学桥梁研究（brAIN），计算机视觉中心，西班牙贝拉特拉；⁹计算机科学系，巴塞罗那自治大学，西班牙贝拉特拉；¹⁰法兰克福莱茵-美因癫痫中心，神经病学系，歌德大学法兰克福，德国法兰克福；¹¹Circulant Labs，德国本斯海姆。 *通讯作者：[email protected]

**摘要**。一个智能体如何仅从源源不断的原始感觉输入及其自身运动中构建出世界的结构化地图，尤其是在自然变异导致精确感觉模式很少重复的情况下？克隆结构因果图算法（CSCG），一种规范的海马体模型，展示了如何从混叠的观测中学习可解释的地图。然而，CSCG 需要一个预定义的离散字母表，并且其期望最大化公式不易与现有神经网络模块结合，从而阻碍了原始图像序列的端到端处理。我们通过将 CSCG 重新表述为一个单一的、完全可微的模块 gradCSCG，并将其与一个学习的矢量量化变分自编码器（VQ-VAE）感知前端耦合，消除了这一障碍。一种软发射前向传递允许地图学习目标流回感知，同时一组损失平衡机制减轻了联合训练期间的模块崩溃。我们首先证明，梯度训练通过在高度混叠的观测中恢复房间拓扑，在原始符号网格世界上复现了 CSCG 的结果。其次，我们展示了在 MNIST 图像序列上地图恢复仍然稳健，其中每次访问一个位置都会得到该位置对应数字的新采样图像。在四个高度混叠的环境中，端到端流水线成功地从视觉输入中直接发现了底层的邻接图，具有高边精确度和召回率。这项工作提供了一个原理证明，即 CSCG 可以作为深度学习架构中的一个可组合构建块。

**关键词**：认知地图、克隆结构认知图（CSCG）、可微序列模型、矢量量化表示学习、海马体、NeuroAI、拓扑恢复。

## 1 引言

良好结构化的内部表示使得生物和人工智能体能够在从未走过的路线上找到捷径，并在试错学习可能致命的情况下提供指导。因此，理解和重新设计这类表示的出现，是神经科学 [1–5] 和人工智能（AI）[6–8] 领域的一个基本研究前沿。

克隆结构因果图（CSCG）算法 [9, 10]，一种规范的海马体模型，解释了良好结构化的表示如何从经验中涌现。从技术上讲，CSCG 是一个过完备隐马尔可夫模型（HMM），具有一个固定的发射矩阵，该矩阵从序列观测中创建一个统计模型。将一系列观测-动作对压缩成更高阶的环境表示，CSCG 学习从混叠的观测中去混淆上下文。满足马尔可夫性质——即任何后续状态仅依赖于当前状态——CSCG 被迫通过在其隐藏节点中创建一个新的克隆来表示一个新的上下文。从这个克隆操作中产生的图提供了一个环境的压缩地图，并适用于规划、巩固和抽象。然而，虽然这种模型在静态且相对较小的环境中优雅地创建了良好结构化的表示，但如何将此方法扩展到更丰富、感知上更复杂的观测尚不清楚。瓶颈在于 CSCG 是通过在固定的离散观测字母表上进行期望最大化训练的，这阻碍了与梯度训练的神经模块的无缝组合。

我们通过将 CSCG 重新实现为一个单一的、可微的、基于梯度训练的计算 [11]（在 TensorFlow 中）——我们称之为 gradCSCG 模型——来解决这个问题。这一创新使我们能够将 gradCSCG 与一个矢量量化变分自编码器（VQ-VAE）共同训练。我们证明，这种方法保留了 CSCG 的表达能力，同时使其能够处理由 MNIST 数字组成的环境中的感觉变异性和复杂性。

**贡献**

1. 1. 一个**端到端可训练流水线**，由一个基于梯度的 CSCG 和一个 VQ-VAE 组成，能够从图像序列创建拓扑图。
2. 2. 用于稳定联合训练的**损失平衡**——长度归一化、权重退火、多样性惩罚和防崩溃措施（第 3.6 节）。
3. 3. 一个正式的、可复用的**拓扑恢复评估套件**（第 3.10 节）。
4. 4. 对四个具有强混叠性的 MNIST 网格世界环境的**实证研究**（第 4–5 节）。

### 1.1 相关工作

**CSCG 工具包。** 一个具有软观测接口和编码器梯度流的可微克隆-HMM 前向传递也已在并发、独立的开源工作中得到开发 [21]，其图像实验将一个卷积网络与一个基于梯度的 CSCG 耦合用于监督数字分类。相对于这项工作，我们的具体贡献是学习的 VQ-VAE 离散化器与序列模型联合训练、保持联合训练稳定的损失平衡，以及——至关重要的是——使用这种基于梯度的设置来**实际恢复和评估环境的地图**：我们从学习的转移中推导出物理邻接图，并根据真实拓扑对其进行评分（第 5 节）。相比之下，他们的图像实验仅将卷积-CSCG 耦合用于监督数字分类，他们没有尝试恢复拓扑，而这正是我们流水线旨在解决的核心问题。

**用于认知地图的神经序列模型。** 一条互补的研究路线从神经预测器的潜在编码中学习地图，而不是从克隆 HMM 中学习。Dedieu 等人 [22] 训练一个带有离散瓶颈的 Transformer 执行下一观测预测，并读取其瓶颈索引中的可解释认知地图，用于部分观测环境中的规划。与我们的流水线一样，神经表示被离散化，并从观测-动作流中恢复地图；不同之处在于**地图所在位置**。在他们的方法中，认知地图作为单独的事后分析从 Transformer 的瓶颈编码中解码出来，并交给外部求解器进行规划，因此网络自身的工作表示保持密集，本身并不是一个可解释的、可直接使用的地图。在我们的方法中，克隆图序列模型与感知离散化器共同训练，其转移矩阵**就是**地图：一个内在可解释的、在线的结构，模型用它进行计算，并且可以直接查询以进行规划和理解环境。

**动态和扩展地图。** 更接近我们在第 7 节中提出的开放问题，de Tinguy 等人 [23] 在线增长一个认知地图——在主动推理智能体中根据预测姿态动态扩展——并在网格环境上与 CSCG 进行基准测试。他们专注于地图在智能体探索时**扩展**的能力，正是我们的固定克隆预算所缺乏的动态分配能力。然而，他们的模型通过不可微的离散计数操作进行学习，这使得难以与深度神经网络无缝耦合以进行视觉处理。相比之下，我们的完全可微公式允许地图与神经前端共同训练，直接从原始像素端到端地学习离散化。

## 2 方法

### 2.1 问题形式化

一个智能体产生一个长度为 T 的回合：观测 x_{1:T}，其中 x_t ∈ X ⊂ R^{H×W×C}，以及动作 a_{1:T-1}，其中 a_t ∈ A = {1, ..., A}，a_t 在时间步 t 和 t+1 之间发生。每个观测都在一个底层物理位置 g_t ∈ G 处发射；环境有一个真实的无向邻接图 M = (G, E)。位置 g_t 和边 E **仅用于评估**，在训练期间从未见过。目标是从 (x_{1:T}, a_{1:T-1}) 单独学习一个潜在模型，其转移结构能够恢复 M。

该流水线（图 1）包含两个模块：一个 VQ-VAE，它将每个图像映射为一个离散标记；以及一个作用在这些标记上的、以动作为条件的克隆 HMM，其潜在图即是学习到的地图。

图像 x_t → 编码器 E_φ → 码本 {e_k} → 解码器 D_ψ → L_rec → 软后验 ρ_t (对 K 个码字) → gradCSCG 前向（软） → L_gradCSCG → z_t → q̃_t → log ρ_t

图 1：gradCSCG 流水线。**实线箭头**：前向计算。**虚线箭头**：梯度流。编码器同时馈送一个重建分支（硬量化 q̃_t + 解码器）和一个序列分支（软码本后验 ρ_t + 可微克隆-HMM 前向传递）。由于 gradCSCG 似然在 ρ_t 上是可微的，拓扑目标 L_gradCSCG 塑造了编码器。码本本身通过指数移动平均（EMA）更新，而不是通过梯度。

### 2.2 感知前端：VQ-VAE

**编码器和量化。** 一个卷积编码器 E_φ: X → R^D 将每个图像映射为潜在向量 z_t = E_φ(x_t)。一个码本 {e_k}_{k=1}^K，e_k ∈ R^D，其中 K 表示离散潜在码字的数量，D 是每个码本嵌入的维度，通过最近邻分配定义一个离散标记，

k_t = arg min_{k∈{1,...,K}} ∥z_t - e_k∥_2^2, q_t = e_{k_t},

其中平方距离计算为 ∥z - e_k∥_2^2 = ∥z∥_2^2 - 2⟨z, e_k⟩ + ∥e_k∥_2^2。梯度通过直通估计器 [18] 穿越非可微的 arg min，

q̃_t = z_t + sg(q_t - z_t),

其中 sg(·) 是停止梯度算子，因此前向值为 q_t，而 ∂q̃_t/∂z_t = I。一个解码器 D_ψ 重建 x̂_t = D_ψ(q̃_t)。

**损失。** 在一个小批次 B 上，

L_rec = (1/|B|) ∑_t ∥x_t - x̂_t∥_2^2,
L_commit = (1/|B|) ∑_t ∥sg(q_t) - z_t∥_2^2,

第二个损失将编码器输出拉向其分配的码字。

**EMA 码本更新。** 码本**不**通过梯度下降训练。使用衰减因子 γ ∈ (0,1)，每个批次我们为每个码字 k 累积聚类大小 n_k 和向量和 m_k，

n_k ← γ n_k + (1-γ) ∑_t 1[k_t = k],
m_k ← γ m_k + (1-γ) ∑_t 1[k_t = k] z_t,

并设置 e_k ← m_k / n̂_k，其中拉普拉斯平滑后的大小为

n̂_k = (n_k + ε) / (∑_{k'} n_{k'} + Kε) * (∑_{k'} n_{k'})。

**软码本后验。** 为了实现可微耦合（第 3.5 节），编码器还发射一个受温度控制的后验分布，作用于码本：

log ρ_t(k) = log softmax_k( -∥z_t - e_k∥_2^2 / τ ),

这在 z_t 上是可微的。当 τ → 0 时，ρ_t 集中在 k_t 上，这恢复了硬分配。

### 2.3 序列模型：以动作为条件的克隆 HMM

**状态空间和克隆结构。** 每个标记 k 被分配 C_k ≥ 1 个克隆——这些潜在状态都发射标记 k，但参与不同的转移上下文。加上一个尾随的汇点状态 ⊥，状态空间为 S = {1, ..., N}，其中 N = 1 + ∑_{k=1}^K C_k。汇点状态是一个特殊的终止状态，不对应于任何视觉标记；它用于在序列结束时吸收概率质量，并使序列终止在 HMM 公式中显式化。一个固定映射 ω: S \ {⊥} → {1, ..., K} 给出每个状态发射的标记；在均匀情况 C_k ≡ C 下，ω(s) = ⌈s/C⌉。发射是**确定性的**：

B_{s,o} = 1[ω(s) = o], log B_{s,o} = {0, ω(s)=o; -∞, 否则},

且汇点不发射真实标记。克隆正是用于消解混叠的机制：在两个位置观测到的同一个标记由两个具有不同转移行的克隆来解释。

**参数。** 模型具有初始状态 logits π ∈ R^N 和以动作为条件的转移 logits Θ ∈ R^{A×N×N}，产生

π̄ = softmax(π), T_{a,i,j} = softmax_j(Θ_{a,i,·})。

所有学习都驻留在 (π, Θ) 中；发射是固定的。

**前向似然。** 记 logsumexp_i u_i = log ∑_i e^{u_i}，一个回合 (o_{1:T}, a_{1:T-1}) 的对数前向消息满足

log α_1(j) = log π̄_j + log B_{j, o_1},
log α_{t+1}(j) = logsumexp_i [ log α_t(i) + log T_{a_t, i, j} ] + log B_{j, o_{t+1}},

并且回合的对数似然为 l(o_{1:T} | a_{1:T-1}) = logsumexp_j log α_T(j)。训练损失是平均负对数似然（