MobiLlama，面向资源受限设备的轻量级全透明GPT模型

人工智能咨询培训老师叶梓 转载标明出处

如何在资源受限的设备上实现与大模型相媲美的语言处理能力呢？这些设备通常需要低能耗、小内存占用以及快速响应，但目前大模型往往因为其庞大的规模和高计算需求而难以部署。针对这一挑战，来自阿联酋人工智能大学、澳大利亚国立大学、阿尔托大学、墨尔本大学和林雪平大学的研究人员提出了一个名为MobiLlama的开源小型语言模型，该模型包含5亿参数，专注于减少资源需求的同时提升性能。MobiLlama的设计始于一个更大型的模型，并应用了精心设计的参数共享方案，以降低预训练和部署的成本。该模型采用了一种共享的多层感知器（MLP）设计，与传统的每个transformer块配备专用FFN的设计不同，这种设计显著减少了整体可训练参数的数量，从而提高了效率。

图1比较了 MobiLlama 0.5B 和 0.8B 模型与近期的 OLMo-1.17B 和 TinyLlama-1.1B 模型在预训练令牌、预训练时间、内存、模型参数、九种基准测试的总体准确性以及设备上的效率（平均电池消耗和在 RTX2080Ti PC 上的平均令牌/秒）方面的表现。明显可以看到，在保持较高准确性的同时MobiLlama 需要的预训练数据更少、预训练时间和 GPU 内存更少，并且在资源受限的设备上部署时更加高效。

方法

基线SLM的设计灵感来源于TinyLlama和Llama-2，是一个0.5B参数规模的模型。这个模型由N层组成，每层有M大小的隐藏维度和5632大小的中间尺寸（MLPs）。词汇量为32K，最大上下文长度为C。设计中考虑了两种方案：

• baseline1：含有22层，每层的隐藏尺寸为1024。
• baseline2：含有8层，每层的隐藏尺寸为2048。

这两种设计各有优势但也面临挑战：较小的隐藏尺寸可以提高计算效率，但可能限制了模型捕捉数据复杂性的能力；而较少的层数会影响模型学习语言层次结构的能力。为了在深度语言理解和上下文分析任务中取得更好的性能，需要结合这两种设计的优势。

但将这两种基线模型的容量合并到一个单一模型（即22层和2048的隐藏尺寸）会导致模型参数规模显著增加到1.2B，训练成本也会相应增加。这个更大的模型被命名为largebase。

MobiLlama模型的设计从一个更大的模型largebase开始，通过精心设计的参数共享方案来减少模型大小，从而显著降低训练成本。在传统的SLM和LLM中，每个transformer块内通常都有一个包含多个前馈网络（FFN）层的专用MLP块。在largebase模型中，FFN层占据了总可训练参数的65%，而注意力机制和头分别贡献了30%和5%。

为了解决FFN层在预训练期间带来的计算成本和收敛速度的挑战，MobiLlama提出了一种参数共享方案，即所有transformer层内的FFN参数都是共享的。这种设计使得MobiLlama的总可训练参数相比largebase减少了60%，同时还能增加模型的层数和隐藏尺寸，而不会显著增加训练成本。

图2展示了MobiLlama与两个基线模型的对比。在基线模型中，每个transformer层都有一个专用的MLP块，包含三个FFN层。相比之下，MobiLlama使用了一个单一的MLP块，该块在SLM的不同transformer层之间共享。这种设计增加了网络的容量，同时没有显著增加总的可训练参数数量。

为了营造一个更具包容性、数据和模型来源清晰、可复现协作的SLM研究开发环境，研究团队致力于开发一个完全透明的开源SLM。为此，他们提供了预训练数据集的详细信息、架构设计配置、训练细节、评估基准和指标。他们还计划公开发布完整的训练和评估代码以及中间模型检查点。

表2（Table 2）展示了Amber数据集中的数据混合情况，这些数据来自不同的来源，包括Arxiv、书籍、C4、精炼的网络数据、StarCoder、StackExchange和Wikipedia，总共包含1259.13亿个tokens。这些数据为模型提供了丰富多样的语言环境。

表3（Table 3）列出了MobiLlama模型的架构和超参数，包括0.5B的参数数量、2048的隐藏尺寸、5632的MLP中间尺寸、32个注意力头、22个隐藏层、1e−6的RMSNorm精度以及最大序列长度和词汇量。

MobiLlama的预训练使用了来自LLM360 Amber数据集的1.2万亿个tokens。Amber数据集提供了丰富多样的语言景观，包含不同类型的文本、主题和风格。数据混合如下：

Arxiv（300亿tokens）：来自科学论文的存储库，提供了复杂、特定领域的语言和术语，丰富了对学术散文的理解。

书籍（289亿tokens）：包含来自广泛文学的tokens，具有多样化的叙事风格、文化背景和丰富的词汇，加深了对故事叙述和语言细微差别的理解。

C4（1977亿tokens）：即Colossal Clean Crawled Corpus (C4)，提供了大量且经过清洗的网络文本，涵盖了各种语域、风格和主题。

精炼网络（665亿tokens）：经过策划的网络爬取，为模型提供了对现代、非正式和多样化的互联网语言的曝光，增强了与现代通信的相关性和适用性。

StarCoder（2919亿tokens）：用于代码理解的大型集合，涵盖了86种编程语言的783GB代码。包括GitHub问题、Jupyter笔记本和提交，总计约2500亿tokens，这些都经过了精心清洗和去重，以提高训练效率。

