移动设备上的大模型优化：MobileLLM
移动设备上的大模型优化：MobileLLM
现有的大模型因为过于庞大难以直接部署在资源受限的移动设备上。针对这一问题，来自Meta公司和PyTorch的研究团队提出了一种新的优化方法，旨在为移动设备设计高效的、参数少于十亿的高质量大模型，即MobileLLM。
图1 为小于10亿参数的大模型（LLMs）在零样本常识任务上的平均得分。图中每个气泡的面积与模型家族中变体的模型大小成正比。为了确保评估过程的一致性，使用开源的Hugging Face模型来评估先前方法的结果。任务的完整列表在表3中。
在小于10亿参数的模型中，MobileLLM在零样本常识任务上表现出色，相较于其他同规模的模型，它在准确度上有显著提升。这表明MobileLLM在模型架构和训练策略上的有效性，使其能够在有限的参数规模下实现更好的性能。
方法
研究团队通过以下方法来构建MobileLLM：
- 深层瘦架构：与普遍观点相反，对于小型LLMs，增加深度比增加宽度更有助于提高性能。
- 嵌入共享：重新审视嵌入共享方法，通过共享输入和输出嵌入层的权重，减少了模型的参数数量，同时保持了模型性能。
- 分组查询注意力：利用分组查询注意力机制，减少了键值头的数量，同时提高了模型的效率。
- 块状权重共享：提出了一种新的权重共享方法，通过在相邻块之间共享权重，减少了模型大小，同时降低了延迟。
研究者针对125M和350M两种模型规模进行了探索，并在这两种情况下都实现了一致的性能提升。
图3 呈现了十亿规模变换器模型的设计发展路线图。图中的前景色和背景色条形分别代表了125M和350M模型在零样本常识推理任务上的平均准确率。起始的125M模型是一个12层768维的结构，通过改善前馈网络设计、网络深度调整和权重共享策略得到了增强。
实验的配置使用了32个A100 GPU，每个GPU的批量大小为32。他们先进行了探索性实验，使用0.25T tokens进行了120k次迭代。然后在1T tokens上对表3和表4中报告的最佳模型进行了480k次迭代的训练。这些模型在零样本常识推理任务上进行了评估，包括多种任务和数据集，如ARC-easy、ARC-challenge、BoolQ、PIQA、SIQA、HellaSwag、OBQA、WinoGrande，以及使用TQA和RACE数据集的问答和阅读理解任务。
研究者们首先调查了通常在前馈网络（FFNs）中使用的激活函数，并发现SwiGLU对小模型也是有益的。通过将普通的FFN（FC → ReLU → FC）改为SwiGLU，125M模型在零样本推理任务上的平均表现从42.6提高到了43.9。因此，后续实验中FFN采用了SwiGLU。
一种普遍的观点认为即变换器模型的性能主要由参数数量、训练数据集的大小和训练迭代次数决定。研究者的实验结果表明，对于小模型来说，增加深度比增加宽度更关键。他们训练了19个模型，包括9个约125M参数和10个约350M参数的模型，这些模型在大小上相似但在深度和宽度上有所不同。实验结果一致表明，更深更瘦的模型比更浅更宽的模型表现更好。图4（a）和（b）展示了在大多数零样本推理任务中更深网络的优越性能，包括ARC-easy、ARC-challenge、PIQA、HellaSwag、OBQA、WinoGrande。特别是在TQA和RACE数据集上，这种趋势更加明显，如图4（c）-（f）所示。
在十亿规模以下语言模型中，嵌入层占据了参数数量的相当一部分。例如，在嵌入维度为512、词汇表大小为32k的情况下，输入和输出嵌入层各有1600万个参数。这些嵌入层加起来占到了125M参数模型总参数的20%以上。与此相反，在更大的语言模型中，这一比例要低得多。例如，在LLaMA 7B模型中，输入和输出嵌入仅占总参数的3.7%，在LLaMA-70B模型中仅占0.7%。研究者重新考虑了输入输出嵌入共享的概念。通过共享嵌入，即重用输入嵌入权重作为输出全连接层的权重，构建了更有效和紧凑的模型架构。表1展示了使用不同嵌入共享策略对30层模型在零样本常识推理任务上性能的影响，包括不共享嵌入、共享嵌入以及共享嵌入加上增加深度（从30层增加到32层）的情况。通过这些实验，研究者们发现共享嵌入可以显著减少模型的参数数量，并且通过增加模型深度可以恢复或提高准确率。
研究者调查了小变换器模型的最优头大小。在头尺寸和多个头的非线性组合之间进行权衡是选择头大小的关键考虑因素。另外大多数先前的研究通常在十亿参数以下语言模型中使用相同数量的键值头和查询头。相反，研究者们发现分组查询注意力，不仅可以减少键值缓存大小，也可以有效地减少小语言模型中键值头的冗余。分组查询注意力可以看作是一种重量共享的形式，用于重量重用，其中键值头的数量是查询头数量的1/n，并且kv头在计算注意力分数和输出时与查询头一起重复n次。在这里，n ∈ Z+表示查询头数量可以整除的正整数。
对于如何通过简单地复制变换器块来实现准确性提升，而不需要进行架构修改或增加模型大小。研究者们进一步检查了三种不同的权重共享策略，并发现在内存受限的语言模型解码过程中，立即块状重复层共享策略在性能上优于其他策略，并且由于共享权重可以立即在缓存中计算两次，因此可以避免在SRAM和DRAM之间传输权重，从而提高了自动回归推理的总体执行速度。这些发现促使研究者们采用立即块状共享策略在模型设计中，并将采用层共享的模型命名为MobileLLM-LS。
实验
实验使用Adam优化器从头开始训练MobileLLM，权重衰减设置为0.1。实验在32个A100 GPU上进行，每个GPU的批量大小为32。初始学习率设置为2e-3，并采用余弦衰减策略。
对0.25T tokens进行了120k次迭代的快速探索性实验，并在1T tokens上对表3和表4中报告的最佳模型进行了480k次迭代的训练。
在零样本常识推理任务、问答和阅读理解任务上比较了MobileLLM模型与现有的十亿参数以下最先进模型的性能。使用开源Hugging Face模型评估基线方法，以确保评估过程的一致性。
表2展示了层共享策略对零样本常识推理任务的影响。例如，对于125M基线模型，使用即时块状共享策略可以将平均准确率从44.6%提高到45.0%。
表3显示了MobileLLM在常识推理任务上的性能，与Cerebras-GPT、OPT、BLOOM等模型相比，MobileLLM在125M模型尺寸下显著优于之前的模型，并且在350M模型尺寸下也超过了之前的最先进模型。
为了验证十亿参数以下模型在移动设备应用中的有效性，研究者们评估了它们在聊天和API调用这两个关键移动设备任务上的性能。证明了即使在参数数量有限的情况下，通过精心设计的模型架构和训练策略，MobileLLM也能够实现与大型模型相媲美的性能。
研究者对MobileLLM模型以及先前的最先进模型进行了微调，并在AlpacaEval和MT-Bench这两个聊天基准上进行了评估。MobileLLM模型显著优于之前的十亿参数以下模型，甚至超过了一些具有10亿参数的模型。
MobileLLM模型和之前最先进模型在AlpacaEval和MT-Bench聊天基准测试上的结果
API调用是移动设备上的常见应用，尤其是在与语音转文字模型配合使用时。研究者们创建了一个包含5000个训练样本和2500个测试样本的合成数据集，并对模型进行了微调。
表6显示MobileLLM-350M在API调用的意图和结构精确匹配得分上与LLaMA-v2 7B相当，表明较小的模型如MobileLLM-350M也能很好地处理常见的移动设备应用场景。
研究者对125M和350M模型尺寸的MobileLLM和MobileLLM-LS进行了量化实验，发现使用W8A8后训练量化(PTQ)仅导致不到0.5个百分点的准确率降低，并且与层共享兼容。
研究者使用LLaMA-v2 7B作为教师模型对125M和350M模型进行知识蒸馏。尽管知识蒸馏可以提高训练时间，但其准确性与基于标签的训练相当或更低。
研究者在iPhone 13上使用ExecuTorch对MobileLLM-125M和MobileLLM-LS-125M FP16模型进行了延迟测量。
表7显示，通过权重共享和翻倍的层数，MobileLLM-LS在加载和初始化时间上仅比MobileLLM增加2.2%，在执行时间上增加2.6%，这得益于数据局部性。
这些实验结果表明，MobileLLM模型在资源受限的移动设备上具有良好的应用潜力，并且通过各种优化技术，实现了与更大模型相当的性能。