[论文解读] Combiner: Full Attention Transformer with Sparse Computation Cost

Hongyu Ren, Hanjun Dai+5 more|arXiv (Cornell University)|Jul 12, 2021

Topic Modeling参考文献 44被引用 28

一句话总结

Combiner 通过将注意力建模为结构化条件分布，在子二次成本下提供完整注意力，在长序列文本和图像任务中取得强结果，并在若干基准测试中达到或超越基线。

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 the same sub-quadratic cost ($\mathcal{O}(L\log(L))$ or $\mathcal{O}(L\sqrt{L})$). Combiner is a drop-in replacement for attention layers in existing transformers and can be easily implemented in common frameworks. An experimental evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach, yielding state-of-the-art results on several image and text modeling tasks.

研究动机与目标

• 实现可扩展的全注意力序列建模，适用于极长序列。
• 将 Combiner 作为保持在子二次成本内的结构化条件期望方法引入，以实现完整注意力。
• 展示如何从现有稀疏注意力模式实例化 Combiner，以保持表达能力。
• 在图像密度估计和长距离序列任务中展示最先进的性能，同时保持易于实现性。

提出的方法

• 将注意力重新解释为给定查询位置的值嵌入的条件期望。
• 将条件分布分解为直接注意力项和多个局部（基于抽象）的注意力项。
• 使用受现有稀疏模式启发的亚二次局部因式分解方案，以实现全注意力并使成本达到 O(L log L) 或 O(L sqrt L)。
• 用缩放点乘和简单抽象（最大池化或 DeepSets）对局部区域的概率进行参数化。
• 提供若干 Combiner 实例（Fixed、Logsparse、Axial、Learnable），将现有稀疏模式转换为全注意力变体。
• 可选的混合分量以提高局部因式分解的表达能力。

实验结果

研究问题

• RQ1 是否可以通过对注意力分布的结构化因式分解，在计算和内存复杂度亚线性地实现全注意力？
• RQ2 如何将现有稀疏注意力模式重新用于产生全注意力，同时不牺牲效率？
• RQ3 Combiner 在自回归和双向任务（文本与图像域）上的经验性性能提升有哪些？

主要发现

• Combiner 在每个头内实现全注意力且成本为亚二次，与稀疏变体的渐进复杂度（如 O(L log L) 或 O(L sqrt L)）相匹配。
• 在 CIFAR-10 上，Combiner-Axial 达到 2.77 bits/dim，优于若干基线。
• 在 ImageNet-64 上，Combiner-Axial 达到 3.42 bits/dim，与此前的稀疏/Transformer 基线相比具有竞争力或更优。
• 在 Wiki-40B 的语言建模中，Combiner-Fixed-2k 与 Combiner-Axial-2k 变体的困惑度低于若干基线，使用 8k 序列时进一步提升。
• 在 Long-Range Arena 基准测试中，Combiner 变体（Fixed、Axial）在多任务中与 Vanilla Transformer 的性能相匹配或超越。
• Combiner 在密度估计与自回归/双向建模表现上表现出强劲的能力，同时可在标准框架中实现。