算力直降97%，GPT-3存儲只用20MB？！這篇直接在1.58-bit下訓練模型的新論文火了

一水 發自 凹非寺

量子位 | 公眾號 QbitAI

好傢伙！1750億參數的GPT-3只需20MB存儲空間了？！

基於1.58-bit訓練，在不損失精度的情況下，大幅節省算力（↓97%）和存儲（↓90%）。

最近，從事機器學習的Will小哥發了一篇論文，直接引來幾十萬網民or同行圍觀。

他提出了一項名為「noise_step」的新技術，允許模型直接在1.58-bit低精度下訓練，且無需反向傳播或動量（Momentum）加速，從而降低算力和存儲消耗。

對此，網民們也紛紛發來祝賀，表示很高興看到模型越來越具有性價比。

話不多說，來看論文具體內容。

反向傳播不需要了

簡單說，noise_step的目標是通過降低模型訓練的精度要求，來減少算力和存儲消耗。

一開始，作者提到前人研究已經表明，大語言模型（LLM）的推理可以在1.58-bit精度下進行，且不會有任何性能損失。

比如下面這篇論文，有人引入了1-bit的LLM變體（即BitNet b1.58），其中LLM的每個參數或權重都是三元的{-1, 0, 1}。

它在困惑度（perplexity）和最終任務性能上與全精度（FP16或BF16）的Transformer LLM相匹配，同時在延遲、存儲、吞吐量和算力消耗方面成本更低。

然而，上述變體是在推理時使用低精度，而在訓練時仍需高精度權重。

因此，noise_step的一個核心區別是：無需反向傳播。

允許模型直接在1.58-bit（三元）精度下進行訓練，而不需要傳統的反向傳播（從後向前檢查每一層）或動量方法。

註：反向傳播(Backpropagation)是訓練神經網絡的核心算法，它通過反向逐層計算損失函數對每個權重的梯度，來反向逐層更新網絡的權重，從而最小化損失函數。

具體而言，will小哥參考了《Gradients without Backpropagation》這篇論文，其中介紹了雅可比向量積（Jacobian Vector Product，JVP）這種不依賴反向傳播的梯度估計方法。

簡單說，通過在前向傳播中引入隨機性，可以生成一個隨機向量。這個隨機向量與目標函數的梯度之間的對齊可以通過計算JVP來評估。

通過在多個隨機方向上重覆JVP計算，可以收集足夠的信息來估計整個梯度向量，從而實現不依賴於反向傳播的梯度估計。

will小哥的具體方法如下：

訓練成本打下來了

上述方式意味著，noise_step允許使用更稀疏的隨機向量和簡單的對齊值。

要知道傳統的梯度計算需要大量計算資源，而noise_step由於不需要存儲或傳輸大量數據，從而減少了存儲使用。

此外，由於noise_step使用偽隨機噪聲，它只需要一個種子（初始值）就能複現整個訓練過程，這意味著不需要存儲大量的擾動向量，從而進一步減少了存儲需求。

而且使用noise_step訓練的模型可以存儲訓練步驟而非權重，這可能會大幅縮小模型尺寸，從而更快地下載模型。

按照will小哥的說法，也許今後一秒鍾內下載一個SOTA模型？

同時由於上面提到的偽隨機噪聲方法，這種特性允許恢復權重的完整歷史，因為每個步驟都是確定性的，並且可以獨立於其他步驟進行計算。

因此微調將變得更加高效，甚至可能允許對過去的訓練步驟進行編輯，例如翻轉（negation）或屏蔽（masking）。

舉個例子，如果發現某個訓練步驟對模型性能產生了負面影響，可以對其進行調整而不必重新訓練整個模型。

也就是說，人們在訓練過程中能進行更精細的控制和調整了。

最後，作者認為這種方式尤為適合分佈式訓練。

在分佈式訓練中，通常需要在不同的計算節點之間同步梯度和優化器狀態，這會限制訓練的速度。而noise_step通過減少每個擾動所需的位數，顯著降低了通信量，從而提高了分佈式訓練的效率。

不過這也導致模型泄露變得更加容易，因為整個模型可以通過幾個字節的訓練步驟來傳輸。

對了，will小哥表示JVP可以和正常推理並行運行，幾乎不增加成本。

除了論文，他也提供了一個CPU實現過程：

One More Thing

BTW，小哥在𝕏上分享完論文後，還順帶「抽水」arXiv不給理由就拒絕了這篇論文。

後來他才補充，原來是卡在了背書（endorse）這一項，也就是需要現有用戶的推薦或通過其他方式獲得背書。

無奈之下，這篇論文目前被小哥放在了GitHub上。

感興趣的童鞋可以進一步查看。

論文：

https://github.com/wbrickner/noise_step?tab=readme-ov-file

CPU實現過程：

https://colab.research.google.com/drive/1hXzf5xB4INzMUNTlAB8CI1V10-JV7zyg?usp=sharing

參考鏈接：

https://x.com/_brickner/status/1871348156786704657