[论文解读] Attention Approximates Sparse Distributed Memory

Trenton Bricken, Cengiz Pehlevan|arXiv (Cornell University)|Nov 10, 2021

Ferroelectric and Negative Capacitance Devices参考文献 64被引用 6

一句话总结

本文表明，Transformer 中的自注意力机制在数学上近似于 Kanerva 的稀疏分布式记忆（SDM）模型，这是一种具有生物学合理性的关联记忆模型。在特定数据条件下——尤其是微调后的 GPT-2 模型中——注意力机制的 Softmax 操作与 SDM 的超球面交集机制高度一致，为注意力机制提供了一种新颖的机制性与生物学解释，即其本质是一种分布式记忆检索形式。

ABSTRACT

While Attention has come to be an important mechanism in deep learning, there remains limited intuition for why it works so well. Here, we show that Transformer Attention can be closely related under certain data conditions to Kanerva's Sparse Distributed Memory (SDM), a biologically plausible associative memory model. We confirm that these conditions are satisfied in pre-trained GPT2 Transformer models. We discuss the implications of the Attention-SDM map and provide new computational and biological interpretations of Attention.

研究动机与目标

• 建立 Transformer 中自注意力机制与 Kanerva 的稀疏分布式记忆（SDM）之间形式化的数学联系，SDM 是一种具有生物学合理性的关联记忆模型。
• 研究在真实数据条件下，如 GPT-2 等预训练模型中的注意力结构与功能特性是否与 SDM 的原理一致。
• 通过将注意力机制视为一种高维近似检索机制，为注意力机制提供一种新的计算与生物学解释。
• 探讨该映射关系对理解注意力机制的成功、可解释性及其在大脑中潜在神经实现的启示。

提出的方法

• 推导 Softmax 注意力机制与 SDM 的读取操作之间的数学等价性，表明注意力中指数函数的求和与指数的比值近似于 SDM 中高维超球面的交集。
• 使用查询-键归一化的 Transformer 变体，在预训练的 GPT-2 模型中实证验证 SDM 与注意力之间的对应关系。
• 通过改变神经元数量与汉明半径，模拟 SDM，测试其在噪声下的收敛性与鲁棒性，并比较二值与连续 SDM 变体。
• 训练投影矩阵以提升 SDM 在 MNIST 与 CIFAR-10 数据集上的临界距离，评估其泛化能力与性能提升。
• 通过比较 MNIST（低复杂度）与 CIFAR-10（高复杂度）数据集在 SDM 收敛行为中的模式相关性与数据复杂度，分析其影响。
• 使用余弦相似度与汉明距离度量，评估在合成与真实图像数据集中的检索性能及对噪声的鲁棒性。

实验结果

研究问题

• RQ1 Transformer 中的注意力机制在多大程度上数学上近似于稀疏分布式记忆（SDM）中的检索过程？
• RQ2 在真实预训练的 Transformer 模型（如 GPT-2）中，是否满足 SDM-注意力对应关系所要求的结构与统计条件？
• RQ3 在 MNIST 与 CIFAR-10 数据集中，SDM 基础的注意力近似性能如何随数据复杂度变化？
• RQ4 学习到的投影矩阵是否能提升 SDM 的临界距离与抗噪鲁棒性？该方法在不同数据集上是否具有泛化能力？
• RQ5 将注意力解释为一种分布式记忆检索形式，其在生物学与计算层面具有哪些启示？

主要发现

• 当查询与键被归一化时，Softmax 注意力机制与 SDM 中的超球面交集操作高度近似，验证了注意力与关联记忆之间数学联系的成立。
• 在预训练的 GPT-2 模型中，满足 SDM-注意力对应关系所必需的条件，证实该理论映射在实践中成立。
• MNIST 数据因复杂度较低，其在 SDM 基础检索中表现出更好的收敛性与更高的鲁棒性；而 CIFAR-10 因模式相关性更高，易出现过拟合。
• 学习投影矩阵可提升 MNIST 与 CIFAR-10 中 SDM 的临界距离，但泛化能力有限，表明对训练数据存在过拟合。
• 神经元数量越少，性能越差；连续 SDM 变体对神经元数量有限的情况更敏感，因其球面交集更小。
• SDM 框架为整个 Transformer 模型提供了更深层次的解释，表明注意力机制本质上是一种高维关联记忆检索机制。