[论文解读] MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases

Zechun Liu, Changsheng Zhao+10 more|arXiv (Cornell University)|Feb 22, 2024

Context-Aware Activity Recognition Systems被引用 9

一句话总结

MobileLLM 表明深度和权重共享技术使子十亿参数的 LLMs 在设备端任务方面达到接近 state-of-the-art 的性能，包括零-shot 推理、问答、聊天和 API 调用，具备高效的设备端延迟和内存。

ABSTRACT

This paper addresses the growing need for efficient large language models (LLMs) on mobile devices, driven by increasing cloud costs and latency concerns. We focus on designing top-quality LLMs with fewer than a billion parameters, a practical choice for mobile deployment. Contrary to prevailing belief emphasizing the pivotal role of data and parameter quantity in determining model quality, our investigation underscores the significance of model architecture for sub-billion scale LLMs. Leveraging deep and thin architectures, coupled with embedding sharing and grouped-query attention mechanisms, we establish a strong baseline network denoted as MobileLLM, which attains a remarkable 2.7%/4.3% accuracy boost over preceding 125M/350M state-of-the-art models. Additionally, we propose an immediate block-wise weight-sharing approach with no increase in model size and only marginal latency overhead. The resultant models, denoted as MobileLLM-LS, demonstrate a further accuracy enhancement of 0.7%/0.8% than MobileLLM 125M/350M. Moreover, MobileLLM model family shows significant improvements compared to previous sub-billion models on chat benchmarks, and demonstrates close correctness to LLaMA-v2 7B in API calling tasks, highlighting the capability of small models for common on-device use cases.

研究动机与目标

• 激励并实现参数量在1B以下、适用于设备端用例的高质量大语言模型，以降低云成本和延迟。
• 深度（不仅是宽度）对小型 LLM 性能至关重要。
• 发展 架构技术：嵌入共享、分组查询注意力、分块权重共享，在有限内存预算下最大化权重利用。
• 通过聊天和 API 调用任务展示设备端的实际能力。

提出的方法

• 设计一个强基线 MobileLLM，使用 SwiGLU FFN、深而瘦的架构、嵌入共享和分组查询注意力。
• 引入立即的分块层共享以在不增加内存成本的情况下提高有效深度。
• 在零-shot 常识任务、问答/阅读理解以及下游设备端任务（聊天和 API 调用）上进行评估。
• 实验深度与宽度、嵌入共享的影响以及分组查询注意力，以确定有效的小模型配置。

实验结果

研究问题

• RQ1 子十亿参数的 LLM 在深度优先架构下，能否实现与竞争者相当的零-shot 常识和问答性能？
• RQ2 在严格的参数预算下，嵌入共享和分组查询注意力是否提升性能？
• RQ3 立即分块层共享是否能在不增加模型大小的情况下提升准确性或降低延迟？
• RQ4 MobileLLM 变体在设备端聊天和 API 调用任务中的表现，与更大模型或先前的子十亿模型相比如何？

主要发现

• 在约125M和350M尺度上，深而瘦的模型比更宽的模型表现更好，挑战了小型语言模型的传统放缩规律。
• 嵌入共享在参数上约减少11.8%，并且几乎无显著准确率损失，且更深的共享可以回收收益。
• 分组查询注意力（GQA）在优化的注意力头配置下为小型模型带来准确性提升；使用16个查询头和4个KV头显示出有利的结果。
• 即时分块层共享（MobileLLM-LS）在延迟几乎不增的情况下，准确性优于 MobileLLM，且没有额外的内存成本。
• MobileLLM-125M 与 MobileLLM-LS-125M 在零-shot 得分高于许多 125M 竞争对手；MobileLLM-350M 与 MobileLLM-LS-350M 比前代 350M 模型领先显著。
• 在聊天基准测试（AlpacaEval、MT-Bench）和 API 调用任务中，MobileLLM-LS-350M 的表现与更大的模型相当甚至更优（例如在 API 调用的精确匹配方面接近 LLaMA-v2 7B）。
• API 调用结果显示 MobileLLM-LS-350M 达到 65.3% 的意图 EM 和 48.8% 的结构 EM，在某些指标上与 7B 模型相当。
• 量化（W8A8 PTQ）与层共享兼容，且准确率损失极小。