LLaMA-NAS: Efficient Neural Architecture Search for Large Language Models

📄 arXiv: 2405.18377v1 📥 PDF

作者: Anthony Sarah, Sharath Nittur Sridhar, Maciej Szankin, Sairam Sundaresan

分类: cs.AI

发布日期: 2024-05-28

💡 一句话要点

LLaMA-NAS：面向大语言模型的高效神经架构搜索

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 神经架构搜索 大型语言模型 模型压缩 遗传算法 量化 LLaMA2 一次性NAS

📋 核心要点

  1. 大型语言模型计算成本高昂，限制了其在资源受限设备上的部署。
  2. 提出LLaMA-NAS，通过一次性神经架构搜索，寻找更小、更高效的LLM架构。
  3. 实验表明，该方法能在保证性能的前提下，显著降低模型大小和提高吞吐量。

📝 摘要（中文）

现代大型语言模型（LLMs）在自然语言处理、复杂推理、情感分析等任务中表现出色，应用广泛。然而，它们的高内存和计算成本限制了在多数硬件平台上的使用。为解决此问题，我们提出了一种基于单次NAS的有效方法，用于寻找基于LLaMA2-7B的Pareto最优网络架构。具体而言，我们仅对LLaMA2-7B进行一次微调，然后应用基于遗传算法的搜索来寻找更小、计算复杂度更低的网络架构。实验表明，对于某些标准基准任务，预训练的LLaMA2-7B网络过于庞大和复杂。我们实现了模型大小减少1.5倍，吞吐量提高1.3倍，而准确率几乎没有下降。此外，我们的方法比某些剪枝或稀疏化技术更有效。最后，我们展示了量化与我们的方法的互补性，并且可以使用量化进一步降低我们找到的网络的规模和复杂性。我们相信，我们的工作提供了一种自动创建LLM的方法，可以在成本较低且更容易获得的硬件平台上使用。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）计算和内存成本过高的问题，这使得它们难以在资源受限的硬件平台上部署。现有方法，如模型剪枝和稀疏化，虽然可以减小模型大小，但往往效率较低，且可能导致显著的性能下降。

核心思路：论文的核心思路是利用神经架构搜索（NAS）自动寻找更优的LLM架构，该架构在保持性能的同时，具有更小的模型尺寸和更低的计算复杂度。通过一次性NAS，避免了对每个候选架构进行独立训练，从而显著提高了搜索效率。

技术框架：LLaMA-NAS的技术框架主要包含两个阶段：首先，对LLaMA2-7B进行一次微调，得到一个“超级网络”（supernet），该网络包含多个候选架构。然后，利用基于遗传算法的搜索策略，在超级网络中寻找Pareto最优的子网络架构。搜索过程以模型大小和性能为优化目标，最终得到一系列在大小和性能之间取得良好平衡的LLM架构。

关键创新：该论文的关键创新在于将一次性NAS应用于大型语言模型，并结合遗传算法进行高效的架构搜索。与传统的NAS方法相比，一次性NAS避免了对每个候选架构进行独立训练，从而大大降低了搜索成本。此外，遗传算法能够有效地探索架构空间，找到性能优异的子网络。

关键设计：论文的关键设计包括：1) 使用LLaMA2-7B作为基础模型，保证了搜索结果的实用性；2) 采用基于遗传算法的搜索策略，平衡了搜索效率和搜索质量；3) 将模型大小和性能作为多目标优化问题，寻找Pareto最优的架构；4) 探索了量化与NAS的结合，进一步降低模型大小。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LLaMA-NAS能够在保证性能的前提下，显著降低LLM的模型大小和提高吞吐量。具体而言，对于某些标准基准任务，该方法实现了1.5倍的模型大小缩减和1.3倍的吞吐量提升，而准确率几乎没有下降。此外，实验还证明了LLaMA-NAS比传统的剪枝和稀疏化方法更有效。

🎯 应用场景

LLaMA-NAS的应用场景广泛，包括移动设备、边缘计算设备等资源受限平台上的LLM部署。该方法能够自动生成适用于特定硬件平台的定制化LLM，降低部署成本，提高推理效率。此外，该研究为开发更高效、更绿色的LLM提供了新的思路，有助于推动人工智能的可持续发展。

📄 摘要（原文）

The abilities of modern large language models (LLMs) in solving natural language processing, complex reasoning, sentiment analysis and other tasks have been extraordinary which has prompted their extensive adoption. Unfortunately, these abilities come with very high memory and computational costs which precludes the use of LLMs on most hardware platforms. To mitigate this, we propose an effective method of finding Pareto-optimal network architectures based on LLaMA2-7B using one-shot NAS. In particular, we fine-tune LLaMA2-7B only once and then apply genetic algorithm-based search to find smaller, less computationally complex network architectures. We show that, for certain standard benchmark tasks, the pre-trained LLaMA2-7B network is unnecessarily large and complex. More specifically, we demonstrate a 1.5x reduction in model size and 1.3x speedup in throughput for certain tasks with negligible drop in accuracy. In addition to finding smaller, higher-performing network architectures, our method does so more effectively and efficiently than certain pruning or sparsification techniques. Finally, we demonstrate how quantization is complementary to our method and that the size and complexity of the networks we find can be further decreased using quantization. We believe that our work provides a way to automatically create LLMs which can be used on less expensive and more readily available hardware platforms.