MobileLLM：“苗条”的模型比较好！ 原创 精华

​1.MobileLLM

Meta AI研究人员推出了MobileLLM，从arXiv论文时间上看2月份发布的，但是在六月底又更新了一版。MobileLLM利用四种先进技术，包括带有SwiGLU的FFN、深度和薄架构、嵌入共享和GQA，构建了一个强大的基线小型LLM，MobileLLM。它是专为智能手机和其他资源受限的设备而设计。

研究人员将其训练出的MobileLLM 125M/350M和参数量相当的State of the Art（SOTA）模型如Cerebras、OPT、BLOOM等进行对话标杆测试。

在零样本（zero-shot）测试中MobileLLM 125M/350M比SOTA模型的平均准确率高出2.7%/4.3%。使用权重共享方法训练出的MobileLLM 125M/350M版本，名为MobileLLM LS，还可以将准确率分别再向上提升0.7%及0.8%。

模型具体的特点在于：

• 采取「深而精简」（deep- and-thin）的模型架构更能以极高效能萃取抽象知识
• 利用先进的权重共享技术，包括嵌入共享、分组查询注意力和块权重共享，在存储受限的场景中实现了权重利用率的显著的提高。

2.苗条架构

大模型领域有一种普遍观点，Transformer模型的性能主要取决于参数数量、训练数据集的大小和训练迭代次数。这种观点认为，架构设计对Transformer模型的性能影响微乎其微。

但是本次实验研究表明，对于较小的模型，为了提高性能，更深比广泛更为重要。下面的研究涉及 19 个模型的训练，其中包括 9个具有∼125M个参数的模型和10个具有∼350M个参数的模型。每个模型的设计大小相似，但深度和宽度有所不同。

上图对八个零样本常识推理任务以及问答和阅读理解基准进行了实验。研究结果一致表明，更深更薄的模型优于更浅更宽的模型。

黑色部分是最佳表现，可以看到基本上层数都较深。同时可以观察到更深的网络在大多数零样本推理任务中的卓越性能，包括 ARC-easy、ARC-challenge、PIQA、HellaSwag、OBQA、WinoGrande，尤其是在TQA和RACE数据集上更加明显。

<划重点！>对于大小约为 125M的transformer 模型，具有30层甚至42层的模型的性能明显优于具有12层的模型。要知道之前大多数的125M模型的层数都限制为12！！

此时细心的读者会发现第一行是升序，第二行是降序，第二行就是每个Transormer块中的头数<温习点>。当模型的层数增加的时候，头数一直在降低，也就是所谓的深而精简！

研究人员最初设计用于减少LLM中的键值缓存大小的GQA，也可以有效减少小型LM中键值头的冗余。GQA可以看作是权重重用的另一种权重共享形式，其中键值头的数量是查询头的1/n，并且kv头在计算注意分数和与查询一起输出时重复n次。

上表的结果表明，使用16个查询头可产生最佳结果。将kv-heads的数量从16个减少到4个，可使125M模型的准确率相当，而350M模型的准确率仅下降0.2个百分点，模型大小减少了近10%。通过采用分组查询注意机制 (GQA)，同时增加嵌入维度以保持模型大小，125M的准确率进一步提高了0.4个百分点，表明GQA是进一步挖掘小模型潜力的有利方法。

3.共享技术

Topic1 嵌入共享：

在1B以下的自然语言模型中，嵌入层构成了参数数量的很大一部分。例如，Embedding<温习点>的维度为 512，词汇量为32k，输入和输出嵌入层就需要16M参数。这些嵌入层合在一起占125M参数模型总参数的20%以上。在大语言模型中，这一比例被稀释了，例如输入和输出嵌入仅占LLaMA-7B模型总参数数量3.7%，在 LLaMA-70B模型中仅占0.7%。

LLM 模型中的输入嵌入将词汇表中的token ID映射到相应的token嵌入，维度为 (vocab_size, embedding_dim)。相反，输出全连接层​<温习点>将嵌入维度映射回词汇表中的logits预测<温习点>，权重大小为 (vocab_size, embedding_dim)。通过共享嵌入，将“输入-中间-输出”的两个转化矩阵参数共用，从而产生更高效、更紧凑的模型架构。

本次研究在30层125M模型上进行了实验。上面的统计证明共享输入和输出嵌入可将参数数量减少16M，约占总参数的11.8%，平均准确率下降0.2个百分点。这时候通过重新分配已保存的参数以添加更多层，可以轻松恢复边际准确率下降。与原始 135M模型相比，将深度增加到32层可使准确率提高 0.4 个百分点，同时仍保持 10M更少的参数。在有限的资源预算下，嵌入共享是一种有价值的技术，可以最大限度地提高权重利用率并优化模型性能。

Topic 2 层共享：

关于层深度与宽度影响的发现表明，对于小型 Transformer 模型来说，更深的层是有利的。这促使我们研究层共享作为一种增加隐藏层数量而无需额外模型存储成本的策略。这种方法在模型大小是主要制约因素的设备上场景中特别有用。

上图(a)代表没有层共享的基线模型；(b)相邻块之间的共享；(c) 重复全部共享；(d) 反向共享。这里每个块就是Transformer块，包含多头自注意力 (MHSA) 和前馈网络 (FFN)。

bcd说白了就是按照特殊规律，相同颜色的块共同拥有相同的参数。虽然c策略的共享性能略高，但b策略可以最好地利用缓存，因为共享权重可以保留在缓存中并立即计算两次。

实验结果表明，只需复制Transformer块即可提高准确度，而无需进行架构修改或扩大模型大小。通过实验研究了三种不同的权重共享策略的效果如下表：<从上往下分别为abcd四种策略！>

结果表明，重复层共享策略C产生最佳性能。但是考虑到硬件内存层次结构<温习点>，用于计算的SRAM通常限制在20MB左右。这个容量通常只足以容纳一个 Transformer 块。因此，b策略可以避免在SRAM和DRAM之间反复传输权重，从而提高自回归推理的整体执行速度。因此在模型设计中选择了B策略，并将采用这个策略的模型称呼为MobileLLM-LS。

4.消融和性能

本次也进行了消融研究，消融（Ablataion）就是，一个全新模型M引入N种优化的方法。为了搞清楚哪种方法对于M的提升最大，于是就在基础模型上面单独添加这项技能进行对比测试，也就是下图的+号的由来）。会发现layer-share和GQA的影响力还是挺大的。

按照上面的优化策略进行模型层数的堆叠，继续将小的模型扩大到1.5B然后，在各项基准的评测下，性能还是不输给同量级的竞争对手。

MobileLLM家族

MobileLLM到目前为止还没公开使用，但Meta已经开源了预训练代码，允许其他研究人员在他们的工作基础上进行开发。小模型技术（尤其是移动手机小模型）的发展会打开个人设备人工智能服务的便利性。

本文转载自​​鲁班模锤​​，作者： 庞德公 ​​​​