LEARNING ASSOCIATIVE INFERENCE USING FAST WEIGHT MEMORY

文章标题：使用快速权重记忆学习联想推理
作者：Imanol Schlag, Tsendsuren Munkhdalai, Jurgen Schmidhuber
机构：The Swiss AI Lab IDSIA / USI / SUPSI, Microsoft Research

A1 主要贡献

核心问题：
现代深度神经网络（NNs）尽管在许多人工智能问题上取得了成功，但在需要组合不同经验中提取的特征并进行关联（即组合泛化）的情境中表现不佳。例如，如果在训练中从未见过某个男人和某个幼儿的关联，但在测试时需要根据上下文（男人抱着幼儿）推断他们可能是父子关系，传统的神经网络模型往往会遇到困难。这种在生产性和系统性方面的缺乏是长期以来反对将神经网络作为人工智能认知架构基础的论点。

研究目标：
本文旨在解决神经网络在组合推理方面的弱点，特别是多步联想推理。研究目标是构建一个能够学习事实的状态表示，并通过组合这些表示来执行联想推理的神经网络模型。为了实现这一点，模型需要能够系统地构建所有可能的状态，以便泛化到前所未见的刺激组合。

创新点/主要贡献：
本文提出了一种新颖的神经网络模型，通过为长短期记忆（LSTM）模型增加一个称为快速权重记忆（Fast Weight Memory, FWM）的联想记忆模块来实现联想推理。

1. 快速权重记忆（FWM）：模型的核心是一个由LSTM控制的、快速变化的FWM。FWM的权重不是在训练过程中缓慢更新的常规权重，而是在输入序列的每一步都由LSTM通过一个类似赫布学习（Hebb-like）的可微机制进行更新。这种快速权重机制使得FWM能够存储和维护组合关联。
2. 组合状态表示：FWM的设计受到了张量积表示（Tensor Product Representation, TPR; 【Smolensky, 1990, Tensor product variable binding and the representation of symbolic structures in connectionist systems. Artif. Intell.】）的启发。通过使用键（key）向量的张量积来形成关联，模型可以为新颖的键组合生成唯一的表示，从而实现对未见过组合的泛化。
3. 多步推理：为了执行更复杂的推理（如传递推理 a→b, b→c, 则 a→c），FWM支持连续多次查询。每次查询检索到的值可以作为下一次查询的键，从而将多个独立的关联链接在一起。
4. 端到端训练与广泛应用：整个模型可以通过梯度下降进行端到端的训练。本文通过在三个不同领域的实验证明了该方法的通用性：
   • 监督学习：在为组合语言推理设计的更具挑战性的bAbI数据集变体（catbAbI）上取得了优异的性能。
   • 元强化学习：在部分可观察马尔可夫决策过程（POMDPs）中，模型能够通过联想推理学会探索环境。
   • 自监督学习：在小规模的词级语言建模任务上，其性能可与最先进的模型相媲美。

A2 方法细节

我们的FWM是一个由变化缓慢的非线性LSTM控制的、快速变化的多线性映射。LSTM的慢速权重是常规的神经网络权重，在训练期间通过梯度下降进行更新。相比之下，FWM的快速权重在输入序列的每一步都由LSTM通过一个类似赫布学习的可微机制进行更新。这使得FWM函数即使在测试期间也能迅速改变——因此得名“快速权重”。除了更新快速权重，LSTM还生成一个内存查询，用于检索先前存储的信息。检索到的信息随后成为模型输出的一部分。

3.1 快速权重记忆（THE FAST WEIGHT MEMORY）

任务定义与模型结构。给定一个来自词汇表V的词元序列 $x = (x_1, ..., x_T)$，语言建模的任务是训练一个模型来最大化联合概率 $p(x)$，我们将其自回归地分解为 $p(x_{1:T}) = \prod_{t=1}^{T} p(x_t|x_{0:t-1})$，其中 $x_0$ 是一个人工的起始词元。在这项工作中，我们训练一个RNN模型，将输入序列 $x_{1:t}$ 编码为LSTM的隐藏状态 $h_t$ 和FWM的快速权重张量 $F_t$，以最大化下一个词元 $x_{t+1}$ 的概率。

模型计算流程。在输入序列的第 $t$ 步，输入词元 $x_t$ 通过一个查找表嵌入到一个 $d_E$ 维的向量空间中，得到 $e_t = \text{embedding}(x_t)$。一个拥有 $d_{LSTM}$ 个隐藏单元的LSTM将嵌入词元的序列编码成一个固定大小的向量表示 $h_t = \text{LSTM}(e_t, h_{t-1})$。下一个词元的概率分布为 $\hat{x}_{t+1} = \text{softmax}(W^{(s)}(h_t + \text{FWM}(h_t, F_t)))$，其中 $F_t \in \mathbb{R}^{d_{FWM} \times d_{FWM}^2}$ 是FWM在第 $t$ 步的快速权重，$W^{(s)} \in \mathbb{R}^{|V| \times d_{LSTM}}$。需要注意的是，快速权重矩阵 $F_t$ 是一个重塑后的三阶张量 $F_t \in \mathbb{R}^{d_{FWM} \times d_{FWM} \times d_{FWM}}$。这使我们能够使用矩阵乘法来描述三阶张量操作。接下来我们将详细描述FWM函数及其快速权重的更新方式。

