好家伙！1750億參數(shù)的GPT-3只需20MB存儲(chǔ)空間了？！

基于1.58-bit訓(xùn)練，在不損失精度的情況下，大幅節(jié)省算力（↓97%）和存儲(chǔ)（↓90%）。

最近，從事機(jī)器學(xué)習(xí)的Will小哥發(fā)了一篇論文，直接引來(lái)幾十萬(wàn)網(wǎng)友or同行圍觀。

他提出了一項(xiàng)名為“noise_step”的新技術(shù)，允許模型直接在1.58-bit低精度下訓(xùn)練，且無(wú)需反向傳播或動(dòng)量（Momentum）加速，從而降低算力和存儲(chǔ)消耗。

對(duì)此，網(wǎng)友們也紛紛發(fā)來(lái)祝賀，表示很高興看到模型越來(lái)越具有性?xún)r(jià)比。

話不多說(shuō)，來(lái)看論文具體內(nèi)容。

反向傳播不需要了

簡(jiǎn)單說(shuō)，noise_step的目標(biāo)是通過(guò)降低模型訓(xùn)練的精度要求，來(lái)減少算力和存儲(chǔ)消耗。

一開(kāi)始，作者提到前人研究已經(jīng)表明，大語(yǔ)言模型（LLM）的推理可以在1.58-bit精度下進(jìn)行，且不會(huì)有任何性能損失。

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

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

然而，上述變體是在推理時(shí)使用低精度，而在訓(xùn)練時(shí)仍需高精度權(quán)重。

因此，noise_step的一個(gè)核心區(qū)別是：無(wú)需反向傳播。

允許模型直接在1.58-bit（三元）精度下進(jìn)行訓(xùn)練，而不需要傳統(tǒng)的反向傳播（從后向前檢查每一層）或動(dòng)量方法。

注：反向傳播(Backpropagation)是訓(xùn)練神經(jīng)網(wǎng)絡(luò)的核心算法，它通過(guò)反向逐層計(jì)算損失函數(shù)對(duì)每個(gè)權(quán)重的梯度，來(lái)反向逐層更新網(wǎng)絡(luò)的權(quán)重，從而最小化損失函數(shù)。

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

簡(jiǎn)單說(shuō)，通過(guò)在前向傳播中引入隨機(jī)性，可以生成一個(gè)隨機(jī)向量。這個(gè)隨機(jī)向量與目標(biāo)函數(shù)的梯度之間的對(duì)齊可以通過(guò)計(jì)算JVP來(lái)評(píng)估。

通過(guò)在多個(gè)隨機(jī)方向上重復(fù)JVP計(jì)算，可以收集足夠的信息來(lái)估計(jì)整個(gè)梯度向量，從而實(shí)現(xiàn)不依賴(lài)于反向傳播的梯度估計(jì)。

will小哥的具體方法如下：

圖片

訓(xùn)練成本打下來(lái)了

上述方式意味著，noise_step允許使用更稀疏的隨機(jī)向量和簡(jiǎn)單的對(duì)齊值。

要知道傳統(tǒng)的梯度計(jì)算需要大量計(jì)算資源，而noise_step由于不需要存儲(chǔ)或傳輸大量數(shù)據(jù)，從而減少了存儲(chǔ)使用。

此外，由于noise_step使用偽隨機(jī)噪聲，它只需要一個(gè)種子（初始值）就能復(fù)現(xiàn)整個(gè)訓(xùn)練過(guò)程，這意味著不需要存儲(chǔ)大量的擾動(dòng)向量，從而進(jìn)一步減少了存儲(chǔ)需求。

而且使用noise_step訓(xùn)練的模型可以存儲(chǔ)訓(xùn)練步驟而非權(quán)重，這可能會(huì)大幅縮小模型尺寸，從而更快地下載模型。

按照will小哥的說(shuō)法，也許今后一秒鐘內(nèi)下載一個(gè)SOTA模型？

同時(shí)由于上面提到的偽隨機(jī)噪聲方法，這種特性允許恢復(fù)權(quán)重的完整歷史，因?yàn)槊總€(gè)步驟都是確定性的，并且可以獨(dú)立于其他步驟進(jìn)行計(jì)算。

因此微調(diào)將變得更加高效，甚至可能允許對(duì)過(guò)去的訓(xùn)練步驟進(jìn)行編輯，例如翻轉(zhuǎn)（negation）或屏蔽（masking）。

舉個(gè)例子，如果發(fā)現(xiàn)某個(gè)訓(xùn)練步驟對(duì)模型性能產(chǎn)生了負(fù)面影響，可以對(duì)其進(jìn)行調(diào)整而不必重新訓(xùn)練整個(gè)模型。

也就是說(shuō)，人們?cè)谟?xùn)練過(guò)程中能進(jìn)行更精細(xì)的控制和調(diào)整了。

最后，作者認(rèn)為這種方式尤為適合分布式訓(xùn)練。

在分布式訓(xùn)練中，通常需要在不同的計(jì)算節(jié)點(diǎn)之間同步梯度和優(yōu)化器狀態(tài)，這會(huì)限制訓(xùn)練的速度。而noise_step通過(guò)減少每個(gè)擾動(dòng)所需的位數(shù)，顯著降低了通信量，從而提高了分布式訓(xùn)練的效率。

不過(guò)這也導(dǎo)致模型泄露變得更加容易，因?yàn)檎麄€(gè)模型可以通過(guò)幾個(gè)字節(jié)的訓(xùn)練步驟來(lái)傳輸。

對(duì)了，will小哥表示JVP可以和正常推理并行運(yùn)行，幾乎不增加成本。

除了論文，他也提供了一個(gè)CPU實(shí)現(xiàn)過(guò)程：

One More Thing

BTW，小哥在??上分享完論文后，還順帶吐槽arXiv不給理由就拒絕了這篇論文。

后來(lái)他才補(bǔ)充，原來(lái)是卡在了背書(shū)（endorse）這一項(xiàng)，也就是需要現(xiàn)有用戶(hù)的推薦或通過(guò)其他方式獲得背書(shū)。

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

感興趣的童鞋可以進(jìn)一步查看。

論文：https://github.com/wbrickner/noise_step?tab=readme-ov-file
CPU實(shí)現(xiàn)過(guò)程：https://colab.research.google.com/drive/1hXzf5xB4INzMUNTlAB8CI1V10-JV7zyg?usp=sharing