Scatterbrain：统一稀疏与低秩注意力近似
机器学习
2021-10-29 v1
摘要
高效 Transformer 的近期进展利用了注意力矩阵的稀疏性或低秩性质，以减少长序列建模中的计算和内存瓶颈。然而，在不同任务中实现模型质量与效率的权衡仍具挑战，难以用一种通用近似兼顾所有情况。为更好理解该权衡，我们观察到稀疏与低秩近似在不同 regime 下表现更优，这取决于注意力中的 softmax 温度，且稀疏 + 低秩可优于各自单独使用。受经典用于稀疏与低秩分解的鲁棒 PCA 算法启发，我们提出 Scatterbrain，一种统一稀疏（通过局部敏感哈希）与低秩（通过核特征映射）注意力以实现准确高效近似的新方法。该估计无偏且误差可证明较低。我们实证表明，在 BigGAN 图像生成和预训练 T2T-ViT 中作为直接替换时，Scatterbrain 可比基线实现 2.1 倍更低误差。在预训练 T2T 视觉 Transformer 上，即使不微调，Scatterbrain 也能以减少 98% 注意力内存为代价仅损失 1% 准确率。我们展示了 Scatterbrain 的端到端训练，在语言建模和 long-range-arena 任务上相比稀疏或低秩高效 Transformer 具有最高 4 点更好困惑度和 5 点更好平均准确率。
引用
@article{arxiv.2110.15343,
title = {Scatterbrain: Unifying Sparse and Low-rank Attention Approximation},
author = {Beidi Chen and Tri Dao and Eric Winsor and Zhao Song and Atri Rudra and Christopher Ré},
journal= {arXiv preprint arXiv:2110.15343},
year = {2021}
}
备注
NeurIPS 2021