3.1.1 写入（WRITING）

写入机制概述。FWM在每个时间步 $t$ 使用本节描述的写入机制进行更新。为此，我们从隐藏状态 $h_t$ 中提取：写入强度 $\beta$（一个使用sigmoid函数 $\sigma$ 限制在0和1之间的标量）、两个键向量 $k_1, k_2$ 和新的值向量 $v$。

写入目的与张量积优势。写入内存的目的是学习输入模式 $k_1 \otimes k_2$ 和输出模式 $v$ 之间特定于上下文的关联。在输入模式中使用张量积可以分解表示空间，从而保证为新颖的键对提供唯一的正交向量表示。Schlag & Schmidhuber (2018) 【Imanol Schlag and Jurgen Schmidhuber. Learning to reason with third order tensor products. In Advances in Neural Information Processing Systems (NeurIPS), 2018.】给出了一个具体例子并进行了演示，其中第一个键学习表示一个实体，第二个键学习表示一个特定的动作，从而学习到一个能够泛化到从未见过的实体和动作组合的表示空间。

更新规则。与TPR-RNN复杂的内存操作形成鲜明对比，我们采用了一个单一、简单且词级的操作，该操作与感知机学习规则（perceptron learning rule）【Rosenblatt, 1958, The perceptron: a probabilistic model for information storage and organization in the brain. Psychological review】密切相关。它允许模型用旧值和新值的凸组合 $\beta v + (1 - \beta)v_{old}$ 来替换先前的关联 $v_{old}$。通过标量 $\beta$，LSTM可以控制新的关联是完全替换旧值（$\beta=1$），还是将两者的信息混合在一起。我们的快速权重更新工作如下：首先，检索与 $k_1 \otimes k_2$ 关联的当前值 $v_{old}$。其次，通过从内存中减去 $\text{vec}(k_1 \otimes k_2) \otimes v_{old}$ 来移除旧的关联，其中 vec 将矩阵向量化。第三，我们加上 $\text{vec}(k_1 \otimes k_2) \otimes (\beta v + (1 - \beta)v_{old})$。这三个步骤可以通过以下更新规则一次性完成（证明见附录B部分）：

权重归一化。为防止快速权重可能无限制地增长，当 $||F'_{t}||_2 > 1$ 时，我们对快速权重进行缩减。这是通过以下逐元素的缩放实现的。

3.1.2 读取（READING）

单步读取。在输入序列的每一步，模型都会查询内存以检索先前存储的值。由于键和值是分开生成的，网络可以检索到与其键在信息上独立的值。

多步读取实现复杂推理。为了执行更复杂的联想推理，例如传递推理（a → b, b → c, 因此 a → c），我们采用了多次读取，其中我们将检索到的值用作下一次查询的键之一（见公式7）。

读取公式详解。在这里，LN 指的是没有学习的逐元素仿射变换的层归一化（layernorm）【Ba et al., 2016b, Layer normalization. arXiv preprint】，「vec」将矩阵重塑为向量，$\phi$ 是双曲正切函数。矩阵 $W_n, W^{(i)}_e \in \mathbb{R}^{d_{FWM} \times d_{LSTM}}$，其中 $i \in \{1..N_r\}$，以及 $W_o \in \mathbb{R}^{d_{LSTM} \times d_{FWM}}$ 是通过梯度下降训练的常规慢速权重，这使我们能够将LSTM的维度与FWM的维度解耦。在公式7中，$F_t$ 是我们用来查询的多线性映射，查询时使用由LSTM生成的“输入”$e^{(i)}$ 和上一次检索的结果 $n^{(i-1)}$（除了第一次查询时两个键都由LSTM生成）。

A3 背景知识/关键Observation/设计原则

RNN与记忆增强。理论上，像长短期记忆（LSTM）【Hochreiter & Schmidhuber, 1997, Long Short-Term Memory. Neural Computation; Gers et al., 2000, Learning to forget: Continual prediction with LSTM. Neural Computation】这样的RNN能够实现任何算法【Siegelmann & Sontag, 1991, Turing computability with neural nets. Applied Mathematics Letters】。然而，全连接RNN隐藏状态的线性增长会导致可训练权重数量的二次增长。早期的工作通过使用额外的记忆【Das et al., 1992, Learning context-free grammars: Capabilities and limitations of a neural network with an external stack memory; Mozer & Das, 1993, A connectionist symbol manipulator that discovers the structure of context-free languages】和可微的快速权重【Schmidhuber, 1992, Learning to control fast-weight memories: An alternative to recurrent nets; 1993, On decreasing the ratio between learning complexity and number of time-varying variables in fully recurrent nets】来解决这个问题。近年来，记忆增强的神经网络已经解决了一些算法类的玩具问题【Graves et al., 2014, Neural turing machines; 2016, Hybrid computing using a neural network with dynamic external memory】以及在合成和自然语言中的推理问题【Weston et al., 2015b, Memory networks; Xiong et al., 2016, Dynamic memory networks for visual and textual question answering】。

