只需激活60%的參數，就能實現與全激活稠密模型相當的性能。

微軟亞洲研究院的一項新研究，實現了模型的完全稀疏激活，讓推理成本大幅下降。

而且適用范圍廣泛，無論是從頭訓練、繼續(xù)訓練還是微調，都能提供有效支持。

該方法名為Q-Sparse，在神經元級別上實現了模型稀疏化，相比于其他方式粒度更細，在相同推理開銷下，無論性能還是稀疏率都更好。

名稱之中，Q指的是量化（Quantization），意味著它除了普通模型之外，也兼容量化技術，適用于各種量化方式的模型。

作者進一步表示，如果把Q-Sparse與模型量化技術結合，還可以實現更大程度的降本增效。

另外在研究Q-Sparse的同時，團隊也對參數規(guī)模、稀疏率和模型性能三者之間的關系進行了深入探尋，并發(fā)現了適用于模型推理優(yōu)化的“Scaling Law”。

有網友認為，這項技術確實不錯，而且比ReLU要更好。

還有人開啟了許愿模式，表示如果（AMD的）ROCm能比英偉達更快支持這項技術就好了。

用Top-K函數實現稀疏化

Q-Sparse所做的最核心的操作，是對輸入的張量應用Top-K稀疏化函數。

具體來說，Transformer架構在注意力層和前饋層中都使用nn.Linear線性層（矩陣乘法）進行投影，可以表示為Y=X·W^T。（其中X就是輸入張量，W代表其權重，Y為輸出張量）

Q-Sparse中，對于一個輸入激活張量X，首先會計算其絕對值|X|并進行排序，找出其中絕對值最大的K個元素。

這里的K是預先設定的超參數，決定了稀疏化的程度。

之后Q-Sparse會創(chuàng)建一個與X形狀相同的二進制掩碼張量M，對于一系列|X|中絕對值最大的K個元素對應的位置，將M中的相應位置設置為1，其余位置設置為0。

接著，將輸入張量X與掩碼張量M進行Hadamard積（逐元素相乘）運算，就得到了稀疏化的張量X_sparse。

在前向傳播過程中，稀疏化后的張量X_sparse將代替原始的輸入張量X參與后續(xù)的計算（如矩陣乘法）。

由于X_sparse中大部分元素已經被設置為零，因此可以顯著減少計算量和內存帶寬需求。

在反向傳播過程中，Q-Sparse使用了直通估計器（Straight-Through Estimator，STE）來計算Top-K函數的梯度。

傳統(tǒng)的訓練方式中，通常需要計算損失函數對網絡參數的梯度，并使用梯度下降法更新參數以最小化損失。

但當網絡中存在量化、Top-K等一些不可微的操作時，梯度的計算就會遇到問題，因為這些操作的輸出對輸入的梯度在大多數點上都是0，導致梯度無法有效傳播。

STE通過直接將梯度傳遞給稀疏化之前的張量，避免了梯度消失的問題。

一般的反向傳播中，損失函數L對x的梯度?L/?x=?L/?y??y/?x，但由于不可微分無法直接計算。

STE的解決方案是只計算損失函數對稀疏化張量y的梯度，然后將其直接復制給原始張量x，也就是直接將?L/?y作為?L/?x的估計。

△有/無STE時的梯度比較

對于前饋層，Q-Sparse使用平方ReLU函數代替常規(guī)的ReLU激活函數，平方運算可以進一步提高激活的稀疏性（⊙表示Hadamard積）。

另外，為了適配量化模型，Q-Sparse在應用Top-K稀疏化之前，會先對輸入張量進行量化，以確保稀疏化操作與量化表示兼容，其函數表示如下：

其中，ε是一個小常數，用于避免出現分母為零的情況。

特別的，對于1-bit量化的權重，Q-Sparse使用以下量化函數，其中α是權重張量W的平均絕對值。

