[论文解读] LiteTransformerSearch: Training-free Neural Architecture Search for Efficient Language Models

Mojan Javaheripi, Gustavo Henrique de Rosa+7 more|arXiv (Cornell University)|Mar 4, 2022

Topic Modeling被引用 5

一句话总结

LiteTransformerSearch 提出了一种无需训练的神经架构搜索（NAS）方法，用于高效自回归 Transformer 模型，通过将解码器参数量作为困惑度的高相关代理，实现在目标设备上的快速、硬件感知的模型搜索。该方法在保持或超越 14 项任务性能的同时，相比基线模型实现了最高 2.5 倍的推理速度提升和 2.0 倍的内存使用降低。

ABSTRACT

The Transformer architecture is ubiquitously used as the building block of large-scale autoregressive language models. However, finding architectures with the optimal trade-off between task performance (perplexity) and hardware constraints like peak memory utilization and latency is non-trivial. This is exacerbated by the proliferation of various hardware. We leverage the somewhat surprising empirical observation that the number of decoder parameters in autoregressive Transformers has a high rank correlation with task performance, irrespective of the architecture topology. This observation organically induces a simple Neural Architecture Search (NAS) algorithm that uses decoder parameters as a proxy for perplexity without need for any model training. The search phase of our training-free algorithm, dubbed Lightweight Transformer Search (LTS), can be run directly on target devices since it does not require GPUs. Using on-target-device measurements, LTS extracts the Pareto-frontier of perplexity versus any hardware performance cost. We evaluate LTS on diverse devices from ARM CPUs to NVIDIA GPUs and two popular autoregressive Transformer backbones: GPT-2 and Transformer-XL. Results show that the perplexity of 16-layer GPT-2 and Transformer-XL can be achieved with up to 1.5x, 2.5x faster runtime and 1.2x, 2.0x lower peak memory utilization. When evaluated in zero and one-shot settings, LTS Pareto-frontier models achieve higher average accuracy compared to the 350M parameter OPT across 14 tasks, with up to 1.6x lower latency. LTS extracts the Pareto-frontier in under 3 hours while running on a commodity laptop. We effectively remove the carbon footprint of hundreds of GPU hours of training during search, offering a strong simple baseline for future NAS methods in autoregressive language modeling.

研究动机与目标

• 为解决在边缘计算和实时应用中寻找在困惑度、延迟和内存使用之间取得平衡的高效 Transformer 架构的挑战。
• 通过识别一种无需训练的模型性能代理，消除 NAS 过程中昂贵的训练需求。
• 实现在设备端的多目标 NAS，同时优化任务性能和硬件约束。
• 通过避免数百小时的 GPU 训练，减少 NAS 的碳足迹。
• 为未来自回归语言建模中的 NAS 提供一种简单、可扩展的基线。

提出的方法

• 基于经验观察：在多种 Transformer 架构中，解码器参数量与困惑度高度相关，因此可作为模型排序的零成本代理。
• 使用进化算法执行 NAS，结合参数量代理与目标设备上的延迟和峰值内存使用量的直接测量。
• 在通用硬件（如笔记本电脑）上执行整个搜索流程，无需 GPU 或超网络训练。
• 构建一个同时优化困惑度、延迟和内存利用率的帕累托前沿模型集合。
• 支持异构搜索空间，其中每个 Transformer 层可具有不同的超参数（如头数、d_model、d_ff）。
• 在搜索过程中使用真实硬件测量，确保目标设备的准确性，避免仿真不准确的问题。

实验结果

研究问题

• RQ1 解码器参数量能否作为自回归 Transformer 中预测困惑度的可靠、无需训练的代理？
• RQ2 参数量-困惑度相关性在异构和同构 Transformer 架构中在多大程度上保持成立？
• RQ3 无需训练的 NAS 方法能否在降低搜索成本和碳足迹的同时实现有竞争力的性能？
• RQ4 在设备端进行的硬件感知 NAS 在识别适用于实际部署的帕累托最优模型方面有多高效？
• RQ5 所提出的方法能否在多种硬件平台上，于延迟、内存和准确率方面超越现有基线？

主要发现

• 在同构和异构搜索空间中，解码器参数量与困惑度之间的等级相关性始终很高（r > 0.9），验证了该代理的可靠性。
• LiteTransformerSearch 在普通笔记本电脑上，无需 GPU 或训练，3 小时内即可提取出困惑度、延迟和内存的完整 3D 帕累托前沿。
• 在 16 层 GPT-2 和 Transformer-XL 上，该方法在保持相同困惑度的同时，实现了最高 2.5 倍的运行速度提升和 2.0 倍的峰值内存使用降低。
• 在零样本和少样本设置下，LTS 模型在 14 项任务上优于 3.5 亿参数的 OPT 模型，平均准确率更高，延迟最高降低 1.6 倍。
• 该方法通过消除数百小时的 GPU 训练，显著降低了 NAS 的碳足迹，为基于训练的方法提供了一种可持续的替代方案。
• 该方法原生可部署于目标设备，实现无需外部计算的实时、硬件感知架构搜索。