Combiner: Full Attention Transformer with Sparse Computation Cost
∗Hongyu Ren†, Hanjun Dai⋄, ∗Zihang Dai⋄ Mengjiao Yang⋄, Jure Leskovec†, Dale Schuurmans⋄,‡, Bo Dai⋄ †Stanford University, ⋄Google Research, Brain Team, ‡University of Alberta indicates equal contribution. The work was completed during HR’s internship at Google Brain.
Abstract
Transformers provide a class of expressive architectures that are extremely effective for sequence modeling. However, the key limitation of transformers is their quadratic memory and time complexity $\mathcal{O}(L^{2})$ with respect to the sequence length in attention layers, which restricts application in extremely long sequences. Most existing approaches leverage sparsity or low-rank assumptions in the attention matrix to reduce cost, but sacrifice expressiveness. Instead, we propose Combiner, which provides full attention capability in each attention head while maintaining low computation and memory complexity. The key idea is to treat the self-attention mechanism as a conditional expectation over embeddings at each location, and approximate the conditional distribution with a structured factorization. Each location can attend to all other locations, either via direct attention, or through indirect attention to abstractions, which are again conditional expectations of embeddings from corresponding local regions. We show that most sparse attention patterns used in existing sparse transformers are able to inspire the design of such factorization for full attention, resulting in th
中文速览
现有Transformer模型在处理超长序列时面临注意力机制计算量随序列长度平方级增长的瓶颈，稀疏注意力虽能降低成本但会牺牲每个位置关注所有其他位置的能力。研究者提出了Combiner方法，将自注意力重新理解为条件期望，并通过对概率空间进行结构化分解来近似这一条件分布——每个位置既可以直接关注局部邻居，也可以通过"抽象键"间接关注更远区域的分组表示，从而在每个注意力头内保留完整的全局注意力能力。这种分解方式能够将已有的稀疏注意力模式（如局部窗口、步幅跳跃、对数稀疏等）"升级"为完整注意力，同时将计算和内存复杂度控制在O(L log L)或O(L√L)的次二次方级别。实验表明，Combiner在图像密度估计（CIFAR-10达到2.77 bits/dim、ImageNet-64达到3.42 bits/dim）和长序列理解基准Long-Range Arena上均达到当时最优性能，且作为即插即用模块可直接替换现有Transformer的注意力层，无需专用硬件，对推动长序列建模的实用化具有重要意义。
原文 arXiv:2107.05768;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2107.05768v2