[论文解读] ABC: Attention with Bounded-memory Control

Hao Peng, Jungo Kasai+6 more|arXiv (Cornell University)|Oct 6, 2021

Topic Modeling被引用 7

一句话总结

本文提出 ABC（带有限制记忆控制的注意力机制），一种统一框架，通过约束记忆大小来泛化高效注意力机制，实现线性时间复杂度。该工作提出 ABC$_{\text{MLP}}$，一种可学习的、上下文感知的记忆控制机制，在长序列上显著提升了准确率与效率，且与标准 Transformer 相比准确率损失极小。

ABSTRACT

Transformer architectures have achieved state-of-the-art results on a variety of sequence modeling tasks. However, their attention mechanism comes with a quadratic complexity in sequence lengths, making the computational overhead prohibitive, especially for long sequences. Attention context can be seen as a random-access memory with each token taking a slot. Under this perspective, the memory size grows linearly with the sequence length, and so does the overhead of reading from it. One way to improve the efficiency is to bound the memory size. We show that disparate approaches can be subsumed into one abstraction, attention with bounded-memory control (ABC), and they vary in their organization of the memory. ABC reveals new, unexplored possibilities. First, it connects several efficient attention variants that would otherwise seem apart. Second, this abstraction gives new insights--an established approach (Wang et al., 2020b) previously thought to be not applicable in causal attention, actually is. Last, we present a new instance of ABC, which draws inspiration from existing ABC approaches, but replaces their heuristic memory-organizing functions with a learned, contextualized one. Our experiments on language modeling, machine translation, and masked language model finetuning show that our approach outperforms previous efficient attention models; compared to the strong transformer baselines, it significantly improves the inference time and space efficiency with no or negligible accuracy loss.

研究动机与目标

• 将多种高效的注意力机制统一于单一抽象框架下，揭示看似迥异的方法之间的隐藏关联。
• 通过限制记忆大小来缓解标准 Transformer 在长序列中的二次方计算开销，同时保持性能。
• 开发一种可学习的、上下文感知的记忆控制策略，以在效率与准确率之间实现优于启发式方法的权衡。
• 证明此前被认为不适用于因果注意力的方法（如 Linformer）可通过 ABC 框架实现适配。
• 提供一种实用的、可直接替换标准自注意力机制的方案，使模型在不损失准确率的前提下显著提升推理速度与内存效率。

提出的方法

• 将注意力建模为一个有限记忆系统，其中固定大小的记忆存储键值对，查询访问由记忆组织策略控制。
• 通过将不同方法（如稀疏注意力、局部注意力、Linformer）的记忆组织机制视为不同的控制机制，该框架抽象出现有方法的共性。
• ABC$_{\text{MLP}}$ 引入一种基于神经网络的可学习控制器，根据上下文输入特征决定存储哪些 token 以及存储位置。
• 控制器使用多层感知机（MLP）计算记忆写入优先级，实现动态、上下文敏感的记忆管理。
• ABC 支持因果与非因果注意力，支持线性复杂度下的高效自回归解码与序列编码。
• 注意力计算通过在有限记忆上进行 softmax 归一化读取操作实现，确保恒定时间与恒定空间的访问开销。

实验结果

研究问题

• RQ1 能否将多种高效的注意力机制统一于单一抽象框架下，揭示其共享的结构原理？
• RQ2 可学习的、上下文感知的记忆控制策略是否在准确率与效率上均优于启发式记忆组织方法？
• RQ3 通过 ABC 框架，此前被认为不适用于因果注意力的成熟方法（如 Linformer）是否能够实现兼容？
• RQ4 记忆大小如何影响不同注意力类型与任务下的模型性能？
• RQ5 ABC$_{\text{MLP}}$ 在减少推理时间与内存使用的同时，能在多大程度上匹配或超越标准 Transformer 的性能？

主要发现

• ABC$_{\text{MLP}}$ 的解码速度相比 RFA 最快提升 1.25 倍，相比 T2R 快 20%，且内存开销显著降低。
• 当记忆大小为 8-8（交叉注意力与因果注意力）时，ABC$_{\text{MLP}}$ 仅比最佳配置低 1.1 BLEU 分，展现出极佳的效率-准确率权衡。
• 在文本编码任务中，ABC$_{\text{MLP}}$ 将内存使用减少 50%，推理速度相比标准 Transformer 基础模型提升 1.5 倍。
• 此前被认为不适用于因果注意力的 Linformer，通过 ABC 抽象被证明可兼容，并在机器翻译任务中取得具有竞争力的性能。
• ABC$_{\text{MLP}}$ 在所有 ABC 变体中准确率最高，且使用更小的记忆规模，证实了上下文感知记忆控制的优势。
• ABC$_{\text{MLP}}$ 的效率增益在长序列上最为显著，表明其在长上下文应用场景中具有显著的渐近优势。