Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. “Reformer: The efficient transformer.” arXiv preprint arXiv:2001.04451 (2020).
背景与动机
这篇论文发表较早,主要关注Transformer的效率问题。标准的Transformer模型在许多自然语言处理任务上取得了最先进的结果,但它在长序列上的训练和推理代价非常大。主要的计算和内存瓶颈在于自注意力机制中的点积注意力,其复杂度为
O
(
L
2
)
\Omicron(L^2)
O(L2),其中
L
L
L是序列长度。随着
L
L
L的增大,计算和内存需求急剧增加。因此,Transformer难以扩展到处理长序列的任务。
模型与方法
引入基于局部敏感哈希(LSH)的近似注意力机制,将自注意力的复杂度从
O
(
L
2
)
\Omicron(L^2)
O(L2)降低到
O
(
L
log
L
)
\Omicron(L\log L)
O(LlogL),大大减少了内存和计算需求。具体来说,文章首先使用随机投影作为敏感哈希函数。相似的query和key通过投影转换后可以映射到相同的哈希桶中。然后根据query和key的哈希值对序列进行排序。相似的query和key会聚集在一起。在排序后的序列上,将每个query块只与相邻的几个query块计算注意力。这样可以大约保证每个块内的query可以attend到相似的key。最后使用多轮不同的哈希函数和注意力计算,综合多个注意力输出,可以降低哈希误差。