在过去的几年里,Transformer 几乎统治了整个 NLP 领域,还跨界到计算机视觉等其他领域。但它也有弱点,比如不擅长处理较长上下文,因为计算复杂度会随着上下文长度的增长而增长,这使其难以有效建模长期记忆。为了缓解这一问题,人们提出了多种 Transformer 变体,但它们的记忆容量都是有限的,不得不抛弃较早的信息。
在一篇论文中,来自 DeepMind 等机构的研究者提出了一种名为 ∞-former 的模型,它是一种具备无限长期记忆(LTM)的 Transformer 模型,可以处理任意长度的上下文。
论文链接:https://arxiv.org/pdf/2109.00301.pdf
通过利用连续空间注意力机制来处理长期记忆,∞-former 的注意力复杂度可以独立于上下文长度。因此,它能够借助一个固定的算力开销建模任意长度的上下文并保持「粘性记忆(sticky memories)」。
在一个综合排序任务上进行的实验证明了∞-former 能够保留来自长序列的信息。此外,研究者还进行了语言建模的实验,包括从头开始训练一个模型以及对一个预训练的语言模型进行微调,这些实验显示了无限长期记忆的优势。
不过,和其他很多 Transformer 变体的论文一样,这篇论文的标题也引发了一些吐槽:
∞-former:一种拥有无限记忆的 Transformer
为了使模型能够处理长程上下文,研究者提出用一个连续 LTM 来扩展原始 transformer,这个 LTM 存储前面步骤的输入嵌入和隐藏状态。他们还考虑了有两种记忆的可能性:LTM 和 STM(短期记忆),类似于 transformer-XL 的记忆。∞-former 的整体架构如下图 1 所示。
为了让新模型的 LTM 达到无限的程度,研究者使用了一个连续空间注意力框架(参见《 Sparse and Continuous Attention Mechanisms 》),它在适用于记忆的信息单元数量(基函数)和这些单元的表示粒度之间进行了权衡。在这一框架中,输入序列被表征为一个连续信号,表示为径向基函数的一个线性组合。这种表征有两个显著的优势:1)上下文可以用 N 个基函数来表示,N 小于上下文中 token 的数量,降低了注意力复杂度;2)N 可以是固定的,这使得在记忆中表征无限长度的上下文成为可能(如图 2 所示),代价是损失 resolution,但不增加其注意力复杂度,O(L^2 + L × N),其中的 L 对应 transformer 序列长度。
为了缓解损失较早记忆 resolution 的问题。研究者引入了「粘性记忆」的概念,将 LTM 新信号中的较大空间归于之前记忆信号的相关区域。这是一种强制重要信息在 LTM 中持续存在的过程,使得模型可以在不损失相关信息的情况下更好地捕捉长上下文,类似于大脑中的长时程增强和突触可塑性。
实验结果
为了检验∞-former 能否建模长上下文,研究者首先针对一个综合任务进行了实验,包括把 token 按其在一个长序列中的频率进行排序,结果如下:
从图中可以看出,在序列长度为 4000 的时候,transformerXL 的准确率要略高于 compressive transformer 和 ∞-former。这是因为 transformerXL 几乎可以在记忆中保存整个完整序列。但随着序列长度的增加,transformerXL 的准确率迅速下降,compressive transformer 和 ∞-former 变化较小。这表明∞-former 更擅长建模长序列。
接下来,他们又进行了语言建模实验,包括:1)从头训练一个模型;2)微调一个预训练语言模型。
第一个语言建模实验的结果如下表 1 所示。从中可以看出,利用长期记忆扩展模型确实会带来更好的困惑度结果,而且使用粘性记忆也可以在一定程度上降低困惑度。
第二个语言建模实验的结果如下表 2 所示。该结果表明,通过简单地将长期记忆添加至 GPT-2 并进行微调,模型在 Wikitext-103 和 PG19 上的困惑度都会降低。这表明∞-former 具有多种用途:既可以从头开始训练模型,也可以用于改进预训练模型。