动态路由条件计算-万亿级参数超大模型关键技术与MindSpore支持（1）

转载于：知乎
作者：金雪锋

最近项目组分析了类Switch Transformer的动态路由条件计算的模型，特分享出来，供大家参考。

主要分两部分，本次先介绍第一部分：

1. 动态路由条件计算-技术介绍
2. 动态路由条件计算-算法实现

动态路由条件计算-技术介绍

• 前言

近来，增大模型规模成为了提升模型性能的主要手段。特别是NLP领域的自监督预训练语言模型，规模越来越大，从GPT3的1750亿参数，到Switch Transformer的16000亿参数，又是一个数量级的增加。

模型规模的数量级的增大，虽然取得了一定程度的性能提升，甚至产生了某些意想不到的“神奇”效果（如GPT3），但其背后的计算开销成了最大的问题，比如GPT3训练使用了万级的GPU和数周的训练时间。

如何既能利用超大规模的参数来提升模型表达和性能，又能控制计算量的较小的增加，成为了最主要的挑战之一。

以MoE为代表的动态神经网络技术被重点引入。

大脑是典型的低能耗高效率的计算模式，稀疏激活是最重要的特性。

除了巨型模型在训练推理特别是训练时的计算效能挑战外，当前巨型模型的训练优化算法另一个更大的挑战是（不在此处讨论），BP算法是当前最为可用的深度网络优化，但更理想的优化算法需要高并行、优化过程非对称、并能够在时空维度通过局部持续优化完成整体优化。

1、传统的神经网络模型，前馈的时候，输入的batch中，每一个样本的处理，都将激活网络中的每一个参数参与计算。

2、条件计算最宽松的定义，指仅激活网络中某些部分的一类算法。Conditional Computation refers to a class of algorithms that activate only some of the different parts in a network. 在具体某类条件计算实现中，条件选择模式，可能按照输入的batch中每sample独立激活网络不同部分，可能按照输入数据空间上不同的部分（比如image不同区域或者channel），可能按照输入数据时间上不同的部分（比如time series的不同slide window或者video的不同的frame。），可能按照目标任务的不同每task独立的，可能按照非可学习的固定的随机分配不同的子网独立计算。

3、对不同的输入（原始或者前层），按照一定条件，选择性的执行后续部分网络的计算，这个技术下，有一些近似或相关的技术，如：dynamic neural network(s), conditional computing, conditional activation, sparse activating, selective execution, mixture of experts (MoE), dynamic routing, …；强相关的一些模型比如 Switch Transformer等。

• 条件计算的分类（广义）：

1、按照routing是否可学习可以分为：learnable routing conditional computation和 unlearnable routing conditional computation.

2、按照activation是否不执行non-activation计算，可以分为：hard conditional computation和soft conditional computation。对于hard-mode的条件计算，通过tensor挑选切分等操作，无论何种条件选择模式，不需要激活的数据将完全不参与不激活的网络部分的计算；soft-mode的条件计算，可能仅采取将相关数据置零等方式来避免产生计算效果，但还是和不需要激活网路部分实际执行计算过程。

• 条件计算的主要优势

1. 计算有效，降低能耗：通过部分激活部分计算，以每样本条件激活的条件计算为例，单个样本只需要经过整个SuperNet的一部分参与计算。
2. 更大网络，表达更强：由于一处到多处的Route，各处（层）的Input被路由到不同的子网独立计算，不同的输入的相互在各层的表达相对独立没有影响，表达能力更强，网络可以更大，但表达效率降低了。

• 条件计算的网络和计算形式

条件计算的网络和计算形式比较灵活，部分构建形式如：（此处省略具体模型和论文引用，参见https://intellabs.github.io/distiller/conditional\_computation.html）

1. 按照CV等task的特点，用多个独立的CNN作为expert网络，按照task来独立路由，尾部组合后给一个大网络。
2. 使用更复杂的cascading等形式组合不同层级的不同的expert网络。
3. 通过决策树等方法做数据变换实现路由。
4. 通过可学习的网络来选择路由。其中策略学习的损失有多种构建形式：直接使用分类等任务的主损失，对不同专家的重要性和负载构建损失作为辅助损失等等。

• 条件计算的路由策略

1. non-learnable/hard-mode，通过某种确定性策略，如LSH等方式计算路由。
2. learnable-mode，通过可学习网络计算路由。 网络规模可大可小，简单的可学习路由为单层权重：G(x) = P(X*W)，G(x)为路由Gate函数，X为输入， W为通损失函数来度量的可学习路由权重，P为某种挑选函数（如topk, sort等），在实际实现中，X*W的输入与权重计算结果可能作为后续网络的输入信息的一部分，不仅仅利用G(x)来选择路由，则需要对X*W的结果做归一化，更典型的形式则为：G(x)=P(N(X*W))，其中N为表达Normalization函数，如Softmax。

