作者介紹：論文一作Liqun Ma目前是MBZUAI機(jī)器學(xué)習(xí)系的博士生，導(dǎo)師為Zhiqiang Shen助理教授，同時(shí)也是該論文的最后作者，其在加入MBZUAI之前為CMU博士后，研究領(lǐng)域主要為機(jī)器學(xué)習(xí)，基礎(chǔ)大模型等等。Liqun的主要研究方向?yàn)楦咝У拇竽Ｐ皖A(yù)訓(xùn)練和微調(diào)，他本科畢業(yè)于天津大學(xué)。論文二作Mingjie Sun目前為CMU計(jì)算機(jī)系在讀博士，導(dǎo)師為Zico Kolter教授。

自回歸訓(xùn)練方式已經(jīng)成為了大語(yǔ)言模型（LLMs）訓(xùn)練的標(biāo)準(zhǔn)模式， 今天介紹一篇來(lái)自阿聯(lián)酋世界第一所人工智能大學(xué)MBZUAI的VILA實(shí)驗(yàn)室和CMU計(jì)算機(jī)系合作的論文，題為《FBI-LLM: Scaling Up Fully Binarized LLMs from Scratch via Autoregressive Distillation》，該論文首次提出了采用自回歸蒸餾的優(yōu)化方式從頭訓(xùn)練二值化的大語(yǔ)言模型，性能可以匹配或者接近FP16或者BF16訓(xùn)練的LLMs，同時(shí)效果遠(yuǎn)超之前所有二值化大語(yǔ)言模型將近十個(gè)點(diǎn)。目前該工作的訓(xùn)練代碼，數(shù)據(jù)和模型權(quán)重已全部開(kāi)源。

• 論文：https://arxiv.org/abs/2407.07093
• 代碼：https://github.com/LiqunMa/FBI-LLM

核心結(jié)論和貢獻(xiàn)

• 相比之前的二值化大語(yǔ)言模型，這是第一個(gè)從頭開(kāi)始訓(xùn)練，不使用任何預(yù)訓(xùn)練參數(shù)的二值化大語(yǔ)言模型。
• 訓(xùn)練過(guò)程僅僅使用自回歸蒸餾損失，沒(méi)有加入其他損失函數(shù)。
• 該工作是一個(gè)全量二值化模型，而不是之前一些方法采用的局部二值化或者三值化大模型。

背景介紹

最近幾年受益于巨大的參數(shù)規(guī)模和海量的訓(xùn)練語(yǔ)料，基于Transformer的大型語(yǔ)言模型（LLMs），如ChatGPT和LLaMA系列，在特定領(lǐng)域知識(shí)的生成和復(fù)雜推理任務(wù)中都表現(xiàn)出色的性能。

此外，LLMs的能力隨著參數(shù)規(guī)模的擴(kuò)大而繼續(xù)增強(qiáng)，給人們?cè)谕ㄍ鵄GI的道路上以無(wú)限遐想。然而，巨大的參數(shù)規(guī)模導(dǎo)致了模型需要巨大的存儲(chǔ)和計(jì)算需求，這大大限制了LLMs的廣泛應(yīng)用和發(fā)展。量化技術(shù)通過(guò)將32位參數(shù)映射到更小的位數(shù)，有效地緩解了這些限制，該技術(shù)可以顯著減少存儲(chǔ)需求，并在推理過(guò)程中提升了計(jì)算速度和能源效率。

作為量化的極端情況，模型二值化僅用{-1, 1}來(lái)表示每個(gè)參數(shù)。它最大限度地實(shí)現(xiàn)了壓縮和推理效率，但代價(jià)是犧牲一定程度的準(zhǔn)確性。以往維持二值化LLMs性能的研究包括如何保留模型中重要參數(shù)或使用接近一位（部分二值化或者三值化）的表達(dá)方式來(lái)表示每個(gè)參數(shù)。

