为了更好地处理长序列数据,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/