Unlimiformer:用搜索的思路来处理注意力机制长度限制

出品人：Towhee 技术团队

作者：王翔宇

Transformer 模型是一种非常强大的序列建模工具，但由于其注意力操作的复杂性，它在处理长序列时受到了限制。在传统的 Transformer 中，每个输入标记都需要与其他标记计算注意力权重，这导致了注意力操作的时间复杂度为输入标记数量的平方。因此，当输入序列非常长时，这种计算开销会变得非常昂贵，甚至无法处理。在 attention 操作中，内积权重占比较大的 value 向量会对结果 embedding 产生更大的影响，那么是否能够在很长的 token 序列中重点关注内积较大的部分而不是全体 token 呢？

｜在这个例子中，一个特定的语言模型编码器的最大输入长度为2个token。一个包含6个token的输入被分块编码，然后被索引。在每个解码器层之前将Unlimiformer注入交叉注意力。在Unlimiformer中，我们执行kNN搜索，从索引中为每个注意力头选择一个 2 个 token 的上下文。这使得交叉注意力可以关注整个输入序列中的 token，而不添加参数，也不改变给定的语言模型的架构。

研究人员了一种名为 Unlimiformer 的基于检索的方法，用于增强预训练语言模型，使其能够在测试时接受无限长度的输入序列。Unlimiformer 可以注入到任何现有的编码-解码 Transformer 中，以实现无限长度的输入。在标准的交叉注意力中，Transformer 解码器会关注编码器的顶层隐藏状态，而编码器通常会截断输入并只对输入序列的前 k 个标记进行编码。作者不再只关注输入的这 k 个标记前缀，而是从 kNN 索引中检索每个 cross-attention 头部的 top-k 隐藏状态，并只关注这些 top-k 隐藏状态。这样就可以从整个输入序列进行检索，而不需要截断。该方法在计算和 GPU 内存方面比关注所有输入标记更加高效；由于 softmax 函数对最大值敏感，检索出最具注意力的标记可以保留绝大部分的注意力集中。

这是注意力机制的公式，一种简单的做法是针对被一个与权重向量 We 相乘后的 K = heWk 作索引，但是权重矩阵每一层每一个注意力头都不一样，所以需要 2xL(层数)xH(注意力头数)个索引,作者做了如下公式改进。这样只需索引编码器的输出 he 即可。

作者做了若干实验，这是在 BookSum 数据集上做的摘要实验，平均有 143 k 个 token，baseline 是 Hierachical，也是一个相当强劲的 baseline，可以看出 unlimiformer 拥有不错的提升。

参考链接：

1. https://arxiv.org/abs/2305.01625
2. https://github.com/abertsch72/unlimiformer

本文由mdnice多平台发布