作者介紹：論文一作Liqun Ma目前是MBZUAI機器學習系的博士生，導師為Zhiqiang Shen助理教授，同時也是該論文的最后作者，其在加入MBZUAI之前為CMU博士后，研究領域主要為機器學習，基礎大模型等等。Liqun的主要研究方向為高效的大模型預訓練和微調，他本科畢業于天津大學。論文二作Mingjie Sun目前為CMU計算機系在讀博士，導師為Zico Kolter教授。

自回歸訓練方式已經成為了大語言模型（LLMs）訓練的標準模式， 今天介紹一篇來自阿聯酋世界第一所人工智能大學MBZUAI的VILA實驗室和CMU計算機系合作的論文，題為《FBI-LLM: Scaling Up Fully Binarized LLMs from Scratch via Autoregressive Distillation》，該論文首次提出了采用自回歸蒸餾的優化方式從頭訓練二值化的大語言模型，性能可以匹配或者接近FP16或者BF16訓練的LLMs，同時效果遠超之前所有二值化大語言模型將近十個點。目前該工作的訓練代碼，數據和模型權重已全部開源。

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

核心結論和貢獻

• 相比之前的二值化大語言模型，這是第一個從頭開始訓練，不使用任何預訓練參數的二值化大語言模型。
• 訓練過程僅僅使用自回歸蒸餾損失，沒有加入其他損失函數。
• 該工作是一個全量二值化模型，而不是之前一些方法采用的局部二值化或者三值化大模型。

背景介紹

最近幾年受益于巨大的參數規模和海量的訓練語料，基于Transformer的大型語言模型（LLMs），如ChatGPT和LLaMA系列，在特定領域知識的生成和復雜推理任務中都表現出色的性能。

此外，LLMs的能力隨著參數規模的擴大而繼續增強，給人們在通往AGI的道路上以無限遐想。然而，巨大的參數規模導致了模型需要巨大的存儲和計算需求，這大大限制了LLMs的廣泛應用和發展。量化技術通過將32位參數映射到更小的位數，有效地緩解了這些限制，該技術可以顯著減少存儲需求，并在推理過程中提升了計算速度和能源效率。

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

雖然這些方法展現出了不錯的性能和潛力，但它們在存儲和效率方面仍有優化的空間，并且額外的全精度參數或采用非2的冪來表示參數編碼在適配特定硬件時會產生額外的開銷。某些全二值化LLMs的研究基于最小化層級損失的優化目標，或使用預訓練的全精度LLM在其基礎上繼續訓練，然后用少量訓練數據進行二值化模型參數糾正，這些方法面臨如下幾個問題：

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

本文作者提出了一種從頭開始訓練的全二值化LLMs（FBI-LLM）。為了實現從頭開始穩定地訓練二值化LLMs，文章提出了一種基于全精度教師模型的自回歸蒸餾的新型訓練模式。具體來說，在訓練過程中，作者逐步從全精度教師模型中生成蒸餾使用的軟標簽，并采用基于自回歸蒸餾的方案來匹配教師模型在每個token位置的預測概率。

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

作者對框架FBI-LLM的有效性進行了詳細評估，訓練了從130M、1.3B到7B規模的不同模型。作者使用廣泛使用的Transformer架構進行LLMs的訓練，結果表明從頭開始訓練全二值化的LLMs是完全可行的，其性能與全精度模型相比只有很小的差距。

與其他基準線方法相比，訓練過程在困惑度和多個下游任務上表現更為出色。這些結果表明自回歸蒸餾是訓練二值化LLMs的關鍵，此外，通過對預訓練更加深入的研究（如權重翻轉比和梯度范數）的分析表明，從全精度LLMs繼承權重與從頭開始訓練二值化LLMs之間沒有顯著差異。

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

下面介紹一下文章具體細節。

模型結構

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

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

模型訓練：自回歸蒸餾（Autoregressive Distillation）

給定一個訓練語料，標準的自回歸語言模型的目標函數是最大化如下似然函數：

其中k表示上下文窗口的大小，條件概率p通過參數為θ的神經網絡建模。不同于一般的自回歸語言模型，本文使用自回歸蒸餾訓練 FBI-LLM。在訓練過程中，一個全精度預訓練 LLM 被用作教師模型，二值化目標模型作為學生模型。假設每個訓練數據實例由輸入token序列x^1,…x^m組成，教師模型對下一個標記的預測概率可以表示為：