雖然這些方法展現(xiàn)出了不錯(cuò)的性能和潛力，但它們?cè)诖鎯?chǔ)和效率方面仍有優(yōu)化的空間，并且額外的全精度參數(shù)或采用非2的冪來(lái)表示參數(shù)編碼在適配特定硬件時(shí)會(huì)產(chǎn)生額外的開(kāi)銷(xiāo)。某些全二值化LLMs的研究基于最小化層級(jí)損失的優(yōu)化目標(biāo)，或使用預(yù)訓(xùn)練的全精度LLM在其基礎(chǔ)上繼續(xù)訓(xùn)練，然后用少量訓(xùn)練數(shù)據(jù)進(jìn)行二值化模型參數(shù)糾正，這些方法面臨如下幾個(gè)問(wèn)題：

• 之前某些方法借助預(yù)訓(xùn)練的全精度模型參數(shù)來(lái)減少訓(xùn)練計(jì)算量和優(yōu)化步驟， 然而二值化過(guò)程會(huì)極大地壓縮原始模型的參數(shù)空間，損害全精度模型中存儲(chǔ)的知識(shí)，因此依然需要足夠的訓(xùn)練數(shù)據(jù)來(lái)讓二值化模型重新學(xué)習(xí)這些知識(shí)并適應(yīng)二值化參數(shù)的模式；
• 從現(xiàn)有預(yù)訓(xùn)練模型中衍生二值化模型 (使用預(yù)訓(xùn)練權(quán)重) 的方案不允許選擇不同的參數(shù)規(guī)模或詞匯表大小，從而限制了模型結(jié)構(gòu)的靈活性和實(shí)際應(yīng)用。

本文作者提出了一種從頭開(kāi)始訓(xùn)練的全二值化LLMs（FBI-LLM）。為了實(shí)現(xiàn)從頭開(kāi)始穩(wěn)定地訓(xùn)練二值化LLMs，文章提出了一種基于全精度教師模型的自回歸蒸餾的新型訓(xùn)練模式。具體來(lái)說(shuō)，在訓(xùn)練過(guò)程中，作者逐步從全精度教師模型中生成蒸餾使用的軟標(biāo)簽，并采用基于自回歸蒸餾的方案來(lái)匹配教師模型在每個(gè)token位置的預(yù)測(cè)概率。

通過(guò)這種簡(jiǎn)單直接的自回歸蒸餾損失，可以成功地從隨機(jī)初始化中訓(xùn)練二值化LLMs。由于該方法相比一般LLM訓(xùn)練的改動(dòng)主要集中在損失函數(shù)上，F(xiàn)BI-LLM可以輕松地融入現(xiàn)有的LLM預(yù)訓(xùn)練過(guò)程。此外，這種方法中的二值化操作與模型訓(xùn)練是分離的，因此任何增強(qiáng)LLM訓(xùn)練效率或者性能的技術(shù)都可以直接應(yīng)用于本文提出的FBI-LLM。

作者對(duì)框架FBI-LLM的有效性進(jìn)行了詳細(xì)評(píng)估，訓(xùn)練了從130M、1.3B到7B規(guī)模的不同模型。作者使用廣泛使用的Transformer架構(gòu)進(jìn)行LLMs的訓(xùn)練，結(jié)果表明從頭開(kāi)始訓(xùn)練全二值化的LLMs是完全可行的，其性能與全精度模型相比只有很小的差距。

與其他基準(zhǔn)線方法相比，訓(xùn)練過(guò)程在困惑度和多個(gè)下游任務(wù)上表現(xiàn)更為出色。這些結(jié)果表明自回歸蒸餾是訓(xùn)練二值化LLMs的關(guān)鍵，此外，通過(guò)對(duì)預(yù)訓(xùn)練更加深入的研究（如權(quán)重翻轉(zhuǎn)比和梯度范數(shù)）的分析表明，從全精度LLMs繼承權(quán)重與從頭開(kāi)始訓(xùn)練二值化LLMs之間沒(méi)有顯著差異。

本文的貢獻(xiàn)可以總結(jié)如下：首先，該論文首次證明可以成功地從頭開(kāi)始訓(xùn)練具有二值權(quán)重的LLMs；其次，本文提出了一種新的蒸餾損失函數(shù)，以穩(wěn)定二值化LLMs的訓(xùn)練，其采用自回歸蒸餾來(lái)匹配教師模型的概率分布；第三，本文進(jìn)行了廣泛的實(shí)驗(yàn)和分析，以更好地理解所提出的方法的有效性。

下面介紹一下文章具體細(xì)節(jié)。

