20-ICLR-LITE TRANSFORMER WITH LONG-SHORT RANGE ATTENTION

[wangqiangneu]

简介

针对Transformer在 mobile device上的改进（CPU）。做法很简单，干了两个事: flaten bottleneck，其实就是不把ffn hidden size阔成4倍，保持一样（这里的做法没啥说的，但是视角很好，参考figure 2(a)，把FFN hidden 阔成4倍那块看作输入，这样就可以把整个网络看作用了bottleneck缩减维度）；LSRA，就是简单的把SAN和DynamicConv合并在一起（作为两个branch）。对于SAN而言，Mult-Add是O(N^2*d)，注意之前的一些paper里SAN的计算复杂度是O(N^2d)也是说的Mult-Add这个意思，否则element-wise的看的话是`O(N^2d^2)，别混淆了。对于Transformer encoder而言，一层的Multi-Add是O(4Nd^2 + N^2d)(4个linear projecion + self attention) +O(2x4Nd^2)`（2个linear projection in FFN）

有意思的点

• CPU上ffn开销最大，跟我之前的认知一样
• 模型整体的Multi-Add是在seq=30上计算的
• 所谓mobile setting，一般指Multi-Add 500M
• 跟pruning、quantization互补

论文信息

• Author: MIT
• Paper

总结

• 整体逻辑：CPU想加速 -> 直接减少embed dim -> 对attn有伤害，而ffn其实消耗最高 -> flatten bottleneck 只减少ffn hidden dim -> 再此基础上再弄一下attn (把conv引进来) -> 用branch的形式，对于attn和conv而言，输入直接half了，单独的计算量自然也就是下来了
• 做法感觉一般般，强在写作、分析、实验~