用神经架构搜索给 LLM 瘦身，同等准确度可让 LLaMA2-7B 模型大小降 2 倍。
大型语言模型（LLM）的一个主要特点是「大」，也因此其训练和部署成本都相当高，如何在保证 LLM 准确度的同时让其变小就成了非常重要且有价值的研究课题。
对此，研究社区已经提出了多种多样的方法，比如剪枝、稀疏化、量化等，它们的效果也各不一样。
近日，Intel Labs 发布了一项研究成果，宣称可使用神经架构搜索（NAS）高效地为 LLM 「瘦身」。他们基于 LLaMA2-7B 模型的实验表明，该技术不仅能降低模型大小，有时甚至还能让模型的准确度获得提升！
- 论文标题：LLaMA-NAS: Efficient Neural Architecture Search for Large Language Models
- 论文地址：https://arxiv.org/pdf/2405.18377
神经架构搜索（NAS）是一种让算法根据样本集自动设计神经网络架构的技术。之前曾被用来探索各式的新型神经网络架构，甚至可能发现人类从未构建过的网络架构，比如谷歌就曾通过 NAS 发现了一个新的特征金字塔架构，参阅《谷歌大脑提出 NAS-FPN：这是一种学会自动架构搜索的特征金字塔网络》
Intel Labs 的这个研究团队猜想：也许可以通过神经架构搜索（NAS）来降低模型的大小和复杂性。单次 NAS 方法往往需要一次训练一个超级网络（super-network）并允许使用权重共享原理进行子网络采样。但是，从头开始训练这些超级网络成本高昂，因此 LLM 领域少有探索。InstaTune 等技术则是利用现成的预训练权重，在微调阶段创造一个超级网络。
为了缓解与 LLM 相关的大小和复杂度问题，该团队做了一番研究探索：
- 该团队首次高效地使用单次 NAS 来降低 LLM 的大小和计算复杂度。他们的研究表明：对于某些标准的基准评测任务而言，LLaMA2-7B 庞大且复杂得没有必要。
- 他们发现，一种使用单次 NAS 的 LLM 压缩方法不仅性能优于剪枝和稀疏化技术，而且还无需剪枝和稀疏化通常需要的额外恢复微调步骤。
- 他们分析了网络参数，以便更好地理解新方法发现的网络的架构特征。他们证明，没有一套单一的架构启发式方法可以应用于多个标准基准测试任务。
- 他们提出了一种框架，无需专门的软件和或硬件就能直接用于产生压缩版 LLM。他们还发现，使用不做任何修改的标准版量化技术就能进一步压缩这些网络。
方法
搜索方法
为了让大型语言模型（LLM）适应多样化的硬件平台和性能要求，该团队的选择是使用 InstaTune 方法，这是一种全新的单次神经架构搜索（NAS）范式，参阅论文《Instatune: Instantaneous neural architecture search during fine-tuning》。
InstaTune 对传统的微调过程进行了扩展，其做法是让模型的架构具有弹性，使其可以探索固定结构之外更宽广的设计空间。这种灵活性让模型可以适应不同的任务、数据集和计算资源。InstaTune 增强 NAS 的方式是利用已有的模型权重，从而无需再从头开始训练超级网络。通过将 NAS 过程嵌入到微调阶段，InstaTune 能在节省计算资源的同时确保派生出的子网络针对目标任务进行了专门优化，而不是针对预训练目标进行了优化。
在这项研究中，研究团队使用 InstaTune 中提出的技术和 Alpaca 数据集对预训练后的 LLaMA2-7B 模型进行了微调，得到了一个超级网络和一个搜索空间 —— 它们可用于搜索过程。
但是，不同于 InstaTune，该团队并未在 LLaMA2-7B 上执行基于强大教师模型或超级网络的知识蒸馏，原因则主要是计算和内存的限制。
此外，虽然 InstaTune 会在同一次迭代中为超级网络和随机采样的子网络同时计算损失，但该团队为了降低内存使用量，选择的做法是交替地为超级网络和随机采样的子网络计算损失。在这个具体案例中，超级网络就是基础 LLaMA2-7B 模型，没有任何额外的层，MLP 模块的中间层大小也没有扩展。
遵循 InstaTune 的方法，该团队也采用了一种进化搜索框架来在多目标设置中对架构进行优化，从而在给定下游任务上对模型大小和准确度进行优化。
具体来说，他们采用的算法是轻量迭代式神经架构搜索（LINAS）算法。LINAS 组合了 NSGA-II 搜索和网络性能预测器，可以高效地找到帕累托最优的网络配置。它能在真实数据上迭代式地评估子网络，从而为性能预测器提供信息，然后它们就能预测大量子网络的性能。之后，将最有潜力的子网络选出来，再在下一轮迭代中进行评估。这个过程持续进行，直到在真实数据上完成一定的评估次数。
搜索空间
如前所述，执行最优子网络搜索的算法是 LINAS。使用 InstaTune 在 Alpaca 数据集上对预训练 LLaMA2-7B 模型进行微调之后，该团队定义了一组允许参数值，以便在搜索最优子网络时使用。这个搜索空间的创建方式是改变整个网络的层数以及每个 MLP 模块的中间大小。
表 1 给出了这些参数值，最终得到的搜索空间大小约为 1.3 × 10^10。
评估实验和结果
任务
该团队选用了以下任务来评估新方法的表现：
- AI2 推理挑战（ARC）
- 大规模多任务语言理解（MMLU）
- TruthfulQA
- WinoGrande
搜索分析
如前所述，该团队使用了 LINAS 算法来搜索 LLaMA2-7B 的最优子网络，使用的搜索空间见表 1。每个任务的搜索过程都执行 250 次评估。
AI2 推理挑战
针对 AI2 推理挑战（ARC）任务，是在模型大小 / ARC-c 准确度和模型大小 / ARC-e 准确度目标空间中执行的搜索。图 1 展示了在这两个目标空间中的帕累托边界。
执行完搜索后，该团队找到了一些能提供更高准确度和 / 或更小规模的子网络架构。比如某些子网络的 ARC-c 准确度比预训练的 LLaMA2-7B 更高，如图中红点所示。举例来说：一个子网络的准确度与预训练 LLaMA2-7B 网络一样，都规模却小 1.1 倍；另一个子网络的大小一样，但准确度高 1.9%。
大规模多任务语言理解
该团队也使用表 1 中的参数搜索了在大规模多任务语言理解（MMLU）任务上的帕累托最优子网络。该搜索的目标空间是模型大小 / MMLU 准确度。此外，为了了解在推理速度方面的增益，该团队还评估了搜索得到的帕累托最优子网络的吞吐量。
图 2 展示了在模型大小 / MMLU 准确度以及吞吐量 / MMLU 准确度这两个目标空间中的帕累托边界。
可以看到，在这些帕累托最优子网络中，某些架构在模型大小和吞吐量这两方面都有优势。比如某些子网络的 MMLU 准确度比预训练的 LLaMA2-7B 更高，如图中红点所示。在一个案例中，子网络的准确度比预训练 LLaMA2-7B 高 1.1%，同时又比它小 1.5 倍，快 1.3 倍。
TruthfulQA
对于 TruthfulQA 任务，该团队发现对于 TruthfulQA MC1，预训练 LLaMA2-7B 明显是过度参数化了。如图 3 所示，新方法发现的许多网络架构都明显优于预训练 LLaMA2-7B，同时大小也小得多。
WinoGrande
在 WinoGrande 上得到的结论与在 ARC-c 上得到的结果很相近。图 4 给出了在模型大小 / WinoGrande 搜索空间中的帕累托边界。
搜索找到了一些有更高准确度和 / 或更小 guim 的子网络架构。举例来说：一个子网络的准确度与预训练 LLaMA2-7B 网络一样，都规模却小 1.1 倍；另一个子网络的大小一样，但准确度高 1.4%。
标准基准性能总结
表 2 给出了各任务的一些比较突出的子网络，并总结了它们的标准基准性能。
剪枝和稀疏化性能比较
在通过新方法找到帕累托最优子网络之后，该团队选出了那些规模同等或更小的子网络（它们是通过剪枝或稀疏化得到的），并评估了它们在不同任务上的性能。然后，该团队比较了它们与 LLM-Pruner 和 SliceGPT 的标准基准性能，结果见表 3 和 4。
在表 3 中，ARC-c 和 MMLU 的某些子网络的准确度高于新方法找到的更大子网络。
原因是网络架构中不同中间大小的位置可能不会影响模型大小，但会影响准确度。因此，较小的子网络也可能取得更高的准确度，因为它选择了更好的中间大小。对于新方法找到的大小相近的子网络，在每个任务上的准确度均高于 LLM-Pruner 和 SliceGPT。
量化
为了进一步提升新方法找到的帕累托最优子网络的性能，该团队使用了定点（INT8）量化，使用的工具是 bitsandbytes。具体来说，首先量化搜索找到的帕累托最优子网络，然后在模型大小 / 准确度目标空间中再次评估它们。解码器层中的所有线性运算都被量化到了 INT8，而剩余层（比如嵌入层）依然是 FP16。
图 5 展示了量化前后在模型大小 / 准确度目标空间中的帕累托边界。
在为帕累托最优子网络使用了定点（INT8）量化之后，在目标空间中的帕累托边界「向左移了」。很显然，量化对准确度的影响很小，但却能显著降低模型大小。事实上，取得了高 1.1% MMLU 准确度且模型大小降低了 1.5 倍的非量化子网络在量化后不仅维持着原来的准确度，而且大小还降低了 2.5 倍。
表 5 比较了表 2 中非量化版子网络与其定点（INT8）量化后的版本。
可以看到，对于 MMLU，相比于非量化版的预训练 LLaMA2-7B 网络，量化版子网络的准确度几乎没有损失，而模型大小却分别降低了 2.0 和 2.5 倍。对于 ARC-e，量化版网络的准确度还比非量化版更高一点。而对于 WinoGrande，量化版子网络的大小不仅分别降低了 1.8 和 2.0 倍，同时还取得了 0.2% 的少量准确度增益。
搜索空间分析
该团队最后分析了层数和中间大小方面的性质。在层数方面，图 6 展示了在不同的任务上，一个被选择的网络对于不同的准确度百分位数具有层数 l ∈ {24, 28, 32} 的概率。
可以看到，层数越多，准确度往往就越高。但是，对于 ARC-c、MMLU 和 TruthfulQA 等某些任务而言，情况却并非如此，这说明 LLaMA2-7B 在这些任务有过度参数化的现象。
在中间大小方面，图 7 给出了在 ARC-c 任务上，选择出的所有 32 层网络架构，一个特定层的中间大小为 s ∈ {5504, 11008} 的概率。
可以看到，取得高 ARC-c 准确度的网络只在某些层有较高概率选择 5504 的中间大小。
具体来说，在 3、4、18、25、26、28、29 和 32 层，中间大小 s = 5504 的概率明显高于 s = 11008。评估 3、4、18、25、26、28、29 和 32 层的中间大小为 s = 5504 且其它层的中间大小为 s = 11008 的网络架构，得到的 ARC-c 准确度为 45.9%，与预训练 LLaMA2-7B 的准确度相当，从而验证了图 7 的结果。
另外，该团队还分析了在 MMLU 任务上评估的 24 层网络架构选择以及在 WinoGrande 任务上评估的 32 层网络架构选择，结果见图 8 和 9。
总体而言，可以得出结论：对于某些任务来说，没有必要使用完整的网络大小。
而这些结果的差异又表明，并不存在一种适用于所有任务的单一架构启发式方法。