基于插槽的记忆。受计算机架构中随机存取存储器的启发，一种常见的方法是将一个软性的、可微分的查找表集成到神经网络模型中。这类基于插槽（slot-based）的记忆矩阵被证明难以训练【Munkhdalai & Yu, 2017b, Neural semantic encoders】，并且需要复杂的机制来进行内存的分配和释放【Csordas & Schmidhuber, 2019, Improving differentiable neural computers through memory masking, de-allocation, and link distribution sharpness control】。Transformer-XL（TXL）【Dai et al., 2019, Transformer-xl: Attentive language models beyond a fixed-length context】，作为Transformer【Vaswani et al., 2017, Attention is all you need】的一种自回归语言模型变体，可以被理解为一个基于插槽的记忆增强RNN，其中每个新状态被推入一个有限大小的不可变队列。尽管它是循环的，但Transformer架构的各层被严格限制只能使用来自较低层的输入，这限制了其通用性。尽管如此，一个足够深且正则化良好的TXL模型已在大型语言建模任务中取得了最先进的性能。

快速权重。作为增加神经网络记忆容量的一种更具生物学合理性的替代方案是快速变化的权重，即可以根据输入变化的、有状态的权重。非可微的快速权重或“动态链接”自1981年以来就有发表【von der Malsburg, 1981, The correlation theory of brain function; Feldman, 1982, Dynamic connections in neural networks; Hinton & Plaut, 1987, Using fast weights to deblur old memories】。随后的工作表明，一个常规网络可以通过梯度下降来训练，以端到端可微的方式控制另一个独立网络的快速权重【Schmidhuber, 1992, Learning to control fast-weight memories: An alternative to recurrent nets】或其自身的快速权重【Schmidhuber, 1993, On decreasing the ratio between learning complexity and number of time-varying variables in fully recurrent nets】。最近，快速权重重新受到关注，并在常规神经网络表现不佳的小型玩具问题上取得了良好结果【Ba et al., 2016a, Using fast weights to attend to the recent past; Schlag & Schmidhuber, 2017, Gated fast weights for on-the-fly neural program generation; Munkhdalai & Yu, 2017a, Meta networks; Pritzel et al., 2017, Neural episodic control; Ha et al., 2017, Hypernetworks; Zhang & Zhou, 2017, Learning to update auto-associative memory in recurrent neural networks for improving sequence memorization; Miconi et al., 2018, Differentiable plasticity: training plastic neural networks with backpropagation; 2019, Backpropamine: training self-modifying neural networks with differentiable neuromodulated plasticity; Schlag & Schmidhuber, 2018, Learning to reason with third order tensor products; Munkhdalai et al., 2019, Metalearned neural memory; Bartunov et al., 2020, Meta-learning deep energybased memory models】。

联想记忆类型。大多数记忆增强的神经网络都基于基于内容或基于键的查找机制。除了将模式存储在查找表中，另一种思路是通过隐式迭代最小化能量函数来重构模式，例如经典的Hopfield网络【Steinbuch, 1961, Die lernmatrix; Willshaw et al., 1969, Non-holographic associative memory; Hopfield, 1982, Neural networks and physical systems with emergent collective computational abilities; Kanerva, 1988, Sparse distributed memory】或现代Hopfield网络【Krotov & Hopfield, 2016, Dense associative memory for pattern recognition; Demircigil et al., 2017, On a model of associative memory with huge storage capacity; Ramsauer et al., 2020, Hopfield networks is all you need】。这通常被描述为一种自联想（auto-associative）类型的记忆，因为它会重构一个与当前模式最相似的先前存储的模式。一个研究较少的变体是异联想（hetero-associative）记忆（例如【Kosko, 1988, Bidirectional associative memories】），其中检索到的模式与输入模式不同。这与我们的用例更相关。我们的目标是训练一个LSTM来构建、维护和编辑其联想记忆。部分编辑Hopfield网络的能力研究得并不充分。因此，我们采用一个简单的（多）线性异联想记忆，因为它与TPR理论（其操作方式已广为人知）关系更密切，并且关联可以在单步内被检索。