模型結(jié)構(gòu)

模型結(jié)構(gòu)主要基于LLaMA的結(jié)構(gòu)。如上左圖，首先是LLM 模塊，其中包含使用可學(xué)習(xí)的α和β組成的 FBI-Linear層。右圖為自回歸蒸餾和模型訓(xùn)練相關(guān)過(guò)程。具體而言，由于在LLM 中，大多數(shù)參數(shù)都位于線性模塊中。FBI-LM 將除causal head以外的所有線性模塊替換為 FBI-linear層。由于causal head直接影響每個(gè)步驟中的輸出token分布，因此對(duì)其參數(shù)進(jìn)行二值化將顯著影響模型輸出的準(zhǔn)確性，因此本文選擇保留其精度。

此外，LLM 的另外兩個(gè)核心模塊（embedding和Layer Norm）中的參數(shù)也需要保持全精確。這是因?yàn)閑mbedding模塊包含有關(guān)所有標(biāo)記的語(yǔ)義信息，并且作為模型輸入的第一層，需要用來(lái)確定文本的初始表示形式。另一方面，Layer Norm 直接縮放激活值，二值化其參數(shù)將顯著降低每層激活值的語(yǔ)義表達(dá)能力，之前其他關(guān)于LLM二值化的工作和研究也選擇采用類(lèi)似的設(shè)置和做法。

模型訓(xùn)練：自回歸蒸餾（Autoregressive Distillation）

給定一個(gè)訓(xùn)練語(yǔ)料，標(biāo)準(zhǔn)的自回歸語(yǔ)言模型的目標(biāo)函數(shù)是最大化如下似然函數(shù)：

其中k表示上下文窗口的大小，條件概率p通過(guò)參數(shù)為θ的神經(jīng)網(wǎng)絡(luò)建模。不同于一般的自回歸語(yǔ)言模型，本文使用自回歸蒸餾訓(xùn)練 FBI-LLM。在訓(xùn)練過(guò)程中，一個(gè)全精度預(yù)訓(xùn)練 LLM 被用作教師模型，二值化目標(biāo)模型作為學(xué)生模型。假設(shè)每個(gè)訓(xùn)練數(shù)據(jù)實(shí)例由輸入token序列x^1,…x^m組成，教師模型對(duì)下一個(gè)標(biāo)記的預(yù)測(cè)概率可以表示為：

其中表示最后一層 transformer 模塊的激活，表示用于預(yù)測(cè)下一個(gè)token概率的線性輸出層的參數(shù)。

學(xué)生模型與教師模型輸出之間的交叉熵被計(jì)算為每一步預(yù)測(cè)下一個(gè)token時(shí)的最終損失函數(shù)。它可以表示為：

其中n表示輸入標(biāo)記的數(shù)量。表示教師模型預(yù)測(cè)的第i步詞匯表上的標(biāo)記分布，而是學(xué)生模型的相應(yīng)預(yù)測(cè)分布。

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

本文使用的訓(xùn)練數(shù)據(jù)集跟一般的LLM訓(xùn)練相似，包含 Refined-Web 、StarCoder  和 RedPajama-v1的混合數(shù)據(jù)集，總共包含 1.26T tokens。

實(shí)驗(yàn)結(jié)果

如下圖所示，首先是對(duì)現(xiàn)有的二值化 LLM 和 FBI-LLM 在 Wikitext2 的困惑度 (Perplexity) 比較。與其他二值化 LLM 相比，F(xiàn)BI-LLM 在相同規(guī)模大小的模型上獲得相似或更低的困惑程度。

其次是在下游任務(wù)上的性能表現(xiàn)，如下表所示，由于 130M 大小的FBI-LLM沒(méi)有對(duì)應(yīng)的基準(zhǔn)模型，本文將 130M 模型與之前 700M 規(guī)模的 BitNetb1.58 進(jìn)行比較。盡管模型規(guī)模相差五倍，權(quán)重量化程度也存在較大的差異，但FBI的模型在 BoolQA 和 OpenbookQA 上的表現(xiàn)仍然優(yōu)于 BitNet b1.58。

