LiteTransformerSearch:面向高效语言模型的免训练神经架构搜索
机器学习
2022-10-19 v2 计算与语言
摘要
Transformer架构被普遍用作大规模自回归语言模型的构建模块。然而,寻找在任务性能(困惑度)与峰值内存占用、延迟等硬件约束之间具有最优权衡的架构并非易事。各类硬件的激增使该问题更趋严重。我们利用了一个略显意外的经验观察:自回归Transformer中解码器参数数量与任务性能具有高通秩相关,而与架构拓扑无关。该观察自然地引出了一个简单的神经架构搜索(NAS)算法,其以解码器参数作为困惑度的代理而无需任何模型训练。我们这一免训练算法的搜索阶段被称为轻量Transformer搜索(LTS),可直接在目标设备上运行,因其无需GPU。利用设备上实测,LTS提取困惑度对任意硬件性能代价的帕累托前沿。我们在从ARM CPU到NVIDIA GPU的多种设备以及GPT-2与Transformer-XL两种流行自回归Transformer骨干上评估LTS。结果显示,16层GPT-2与Transformer-XL的困惑度可分别以最高1.5倍、2.5倍更快运行时间及1.2倍、2.0倍更低峰值内存占用实现。在零样本与单样本设定下,LTS帕累托前沿模型在14项任务上相比3.5亿参数OPT取得更高平均准确率,且延迟最高降低1.6倍。LTS在商用笔记本上运行可在3小时内提取帕累托前沿。我们有效消除了搜索期间数百GPU小时训练的碳足迹,为未来自回归语言建模的NAS方法提供了强劲而简单的基线。
引用
@article{arxiv.2203.02094,
title = {LiteTransformerSearch: Training-free Neural Architecture Search for Efficient Language Models},
author = {Mojan Javaheripi and Gustavo H. de Rosa and Subhabrata Mukherjee and Shital Shah and Tomasz L. Religa and Caio C. T. Mendes and Sebastien Bubeck and Farinaz Koushanfar and Debadeepta Dey},
journal= {arXiv preprint arXiv:2203.02094},
year = {2022}
}