与近期工作的关系。我们的工作直接建立在两个可微快速权重记忆的例子之上：Schlag & Schmidhuber (2018) 的 TPR-RNN【Learning to reason with third order tensor products】和 Munkhdalai et al. (2019) 的元学习神经记忆（MNM）【Metalearned neural memory】。TPR-RNN 是一个用于文本推理的句子级模型。它在常规 bAbI 任务上取得了优异的结果，但在词级 bAbI【Schlag et al., 2019, Enhancing the transformer with explicit relational encoding for math problem solving】或算法玩具问题【Le et al., 2020, Self-attentive associative memory】上表现不佳。相比之下，MNM 是一个词级模型，它通过一个全连接的多层前馈网络来增强 LSTM 作为其记忆，并使用元学习目标进行训练。MNM 和 TPR-RNN 都是在常规 bAbI 数据集上开发的，该数据集只包含短序列，并且不要求模型从其记忆中移除过时的关联。在这项工作中，我们在一个无限的 bAbI 故事序列上进行训练，我们的 FWM 取得了优异的性能，并优于 MNM。我们进一步在小规模语言建模和元强化学习中展示了强大的性能，证明了我们贡献的通用性。

A4 实验环境

数据集：
1. concatenated-bAbI (catbAbI)：这是一个基于 bAbI 数据集【Weston et al., 2015a, Towards ai-complete question answering: A set of prerequisite toy tasks】的更具挑战性的变体。它将 bAbI 故事拼接成一个无限序列，用于自回归方式的训练，模拟闭卷问答。它被用于监督学习实验。
2. 部分可观察马尔可夫决策过程 (POMDPs)：这是一个为元强化学习设计的任务集。每个环境是一个由5个状态和3个动作组成的稀疏有向图。代理需要在训练中未见过的新图中最大化奖励。
3. Penn Treebank (PTB) 和 WikiText-2 (WT2)：这是两个流行的小规模词级语言建模基准数据集，用于自监督学习实验。

模型架构：
* FWM：核心模型。在 catbAbI 实验中，LSTM 隐藏层大小 $d_{LSTM}=256$，FWM 维度 $d_{FWM}=32$，读取次数 $N_r=3$。在语言建模实验中，FWM 使用了与 AWD-LSTM 相同的三层 LSTM 作为其慢速 RNN 部分。
* 基线模型：
* 正则化 LSTM：一个4层、带残差连接的 LSTM。
* Transformer-XL (TXL)：一个3层、注意力窗口为1400个词元的 TXL 模型。
* Metalearned Neural Memory (MNM)：当前词级 bAbI 的 SOTA 模型。
* Ba's Fast Weights (JBFW)：一种早期的快速权重模型，在实验中未能收敛。
* AWD-LSTM：在语言建模实验中作为基线。
* AWD-TXL：在语言建模实验中作为基线。

硬件配置：
* GPU：实验被限制在约 16GB 显存的 GPU 上运行。
* 致谢中提及：NVIDIA 捐赠了多台 DGX 服务器，IBM 捐赠了一台 Minsky 服务器。

软件配置：
* 优化器：Adam 优化器。在语言建模实验中，后期切换到平均随机梯度下降（ASGD）。
* 训练细节 (catbAbI)：所有模型使用截断的反向传播（tBPTT），截断长度为200个词元。
* 训练细节 (Meta-RL)：使用优势演员-评论家（A2C）算法进行优化。
* 代码库：语言建模实验代码基于 Uber AI Lab 和 Salesforce 的 AWD-LSTM-LM 代码库。

A4 实验结果

4.1 CONCATENATED-BABI 实验

• 实验内容：在一个更具挑战性的 bAbI 数据集变体 catbAbI 上，评估 FWM 和其他基线模型的推理能力。实验分为两种模式：仅在回答问题时计算损失的 QA-mode，和对所有词元都计算损失的 LM-mode。
• 实验结果：
  • 如 表1 和 图2 所示，FWM 在 QA-mode 和 LM-mode 下均取得了最高的准确率（QA-mode: 98.6%，LM-mode: 98.4%）和最低的困惑度。
  • FWM 在参数数量上远小于其他基线模型，是最轻量级的模型。
  • 正则化的 LSTM 和 Transformer-XL (TXL) 表现强劲，但 TXL 的长注意力窗口使其更像一个“开卷”考试系统，因为它能直接查找历史信息，而 RNN 模型必须学会将信息保存在隐藏状态中。
  • 之前的 SOTA 模型 MNM 在 catbAbI 上的性能相比于原始 bAbI 有约10%的下降，显示了 catbAbI 的更高难度。
  • 早期的快速权重模型 JBFW 未能收敛。
• 分析结论：
  • FWM 的快速权重机制对于解决需要移除旧关联并建立新关联的连续推理任务至关重要。
  • 图3 可视化了 FWM 如何通过链式查询记忆来执行传递推理，证明了其多步推理能力。FWM 首先将 "gertrude" 和 "afraid" 关联，然后利用检索到的信息进一步关联到 "gertrude is a mouse"，最后关联到 "mice are afraid of wolves"，从而正确推断出答案 "wolf"。

4.2 元强化学习实验

• 实验内容：在一个由随机生成的稀疏有向图构成的部分可观察环境中训练智能体。智能体的目标是到达奖励位置。模型在600个训练图上训练，在600个未见过的新图上测试，以评估其泛化和适应能力。
• 实验结果：
  • 如 图5 所示，基于 FWM 的智能体训练速度更快，并且能很好地泛化到测试图上。
  • 相比之下，基于 LSTM 的智能体需要更多的训练回合、更大的网络，并且最终在训练图上过拟合。即使参数量是 FWM 的数千倍，LSTM 在测试环境中的总奖励也远低于 FWM。
