20-Arxiv-Fast Transformer Decoding: One Write-Head is All You Need

[wangqiangneu]

简介

分析了transformer中self-attention的复杂度，包括计算复杂度，空间复杂度，然后根据 空间复杂度/计算复杂度的结果，看内存带宽是否是瓶颈。这个值越小越好，当计算复杂度不变时，如果空间复杂度小，则说明在计算量相同的前提下，对内存的访问越小，则对于GPU来说更友好。标准的SAN在batch处理时，计算复杂度O(bnd^2)，空间复杂度O(bn^2d+nd^2)，比值O(n/d + 1/b)，其中n是序列长，b是bsz，d是输入向量维度。文章提出一个简单的方法，就是multi-head计算时，让Q还是multi-head，K，V则share a head。同时，减少的参数被补到ffn中。e.g. 文章中实现时ffn从4096 ->5440，文章称这种方法为multi-query。这种做法的计算复杂度保持不变，空间复杂度变为O(bnd+bn^2d+nd^2)，比值成为O(1/d+1/b+n/(dxh))，可以看到多了一个/h的系数。 multi-query被应用到所有类型的san中。他也对比了直接减少K,V的维度的方法，就是还是multi-head，但是head的dim变小。

论文信息

• Author: google
• Paper

总结

• 方法很简单，不知实际效果如何。
• 他在seq_len=128的句子上测的速度，猛涨（看表2里的样子是快6倍？？有点扯呢），这是合理的。但是实际应用时seq_len往往没这么大，提高还有这么明显吗？