其中表示最后一層 transformer 模塊的激活，表示用于預測下一個token概率的線性輸出層的參數。

學生模型與教師模型輸出之間的交叉熵被計算為每一步預測下一個token時的最終損失函數。它可以表示為：

其中n表示輸入標記的數量。表示教師模型預測的第i步詞匯表上的標記分布，而是學生模型的相應預測分布。

訓練數據

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

實驗結果

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

其次是在下游任務上的性能表現，如下表所示，由于 130M 大小的FBI-LLM沒有對應的基準模型，本文將 130M 模型與之前 700M 規模的 BitNetb1.58 進行比較。盡管模型規模相差五倍，權重量化程度也存在較大的差異，但FBI的模型在 BoolQA 和 OpenbookQA 上的表現仍然優于 BitNet b1.58。

對于 1.3B 規模的二值化模型，FBI-LLM 在大多數下游任務和困惑度中都取得了最佳性能，甚至接近或超過了某些 7B 規模的二值化模型（如 BiLLM-LLaMA2-7B）的性能。與相同規模的全精度模型相比， FBI-LLM 1.3B 在下游任務中可以達到其 87% 的性能。在 7B 規模中，FBI模型依然顯著優于之前幾乎所有的基準線模型，具體來說，FBI-LLM相比之前最好的方法平均提升了將近十個點。

此外，由于計算資源的限制，FBI-LLM 7B當前匯報的結果并不是最終結果。作者只使用了整個 數據集的 8.6%（31 個塊）。下圖展示了FBI-LLM-7B訓練過程中下游任務準確率和困惑度的變化。顯然，從目前的訓練進度來看，FBI-LLM-7B的性能將持續提高，更進一步的訓練可能會得到更好的效果。

模型分析和可視化

二值化大模型是 從頭開始訓練還是從預訓練的 LLM 接著訓練？

直觀地說，從預訓練的 LLM 繼續訓練可以讓二值化模型從全精度原始模型中繼承知識，從而可能比從頭開始訓練獲得更好的結果。為了論證這一假設，本文進行了全面的消融和分析實驗，以記錄和比較模型在兩種不同訓練模式下的行為。

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

雖然從頭開始訓練的損失在中間階段比繼續訓練的損失略高，但過一段時間后，從頭開始訓練的損失再次與接著訓練的損失相當，甚至變得比繼續訓練的損失更為穩定。值得注意的是，在大約第 1000 步時，如圖 (a) 所示，當從預訓練的 LLM 繼續訓練時，FF ratio 開始出現明顯的波動。同樣，在圖 (b) 所示的第 1700 步左右，訓練損失也遇到了類似的問題。

這些發現挑戰了之前相關工作的一些假設，即從預訓練的 LLM權重開始訓練二值化大模型將賦予二值化LLM繼承之前預訓練的知識，從而提高性能。然而，本文結果和分析暗示了二值化LLM對參數初始化的方式不敏感，或者說隨機參數初始化反而更好。對于原因，本文推測二值化和全精度LLM采用不同的參數組合和配置來編碼語義，這導致其參數空間模式存在實質性差異。為了適應這種模式，通過從預訓練的 LLM 繼續訓練來優化二值化的過程可能需要更大幅度的參數數值調整和變換，從而破壞預訓練參數模式結構，繼而破壞其中保存的知識。這個理論可以部分地解釋為什么在訓練期間與從頭開始訓練相比，繼續預訓練的權重反而讓二值化LLM變得更不穩定。

訓練的穩定性分析

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

存儲效率分析

上表顯示了不同規模的FBI-LLM與相同結構的全精度LLaMA相比所需的理論存儲空間，同時還詳細給出了 FBI-LLM 引入的附加參數（α 和 β）的比例。表中的比較表明，FBI-LLM可以實現高壓縮比，大大減輕了LLMs的存儲負擔。盡管 FBI-LLM 引入的用于縮放和移動的額外參數需要保留全精確，但它們的比例相比整個模型很小，因此其對存儲的影響可以忽略不計。

生成結果示例

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

模型配置和訓練細節

具體的模型配置和訓練細節如下所示：

更多細節歡迎閱讀文論原文。