9 天前
组合器:具有稀疏计算成本的完整注意力Transformer
Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovec, Dale Schuurmans, Bo Dai

摘要
Transformer 架构提供了一类极具表达能力的模型,广泛应用于序列建模任务中。然而,其核心局限在于注意力层的计算复杂度和内存开销随序列长度呈二次方增长,即 $\mathcal{O}(L^2)$,这严重限制了其在超长序列场景中的应用。现有大多数方法通过在注意力矩阵中引入稀疏性或低秩假设来降低计算成本,但往往以牺牲模型表达能力为代价。为此,我们提出 Combiner,该方法在保持每个注意力头完整注意力能力的同时,实现了低计算与内存复杂度。其核心思想是将自注意力机制视为在每个位置上对嵌入向量的条件期望,并通过结构化分解来近似该条件分布。在此框架下,每个位置均可直接关注其他所有位置,或通过关注对应局部区域的抽象表示(即这些局部区域嵌入的条件期望)实现间接关注。我们证明,现有稀疏 Transformer 中广泛使用的多种稀疏注意力模式,均可启发此类结构化分解的设计,从而在保持全注意力能力的前提下,实现亚二次复杂度($\mathcal{O}(L\log L)$ 或 $\mathcal{O}(L\sqrt{L})$)。Combiner 可作为现有 Transformer 中注意力层的即插即用替代方案,且在主流深度学习框架中易于实现。在自回归与双向序列建模任务上的实验评估表明,该方法具有卓越性能,在多个图像与文本建模任务中取得了当前最优结果。