注意力机制¶

大多数 transformer 模型使用全注意力机制，这意味着注意力矩阵是方形的。对于长文本而言，这可能会成为一个巨大的计算瓶颈。为了更高效，Longformer 和 Reformer 等模型使用了稀疏注意力矩阵来加速训练。

局部敏感哈希 (LSH) 注意力¶

Reformer 模型采用了 LSH 注意力机制。在 softmax(QK^t) 中，只有矩阵 QK^t 中最大的元素（在 softmax 维度上）会给出有用的贡献。因此，对于查询矩阵 Q 中的每个查询 q，我们只需要考虑与 q 接近的键 k。通过哈希函数来确定 q 和 k 是否接近。为了防止查询和键相等（即非常相似），除了在第一个位置外，注意力掩码会屏蔽当前的 token。由于哈希函数可能是随机的，实际上会使用多个哈希函数（由 n_rounds 参数确定），然后对结果进行平均。

局部注意力¶

Longformer 模型使用了局部注意力机制：通常情况下，局部上下文（例如，左右两个 token）对于处理给定的 token 已经足够了。通过堆叠具有小窗口的注意力层，最后一层将具有超出窗口 token 的感受野，从而可以构建整个句子的表示。

一些预选的输入 token 可以获得全局注意力：对于这些少数 token，注意力矩阵可以访问所有 token，这个过程是对称的：所有其他 token 也可以访问这些特定 token（除了它们本地窗口中的 token）。如下图所示，这是论文中的一个示例注意力掩码：

使用这些参数较少的注意力矩阵可以让模型处理更长的序列。

其他技巧¶

轴向位置编码¶

Reformer 模型使用了轴向位置编码：在传统的 transformer 模型中，位置编码 E 是一个大小为 l × d 的矩阵，其中 l 是序列长度，d 是隐藏状态的维度。如果文本很长，这个矩阵可能会非常大，占用过多的 GPU 空间。为了缓解这个问题，轴向位置编码通过将大矩阵 E 分解成两个较小的矩阵 E1 和 E2 来实现，其尺寸分别为 l1 × d1 和 l2 × d2，满足 l1 × l2 = l 和 d1 + d2 = d（通过长度相乘，最终会小得多）。时间步 j 在 E 中的嵌入是通过连接 E1 中的时间步 j % l1 和 E2 中的时间步 j // l1 的嵌入获得的。