60%激活參數達到相同效果

對比實驗表明，無論是稀疏率還是模型表現，Q-Sparse都顯著優(yōu)于此前的ReLU方法。

針對Q-Sparse的具體效果，作者對其在從頭訓練、繼續(xù)訓練和微調三項任務上的性能進行了評估。

從頭訓練實驗使用的模型為Llama，結果在700M和7B模型上，使用70% top-K(即40%的整體稀疏率)的Q-Sparse可以達到與密集baseline相當的訓練損失。

繼續(xù)訓練的目的是將稠密模型稀疏化，這里的實驗對象是Mistral-7B。

結果，在激活參數為2.9B和3.8B的情況下，模型在ARC、MMLU等數據集中的得分均未發(fā)生明顯下降。

在微調實驗中，對于Qwen-7B和Mistral-7B兩種模型，Q-Sparse顯示出了與繼續(xù)訓練相似的結果，用60%左右的激活參數實現了與密集模型十分接近的表現。

這些結果意味著，在相同的性能下，與密集模型相比，稀疏激活模型在推理過程中可以顯著減少激活參數，進而降低消耗FLOPS的數量。

對于量化模型，團隊在自研的BitNet b1.58模型上應用了Q-Sparse，并在多個數據集上進行了訓練和評估。

可以看到，在700M和7B兩種規(guī)模下，使用Q-Sparse的量化模型的收斂速度和最終損失函數值與未使用Q-Sparse的量化模型(BitNet b1.58)相當。

這說明Q-Sparse可以無縫集成到量化模型中，而不會顯著影響模型的訓練和收斂。

據此作者認為，將Q-Sparse與量化技術相結合，可以進一步提高大語言模型在推理階段的效率。

發(fā)現推理優(yōu)化新“Scaling Law”

除了測評這些模型采取稀疏激活時的表現，作者也對模型性能、規(guī)模和稀疏率三者之間的關系進行了探究，并有了一些新的發(fā)現。

稀疏激活模型的性能縮放定律: 作者發(fā)現，與密集模型類似，稀疏激活模型的性能也遵循一個冪律縮放關系。

具體來說，給定稀疏率S，模型在收斂時的損失函數值L(N,S)可以用以下公式近似：

其中，N是模型參數的數量；E是一個常數，表示模型在無限大時的損失；A(S)是一個與稀疏率S有關的縮放因子。

這個縮放定律表明，稀疏激活模型的性能隨著模型規(guī)模的增大而提高，但提高的速度會逐漸變慢。

同時作者發(fā)現，模型的性能也會受到稀疏率的影響。

在參數規(guī)模與性能之間關系的部分提到，A(S)是一個與稀疏率S有關的縮放因子，可以用以下公式近似：

其中B和C是常數，β是一個控制指數衰減速度的參數。

這個公式表明，當稀疏率S增大（模型變得更稀疏）時，意味著更高的稀疏率會導致性能的下降，下降的速度是指數級的。

基于上述發(fā)現，作者得出了一個推理最優(yōu)的稀疏率S*，能在預算（推理時的浮點操作數）一定時，實現模型損失函數值的最小化。

對于全精度（FP32）模型，最優(yōu)稀疏率約為45.58%；而低精度（如1.58-bit）模型的最優(yōu)稀疏率則更高，約為61.25%。

作者觀察到，隨著模型規(guī)模的增大，稀疏激活模型與密集模型之間的性能差距逐漸縮小。

這可以從縮放定律中得到解釋：當模型規(guī)模N趨于無窮大時，稀疏激活模型的損失函數值趨于L(∞,S)=E，而密集模型的損失函數值趨于L(∞,0)=E。

這意味著，在極大規(guī)模下，稀疏激活模型有可能達到與密集模型相當的性能，為設計和訓練大規(guī)模稀疏激活模型提供了一個有用的參考。

論文地址：https://arxiv.org/abs/2407.10969