StackExchange（218亿tokens）：来自问答网站网络，这个子集帮助模型学习问答格式以及跨不同主题的技术讨论。

Wikipedia（239亿tokens）：百科全书集合，提供结构化和事实性的内容，帮助模型学习百科知识和正式写作风格。

这些数据子集总计达到1259.13亿tokens。原始数据来源经过分词处理，并使用Huggingface LLaMA分词器进行组织。然后，这些tokens被组织成每个包含2048个tokens的序列，并合并并划分为360个不同的数据段，每个数据段作为一个jsonl文件，包含一个token ID数组以及一个表示原始数据集的源标识符。每个数据样本设计为有2049个tokens。

MobiLlama 0.5B的架构包括2048大小的隐藏尺寸、5632大小的MLP中间尺寸，并在22个隐藏层中使用32个注意力头。它被设计为能够处理长达2048个tokens的序列，并得到32000大小的词汇量的支持。为了获得更稳定的训练，正则化精度由1e−6的RMSNorm保证。MobiLlama采用了RoPE（旋转位置嵌入）来编码位置信息，并采用了Swish和Gated Linear Units的组合作为激活函数。

MobiLlama模型的训练使用AdamW优化器进行，使用了β1 = 0.9, β2 = 0.95的超参数，初始学习率为η = 3e−4，并遵循余弦学习率调度，最终降至η = 3e−5。模型还结合了0.1的权重衰减，并在2000步的预热期内应用了1.0的梯度裁剪。为了适应20个GPU节点的硬件配置，优化了预训练批量大小为800（160×5），在单个GPU上实现了大约14k-15k tokens/秒的吞吐量。在模型预训练期间，每3.3B个tokens后会保存一次中间检查点，并将公开发布。

为了全面评估性能，使用了Open LLM Leaderboard上的九个不同的基准。这些基准测试包括HellaSwag、TruthfulQA、MMLU、ARC_Challenge、CrowsPairs、PIQA、Race、SIQA和Winogrande，覆盖了常识推理、事实理解、跨学科知识、复杂推理、偏见评估、物理常识、阅读理解、社会常识和文本歧义解决等多个方面。评估遵循Analysis-360框架，使用不同的shot设置进行详细评估，验证了模型在不同语言任务中的鲁棒性和适应性。

结果

在对0.5B模型系列的基线评估中，研究团队将所有模型在Amber数据集的相同1000亿tokens上进行预训练，并在HellaSwag、TruthfulQA、MMLU和Arc_C四个基准上报告结果。MobiLlama模型在这些基准上的平均得分为34.4，与两个基线模型相比表现出更优的性能。值得注意的是，这种性能提升并没有显著增加训练成本，这突显了所提出的SLM设计的优势。

表4展示了MobiLlama与其他小于10亿参数模型在九个基准测试上的性能比较。MobiLlama在0.5B模型系列中平均性能上取得了2.4%的显著提升，而0.8B模型的平均得分为46.67。

表5展示了两个基线模型和MobiLlama在四个基准测试上的比较结果。尽管所有模型的参数相同（0.5B），并且都是从Amber数据集的1000亿tokens中预训练的，但MobiLlama在性能上优于两个基线模型。

在表6中，MobiLlama 0.5B与现有参数量小于10亿的SLMs进行了比较，包括在不同低端硬件设备上的效率和资源消耗的评估。这些设备包括搭载RTX-2080Ti GPU的PC、搭载i7 CPU的笔记本和搭载Snapdragon-685处理器的智能手机。MobiLlama在处理速度（每秒平均tokens数）、内存消耗（平均内存消耗量，以MB计）和电池效率（每处理1000个tokens的平均电池消耗量，以mAH计）方面均展现出较好的性能。

表7展示了MobiLlama与现有大型语言模型（LLMs）在参数缩减至30%后的比较。这是通过使用Wiki2Text数据集进行的后训练稀疏化方案实现的。MobiLlama在HellaSwag、TruthfulQA、MMLU、Arc_C、CrowsPairs、piqa、race、siqa和winogrande等九个基准上的综合得分为46.0，优于pythia-410m的43.57。

表8展示了多模态MobiLlama模型（MobiLlama-V 0.8B）在不同基准测试（GQA、SQA、TextQA和MME）上的性能。这个模型通过结合CLIP的视觉编码器和MobiLlama的语言解码器，并在665k视觉-语言指令集上进行端到端的微调，展现了出色的视觉推理能力。

图3和图4展示了MobiLlama在多种任务上的反应示例，包括创意故事讲述、编码练习、经济分析和烹饪指导。这些示例显示了MobiLlama在抽象概念和实践步骤过程上的广泛适用性。

表9中，研究者将MobiLlama的大型基础模型（1.2B参数）与现有完全开源的、参数量小于20亿的SLMs进行了比较。尽管MobiLlama在数据量上（1.2T tokens）少于其他模型（如3T tokens），但其在九个基准测试上的平均得分达到了49.06，显示出优异的性能。

通过这些详细的评估和比较，MobiLlama证明了其在资源受限设备上实现高效、准确语言处理的潜力，同时也展示了其在多模态任务中的应用前景。MobiLlama项目的完整训练数据流程、训练代码、模型权重以及超过300个检查点连同评估代码已在GitHub上公开，以确保完全透明，便于社区访问和进一步研究。

论文链接：https://arxiv.org/abs/2402.16840

项目链接： MobiLlama GitHub Repo