對(duì)于 1.3B 規(guī)模的二值化模型，F(xiàn)BI-LLM 在大多數(shù)下游任務(wù)和困惑度中都取得了最佳性能，甚至接近或超過(guò)了某些 7B 規(guī)模的二值化模型（如 BiLLM-LLaMA2-7B）的性能。與相同規(guī)模的全精度模型相比， FBI-LLM 1.3B 在下游任務(wù)中可以達(dá)到其 87% 的性能。在 7B 規(guī)模中，F(xiàn)BI模型依然顯著優(yōu)于之前幾乎所有的基準(zhǔn)線模型，具體來(lái)說(shuō)，F(xiàn)BI-LLM相比之前最好的方法平均提升了將近十個(gè)點(diǎn)。

此外，由于計(jì)算資源的限制，F(xiàn)BI-LLM 7B當(dāng)前匯報(bào)的結(jié)果并不是最終結(jié)果。作者只使用了整個(gè) 數(shù)據(jù)集的 8.6%（31 個(gè)塊）。下圖展示了FBI-LLM-7B訓(xùn)練過(guò)程中下游任務(wù)準(zhǔn)確率和困惑度的變化。顯然，從目前的訓(xùn)練進(jìn)度來(lái)看，F(xiàn)BI-LLM-7B的性能將持續(xù)提高，更進(jìn)一步的訓(xùn)練可能會(huì)得到更好的效果。

模型分析和可視化

二值化大模型是 從頭開(kāi)始訓(xùn)練還是從預(yù)訓(xùn)練的 LLM 接著訓(xùn)練？

直觀地說(shuō)，從預(yù)訓(xùn)練的 LLM 繼續(xù)訓(xùn)練可以讓二值化模型從全精度原始模型中繼承知識(shí)，從而可能比從頭開(kāi)始訓(xùn)練獲得更好的結(jié)果。為了論證這一假設(shè)，本文進(jìn)行了全面的消融和分析實(shí)驗(yàn)，以記錄和比較模型在兩種不同訓(xùn)練模式下的行為。

從下圖 (a) 中可以觀察到，在訓(xùn)練初期，兩種訓(xùn)練方式的 FF ratio 趨勢(shì)基本保持一致。在整個(gè)訓(xùn)練過(guò)程中，兩種方法的 FF ratio 都處于相似的大小上，并且數(shù)值相對(duì)較小。該結(jié)果表明，兩種不同的參數(shù)初始化方法對(duì)二值化優(yōu)化過(guò)程的影響沒(méi)有顯著差異。下圖 (b) 展示了兩種訓(xùn)練模式下的訓(xùn)練損失變化。在訓(xùn)練的初始階段，兩種方法的訓(xùn)練損失基本相同，表明模型的訓(xùn)練損失不會(huì)顯著受初始化方法的影響。

雖然從頭開(kāi)始訓(xùn)練的損失在中間階段比繼續(xù)訓(xùn)練的損失略高，但過(guò)一段時(shí)間后，從頭開(kāi)始訓(xùn)練的損失再次與接著訓(xùn)練的損失相當(dāng)，甚至變得比繼續(xù)訓(xùn)練的損失更為穩(wěn)定。值得注意的是，在大約第 1000 步時(shí)，如圖 (a) 所示，當(dāng)從預(yù)訓(xùn)練的 LLM 繼續(xù)訓(xùn)練時(shí)，F(xiàn)F ratio 開(kāi)始出現(xiàn)明顯的波動(dòng)。同樣，在圖 (b) 所示的第 1700 步左右，訓(xùn)練損失也遇到了類(lèi)似的問(wèn)題。

這些發(fā)現(xiàn)挑戰(zhàn)了之前相關(guān)工作的一些假設(shè)，即從預(yù)訓(xùn)練的 LLM權(quán)重開(kāi)始訓(xùn)練二值化大模型將賦予二值化LLM繼承之前預(yù)訓(xùn)練的知識(shí)，從而提高性能。然而，本文結(jié)果和分析暗示了二值化LLM對(duì)參數(shù)初始化的方式不敏感，或者說(shuō)隨機(jī)參數(shù)初始化反而更好。對(duì)于原因，本文推測(cè)二值化和全精度LLM采用不同的參數(shù)組合和配置來(lái)編碼語(yǔ)義，這導(dǎo)致其參數(shù)空間模式存在實(shí)質(zhì)性差異。為了適應(yīng)這種模式，通過(guò)從預(yù)訓(xùn)練的 LLM 繼續(xù)訓(xùn)練來(lái)優(yōu)化二值化的過(guò)程可能需要更大幅度的參數(shù)數(shù)值調(diào)整和變換，從而破壞預(yù)訓(xùn)練參數(shù)模式結(jié)構(gòu)，繼而破壞其中保存的知識(shí)。這個(gè)理論可以部分地解釋為什么在訓(xùn)練期間與從頭開(kāi)始訓(xùn)練相比，繼續(xù)預(yù)訓(xùn)練的權(quán)重反而讓二值化LLM變得更不穩(wěn)定。