• 分析结论：
  • FWM 的联想推理能力使智能体能够根据环境反馈（观察和动作）快速构建对当前环境（图结构）的表征，并利用这些表征来规划路径以最大化奖励，从而实现对新环境的快速适应。这表明 catbAbI 中的推理过程与 RL 中的探索和利用过程有相似之处。

4.3 语言建模实验

• 实验内容：在两个标准的词级语言建模基准数据集 PTB 和 WT2 上评估 FWM 的性能，并与强大的基线模型 AWD-LSTM 和 AWD-TXL 进行比较。
• 实验结果：
  • 如 表2 所示，FWM 在 PTB 和 WT2 数据集上的测试困惑度均优于 AWD-LSTM 和 AWD-TXL 基线。
  • 图6 的一个案例分析显示，当文本中反复出现稀有主题词（如人名或公司名）时，FWM 的预测不确定性（损失）比 LSTM 更低。这表明 FWM 能够利用其记忆机制捕捉并利用这些长距离依赖关系。
• 分析结论：
  • 即使在不像 bAbI 那样明确需要符号推理的传统语言建模任务中，FWM 的联想记忆机制也能带来性能提升。
  • 然而，FWM 相对于基线的提升幅度不大，这可能表明 PTB 和 WT2 这类新闻和百科文章数据集从联想推理能力中获益不如 catbAbI 那么显著。

A5 结论

本文提出了一种名为快速权重记忆（FWM）的新型快速权重架构。FWM 能够通过学习合成数据来回答需要多种符号推理技巧的问题。为了提升方法的通用性，本文引入了一个更通用、更困难的 bAbI 数据集变体——catbAbI，并在此数据集上取得了优异的性能，超过了基于最先进语言模型的强基线模型以及之前在词级 bAbI 上的最先进模型。此外，本文还将 FWM 应用于一个具有挑战性的元强化学习环境，在该环境中，智能体通过从观察和行动中学习，成功地泛化到新颖的环境。最后，在一个自监督学习的设定中，本文将 FWM 应用于 PTB 和 WT2 上的词级语言建模，其性能超过了 AWD-LSTM 和 AWD-Transformer-XL 基线。

A6 附录

A 进一步讨论

三阶张量记忆的容量。评估三阶张量记忆容量的一种方法是其秩（类似于矩阵的秩）。然而，对于一个给定的高阶张量 $A \in \mathbb{R}^{I \times J \times K}$，目前没有通用算法来确定其秩。只有一个宽松的上限，即 $rank(A) \leq \min\{IJ, IK, JK\}$【Kruskal, 1989, Rank, decomposition, and uniqueness for 3-way and n-way arrays; Kolda & Bader, 2009, Tensor decompositions and applications】。

张量积与拼接的比较。为了降低空间和时间复杂度，人们可能会尝试用拼接（concatenation）代替输入的外积（outer-product）来简化FWM。然而，在高度组合性的领域中，拼接的输入会遭受记忆间的干扰。考虑一个问题，需要从10个符号集合中关联任意三个符号，由向量 $s, r, t \in \mathbb{R}^{10}$ 表示。在使用拼接的情况下，快速权重记忆的一阶（rank-1）形式为 $[s; r] \otimes t$，其中 $[s; r]$ 是键表示。读取时，读取向量 $s', r' \in \mathbb{R}^{10}$ 被拼接起来并通过矩阵乘法检索先前的关联 $\hat{t} = F[s'; r']$。这里 $[s'; r']$ 是查询表示。由于有10个不同的符号，任意两个都可以作为键表示，因此存在 $10^2 = 100$ 个唯一的键模式。为保证在任何上下文中无噪声检索，键表示的向量必须是正交的。然而，$[s'; r']$ 只是一个20维空间，这意味着某些键表示不能在不产生干扰的情况下同时使用。另一方面，张量积能够实现无噪声检索，因为它将键表示为 $s \otimes r \in \mathbb{R}^{10 \times 10}$，这允许100个正交的键，从而可能实现无噪声检索。我们得出结论，如果问题是高度组合性的，即每个组件都可以与任何其他组件组合，那么张量积将比拼接更合适。我们在E节的实验中评估了拼接键。结果显示，拼接键会导致性能轻微下降（见图8）。作为替代方案，非线性记忆，例如通过使用softmax，可能不需要其键是正交的就能避免干扰，并可能带来更大的存储容量。

B 更新规则的推导

定理 B.1。给定两个键向量 $k_1, k_2 \in \mathbb{R}^d$ 和两个值向量 $v_{old}, v_{new} \in \mathbb{R}^d$，其中 $d \in \mathbb{Z}_{>0}$，混合系数 $\beta \in (0, 1)$，以及快速权重记忆 $F_{old} = \text{vec}(k_1 \otimes k_2) \otimes v_{old}$，其中 vec 指的是高阶张量的向量化。那么，由 $F_{old} + \beta \text{vec}(k_1 \otimes k_2) \otimes (v_{new} - v_{old})$ 给出的（循环）快速权重更新规则将得到 $F_{new} = \text{vec}(k_1 \otimes k_2) \otimes [(1-\beta)v_{old} + \beta v_{new}]$。

