AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
作者介紹：論文一作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展示了進一步擴大模型規模從而達到更高智力水平的二值化模型的潛力，這種型別的大模型 (二值化大模型) 對部署的硬體需求也更加友好。
模型配置和訓練細節
具體的模型配置和訓練細節如下所示：
更多細節歡迎閱讀論原文。