FLAT：一种用于缓解注意力瓶颈的优化数据流
机器学习
2022-09-27 v7 硬件体系结构
摘要
注意力机制最初旨在捕捉词与词之间的两两相关性，现已成为机器学习的支柱，其应用从自然语言处理扩展到其他领域。这种适应性的增长以极高的内存需求和计算复杂度为代价，尤其当输入元素数量较多时。该限制源于固有的有限数据复用机会以及内存占用的二次增长，导致严重的内存受限性和输入元素可扩展性的局限。本工作通过设计一种名为 FLAT 的定制化数据流优化来解决这些挑战，且不改变注意力机制的功能。该数据流通过独特的融合机制处理代价高昂的注意力操作，将内存占用的二次增长转变为仅线性增长。为充分发挥这一定制机制的潜力，我们提出一种分块方法来增强注意力操作间的数据复用。我们的方法既缓解了片外带宽瓶颈，又降低了片上内存需求。与最先进的未采用定制化数据流优化的边缘（云）加速器相比，FLAT 分别实现了 1.94 倍（1.76 倍）的加速以及 49%（42%）的能耗节省。当片上资源稀缺（20 KB–200 KB）时，FLAT 在输入序列长度从 512-token 到 64K-token 的一系列多样化常规基于注意力的模型上，平均实现了 1.5 倍的端到端延迟降低。我们的评估表明，最先进的 DNN 数据流应用于注意力操作时，在输入超过 512 个元素时达到效率极限。相比之下，FLAT 为输入多达 64K 个元素的 transformer 模型解除了限制。
引用
@article{arxiv.2107.06419,
title = {FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks},
author = {Sheng-Chun Kao and Suvinay Subramanian and Gaurav Agrawal and Amir Yazdanbakhsh and Tushar Krishna},
journal= {arXiv preprint arXiv:2107.06419},
year = {2022}
}