神经网络通常被视为复杂的黑盒,但有一种特定的视角将每一层都视为一种关联记忆。在这种框架下,记忆库是一个持有键值对(key-value pairs)集合的函数。当用户提供输入时,系统会根据与存储键的相似度进行查找,并检索相应的数值。该系统面临的主要挑战是干扰:随着存储的记忆越多,它们开始相互重叠并产生干扰,从而降低检索的准确性。

这些记忆库的容量从根本上受限于系统的维度。基础线性层将记忆作为外积进行累积,只能有效地存储与维度成比例的键值对数量($O(d)$)。为了增加这种容量,网络必须抑制噪声和干扰。通过引入非线性因素,例如多层感知机(MLPs)中存在的非线性,容量可扩展到 $O(d^2)$。例如,从标准的线性设置转向基于 ReLU 的 MLP 或门控线性单元(GLU),可以让网络通过有效地“平方”可用容量来存储显著更多的信息。

一只栖息在竹枝上的开亚鸟水墨画
开亚鸟的艺术表现,配合对神经网络中关联记忆理论框架的讨论。alexwang.ai

这张图片是文章探讨神经网络层如何作为关联记忆库的主要视觉素材。

Transformer 架构代表了这一演进的巅峰。由于指数函数可以被视为幂的求和,具有 $\sqrt{d}$ 缩放的自注意力机制实际上充当了一个高容量的记忆库。在 64 维空间中,该机制理论上可以存储多达 100 万个键值对。这一视角表明,从简单的 MLP 到复杂的 Transformer,现代人工智能中的架构选择本质上是管理记忆堆积和相互干扰的不同方法,旨在最大化“信道容量”。然而,虽然这些启发式论点为理解模型如何“记住”事实提供了一条清晰的路径,但特定权重如何表示复杂概念的确切机制仍是当前机械可解释性研究的核心焦点。

这一分析强调了复杂性与容量之间的权衡:当我们从线性层转向自注意力层时,我们获得了存储更多不同记忆的能力,但同时也增加了模型必须管理的干扰模式的复杂性。理解这些限制对于改进模型如何存储和检索世界知识至关重要,尤其是在数据量增长的同时仍需保持精度。

虽然这种以记忆为中心的视角为理解 Transformer 层提供了一个强大的启发式框架,但需要注意的是,它是一个用于解释容量的理论框架,而非对每个权重功能的字面描述。模型的实际性能既取决于其架构的理论容量,也取决于用于填充这些记忆的特定训练数据。

最终,通过关联记忆的视角审视神经网络有助于阐明为何选择某些特定架构。通过将每一层视为记忆库,研究人员可以根据模型维度及其用于抑制噪声的非线性因素,更好地估算模型能够“知道”内容的极限。这一视角在传统的通信工程概念(如信道容量)与现代深度学习架构之间架起了一座桥梁。

目前仍存在的一个不确定因素是,这些理论容量如何转化为现实世界的“推理”而非简单的事实检索。虽然数学上支持 Transformer 在键值存储上的巨大增长,但从存储一个事实到执行多步逻辑推演的质变,仍是机械可解释性领域中活跃的研究课题。