要說(shuō) ChatGPT 拉開(kāi)了大模型競(jìng)賽的序幕，那么 Meta 開(kāi)源 Llama 系列模型則掀起了開(kāi)源領(lǐng)域的熱潮。在這當(dāng)中，蘋(píng)果似乎掀起的水花不是很大。

不過(guò)，蘋(píng)果最新放出的論文，我們看到其在開(kāi)源領(lǐng)域做出的貢獻(xiàn)。

近日，蘋(píng)果發(fā)布了 OpenELM，共四種變體（參數(shù)量分別為 270M、450M、1.1B 和 3B），這是一系列基于公開(kāi)數(shù)據(jù)集進(jìn)行預(yù)訓(xùn)練和微調(diào)的模型。OpenELM 的核心在于逐層縮放，即 OpenELM 中的每個(gè) Transformer 層都有不同的配置（例如，頭數(shù)和前饋網(wǎng)絡(luò)維度），導(dǎo)致模型每層的參數(shù)數(shù)量不同，從而實(shí)現(xiàn)了更有效的跨層參數(shù)分配。

值得一提的是，蘋(píng)果這次發(fā)布了完整的框架，包括數(shù)據(jù)準(zhǔn)備、訓(xùn)練、微調(diào)和評(píng)估程序，以及多個(gè)預(yù)訓(xùn)練的 checkpoint 和訓(xùn)練日志，以促進(jìn)開(kāi)源研究。

• 論文地址：https://arxiv.org/pdf/2404.14619.pdf
• 項(xiàng)目地址：https://github.com/apple/corenet
• 論文標(biāo)題：OpenELM: An Efficient Language Model Family with Open-source Training and Inference Framework

結(jié)果顯示，OpenELM 的性能優(yōu)于使用公開(kāi)數(shù)據(jù)集進(jìn)行預(yù)訓(xùn)練的現(xiàn)有開(kāi)源 LLM（表 1）。例如，具有 11 億個(gè)參數(shù)的 OpenELM 性能優(yōu)于 OLMo。

方法介紹

OpenELM 架構(gòu)

OpenELM 采用只有解碼器的 transformer 架構(gòu)，并遵循以下方式：

（1）不在任何全連接（也稱(chēng)為線性）層中使用可學(xué)習(xí)的偏差參數(shù)；

（2）使用 RMSNorm 進(jìn)行預(yù)標(biāo)準(zhǔn)化，旋轉(zhuǎn)位置嵌入（ROPE）用于編碼位置信息；

（3）使用分組查詢注意力（GQA）代替多頭注意力（MHA）；

（4）用 SwiGLU FFN 替換前饋網(wǎng)絡(luò)（FFN）；

  (5) 使用 flash 注意力來(lái)計(jì)算可縮放的點(diǎn)積注意力；

  (6) 使用與 LLama 相同的分詞器（tokenizer）。

一般來(lái)講，LLM 中每個(gè) transformer 層使用相同的配置，從而實(shí)現(xiàn)跨層參數(shù)的統(tǒng)一分配。與這些模型不同的是，OpenELM 中的每個(gè) Transformer 層都有不同的配置（例如，頭數(shù)和前饋網(wǎng)絡(luò)維度），導(dǎo)致模型每層的參數(shù)數(shù)量不同。這使得 OpenELM 能夠更好地利用可用的參數(shù)預(yù)算來(lái)實(shí)現(xiàn)更高的精度。蘋(píng)果使用逐層縮放（layer-wise scaling）來(lái)實(shí)現(xiàn)跨層參數(shù)的非均勻分配。

逐層縮放：標(biāo)準(zhǔn) Transformer 層由多頭注意力（MHA）和前饋網(wǎng)絡(luò)（FFN）組成。針對(duì) Transformer 層參數(shù)分配不均勻的問(wèn)題，蘋(píng)果對(duì)各個(gè) Transformer 層的注意力頭數(shù)和 FFN 乘法器進(jìn)行了調(diào)整。

蘋(píng)果是這樣做的。設(shè)參數(shù)分配均勻的標(biāo)準(zhǔn) Transformer 模型有 N 層 transformer，假設(shè)每層輸入的維數(shù)為 d_model。MHA 有 n_h 個(gè)頭，每個(gè)頭的維度為，F(xiàn)FN 的隱藏維度為：

蘋(píng)果引入?yún)?shù) α 和 β 兩個(gè)超參數(shù)來(lái)分別縮放每層注意力頭的數(shù)量 n_h 和 m。對(duì)于第 i 層，n_h 和 m 計(jì)算為：

預(yù)訓(xùn)練數(shù)據(jù)

對(duì)于預(yù)訓(xùn)練，蘋(píng)果使用公共數(shù)據(jù)集。具體來(lái)說(shuō)，他們的預(yù)訓(xùn)練數(shù)據(jù)集包含 RefinedWeb、deduplicated PILE、RedPajama 的子集和 Dolma v1.6 的子集，總計(jì)約 1.8 萬(wàn)億個(gè) token 。如下表所示。

訓(xùn)練細(xì)節(jié)

蘋(píng)果使用自家開(kāi)源的 CoreNet 庫(kù)（以前稱(chēng)為 CVNets ，專(zhuān)門(mén)用于訓(xùn)練深度神經(jīng)網(wǎng)絡(luò)）訓(xùn)練 OpenELM 變體，訓(xùn)練過(guò)程迭代了 35 萬(wàn)次。最終訓(xùn)練出了 OpenELM 四種變體（參數(shù)量為 270M、450M、1.1B 和 3B）。

