[论文解读] Scatterbrain: Unifying Sparse and Low-rank Attention Approximation

Beidi Chen, Tri Dao+4 more|arXiv (Cornell University)|Oct 28, 2021

Advanced Neural Network Applications参考文献 66被引用 9

一句话总结

Scatterbrain 提出了一种统一的、理论基础坚实的稀疏与低秩近似方法，用于高效注意力机制的 Transformer 模型，通过局部敏感哈希（LSH）实现稀疏性，通过核特征映射实现低秩结构。该方法在近似误差上相比基线方法降低最多 2.1 倍，同时将注意力显存占用减少 98%，在视觉模型中仅造成 1% 的准确率下降，且在语言与视觉任务上均优于单独使用稀疏或低秩方法的性能。

ABSTRACT

Recent advances in efficient Transformers have exploited either the sparsity or low-rank properties of attention matrices to reduce the computational and memory bottlenecks of modeling long sequences. However, it is still challenging to balance the trade-off between model quality and efficiency to perform a one-size-fits-all approximation for different tasks. To better understand this trade-off, we observe that sparse and low-rank approximations excel in different regimes, determined by the softmax temperature in attention, and sparse + low-rank can outperform each individually. Inspired by the classical robust-PCA algorithm for sparse and low-rank decomposition, we propose Scatterbrain, a novel way to unify sparse (via locality sensitive hashing) and low-rank (via kernel feature map) attention for accurate and efficient approximation. The estimation is unbiased with provably low error. We empirically show that Scatterbrain can achieve 2.1x lower error than baselines when serving as a drop-in replacement in BigGAN image generation and pre-trained T2T-ViT. On a pre-trained T2T Vision transformer, even without fine-tuning, Scatterbrain can reduce 98% of attention memory at the cost of only 1% drop in accuracy. We demonstrate Scatterbrain for end-to-end training with up to 4 points better perplexity and 5 points better average accuracy than sparse or low-rank efficient transformers on language modeling and long-range-arena tasks.

研究动机与目标

• 统一稀疏与低秩注意力近似方法，因为目前二者在不同场景下表现优异，通常被单独使用。
• 通过结合两种近似类型，缓解长序列建模中效率与准确率之间的权衡。
• 开发一种具有理论保证、无偏且高效的近似方法，避免在简单组合中出现重复计数问题。
• 通过实证验证，稀疏+低秩近似在图像生成与语言建模等多样化任务中优于单独使用任一方法。
• 实现全注意力的即插即用替换，仅造成极小的准确率损失并实现显著的显存节省，适用于预训练模型。

提出的方法

• Scatterbrain 使用局部敏感哈希（LSH）识别并近似 softmax 化注意力矩阵中的稀疏部分，而无需显式构建完整矩阵。
• 通过核特征映射近似注意力矩阵的低秩部分，利用随机特征近似实现高效计算。
• 该方法将注意力矩阵分解为稀疏与低秩分量之和，通过受鲁棒主成分分析启发的结构化分解确保无重复计数。
• 提供理论保证，证明其近似误差严格低于仅使用低秩基线方法的误差。
• 该框架设计为标准注意力层的即插即用替代品，与 Reformer 和 Performer 等现有高效 Transformer 架构兼容。
• 根据 softmax 温度与注意力分布的熵，自适应选择稀疏与低秩分量的最佳组合。

实验结果

研究问题

• RQ1 在哪些场景（例如基于 softmax 熵）下，稀疏、低秩或联合近似方法在注意力矩阵上的表现最佳？
• RQ2 是否能通过统一方法结合稀疏与低秩近似，实现低于任一方法单独使用时的近似误差？
• RQ3 是否存在一种理论合理且高效的组合方式，可避免重复计数或引入偏差？
• RQ4 改进的近似误差在语言建模与图像分类等下游任务的端到端性能上如何体现？
• RQ5 Scatterbrain 是否可作为全注意力的即插即用替代品，实现极小的准确率损失与显著的显存节省？

主要发现

• 在预训练的 T2T-ViT 模型中，Scatterbrain 将注意力显存减少 98%，仅造成 1% 的 top-1 准确率下降，优于 SMYRF 与 Local attention 方法。
• 在 BigGAN 图像生成任务中，作为即插即用替换，Scatterbrain 的近似误差比基线方法低 2.1 倍。
• 在语言建模与 Long Range Arena 任务的端到端训练中，Scatterbrain 的困惑度最高可比稀疏或低秩基线方法提升 4 个点，平均准确率提升 5 个点。
• 在 GLUE 基准测试中，当替换 BERT 的注意力层时，Scatterbrain 在 9 项下游任务中的 8 项上优于所有基线方法，包括 Performer 与 Smyrf。
• 可视化分析表明，稀疏+低秩近似在各层中误差最小，尤其在中熵区间表现最优，而单一方法在此区间常失效。
• 该方法在层次化（如语言）与非层次化（如图像）任务中均达到最先进性能，展现出广泛的适用性。