谷歌大脑Quoc发布Primer，从操作原语搜索高效Transformer变体

　　新智元报道

　　来源：arXiv

　　编辑：LRS

　　【新智元导读】调参、改激活函数提高模型性能已经见怪不改了。最近 Google Brain 的首席 Quoc 发布了一个搜索框架，能够自动搜索高效率的 Transformer 变体，并找到一些有效的模型 Primer，其中 ReLU 加个平方竟然能提升最多性能！

　　目前自然语言处理领域发展的红利都来自于大型的、基于 Transformer 的语言模型，但这些语言模型的训练成本、推理成本都高到劝退平民炼金术师。

　　而当模型参数量大到一定程度的时候，研究人员也在考虑如何在缩小模型的情况下，保持性能不变。

　　Google Brain 团队最近在 arxiv 上传了一篇论文，目标是通过寻找更高效的 Transformer 变体来降低训练和推理成本。

　　与之前的方法相比，新提出的方法在更低级别上执行搜索，在 Tensorflow 程序的原语上定义和搜索 Transformer。并提出了一种名为 Primer 的模型架构，训练成本比原始 Transformer 和用于自回归语言建模的其他模型变体要更小。

　　https://arxiv.org/abs/2109.08668

　　论文的作者是大神 Quoc V. Le，在斯坦福读博期间导师是吴恩达教授，目前是谷歌的研究科学家，Google Brain 的创始成员之一；seq2seq 的作者之一；谷歌 AutoML 的奠基人，提出包括神经架构等方法；EfficientNet 的作者等。

　　研究人员使用 TensorFlow（TF）中的操作来构造 Transformer 变体的搜索空间。在这个搜索空间中，每个程序定义了自回归语言模型的可堆叠解码器块。给定输入张量是一个长度为n且嵌入长度为d的序列，程序能够返回相同形状的张量。

　　堆叠时，其输出表示每个序列位置的下一个 token 的预测 embedding，并且程序只指定模型架构，没有其他内容。换句话说，输入和输出 embedding 矩阵本身以及输入预处理和权重优化不在这个程序的任务范围内。

　　解码器模型程序（DNA, decoder model program）定义了一个自回归语言模型，每个 DNA 都有一组子程序，其中子程序 0 是 MAIN 函数的入口点。每个子程序都由指令组成，这些指令被转换为 TensorFlow 代码行。

　　指令操作映射到原语词汇表中的基本 TensorFlow 库函数或父 DNA 子程序之一，原语词汇表由简单的原语 TF 函数组成，如 ADD、LOG、MATMUL 等等，但像 self-attention 这样的高级构建块不是搜索空间中的操作，自注意力可以从低级操作中构建出来的。

　　DNA 的子程序库由附加程序组成，这些程序可以通过指令作为函数执行。每个子程序只能调用子程序库中索引较高的子程序，这样就消除了循环的可能性。

　　使用父指令的参数集填充操作的参数，该参数集包含所有潜在操作参数的值，参数包括 Input 1（ 用作第一个 tensor 输入的隐藏状态的索引）、Input 2（第二个 tensor 输入的隐藏状态的索引）、Constant（实值常数，可以用于 MAX 等函数）、Dimensionsize（用来表示输出维度大小的整数）。特定操作中没有使用的参数就直接被省略掉。

　　研究人员还提出进化搜索（evolutionary search），目标是在搜索空间中找到最有效的模型架构。主要方法是设计一个固定的训练预算（使用 TPUv2 限时 24 小时），并将其适应性指标定义为 Tensor2Tensor 中 One Billion Words Benchmark (LM1B)上的困惑度。

　　这些架构搜索工作的明确目标是在优化效率时减少训练或推理步骤时间，在搜索过程中，可以发现将步长时间增加一倍、采样效率提高三倍是一个不错的修改方案，因为它最终使模型架构的计算效率更高。还可以将 ReLUs 平方化，并在注意力上增加深度卷积，从而增加训练步长时间。

　　这些操作极大地提高了模型的采样效率，通过大幅减少达到目标质量所需的训练步骤数量，减少了达到目标性能所需的总计算量。

　　通过这个搜索程序找到的模型被研究人员命名为 Primer，也就是原语搜索 Transformer（PRIMitives searched transformER）。

　　Primer 的改进主要有平方 ReLU 激活并在自注意力中的每个 Q、K 和 V 投影后添加一个深度卷积层。

　　最有效的修改是将变 Transformer 前馈块中的 ReLU 激活改进为平方 ReLU 激活函数，这也是第一次证明这种整流多项式激活在 Transformer 中有用。并且高阶多项式的有效性也可以在其他 Transfomer 非线性激活函数中观察到，例如 GLU 的各种变体，ReGLU、近似 GELU 等。然而平方 ReLU 与最常用的激活功能相比 ReLU、GELU 和 Swish 具有截然不同的渐近性。

　　平方 ReLU 确实与 ReGLU 有显著重叠，事实上，当 ReGLU 的U和V权重矩阵相同时，平方 ReLU 与 ReLU 是等效的。并且平方 ReLU 在更简单的同时，也能获得 GLU 变体的好处，且无需额外参数，并提供更好的质量。

　　研究人员使用三个 Transformer 变体与 Primer 进行对比：

　　1、Vanilla Transformer: 原始 Transformer，使用 ReLU 激活和 layer normalization。

　　2、Transformer+GELU: Transformer 的常用变体，使用 GELU 近似激活函数

　　3、Transformer++: 使用 RMS 归一化、Swish 激活和 GLU 乘法分支在前馈反向瓶颈（SwiGLU）中。这些修改在 T5 中进行了基准测试，并被表明是有效的。

　　实验表明，随着计算规模的增长，Primer 相对于 Transformer 的收益会增加，并且在最佳模型大小下遵循与质量相关的幂律。

　　研究人员还凭经验验证了 Primer 可以放入不同的代码库，发现可以显著加快训练速度，而无需额外调整。例如，在 500M 的参数大小下，Primer 在 C4 自回归语言建模上改进了原始 T5 架构，将训练成本降低了 4 倍。

　　此外，降低的训练成本意味着 Primer 需要更少的计算来达到目标 one shot 性能。例如，在类似于 GPT-3 XL 的 1.9B 参数配置中，Primer 使用 1/3 的训练计算来实现与 Transformer 相同的一次性性能。

　　研究人员已经开源了模型，以帮助提论文可重复性。

　　参考资料：

　　https://arxiv.org/abs/2109.08668