實(shí)驗(yàn)

本文評(píng)估了 OpenELM 在零樣本和少樣本設(shè)置下的性能，如表 3 所示。研究者將 OpenELM 與公開(kāi)的 LLM 進(jìn)行了比較，其中包括 PyThia 、Cerebras-GPT 、TinyLlama 、OpenLM 、MobiLlama  和 OLMo 。與本文工作較為相關(guān)的是 MobiLlama 和 OLMo。這些模型都是在類(lèi)似的數(shù)據(jù)集上訓(xùn)練的，具有相似或更多的預(yù)訓(xùn)練 token。

圖 1 繪制了 OpenELM 在 7 個(gè)標(biāo)準(zhǔn)零樣本任務(wù)上隨訓(xùn)練迭代次數(shù)的準(zhǔn)確率。可以發(fā)現(xiàn)，在大多數(shù)任務(wù)中，隨著訓(xùn)練持續(xù)時(shí)間的延長(zhǎng)，準(zhǔn)確率在總體上會(huì)有所提高。此外，通過(guò)平均最后五個(gè)檢查點(diǎn)（每 5000 次迭代收集一次）得到的檢查點(diǎn)，在準(zhǔn)確率上與經(jīng)過(guò) 350k 次迭代后得到的最終檢查點(diǎn)相當(dāng)，或略有提高。這種改進(jìn)很可能是由于權(quán)重平均降低了噪聲。因此，在表 4 的主要評(píng)估、表 5 的指令調(diào)優(yōu)實(shí)驗(yàn)和表 6 的參數(shù)效率調(diào)優(yōu)實(shí)驗(yàn)中，研究者使用了平均檢查點(diǎn)。

表 4 中的結(jié)果橫跨各種評(píng)估框架，突出了 OpenELM 相對(duì)于現(xiàn)有方法的有效性。表 4 中的結(jié)果跨越了不同的評(píng)估框架，凸顯了 OpenELM 相對(duì)于現(xiàn)有方法的有效性。例如，與擁有 12 億個(gè)參數(shù)的 OLMo 相比，擁有 11 億個(gè)參數(shù)的 OpenELM 變體的準(zhǔn)確率分別提高了 1.28%（表 4a）、2.36%（表 4b）和 1.72%（表 4c）。值得注意的是，OpenELM 達(dá)成了這樣的準(zhǔn)確率，但是使用的預(yù)訓(xùn)練數(shù)據(jù)比 OLMo 少的多。

如圖 5 所示，在不同的評(píng)估框架中，指令微調(diào)始終能將 OpenELM 的平均準(zhǔn)確率提高 1-2%。

參數(shù)高效微調(diào)（PEFT）結(jié)果。研究者使用常識(shí)推理的訓(xùn)練和評(píng)估設(shè)置。這個(gè)設(shè)置為不同方法提供了 8 個(gè)多項(xiàng)選擇數(shù)據(jù)集的 170k 訓(xùn)練樣本進(jìn)行 PEFT 研究，包括 LoRA 和 DoRA。研究者將 OpenELM 與這些方法整合在一起，并使用 8 個(gè) NVIDIA H100 GPU 對(duì)所生成的模型進(jìn)行了三個(gè)訓(xùn)練周期的微調(diào)。如表 6 所示，PEFT 方法可以應(yīng)用于 OpenELM。在給定的 CommonSense 推理數(shù)據(jù)集上，LoRA 和 DoRA 的平均準(zhǔn)確率相似。

表 7a 和 7b 分別展示了本項(xiàng)工作在 GPU 和 MacBook Pro 上的基準(zhǔn)測(cè)試結(jié)果。盡管 OpenELM 在相似參數(shù)數(shù)量下準(zhǔn)確度更高，但其速度比 OLMo 慢。雖然這項(xiàng)研究的主要關(guān)注點(diǎn)是可復(fù)現(xiàn)性而不是推理性能，但研究者還是進(jìn)行了全面的性能分析來(lái)判斷工作的瓶頸所在。

分析表明，OpenELM 的處理時(shí)間的相當(dāng)部分可歸因于研究者對(duì) RMSNorm 的簡(jiǎn)單實(shí)現(xiàn)（見(jiàn)表 8）。詳細(xì)來(lái)說(shuō)，也就是簡(jiǎn)單的 RMSNorm 實(shí)現(xiàn)導(dǎo)致許多單獨(dú)的內(nèi)核啟動(dòng)，每個(gè)都處理少量輸入，而不是像 LayerNorm 那樣啟動(dòng)單個(gè)融合內(nèi)核。通過(guò)用 Apex 的 RMSNorm 替換簡(jiǎn)單的 RMSNorm，研究者發(fā)現(xiàn) OpenELM 的吞吐量顯著提高。然而，與使用優(yōu)化 LayerNorm 的模型相比，仍有顯著的性能差距，部分原因是（1）OpenELM 有 113 層 RMSNorm，而 OLMo 有 33 層 LayerNorm；（2）Apex 的 RMSNorm 沒(méi)有為小輸入優(yōu)化。為了進(jìn)一步說(shuō)明由 RMSNorm 引起的性能下降，蘋(píng)果用 RMSNorm 替換了 OLMo 中的 LayerNorm，觀察到生成吞吐量顯著下降。在未來(lái)的工作中，研究者計(jì)劃探索優(yōu)化策略以進(jìn)一步提高 OpenELM 的推理效率。

更多詳細(xì)內(nèi)容，請(qǐng)閱讀原論文。