Transformer 遇见随机块模型：具有数据自适应稀疏性与代价的注意力

为克服自注意力的二次代价，近期工作提出了各种稀疏注意力模块，大多归于以下两类之一：1) 手工模式下的稀疏注意力，2) 全注意力后接如 $\alpha$-entmax 的稀疏变体 softmax。遗憾的是，第一类缺乏对数据的适应性，而第二类在训练中仍需要二次代价。本工作中，我们提出 SBM-Transformer，一种通过赋予每个注意力头一个混合成员随机块模型 (SBM) 来解决这两个问题的模型。然后，每个注意力头数据自适应地采样一个二部图，其邻接关系被用作每个输入的注意力掩码。在反向传播中，使用直通估计器将梯度流过离散采样步骤，并基于预测损失调整采样边的概率。因此前向与反向代价对边数呈线性，每个注意力头也可基于输入灵活选择。通过评估图的分布，我们从理论上证明 SBM-Transformer 是任意序列到序列函数在期望下的通用逼近器。在 LRA 和 GLUE 基准下的实证评估表明，我们的模型优于以往高效变体以及具有全注意力的原始 Transformer。我们的实现可在 https://github.com/sc782/SBM-Transformer 找到。