訓(xùn)練的穩(wěn)定性分析

二值化和全精度LLM訓(xùn)練在之前一些工作中都被發(fā)現(xiàn)表現(xiàn)出不穩(wěn)定的訓(xùn)練行為。 FBI-LLM 也有類(lèi)似的問(wèn)題，具體表現(xiàn)為在訓(xùn)練 1.3B 和 7B FBI-LLM 時(shí)訓(xùn)練損失突然激增，有時(shí)在此之后繼續(xù)訓(xùn)練模型也無(wú)法收斂。本文采用類(lèi)似于 PaLM的解決方案：如果損失值不再趨于收斂，模型將恢復(fù)到前一個(gè)檢查點(diǎn)，并跳過(guò)觸發(fā)不穩(wěn)定損失的數(shù)據(jù)塊以繼續(xù)訓(xùn)練。使用此方法，模型在相同的訓(xùn)練步驟中不再遇到問(wèn)題。根據(jù)該論文的觀察，從頭開(kāi)始訓(xùn)練 7B FBI 模型大約有 6% 的概率導(dǎo)致出現(xiàn)損失峰值。對(duì)于 1.3B 模型，由于其模型能力較低，訓(xùn)練更加不穩(wěn)定，損失峰值的出現(xiàn)概率約為 15%。這與在全精度 LLM 中看到的預(yù)訓(xùn)練行為相似，而尖峰的概率明顯更高，這可能與二值化參數(shù)的有限表達(dá)能力有關(guān)。為了解決這個(gè)問(wèn)題，F(xiàn)BI跳過(guò)了發(fā)生損失峰值且沒(méi)法通過(guò)后續(xù)訓(xùn)練恢復(fù)的數(shù)據(jù)塊。

存儲(chǔ)效率分析

上表顯示了不同規(guī)模的FBI-LLM與相同結(jié)構(gòu)的全精度LLaMA相比所需的理論存儲(chǔ)空間，同時(shí)還詳細(xì)給出了 FBI-LLM 引入的附加參數(shù)（α 和 β）的比例。表中的比較表明，F(xiàn)BI-LLM可以實(shí)現(xiàn)高壓縮比，大大減輕了LLMs的存儲(chǔ)負(fù)擔(dān)。盡管 FBI-LLM 引入的用于縮放和移動(dòng)的額外參數(shù)需要保留全精確，但它們的比例相比整個(gè)模型很小，因此其對(duì)存儲(chǔ)的影響可以忽略不計(jì)。

生成結(jié)果示例

如上圖所示，盡管 FBI-LLM 的生成質(zhì)量無(wú)法完全匹配全精度 LLM模型，但 FBI-LLM 仍然可以生成流暢且有意義的內(nèi)容。與具有更高參數(shù)位寬的BitNet b1.58模型相比，F(xiàn)BI-LLM對(duì)提示的理解更好，并且在一些生成的示例中包含了更多的知識(shí)。這表明FBI-LLMs具有更強(qiáng)的生成能力，并包含足夠的知識(shí)。此外，F(xiàn)BI-LLM展示了進(jìn)一步擴(kuò)大模型規(guī)模從而達(dá)到更高智力水平的二值化模型的潛力，這種類(lèi)型的大模型 (二值化大模型) 對(duì)部署的硬件需求也更加友好。

模型配置和訓(xùn)練細(xì)節(jié)

具體的模型配置和訓(xùn)練細(xì)節(jié)如下所示：

更多細(xì)節(jié)歡迎閱讀文論原文。