【论文速读】逐层扩展，《OpenELM: An Efficient Language Model Family with Open Training and Inference Framework》

文章链接：https://arxiv.org/html/2404.14619v2

这篇文章的标题是《OpenELM: An Efficient Language Model Family with Open Training and Inference Framework》，由 Sachin Mehta 等人撰写，来自苹果公司，旨在提供用于在移动设备上部署的小型 LLM。
文章的亮点是逐层扩展（Layer-wise scaling）技术，它允许跨层更有效地分配参数。这种方法在接近输入的 transformer 层中使用较小的潜在维度，然后逐渐扩大接近输出的层的参数预算，通过调整每个 transformer 层中的注意力头数和前馈网络 (FFN) 隐藏层的维度，从而实现了非均匀参数分配。

摘要 (Abstract)

文章强调了大型语言模型 (LLM) 在开放研究中的可复制性和透明度的重要性，这对于确保结果的可信度、调查数据和模型偏差以及潜在风险至关重要。为此，作者发布了 OpenELM，这是一个先进的开放语言模型。OpenELM 采用逐层扩展策略，在每个 transformer 模型层内有效分配参数，从而提高准确性。例如，在大约十亿参数的预算下，OpenELM 与 OLMo 相比在准确性上提高了 2.36%，同时需要的预训练 tokens 减少了 2 倍。

与以往只提供模型权重和推理代码以及在私有数据集上预训练的做法不同，作者的发布包括了在公开可用数据集上训练和评估语言模型的完整框架，包括训练日志、多个检查点和预训练配置。此外，作者还发布了将模型转换为 MLX 库的代码，以便在 Apple 设备上进行推理和微调。这一全面的发布旨在赋予和加强开放研究社区的力量，为未来的开放研究工作铺平道路。

模型主要特点：

1、OpenELM 有 4种相对小且方便的规模：270M、450M、1.08B 和 3.04B
2、对于每种规模，有经过Instruction tuning版本可用，经过拒绝采样或DPO训练
3、OpenELM 的性能略好于 OLMo，尽管它的训练数据量少了 2 倍
4、OpenELM 架构的核心创新在于其逐层扩展方法，这与传统的均匀参数分配方法不同。模型采用了仅解码器的基于 transformer 的架构，关键点包括：(1) 在任何全连接（也称为线性）层中不使用可学习的偏置参数，(2) 应用 RMSNorm 进行预归一化，有助于稳定训练过程，并且同样使用旋转位置嵌入 (RoPE) 编码位置信息，(3) 使用分组查询注意力 (GQA) 代替多头注意力 (MHA)，(4) 用 SwiGLU FFN 替换前馈网络 (FFN)，使用 SwiGLU 激活函数代替传统的 ReLU 或 GELU，可能有助于改善信息流和梯度传播，从而提高模型性能，(5) 使用 Flash attention注意力机制，可以更快、更高效地计算注意力，以及 (6) 使用与 LLama 相同的分词器。
最关键点在于逐层扩展（Layer-wise scaling）：通过调整每个层的注意力头数和 FFN 乘数，OpenELM 能够在不同层之间实现参数的非均匀分配，这允许模型根据每层的需求分配资源，而不是简单地在所有层中均匀分配参数。 设置超参数 α 和 β：这些超参数允许模型在不同层中调整注意力头数和 FFN 宽度，从而实现更精细的控制，以优化模型的性能。

逐层扩展（Layer-wise scaling）

现有的 LLM 在模型中的每个 transformer 层使用相同的配置，导致各层之间参数的均匀分配。与这些模型不同，OpenELM 中的每个 transformer 层有不同的配置（例如头数和前馈网络维度），模型中每层的参数数量不同。这使得 OpenELM 能够更好地利用可用的参数预算，以实现更高的准确率。通过逐层扩展（也称为块扩展）实现了这种非均匀的参数分配。

一个标准的 transformer 层由多头注意力 (MHA) 和前馈网络 (FFN) 组成。为了在 transformer 层中非均匀地分配参数，调整每个 transformer 层的注意力头数和 FFN 乘数。

假设具有均匀参数分配的标准 transformer 模型具有 N 层，每个层的输入维度是 dm⁢o⁢d⁢e⁢l，MHA 有 nh个头，每个头的维度是 dh=dm⁢o⁢d⁢e⁢l/nh，同时，FFN 的隐藏维度是 dFFN=𝑚*dm⁢o⁢d⁢el⁢，其中 𝑚是一个标量乘数。

引入了参数 α 和 β 来分别按层扩展注意力头 nh和 FFN 乘数 𝑚。对于第 𝑖层，nh 和 𝑚计算如下：

这里，αm⁢i⁢n 和 αm⁢a⁢x是允许扩展注意力头的超参数，类似地，βm⁢i⁢n 和 βm⁢a⁢x是允许改变 FFN层宽度的超参数。通过这些超参数改变标准 transformer 层的配置，在模型中实现了非均匀的参数分配。注意，将 αm⁢i⁢n=αm⁢a⁢x=1.0，𝑚𝑖=𝑚将产生标准的均匀 transformer 模型。

总结

OpenELM 的核心是逐层扩展（Layer-wise scaling），它允许跨层更有效地分配参数。这种方法在接近输入的 transformer 层中使用较小的潜在维度，并逐渐加宽接近输出的层。通过调整每个 transformer 层中的注意力头数和前馈网络 (FFN) 乘数，实现了这种非均匀参数分配。

文章还详细描述了 OpenELM 的架构，包括模型架构、预训练数据、训练超参数和评估方法。作者采用了无偏置参数的全连接层、RMSNorm 预归一化、旋转位置嵌入 (RoPE)、分组查询注意力 (GQA)、SwiGLU FFN、Flash attention等技术，并使用了与 LLama 相同的分词器。

在预训练数据方面，作者使用了公共数据集，包括 RefinedWeb、去重的 PILE、RedPajama 的子集和 Dolma v1.6 的子集，总共约 1.8 万亿 tokens。此外，作者还介绍了即时分词和数据过滤的方法，与以往使用预分词数据的方法不同。

在训练细节方面，作者使用了 CoreNet 训练 OpenELM 变体，使用了 AdamW 优化器和余弦学习率调度策略。他们训练了四种不同规模的 OpenELM（270M、450M、1.1B 和 3B）。

最后，在评估细节方面，作者使用了 LM Evaluation Harness 评估 OpenELM 在不同任务上的性能，包括标准零样本任务、OpenLLM 排行榜任务和 LLM360 排行榜任务。这些评估框架允许全面评估 OpenELM 在推理、知识理解和错误信息&偏见方面的表现。

其他

文章比较了 LoRA 和 DoRA的两种高效参数微调，文章表明，这两种方法之间没有明显的区别。

（欢迎关注公众号：老无谈藏书票）