【论文标题】SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruning

【发表会议】HPCA 2021

【作者团队】Hanrui Wang, Zhekai Zhang, Song Han

【作者机构】MIT

【论文链接】https://arxiv.org/abs/2012.09852

【推荐理由】

  • 注意力机制(Attention)现在是自然语言处理领域中的基础方法 (例如Transformer, BERT, GPT-2), 但是注意力模块速度不快. 本文从软件-架构联合设计的角度提出一种加速和节能算法, 叫做SpAtten, 结合剪枝和量化.
    • 剪枝部分: 提出串联 (cascade) token和head pruning, 对不重要的token和head进行剪枝 (注意区别: 传统的剪枝算法, 减掉的是权重, 这里减掉的是token和head). cascade的含义是, 一旦一个token或head被减掉了, 在之后的层中, 这个token或head也都不复存在, 如下图所示 (所以越深的层, 被减掉的token或head就越多).

      • 剪枝中需要确定不同token或者head的相对重要性, 所以需要一个重要性标准.
      • 对于token, 文中提出了一种累积计分策略: 累积不同轮的注意力概率 (概率越大, 输出受当前token的影响就大, 因而更重要)
      • 对于head, 重要性也是累积记分: 一个head输出的幅值(magnitude)作为其重要性的度量, 幅值越大, 越重要.
      • 为了高效对token或者head进行重要性排序, 本文提出一个top-k engine. 功能是从一个序列中找出top k的元素. 传统的排序需要O(n*logn)的时间复杂度, 本文的top-k engine可以做到O(n).
    • 量化部分: 提出渐进式量化 (progressive quantization): 如果过度量化损伤了准确率, 就使用更多的bit.
    • 剪枝得到的结果和量化的比特数(bit width)是会对着输入的不同而改变, 是自适应的(adaptive)的.

如下, 是整个加速器的系统框图:

  • 算法实现: 使用SpinalHDL对SpAtten进行实验, 并编译成RTL. 使用Verilator进行应用模拟. 对于HBM modeling, 使用Ramulator (HBM2模式). 使用Cadence Genus进行SpAtten的模型综合.
  • 实验效果: 比传统的平台 (TITAN Xp GPU, Xeon CPU)要快2个数量级, 比之前的同类加速器A^3 [HPCA, 2020], MNNFast [ISCA, 2019]分别快1.6x, 3.0x.

内容中包含的图片若涉及版权问题，请及时与我们联系删除