证明。

C 关于常规 bAbI 数据集和先前工作的评论

bAbI 数据集概述。bAbI 任务是一个流行的玩具数据集，用于测试具有记忆增强和推理能力的神经网络【Weston et al., 2015a, Towards ai-complete question answering: A set of prerequisite toy tasks】。它由一系列简短的故事组成，问题嵌入在文本中。这些故事通过模拟虚拟环境中的多个实体生成，涵盖了实体自身状态改变或通过互动改变状态的不同情境。每个故事样本属于20个不同任务之一，这些任务被数据集作者认为对智能对话代理很重要。这些任务包含需要推理能力（如推导、共指或计数）的问题。所有任务都需要一定程度的符号推理，最初的神经网络和非神经网络基线在测试数据上表现出较差的泛化性能【Weston et al., 2015a】。除了故事句子、问题和答案，数据集还包括“支持事实”（supporting facts），用于标出与问题相关的句子。故事通常遵循多个平行的情节线，每个新句子都将其中一个情节线推进一个事实。

先前工作的局限性。bAbI 数据集没有严格的实验协议，导致出现了几种略有不同的变体。早期的方法通过依赖支持事实【Weston et al., 2015b, Memory networks; Kumar et al., 2016, Ask me anything: Dynamic memory networks for natural language processing】或其他监督训练信号【例如 Johnson (2017), Learning graphical state transitions; Li et al. (2016), Gated graph sequence neural networks】取得了良好结果。

开卷式 vs. 闭卷式问答。一些研究者通过重新格式化数据，使得问题在故事之前被读取，或者类似地，让模型在读取问题后有能力通过某种注意力机制查找故事的部分内容，从而取得了很好的结果【Sukhbaatar et al., 2015, End-to-end memory networks; Xiong et al., 2016, Dynamic memory networks for visual and textual question answering; Dehghani et al., 2019, Universal transformers】。这类方法在可以访问完整故事的情况下回答问题非常有效。我们认为这类似于“开卷式”（open-book）问答。在这种设置下，模型被激励去查找信息，而不是捕捉它所看到的数据中的有用片段。后者的优势在不同场景下更为明显：想象一个模型正在处理一本书，用户可以随时提问关于内容的问题。开卷式方法将不得不将其所有先前的句子存储在内存中，并将其答案搜索机制应用于所有数据。相反，“闭卷式”（closed-book）方法会存储故事的压缩版本，或者故事中与问题相关的信息。

句子级 vs. 词级模型。必须承认，所有任务的 bAbI 故事中的句子都是简短和简单的。几乎每个句子都只包含一个事实。因此，句子级模型可能比词级模型更有优势。确实，之前的一个句子级模型在词级设置下报告了较差的性能【Schlag & Schmidhuber, 2018, Learning to reason with third order tensor products】。这限制了它们的通用性，因为自然语言中的句子通常不限于单个事实。

数据预处理问题。最后，尽管 bAbI 数据集最初设计时问题是嵌入在故事中的，但几乎所有到目前为止的方法都对数据集进行了预处理，将一个包含四个问题的样本拆分为四个各含一个问题的样本【Weston et al., 2015b】。这无疑简化了问题，因为一旦读到问题，模型就不再需要维护与该问题无关的其他实体的状态。然而，这是否会导致性能下降还有待检验。

D Concatenated-bAbI 详情

catbAbI 定义。Concatenated-bAbI (catbAbI) 是一种预处理和实验流程，用于评估自回归模型在预测需要特定推理技能的单词（这里指问题的答案）方面的能力。在这项工作中，我们只关注每个任务10k样本版本的bAbI，但我们所有的脚本也可以应用于1k版本。我们使用了与常规bAbI相同的数据训练/测试/验证集划分。与先前的工作不同，我们不将故事拆分以仅包含一个问题。我们移除了句子索引，并将句子与问号后的答案连接成一个长单词序列。预处理后的数据是样本的随机列表。每个样本都附带其任务ID以供诊断。所有答案都以问号开头。

数据处理流程。为了确保故事不会重叠和变得模棱两可，我们在连接新故事之前添加了一个特殊的“故事结束”标记。对于每个单词，预处理脚本提供了其任务ID，以衡量在不同任务上的性能。同样，它还提供了一个特殊的“答案”标记，表示当前单词是否是答案。当然，任务ID和答案信息不会作为输入提供给模型。验证和测试数据也同样处理，但为了对各种模型进行适当比较，验证和测试数据只进行一次洗牌。在训练和评估期间，验证和测试故事是确定性地抽取的。

