Transformer变体（Sparse Transformer，Longformer，Switch Transformer）

不知不觉Transformer已经逐步渗透到了各个领域，就其本身也产生了相当多的变体，如上图。博主前一篇类似博文更新了 Transformer变体（Star-Transformer，Transformer-XL），本篇文章想整理一下这两篇很重要的Transformer变体，分别是Sparse Transformer and Switch Transformer。

Explicit Sparse Transformer: : Concentrated Attention Through Explicit Selection
标准Transformer的复杂度为O（n^2），但是否序列内的所有元素都有必要被关注到，是否有方法可以简化这个机制？所以本文的“Sparse”重点就体现在只有少量的token参与attention分布的计算，以提升注意力机制的集中度。即本来一个词只和少量的词有关，但是标准自注意力却会给所有的词都分配权重然后聚合，一种很自然的想法就是通过显式选择，只让模型关注少数几个元素就行。

模型图如上图，最左边是标准计算注意力的路线，中间的是Sparse的实现，可以看到区别就在于中间多了一个手工选择的Sparsification，最右则是它的执行示意图。简单来说就是在算softmax分数之前先进行top-k选择出少数的重要元素即可。具体来说先算内积：$$P=\frac{Q{K}^T}{\sqrt{d}}$$然后人工按内积分数过滤top-k个元素，即下式中的M操作，其他的P则直接置为负无穷，这样强制只让k个元素被关注。$$A=softmax(M(P,k))$$最后再把分数回乘给V：$$C=AV$$通过这种操作，可以使注意力更集中。这种减轻计算量的操作也让GPT-3等模型能够架得更大更暴力也取得了很好的但是玩不起效果。。。

• paper：https://arxiv.org/abs/1912.11637
• code：https://github.com/lancopku/Explicit-Sparse-Transformer

Longformer: The Long-Document Transformer
Longformer也算是一种比较经典的Sparse的方法了吧。一共提出了3种策略：

• Sliding Window：如上图(b)所示，跟CNN很像，给定一个固定的窗口大小w，其两边都有个w/2个token与其做attention。计算复杂度降为 O（n x w），即复杂度与序列长度呈线性关系。而且如果为每一层设置不同的窗口size可以很好地平衡模型效率和表示能力。

• Dilated sliding window：如上图所示，类似扩张CNN，可以在计算复杂度不变的情况下进一步扩大接收域。同样的，如果在多头注意力机制的每个头设置不同的扩张配置，可以关注文章的不同局部上下文，特别是通过这种Dilated可以扩张甚至很远的地方。

• Global Attention ：如图(d)所示，计算全局token可能表征序列的整体特性。比如BERT中的[CLS]这种功能，复杂度降为O（n）

完整内容可以看原文：

• paper：https://arxiv.org/pdf/2004.05150.pdf
• code：https://github.com/allenai/longformer

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
相比起Sparse Attention需要用到稀疏算子而很难发挥GPU、TPU硬件性能的问题。Switch Transformer不需要稀疏算子，可以更好的适应GPU、TPU等稠密硬件。主要的想法是简化稀疏路由。即在自然语言 MoE （Mixture of experts）层中，只将 token 表征发送给单个专家而不是多个的表现会更好。模型架构如上图，中间蓝色部分是比价关键的部分，可以看到每次router都只把信息传给分数p最大的单个FFN。而这个操作可以大大降低计算量。

然后另一方面，其成功的原因是有一个很秀的并行策略，其同时结合了数据并行+模型并行+expert并行。具体如下图：

其实对应模型架构来看，experts并行是算子间的并行，其对应的FFN内部有算子级的模型并行，而整个experts在计算图上又是个多并行FFN分支，这是算子间模型并行，所以能够获得更低的通信开销，提高并行的效率。

• paper：https://arxiv.org/abs/2101.03961
• code：https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/moe.py

下一篇文章继续整理相关内容：

• Transformer变体（Routing Transformer，Linformer，Big Bird）