为了更好地处理长序列数据，Scaling Transformer 进一步配备了来自 Reformer 的 LSH（局部敏感哈希）注意力和 FFN 块循环，从而产生了 Terraformer 模型。
混合专家系统 MoE
专家混合系统 (MoE) 模型是一种专家网络的集合，每个样本仅激活网络的一个子集来获得预测结果。这个想法起源于上世纪九十年代并且与集成方法密切相关。有关如何将 MoE 模块合并到 Transformer 的详细信息，可以查看本文作者之前写的关于大型模型训练技术的帖子和 Fedus 等人关于 MoE 的论文。
使用 MoE 架构，在解码时仅使用部分参数，因此节省了推理成本。每个专家的容量可以通过超参数容量因子 C 进行调整，专家容量定义为：
每个 token 需要选择前 k 个专家。较大的 C 会扩大专家容量，提高性能，但这样做计算成本更高。当 C>1 时，需要增加一个松弛容量；当 C<1 时，路由网络需要忽略一些 token。
路由策略改进
MoE 层有一个路由网络来为每个输入 token 分配一个专家子集。原生 MoE 模型中的路由策略是将每个 token 以不同的方式路由到按自然顺序出现的首选专家。如果路由到的专家已经没有多余的空间，token 将被标记为溢出并被跳过。
V-MoE 将 MoE 层添加到 ViT (Vision Transformer) 中。它与以前的 SOTA 模型的性能相匹配，但只需要一半的推理计算。V-MoE 可以扩展成一千五百万个参数。有研究者在实验中将 k=2、专家需要 32 位，每 2 位专家间放置一层 MoE。
由于每个专家的能力有限，如果某些重要且信息丰富的 token 在预定义的序列顺序（例如句子中的单词顺序或图像中 patch 的顺序）中出现得太晚，则可能不得不丢弃它们。为了避免原生路由方案中的这种缺陷，V-MoE 采用 BPR（批量优先路由）首先将专家分配给具有高优先级分数的 token。BPR 在专家分配之前计算每个 token 的优先级分数（前 k 名路由器得分的最大值或总和），并相应地更改 token 的顺序。这保证了核心的 token 能优先使用专家容量的缓冲区。
图 15. 当 C<1 时，根据优先级分数丢弃图像 patch 的方式。
当 C≤0.5 时，BPR 比普通路由效果更好，此时模型开始丢弃大量 token。这使模型即使在非常低的容量下也能与稠密网络一较高低。
在研究如何解释图像的类别与专家之间的关系时，研究者观察到早期的 MoE 层更通用，而后期的 MoE 层可以专门用于某类图像。
任务级 MoE 将任务信息考虑在内，并且将路由 token 在任务级的视角来处理。研究者以 MNMT（多语言神经机器翻译）为例，根据目标语言或语言对进行翻译任务分组。
Token 级路由是动态的，每个 token 的路由决策是不相交的。因此，在推理时，服务器需要预加载所有专家。相比之下，任务级路由是静态的，甚至是固定的任务，因此一个任务的推理服务器只需要预加载 k 个专家（假设 top-k 才有路由）。根据研究者的实验，与稠密模型的 baseline 相比，任务级 MoE 可以实现与 token MoE 类似的性能增益，峰值吞吐量高 2.6 倍，解码器小 1.6%。
任务级 MoE 本质上是根据预定义的启发式方法对任务分布进行分类，并将此类人类知识纳入路由器。当这种启发式不存在时，任务级 MoE 就难以使用了。
PR MoE 让每个 token 通过一个固定的 MLP 和一个选定的专家。由于靠后的 MoE 更有价值，PR MoE 在靠后的层上设计了更多的出口。DeepSpeed 库实现了灵活的多专家、多数据并行，以支持使用不同数量的专家来训练 PR MoE。
图 16。PR MoE 架构与标准 MoE 的对比图。
内核方面的改进措施
专家网络可以托管在不同的设备上。然而，当 GPU 数量增加时，每个 GPU 上的专家数量就会减少，专家之间的通信成本变得更加昂贵。跨多个 GPU 的专家之间的多对多通信依赖于 NCCL 的 P2P API，这个接口不能占据高速链路所有的带宽，这是因为使用的节点越多，单个 chunk 越小。现有的多对多算法在大规模问题上性能较差，单个 GPU 的工作量不能提升。针对这种情况，有多种内核改进措施来实现更高效的 MoE 计算，例如使多对多通信更便宜 / 更快。
DeepSpeed 库和 TUTEL 都实现了基于树的分层多对多算法，该算法在节点内使用多对多算法处理，然后再在节点间实现多对多。这种算法将通信跳数从 O（G）减少到，其中 G 是 GPU 节点的总数，G_(node) 是每个节点的 GPU 内核数。尽管在这样的实现中通信量增加了一倍，但当批大小较小时 1×1 卷积层存在延迟，因此可以更好地扩展 batch 的规模。
DynaMoE 使用动态再编译使计算资源适应专家之间的动态工作负载。再编译机制需要从头开始编译计算图，并且只在需要时重新分配资源。它会琢磨分配给每个专家的样本数量，并动态调整其容量因子 C，以减少运行时的内存和计算需求。这种方法基于在训练早期对专家和样本的分配关系的观察，在模型收敛后引入样本分配缓存，然后使用再编译算法消除门控网络和专家之间的依赖性。
架构优化
论文《Efficient Transformers: A Survey》回顾了一系列新的 Transformer 架构，并针对提高计算和内存效率进行了一些改进，除此以外，大家还可以阅读这篇文章《The Transformer Family》，以深入了解几种类型的 Transformer 改进。
图 17. 高效 transformer 模型的分类
自注意力机制的二次时间复杂度和内存复杂性问题是提高 transformer 解码效率的主要瓶颈，因此所有高效 transformer 模型都对原本稠密的注意力层应用了某种形式的稀疏化措施。
1. 固定模式：使用预定义的固定模式限制注意力矩阵的感受野：
可以将输入序列分成固定的块；
图像 transformer 使用了局部注意力；
稀疏 transformer 使用了跨线注意力模式；
Longformer 使用了 dilated 注意力窗口；
可以使用 strided 卷积压缩注意力来减少序列长度。
2. 组合模式：对输入的 token 进行排序 / 聚类以实现更优化的序列全局视图，同时保持固定模式的效率优势
稀疏 transformer 结合了跨步和局部注意力；
给定高维输入张量，axial transformer 不会将输入 flattened 后再使用注意力机制，而是使用多注意力机制，一个注意力对应着输入张量的一个轴；
Big Bird 模型设计了一些关键组件，即（1）全局 token，（2）随机注意力（query 向量随机绑定 key 向量）和（3）固定模式（局部滑动窗口）。
3. 可学习模式：通过学习确定最佳注意力模式：
Reformer 使用局部敏感哈希将 token 聚类；
路由 transformer 用 k-means 将 token 聚类；
Sinkhorn 排序网络会对输入序列块的排序算法进行学习。
4. 递归：通过递归连接多个 block/segment：
Transformer-XL 通过在 segment 之间重用隐藏状态来获取更长的上下文；
通用 transformer 将自注意力与 RNN 中的循环机制相结合；
Compressive transformer 是 Transformer-XL 的扩展，具有额外的内存，具有 n_m 个内存槽和 n_(cm) 个压缩内存槽。每当有新的输入段被输入到模型当中时，主内存中最久未更新的前 n_s 个激活函数都会被转移到压缩内存中。
5.Side Memory：使用可以一次访问多个 token 的 Side Memory 模块
Set Transformer 设计了一种受归纳点方法启发的新注意力；
ETC（Extended transformer construction）是 Sparse Transformer 的变体，具有新的全局 - 局部注意力机制；
Longformer 也是 Sparse Transformer 的变体，使用 dilated 滑动窗口。随着模型网络的深入，感受野也会逐渐增加。
6. 节省内存：更改架构以使用更少的内存：
Linformer 将 key 和 value 的代表长度的维度投影到低维表示（N→k），因此内存复杂度从 N×N 降低到 N×k；
Shazeer 等人提出了多 query 注意力，在不同注意力头之间共享 key 和 value，大大减少了这些张量的大小和内存成本。
7. 使用内核：使用内核可以让自注意力机制的公式书写起来更简单。需要注意的使，这里的内核是指内核方法中的内核，而不是 GPU 操作程序。
8. 自适应注意力：让模型学习最佳注意力广度，或决定何时按每个 token 提前退出：
自适应注意力广度训练模型，通过 token 和其他 key 之间的 soft mask 机制，为每个 token、每个注意力头学习最佳的注意力广度；
通用 transformer 结合了循环机制，并使用 ACT（自适应计算时间）来动态决定循环几次；
深度自适应 transformer 和 CALM 使用一些置信度度量方法来学习何时提前退出每个 token 的计算层，这样可以在性能和效率之间找到一种平衡。
原文链接：https://lilianweng.github.io/posts/2023-01-10-inference-optimization/
