You Only Sample (Almost) Once：基于伯努利采样的线性代价自注意力
机器学习
2021-11-19 v1 计算与语言
摘要
基于 Transformer 的模型在自然语言处理（NLP）中被广泛使用。Transformer 模型的核心是自注意力机制，其捕捉输入序列中词对间的交互，且复杂度随序列长度呈二次增长。在更长序列上训练此类模型代价高昂。本文中，我们表明一种基于局部敏感哈希（LSH）的伯努利采样注意力机制，可将此类模型的二次复杂度降至线性。我们将自注意力视为与伯努利随机变量相关联的单个词之和，原则上可通过单次哈希一次性采样（尽管实践中该次数可能为一个小常数），从而规避二次代价。这导出了一种高效的采样方案来估计自注意力，其依赖于对 LSH 的特定修改（以使其在 GPU 架构上可部署）。我们在标准 512 序列长度的 GLUE 基准上评估了我们的算法，相对于标准预训练 Transformer 取得了良好性能。在用于评估长序列性能的 Long Range Arena（LRA）基准上，我们的方法取得了与 softmax 自注意力一致的结果，但具有可观的加速与内存节省，且常优于其他高效自注意力方法。我们的代码见 https://github.com/mlpen/YOSO
引用
@article{arxiv.2111.09714,
title = {You Only Sample (Almost) Once: Linear Cost Self-Attention Via Bernoulli Sampling},
author = {Zhanpeng Zeng and Yunyang Xiong and Sathya N. Ravi and Shailesh Acharya and Glenn Fung and Vikas Singh},
journal= {arXiv preprint arXiv:2111.09714},
year = {2021}
}
备注
Proceedings of the 38th ICML (2021)