[论文解读] Primer: Searching for Efficient Transformers for Language Modeling

David R. So, Wojciech Mańke+4 more|arXiv (Cornell University)|Sep 17, 2021

Topic Modeling参考文献 56被引用 11

一句话总结

本文提出 Primer，一种更高效的自回归语言建模 Transformer 架构，通过直接在 TensorFlow 计算图上进行搜索实现。通过将 ReLU 激活函数平方化，并在 Q/K/V 投影后添加深度可分离卷积，Primer 相较于 T5 将训练计算量减少最多 4.2 倍，相较于 GPT-3 XL 级模型减少 3 倍，同时在各种设置下实现与原模型相当或更优的性能，且计算量更低。

ABSTRACT

Large Transformer models have been central to recent advances in natural language processing. The training and inference costs of these models, however, have grown rapidly and become prohibitively expensive. Here we aim to reduce the costs of Transformers by searching for a more efficient variant. Compared to previous approaches, our search is performed at a lower level, over the primitives that define a Transformer TensorFlow program. We identify an architecture, named Primer, that has a smaller training cost than the original Transformer and other variants for auto-regressive language modeling. Primer's improvements can be mostly attributed to two simple modifications: squaring ReLU activations and adding a depthwise convolution layer after each Q, K, and V projection in self-attention. Experiments show Primer's gains over Transformer increase as compute scale grows and follow a power law with respect to quality at optimal model sizes. We also verify empirically that Primer can be dropped into different codebases to significantly speed up training without additional tuning. For example, at a 500M parameter size, Primer improves the original T5 architecture on C4 auto-regressive language modeling, reducing the training cost by 4X. Furthermore, the reduced training cost means Primer needs much less compute to reach a target one-shot performance. For instance, in a 1.9B parameter configuration similar to GPT-3 XL, Primer uses 1/3 of the training compute to achieve the same one-shot performance as Transformer. We open source our models and several comparisons in T5 to help with reproducibility.

研究动机与目标

• 降低大规模 Transformer 语言模型的高训练与推理成本。
• 通过在 TensorFlow 计算图上进行低层级搜索，发现更高效的 Transformer 变体。
• 开发简单、实用且无需超参数调优的修改方法。
• 在不牺牲模型质量的前提下实现显著的计算量节省，尤其是在大规模场景下。
• 实现标准 Transformer 的即插即用替换，仅需极少代码修改。

提出的方法

• 搜索在定义自回归语言模型解码器模块的 TensorFlow 程序上进行。
• 搜索空间包含低层级操作，如激活函数、归一化层和注意力组件。
• 使用进化算法探索搜索空间，并在固定计算预算下优化验证损失。
• 每个模型均使用来自 T5 和 T2T 等成熟库的固定超参数集进行评估。
• 搜索空间支持组件重排和原语修改（例如改变激活或归一化顺序），不同于结构僵化的模块化搜索空间。
• 识别出两项关键修改：将 ReLU 激活函数平方化，并在 Q、K、V 投影后插入深度可分离卷积层。

实验结果

研究问题

• RQ1 对 Transformer 计算图进行低层级架构修改，是否能在不损失性能的前提下实现显著的训练成本降低？
• RQ2 所提出的修改（平方 ReLU 和 Q/K/V 投影后的深度可分离卷积）在不同模型规模和计算量下如何提升效率与质量？
• RQ3 所发现的架构是否可作为即插即用替换部署在现有代码库中而无需重新调优？
• RQ4 Primer 的效率提升是否随计算量增加而增强？其与模型质量的关系是否符合幂律规律？
• RQ5 这些改进是否在不同数据集、硬件平台和模型族之间具有泛化能力？

主要发现

• 对于参数量为 500M 的 C4 数据集，Primer 实现与原始 T5 模型相当的性能，但训练计算量减少 4.2 倍。
• 对于参数量为 19 亿的模型（类似 GPT-3 XL），Primer 仅用 1/3 的训练计算量，即可达到与 Transformer 相当的一次性推理性能。
• 改进效果在模型规模（2000 万至 19 亿参数）、数据集（LM1B、C4、PG19）、硬件（TPUv2 至 V100）和代码库（T5、Lingvo、T2T）之间均表现稳健。
• 两项关键修改——平方 ReLU 激活函数和在 Q/K/V 投影后插入深度可分离卷积——贡献了绝大部分性能提升。
• 当使用最优规模模型时，Primer 的计算节省随模型规模增加而增强，且与模型质量的关系符合幂律规律。
• Primer-EZ（简化版）在 T5 中表现良好，但完整版 Primer 在 Lingvo 和 T2T 等其他代码库中表现更优，表明其具备更广泛的适用性。