训练流程。在训练期间，我们无放回地均匀抽样故事，并将它们连接成一个长序列。由于问号并不总是故事的结尾，我们通过用一个特殊的“故事结束”标记分隔故事来解决任何歧义。模型在这个长序列上以自回归方式进行训练，并使用截断反向传播。在每个epoch结束时，如果批次中的序列长度不同，我们用填充符号填充批次。

损失计算与状态传递。在LM模式下，我们屏蔽填充标记；在QA模式下，我们屏蔽除输入为问号的步骤之外的所有内容。在epoch结束时，我们将隐藏状态传递到新的epoch。将所有隐藏状态重置为相同或零对最终性能有轻微的负面影响，但未进行深入探讨。为了在验证和测试集上进行评估，使用了第一个批次元素的隐藏状态副本。在训练期间，对验证集的评估使用大批量大小以保持速度。对测试集的评估使用批量大小为1。在验证和测试集上评估期间，样本是顺序选取的，以确保所有模型都在相同的验证和测试bAbI故事序列上进行评估。

E 消融实验

推理步数影响。我们评估了具有不同循环步数（$N_r$）的FWM模型。图7中的实验表明，仅一步（$N_r=1$）就已经达到了超过95%的准确率，但更多的推理步骤有助于解决更罕见但更困难的任务。

键组合方式影响。我们还测试了一个FWM版本，其中读取和查询的键是拼接（concatenated）而不是通过张量积相乘的。在这个版本中，FWM产生一个权重矩阵，其维度为 $\mathbb{R}^{2d_{FWM} \times d_{FWM}}$ 而不是 $\mathbb{R}^{d_{FWM}^2 \times d_{FWM}}$。图8中的结果表明性能有所下降。

F catbAbI 的超参数搜索

通用设置。由于catbAbI是一个连续的故事序列，对所有模型来说，通过时间反向传播（BPTT）是不可行的，因此我们将BPTT截断到最后200个词元。超参数的选择使得模型大约能在一个有16GB内存的GPU上运行。所有模型都使用256维的词元嵌入和Adam优化器。我们专门为QM（问答）设置调整了超参数，并将最佳设置直接用于LM（语言模型）设置。我们进行了网格搜索，涵盖了批次大小、学习率以及各种模型特定的参数，如dropout率或层数，此外还有一些手动选择的设置。出于计算原因，我们进行了两轮网格搜索：第一轮运行3000步，其中表现最好的模型进入第二轮，训练30000或60000步。

F.1 快速权重记忆 (Fast Weight Memory)

搜索参数。我们设置 $d_{LSTM} = 256$, $d_{FWM} = 32$, $N_r = 3$，并对批次大小（64, 128）和学习率（0.0001, 0.00025, 0.0005, 0.001, 0.002）进行了两个种子的实验。

F.2 元学习神经记忆 (Metalearned Neural Memory)

搜索参数。我们只实验了MNM的可塑性版本（plastic version），因为它被报道是最好的。我们使用了Munkhdalai等人(2019)报告的相同的快速权重超参数：3层快速权重，维度为100。我们搜索了批次大小（64, 128）；学习率（0.00025, 0.0005, 0.001, 0.002）；以及元目标系数（reg）（1.0, 2.0）。在最初的3000步中，MNM没有表现出任何不稳定性，但对于更长的运行，MNM有时会产生NaN或变得不稳定。

F.3 Transformer-XL

实现与搜索。我们将官方的Transformer-XL实现移植到我们自己的代码库中，完全重用了模型代码用于我们的catbAbI实验。我们在最初的1000步中采用了线性学习率预热（warm-up）计划，并对批次大小、学习率、层数和记忆长度进行网格搜索，并附加了一些手动选择的参数。我们最佳的设置使用了0.00025的学习率，1200的记忆宽度，隐藏状态大小 $d_{model} = 512$，全连接部分内部维度 $d_{inner} = 2048$，以及3个transformer层。几次长时间的运行可以在图12中看到。我们的实验显示，各种种子最终会变得不稳定并过拟合。某些设置也导致了NaN，我们已从图12中移除。表现最好且最稳定的是具有大记忆和低学习率的3层模型（见图13）。

F.4 LSTM

正则化。我们对一个四层堆叠的残差连接LSTM单元进行了重度正则化，每个单元有512个隐藏单元。受AWD-LSTM【Merity et al., 2018, Regularizing and optimizing LSTM language models】的启发，我们使用四种不同的dropout方式来正则化模型。我们对输入序列的词元、嵌入向量的元素、循环权重矩阵的元素以及LSTM层之间隐藏表示的元素进行dropout。

F.5 注意近期历史的快速权重 (Attention to the Recent Past Fast Weights)

实现与结果。我们评估了我们自己实现的、由Ba等人（2016a）【Using fast weights to attend to the recent past】引入的快速权重模型。他们提出了一种用快速权重增强的RNN，该快速权重使用固定的学习率和衰减率来调节Elman RNN的慢速权重（JBFW）。我们的超参数搜索没有产生任何在测试数据上性能超过15%的模型。

G 各任务上最佳 catbAbI 运行结果分解

