「在移动设备上运行 LLM？可能需要 Meta 的一些技巧。」刚刚，图灵奖得主 Yann LeCun 在个人社交平台表示。
他所宣传的这项研究来自 Meta 最新论文《 MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases 》，在众多作者中也有我们熟悉的来自 Meta FAIR 田渊栋。
田渊栋表示：「我们的 MobileLLM 预训练模型（125M/350M），性能达到 SoTA，特别是在聊天 / API 调用方面表现出色。此外，本工作中的一个有趣研究是跨 Transformer 层的权重共享，这样不仅节省了参数，还减少了推理过程中的延迟。」
论文地址：https://arxiv.org/pdf/2402.14905.pdf
现阶段大语言模型（LLM）已经渗透到人类生活的各个方面，尤其是以 ChatGPT 等为代表的模型，这类研究主要在云环境中运行。
然而领先的模型如 ChatGPT4 的参数量已经超过了 1 万亿。我们设想这样一个场景，这个场景广泛依赖 LLM，不仅用于前端的会话界面，也用于后端操作，如推荐系统，覆盖人类约 5% 的时间。在这一假设场景中，假如以 GPT-4 每秒处理 50 个 token 的速率来计算，则需要部署大约一亿个 H100 GPU，每个 GPU 的计算能力为 60 TFLOPs/s。这种计算规模，还不包括通信和数据传输的开销，就已经与 160 个 Meta 规模的公司相当。随之而来的能源消耗和二氧化碳排放将带来巨大的环境挑战。
因此，最好的解决方案是缩小 LLM 的规模。
此外，在当前的移动技术领域，由于主内存（DRAM）容量的限制，将像 LLaMAv2 7B 这样的 LLM 与 8 位权重整合起来代价过高。移动设备中普遍的内存层结构如图 2 所示。随着 DRAM 容量从 iPhone 15 的 6GB 到 Google Pixel 8 Pro 的 12GB 不等，一个移动应用不应超过 DRAM 的 10%，因为 DRAM 需要与操作系统和其他应用程序共享。这一要求促进了部署小于十亿参数 LLM 更进一步的研究。
基于上述考量，来自 Meta 的研究者专注于设计参数少于十亿的高质量 LLM，这是在移动端部署 LLM 比较好的解决方案。
与强调数据和参数数量在决定模型质量方面的关键作用的普遍观点相反，Meta 强调了模型架构对少于十亿（sub-billion）规模 LLM 的重要性。
基于深而窄的架构，加上嵌入共享和分组查询注意力机制，Meta 建立了一个强大的基线网络，称为 MobileLLM，与之前的 125M/350M 最先进模型相比，其准确率显著提高了 2.7%/4.3% 。这也说明了与缩放定律（scaling law）相反，该研究证明对于小型 LLM 来说深度比宽度更重要，一个深而窄的模型结构在捕获抽象概念方面更为出色。
此外，Meta 还提出了一种及时逐块权重共享（ immediate block-wise weight sharing）方法，该方法不会增加模型大小，所得模型表示为 MobileLLM-LS，其准确率比 MobileLLM 125M/350M 进一步提高了 0.7%/0.8%。此外，在下游任务中，例如 Chat 和 API 调用，MobileLLM 模型家族显著优于同等规模的模型。在 API 调用任务中，与规模较大的 LLaMA-v2 7B 相比，MobileLLM 甚至实现了相媲美的分数。
看到这项研究后，网友纷纷表示「我们应该向 Meta 致敬，很高兴看到这个领域的活跃玩家。该机构通过使用低于 10 亿参数的模型，并且 350M 8 位模型的能源消耗仅为 0.035 J/token ，要是部署在 iPhone 上的话，可以支持用户一整天的会话使用。」
改进十亿以下参数规模的 LLM 设计
研究者介绍了从十亿以下参数规模的基线模型到新的 SOTA 模型的演进之路（如下图 3 所示）。他们分别研究了 125M 和 350M 参数规模的模型，并在这两个规模下展示了一致的改进。对于模型尺寸成为主要制约因素的设备用例而言，如何有效地分配有限的权重参数变得比以往更加重要。
研究者首先通过测试四种有益于十亿以下规模 LLM 的模型设计方法，提出了一个名为MobileLLM 的强大基线模型。这四种模型设计方法包括 1）采用 SwiGLU FFN，2）强制使用深和薄的架构，3）重新审视嵌入共享方法，4）利用分组查询注意力。
接下来，研究者开发了一种直接的逐块层共享方法，基于该方法可以进一步提高准确度，而不产生任何额外的内存开销，并在内存有限的 LM 解码过程中产生很小的延迟开销。他们将具有层共享的模型表示为 MobileLLM-LS。
训练设置
研究者在 32 个 A100 GPU 上进行实验，其中每个 GPU 的批大小为 32。他们在 0.25T 的 tokens 上执行了 120k 次迭代的探索性实验。下文中表 3 和表 4 报告了在 1T 的 tokens 上执行 480k 次迭代训练的 top 模型。
层共享
关于层深度与宽度影响的研究结果表明，更深的层有利于小型 transformer 模型。这促使本文研究层共享作为增加隐藏层数量而不增加存储成本的策略。这种方法在模型大小成为主要制约因素的场景中尤其有用。
令人惊讶的是，实验结果表明，通过简单地复制 transformer 块就可以提高准确度而无需任何架构修改或扩大模型尺寸。研究者进一步探究三种不同的权重共享策略，具体如下图 6 所示。
下表 2 结果表明，重复层共享策略在立即块重复、全面重复（repeat all-over）和反向共享策略中产生了最佳性能。
不过，考虑到硬件内存的层级结构（如图 2），用于计算的 SRAM 通常限制在了 20M 左右。该容量通常仅够容纳单个 transformer 块。因此，将共享权重放入缓存中并立即计算两次则无需在 SRAM 和 DRAM 之间传输权重，提高了自回归推理的整体执行速度。
研究者在模型设计中选择了直接的分块共享策略，并将提出的带有层共享的模型表示为 MobileLLM-LS。
实验结果
该研究进行实验比较了模型在零样本（zero-shot）常识推理任务、问答和阅读理解任务上的性能。
零样本常识推理任务的实验结果如下表 3 所示：
在问答和阅读理解任务上，该研究采用 TQA 问答基准和 RACE 阅读理解基准来评估预训练模型，实验结果如下表 4 所示：
为了验证将模型用于设备上应用程序的有效性，该研究评估了模型在两个关键任务上的性能：聊天和 API 调用。
针对聊天任务，该研究在两个基准上进行了评估实验：AlpacaEval（单轮聊天基准）和 MT-Bench（多轮聊天基准），实验结果如下表 5 所示：
在 API 调用方面，如下表 6 所示，MobileLLM-350M 表现出与 LLaMA-v2 7B 相当的 EM_intent 和 EM_structure，其中 EM_intent 越高，表明模型对用户计划调用 API 的预测就越准确，而 EM_structure 反映了预测 API 函数内内容的熟练程度。
该研究进一步在 MobileLLM 和 MobileLLM-LS 模型上针对每个 token 进行最小 / 最大训练后量化 (PTQ) 实验，模型大小分别为 125M 和 350M，在 0.25T token 上进行训练，实验结果如下图 7 所示：
模型加载、初始化和执行时间如下表 7 所示：
更多技术细节请参阅原论文。