• 条件计算的冗余策略

条件计算的冗余策略，可分为无冗余条件计算和冗余条件计算：

1、无冗余条件计算可通过P(.)函数的实现如topk(k=1,…)来实现；

2、冗余条件计算，可以多种实现形式，可以通过P(.)函数的实现如topk(k=n,…)，n>=2来实现，也可以通过硬冗余模式，整个网络中支持输入的复制和多路计算实现。

• 条件计算的挑战

1、路由算法对模型质量的影响

无论输入和路由权重作用的信息（X*W），是仅作为路由选择并作为后续网络单元的输入，还是直接作为后续网络单元的输入的一部分，路由算法决定了输入信息的处理流向，对模型的整体质量都有很大影响。

2、路由(routing)/门(gate)的稳定性

随机初始化的路由/门的权重，权重自身在不断被训练调整；在前后层的网络持续训练变化，同一样本在训练的不同阶段会被分派到不同的后续网络单元中，这种动态变化过于剧烈，将严重影响整个网络训练过程的稳定性和收敛速度。

3、路由的专家样本重要性和负载的平衡性

训练阶段，每专家和样本批次中样本的关联度重要性，和每批次中样本被均衡分派到不同专家的负载平衡性，这两个指标既相关又冲突。需要分别构建损失函数作为辅助损失，来优化这两个指标。在arxiv:1701.06538《Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer》做了相关讨论。

• 关于条件计算/动态神经网络，更多的信息。

在《Dynamic Neural Networks: A Survey》arxiv:2102.04906 (https://arxiv.org/abs/2102.04906)一文中，作者对广义的动态神经网络，将各种动态网络相关的技术按照实例级、时间级、空间级做了分类。

1、Instance-wise Dynamic NN： 逐实例动态，每样本独立激活不同的网络和参数（MoE为这个方向）。

Dynamic Architecture：Dynamic Depth、Dynamic Width、Dynamic Routing/MoE；

Dynamic Parameter：Parameter Adjustment、Parameter Prediction、Dynamic Feature(s)

2、Spatial-wise Dynamic NN：空间级动态：图像等不同空间位置激活后续不同网络和参数。

(CNN等)：Pixel Level、Region Level、Resolution Level

3、Temporal-wise Dynamic NN：时间级动态：时序数据按时序维切分激活后续不同网络和参数。（video-frames, text-sequence, time-series, stream, ...)

Text-Sequence

Video-Frames

上述为该综述论文对Dynamic NN的总体分类。

从超大规模网络动态网络技术支撑角度，高表达能力，低计算代价为主的来考虑分类，从两个维度对动态网络技术分类:

1、按照在前馈计算时是否部分激活:

Hard-Dynamic：在前馈的时候，部分网络绝对不激活参与计算

Soft-Dynamic：在前馈的时候，部分网络经过softmax等gate/route后，通过张量元素置零等方式，失去表达能力，但会参与计算。

2、按照动态激活判定算法的输入：

逐样本级：（在输入层）按照每样本的实例来决定动态网络的后续激活。

亚样本级：（在输入层）样本内时间/空间级激活不同的后续网络单元。一般深度网络，不仅在输入层会被选择性激活执行，在中间层也类似。

其中，智能平台支持Hard-Dynamic逐样本级的动态神经网络，能比较自然的获得网络结构大颗粒的稀疏激活，在超大模型中能实现训练和推理的搞能效。

动态神经网络相比与静态结构的神经网络，在相关研究中，从效能，表达，泛化、鲁棒，可解释等方面做了大量对比研究。从智能平台通过计算成本尽量低的支持超大规模网络来提升模型性能的角度看，Efficiency和Representation最为重要：

1、Efficiency：静态网络“牵一发而动全身”，每一个样本输入整个网络/所有参数都要响应，这对超大网络来取得领先效果的模型能耗挑战太大。

2、Representation: 参数量更大，表达容量更大；但MoE等结构在深度网络的各层特征的表达上，复用降低，每参数的表达效率更低。

后续，将继续分享探讨算法实现。

最后，感谢岳大师的支持

推荐阅读

• Transformer系列 | 更深、更强、更轻巧的Transformer，DeLighT
• TeraPipe（支持超大模型训练的Token-Level的Pipeline并行）论文分析
• Towards Open World Object Detection -CVPR2021 Oral

更多嵌入式AI技术干货请关注嵌入式AI专栏。