H 语言建模

代码与实现。我们的语言建模实验代码分支自Uber AI Lab的代码库，该代码库本身又分支自Salesforce的语言模型工具包。FWM使用了与Merity等人（2018）【Regularizing and optimizing LSTM language models】所做的相同优化的三层LSTM作为慢速RNN。我们没有探索的一种替代方案是使用多个FWM层，每层包含一个LSTM单元和一个FWM。我们在PTB上训练了1000个epoch，在WT2上训练了1600个epoch。与Merity等人（2018）类似，我们在PTB和WT2模型分别训练了916和1372个epoch后，将优化器从Adam切换到平均随机梯度下降（ASGD）。我们在验证集上调整了dropout参数，并且在训练后，我们还调整了softmax温度（调整softmax温度带来了约1个PPL的提升）。嵌入层是从均匀分布uniform(-0.25, 0.25)中随机初始化的，这在我们的FWM语言模型中至关重要。所有报告结果使用的超参数见表4。

Transformer-XL 设置。Transformer-XL在PTB上的结果是基于作者的官方代码和超参数设置，其中包括AWD风格的正则化、模型平均和softmax调优。WT2的结果是基于相同的代码，使用了Tim Dettmers找到的最佳超参数。

H.1 结果

I 元强化学习

环境设置。元强化学习实验在一个训练POMDP集合中训练智能体，并在一个测试POMDP集合中进行评估。环境是有向图，边上有标签。作为数据生成过程的一部分，根据清单1中的python算法采样新的图。动作和状态是独热编码的。智能体接收一个17维的输入：奖励位置、当前位置、上一个动作、一个固定位、作为当前步数/总步数的分数进度，以及当前的奖励总和。到达奖励位置会得到10的奖励。选择一个无效的动作会得到0.05的惩罚。我们使用0.9的折扣因子和0.1的值系数。A2C的熵系数设置为0.03。

训练细节。智能体和奖励的位置在每回合开始时随机选择。由于只有5个状态，奖励最多在5步内可以达到。如4.2节所述，这种最优行为只有在智能体从经验中学习了图的结构后才可能实现。每当奖励被放置在环境中时，一个重置计时器被设为0。当智能体到达奖励位置，或者在6次不成功的步骤后，重置计时器被设为0，奖励和智能体的位置会被随机重新放置。我们使用600个智能体的批次大小进行训练，并使用Adam优化器优化平均步长损失。

# import numpy as np

def sample_adjacency_matrix(n_actions, n_states, random_state):
    while True:
        A = np.zeros((n_actions, n_states, n_states))
        # every state has to be leavable by at least one action
        for from_state in range(n_states):
            to_state = random_state.choice([i for i in range(n_states) if i != from_state])
            action = random_state.randint(0, n_actions)
            A[action, from_state, to_state] = 1
        # every state has to be reachable by one or more from-states
        for to_state in range(n_states):
            # only select states which don't have any neighbours given an action
            action_list, from_list = np.where(A.sum(2) == 0)
            # remove self from the selection
            options = np.asarray(list(filter(lambda x: x[0] != to_state, zip(from_list, action_list))))
            indecies = np.arange(options.shape[0])
            chosen_idx = random_state.choice(indecies)
            from_state, action = options[chosen_idx]
            A[action, from_state, to_state] = 1
        # reject if they are not all connected
        Q = A.sum(0)
        Q[Q > 0] = 1
        for _ in range(n_states):
            Q = np.matmul(Q,Q)
        if (Q == 0).sum() == 0:
            return A

清单1：用于采样新环境的Python3代码，确保任何状态都可以从任何其他状态到达。

A7 补充细节

计算成本与局限性。一个三阶记忆张量是构建组合式状态表示的一种计算成本高昂的方法。对于在 $\mathbb{R}^n$ 中的向量分量，仅张量积计算就具有 $O(n^3)$ 的空间和时间复杂度。出于实际原因，这迫使FWM相对于慢速神经网络保持较小，从而限制了可以同时维持的关联数量。先前的工作已经提出了以方差最优的方式对此类记忆张量进行近似的方法【Schlag et al., 2019, Enhancing the transformer with explicit relational encoding for math problem solving】。在E节的消融实验中，我们在catbAbI上显示，将键进行拼接会导致性能准确率下降约5%。我们还实验了较少的读取操作（较小的$N_r$），这也导致了性能下降（附录图7）。然而，进一步的改进可能不会来自规模扩大，而是来自更通用的符号操作。我们在附录的A节中从线性异联想记忆的角度讨论了FWM的容量和张量积的必要性。

与工作记忆的类比。最后，我们的快速权重记忆可以被认为是模型的原始“工作记忆”——类似于人脑中的工作记忆【Spalding et al., 2018, Ventromedial prefrontal cortex is necessary for normal associative inference and memory integration】。这一观点得到了最近工作的支持，该工作提出了一个基于此类高阶张量的人脑认知模型【Tresp & Ma, 2017, The tensor memory hypothesis】。