Treeformer：用于高效注意力计算的稠密梯度树
计算与语言
2023-03-20 v2 机器学习
摘要
基于Transformer架构的标准推理和训练随输入序列长度呈二次方扩展。对于包括网页翻译、问答等在内的多种应用而言，这一开销大得难以承受。因此，近期出现了若干方法，通过施加不同的注意力结构（如稀疏性、低秩、使用核近似注意力）来加速注意力计算。在本工作中，我们将注意力计算视为最近邻检索，并使用基于决策树的分层导航，将每个查询词的检索代价从随序列长度线性降为近乎对数。基于此种分层导航，我们设计了Treeformer，它可使用两种高效注意力层之一——TF-Attention和TC-Attention。TF-Attention以细粒度风格计算注意力，而TC-Attention是一种粗粒度注意力层，同时也确保梯度是“稠密”的。为优化此类具挑战性的离散层，我们提出了一种两级自举训练方法。通过在标准NLP基准（尤其是长序列）上的大量实验，我们证明我们的Treeformer架构在注意力层少用30倍FLOPs的情况下，准确率几乎与基线Transformer相当。与Linformer相比，在注意力层使用相似FLOPs时，准确率可高出多达12%。
引用
@article{arxiv.2208.09015,
title = {Treeformer: Dense Gradient Trees for Efficient Attention Computation},
author = {Lovish Madaan and Srinadh Bhojanapalli and Himanshu Jain and Prateek Jain},
journal= {arXiv preprint arXiv:2208.09015},
year = {2023}
}
备注
ICLR 2023