[论文解读] Relating transformers to models and neural representations of the hippocampal formation

James C. R. Whittington, Joseph W. Warren+1 more|arXiv (Cornell University)|Dec 7, 2021

Neural Networks and Applications被引用 59

一句话总结

带有循环位置编码的 Transformer 可以再现海马体空间表征（位置细胞/网格细胞），并且与如 Tolman-Eichenbaum Machine (TEM) 的海马模型密切相关；Transformer 的变体 (TEM-t) 取得显著的性能提升，并为海马体索引与学习的位置信编码提供新的见解。

ABSTRACT

Many deep neural network architectures loosely based on brain networks have recently been shown to replicate neural firing patterns observed in the brain. One of the most exciting and promising novel architectures, the Transformer neural network, was developed without the brain in mind. In this work, we show that transformers, when equipped with recurrent position encodings, replicate the precisely tuned spatial representations of the hippocampal formation; most notably place and grid cells. Furthermore, we show that this result is no surprise since it is closely related to current hippocampal models from neuroscience. We additionally show the transformer version offers dramatic performance gains over the neuroscience version. This work continues to bind computations of artificial and brain networks, offers a novel understanding of the hippocampal-cortical interaction, and suggests how wider cortical areas may perform complex tasks beyond current neuroscience models such as language comprehension.

研究动机与目标

• 证明带有循环位置编码的 Transformer 能再现海马体的空间表征，如网格细胞和位置细胞。
• 在 Transformer 与海马形成物理神经科学模型（TEM）之间建立紧密的数学关系。
• 提出一种新颖的、以海马体索引为灵感的位置信编码与记忆检索的解释。
• 显示以 TEM 为灵感的 Transformer（TEM-t）在学习效率和记忆容量方面优于原始 TEM。
• 讨论海马-皮层相互作用的广泛含义以及在非空间认知领域的潜在应用。

提出的方法

• 描述标准 Transformer 自注意力及其方程，并引入一种修改后的设置，使查询和键专注于位置信息编码，而值取决于刺激。
• 采用对过去位置和刺激有记忆的因果 Transformer，限制对过去时间步的访问（因果掩蔽）。
• 使用循环生成位置信息编码 e_{t+1}=sigma(e_t W_a) 使编码可学习且与序列相关。
• 在一个空间预测任务上训练，其中代理根据在二维环境中的动作预测即将到来的感官观测，要求对抽象空间规则的泛化。
• 将 TEM 模型与 Transformer 自注意力联系起来，展示记忆检索与连结表征映射到注意力机制，并推导 TEM-t 的表达式。
• 将 TEM-t 与 TEM 进行比较，突显样本效率、可扩展性和记忆存储方面的改进。

实验结果

研究问题

• RQ1 在基于空间预测任务训练时，带有循环位置编码的 Transformer 能否再现海马体的空间表征（网格细胞和地点细胞）？
• RQ2 TEM（Tolman-Eichenbaum Machine）与 Transformer 自注意力之间的数学关系是什么，这如何帮助我们理解海马体记忆结合？
• RQ3 TEM-t 是否在数据效率和记忆容量方面优于 TEM？
• RQ4 TEM-t 中学习到的位置信编码如何与海马体样表征相关联，并可能推广到非空间认知领域，如语言？

主要发现

• 具有循环位置编码的 Transformer 学习到的空间表征，类似于在内嗅皮层/海马回路中观察到的网格细胞和带状细胞。
• TEM-t 相比原始 TEM 提供显著的性能提升，包括更高的样本效率和存储/检索更多记忆的能力。
• 数学框架显示 TEM 的记忆检索与 transformer 自注意力类似，且路径积分编码充当学习到的位置编码。
• 将 TEM-t 实例化为一个在神经层面上可信的架构，其中海马体记忆对皮层模式进行索引，与海马体索引理论相一致。
• 记忆神经元中出现类似地点细胞的反应，伴随环境特异性重新映射，与海马体数据及模型预测一致。