為什麼網路比實際需要的還大
當你訓練完一個網路後,它幾乎從來不會用盡自己全部的容量。研究人員一再發現:你可以丟掉很大一部分權重——有時高達 80% 或 90%——只要稍加修補,網路的表現就幾乎和原來一模一樣。這就是本篇所有內容背後的核心事實:訓練好的網路是過度參數化的。它們擁有的旋鈕(權重)遠比任務真正需要的還多。這些多出來的旋鈕在訓練過程中有用——它們讓損失地形更好走——但一旦訓練結束,其中很多就變得多餘、接近零,或只是在做和鄰居一樣的工作。
卷積神經網路示意圖:輸入影像由左到右流經數個空間尺寸逐漸縮小的卷積與池化區塊,最後接到全連接分類器。
一個很有幫助的想像是樂透彩券直覺。把一個大網路想成一大疊樂透彩券,這疊彩券裡藏著一個小小的子網路——一張「中獎彩券」——如果你一開始就單獨訓練它,它能把工作做得幾乎和整疊一樣好。訓練大網路某種程度上就是在尋找那張中獎彩券。誠實地說、不誇大:實務上我們通常得先訓練完大網路才找得到那個中獎子網路,而且要乾淨地把它抽取出來至今仍是熱門的研究題目。但實用的啟示很扎實——你的大網路裡藏著一個更小卻同樣能幹的網路,而壓縮就是把它抽出來的手藝。
把這些抽取小網路的技巧全部涵蓋起來的統稱,就是模型壓縮。它是一個家族,而不是單一招式;本篇會帶你逛過你最常用到的四位成員:剪枝(刪掉不重要的權重或通道)、知識蒸餾(訓練一個小模型去模仿大模型)、低秩分解(把一個又胖的權重矩陣換成兩個瘦的)、以及——預告下一篇——量化(用更少的位元來儲存與運算數字)。每一種都從不同角度去攻擊「太大」這件事,而正如我們在最後會看到的,它們可以疊著用。
剪枝:剪掉無用的權重
網路剪枝最直觀的形式——也是整個主題的入門——就是以權重大小為依據的剪枝。做法簡單到幾乎令人起疑:看過某一層裡的每一個權重,按絕對值大小排序,把最小的那些設成零。背後的道理是:權重會去乘上一個輸入,所以如果權重很小,它對輸出的改變就很小,刪掉它應該也幾乎不痛不癢。這是模型壓縮家族中最簡單的成員,也是幾乎每個人開始的地方。
稀疏度 s,以及以權重大小為準則:把任何絕對值低於門檻 τ 的權重設為零。
我們把兩條式子都拆開來看。稀疏度 s 就是你設成零的權重所佔的比例:N_{\text{zeroed}} 是被設成零的權重個數,N_{\text{total}} 是該層(或整個模型)的權重總數。如果一層有 1,000,000 個權重而你把其中 600,000 個設成零,那麼 s = 600{,}000 / 1{,}000{,}000 = 0.6——也就是「60% 稀疏度」。第二條式子則是決定要把哪些權重設成零的規則:只要權重 w 的絕對值 |w| 低於你選定的門檻 \tau,就把它剪掉(你會把 \tau 設成讓剛好目標比例的權重落在它之下)。為什麼用 |w|?因為一個權重對輸出的貢獻是(權重 × 輸入),而接近零的權重不管輸入是多少,貢獻都幾乎為零——所以它最不值錢、最該被捨棄。關鍵在於:這只是對真正重要性的一種近似——一個小權重若餵進一條非常敏感的路徑,仍可能很重要;一個大權重也可能其實多餘。權重大小是衡量重要性的一個快速、出奇好用的代理指標,但不是保證。
import torch
def magnitude_prune(weight, sparsity):
# weight: one layer's weight tensor
# sparsity: fraction of weights to zero out, e.g. 0.6
k = int(sparsity * weight.numel()) # how many weights to remove
if k == 0:
return weight
# threshold = the k-th smallest |w| in the layer
threshold = weight.abs().flatten().kthvalue(k).values
mask = weight.abs() > threshold # keep weights above threshold
return weight * mask # zeroed weights stay in place (same shape!)一個 ResNet 殘差區塊,包含兩個卷積層與一條跳接;卷積權重格子中有許多被塗灰,表示被設成零的權重。
結構化剪枝:移除整個通道
結構化剪枝正好修補了那個掃興之處。它不是把散落各處的個別權重設成零,而是移除整個結構——整個卷積濾波器(輸出通道)、注意力頭、甚至整個區塊。當你刪掉一個輸出通道,你刪掉的是權重張量中一整片相連的部分,可以實際把它切掉,留下一個真正更小、卻依然稠密的網路。「稠密」意味著一般硬體可以全速執行它,不需要任何特殊的稀疏矩陣技巧。這才是真正能在延遲上得到回報的網路剪枝形式。
我們怎麼決定要剪掉哪些通道?我們替每個通道的重要性打分數,然後丟掉分數最低的。兩種便宜又常用的分數:(1) 通道權重的 L2 範數——一個權重全都很小的濾波器,產生的輸出也小,所以多半可有可無;(2) BatchNorm 的縮放係數 γ——大多數 CNN 通道會通過一個 BatchNorm 層,被一個學到的 γ 相乘,而 γ 接近零的通道,等於是網路自己主動把它調小,這強烈暗示它不重要。這些都是啟發式規則,跟前面的權重大小規則一樣,但它們計算便宜、實務上也好用。
一個卷積層的運算成本(出自第 1 篇)。輸出通道數 C_out 是其中一個直接的乘數。
這正是第 1 篇的FLOPs公式,而它精確地告訴我們為什麼結構化剪枝能換到速度。H_{\text{out}} 和 W_{\text{out}} 是輸出特徵圖的高與寬;C_{\text{in}} 和 C_{\text{out}} 是輸入與輸出通道的數量;K 是核大小(所以 K^2 是一個濾波器的面積)。成本是一個連乘,所以減少輸出通道,就會讓 C_{\text{out}} 這一項成正比地縮小。 拿一個真實的層來看:一個 3\times3 卷積,C_{\text{in}}=128、C_{\text{out}}=256、輸出 56\times56。它是 56\cdot56\cdot128\cdot256\cdot9 \approx 9.2 億次乘加。把它 25% 的輸出通道剪掉(256 \to 192),成本就降到 56\cdot56\cdot128\cdot192\cdot9 \approx 6.9 億——乾淨俐落地砍掉 25%,因為 192/256 = 0.75。更妙的是,這一刀會連鎖:下一層現在看到的 C_{\text{in}} 變成 192 而不是 256,所以它也跟著縮小。移除一個通道,瘦了兩層。
一個卷積層產生的一疊特徵圖切片;其中兩張被打叉刪除,下游的張量也相應地畫得更薄。
知識蒸餾:向老師學習
知識蒸餾從一個完全不同的角度切入模型壓縮:你不是去縮小一個已經訓練好的模型,而是從頭訓練一個小模型——學生——去模仿一個又大又準的模型——老師。巧妙之處在於學生模仿的是什麼。一般的分類器是用硬標籤來訓練的:「這張影像是貓(1),不是狗(0),不是車(0)。」但老師提供的東西豐富得多——它在所有類別上的完整機率分布——而我們要學生去比對的,正是這個。
為什麼老師的完整分布這麼有價值?想想一張貓的照片。一個好老師可能會輸出「貓 0.84、狗 0.11、車 0.0001」。注意這裡藏著的資訊:老師在說,貓看起來有一點像狗、完全不像車。這種對錯誤類別之間的相對排序——常被稱作暗知識——是一堂關於各類別彼此關係的免費課程,而硬標籤「1, 0, 0」把它全丟掉了。不過有個小麻煩:一個有自信的老師,機率分布太尖銳(0.84 對 0.11 對 0.0001),以致那些小數字之間有趣的結構幾乎看不見。解方是溫度:在做 softmax 之前先把 logits 除以一個數 T,這會軟化分布、把那些小機率抬高到學生真的學得到的位置。
帶溫度 T 的 softmax。T 越大分布越平坦;T = 1 就是一般的 softmax。
我們慢慢讀,因為這一節後面都建立在它上面。\sigma(z)_i 是類別 i 的輸出機率。z_j 是logits——網路在 softmax 之前產生的、未經正規化的原始分數。e^{z_i/T} 把每個 logit(先除以溫度 T 之後)取指數,讓所有東西都變正;分母 \sum_j e^{z_j/T} 對所有類別加總,好讓輸出總和為 1。現在用一個真實例子看看 T 做了什麼。設老師對 [貓, 狗, 車] 的 logits 為 z = [4, 2, 1]。在 T = 1(一般 softmax)下我們得到 [0.844,\,0.114,\,0.042]——貓壓倒性地高,狗和車被壓得扁平。在 T = 4 下,我們先相除得到 [1,\,0.5,\,0.25],softmax 現在給出大約 [0.481,\,0.292,\,0.227]。排序一樣,但看看差距縮小了多少:學生現在清楚看到「狗是個合理的第二名,車是個微弱但真實存在的第三名」。這就是暗知識被顯現出來了。當 T \to \infty,分布趨於平坦的均勻分布;當 T \to 1,它又銳化回原樣。T 取 2 到 6 之間是常見的。
蒸餾損失:一個硬標籤項,加上一個軟性的「模仿老師」項,由 α 平衡、並以 T² 重新縮放。
學生要最小化的損失,融合了兩個目標。第一項 \mathrm{CE}(y, \sigma(z_s)) 是對真實硬標籤 y 的一般交叉熵——「持續把真正的答案答對」——其中 z_s 是學生的 logits,這裡的 \sigma 是普通(T=1)的 softmax。第二項是 Kullback–Leibler 散度 \mathrm{KL}(\,\cdot\,\|\,\cdot\,),衡量兩個分布有多不同(相同時為零,越分歧越大),比較的是老師軟化後的分布 \sigma(z_t/T) 與學生軟化後的分布 \sigma(z_s/T),其中 z_t 是老師的 logits。最小化它,會推著學生的軟性輸出去比對老師的——暗知識就是在這裡轉移的。權重 \alpha \in [0,1] 在兩者間取捨:\alpha 大就偏向真實標籤,\alpha 小就偏向模仿(常取 0.1 到 0.5)。那麼為什麼要乘 T^2? 因為把 logits 除以 T 會讓軟性項的梯度大約縮小 1/T^2,所以如果不修正,這一項會隨著你調高溫度而悄悄消失。乘上 T^2 正好抵消這個縮小,讓兩項不管你選什麼 T 都維持差不多的份量。
低秩分解:精打細算的矩陣
低秩分解是模型壓縮工具箱裡最有線性代數味道的一把,而一旦你想通了,它感覺幾乎像作弊。一個全連接層不過就是乘上一個權重矩陣 W。如果 W 是,比方說,1000 \times 1000,那就是要儲存一百萬個數字、執行一百萬次乘加。它的點子是:把那一個又胖的矩陣,換成兩個瘦矩陣的乘積——一個 1000 \times r 的矩陣乘上一個 r \times 1000 的矩陣——並挑選它們,使其乘積接近原本的 W。
為什麼兩個瘦矩陣能頂替一個胖的?答案是 SVD(奇異值分解),它把任何矩陣寫成一堆簡單的秩-1 小塊之和,每一塊都標著一個「奇異值」,用來衡量那一塊有多重要。在真實訓練好的層裡,奇異值衰減得很快:少數幾個方向就帶走了矩陣大部分的「能量」,而長長的尾巴幾乎沒有貢獻。所以如果我們只保留前 r 個方向——一個截斷的 SVD——我們就丟掉了可忽略的尾巴,同時留住了真正在做事的結構。這個保留下來的數目 r,就是我們這個近似的秩。
用一個高的 U 與一個寬的 V(秩為 r)的乘積,去近似一個 m×n 的矩陣 W。參數量從 m·n 降到 r·(m+n)。
逐個符號來看:W 是原本的權重矩陣,有 m 列、n 行,所以它裝著 m \cdot n 個數字。我們把它近似成 U V,其中 U 是 m \times r(又高又瘦),V 是 r \times n(又矮又寬);r 是我們選擇保留的秩。U 裝 m \cdot r 個數字、V 裝 r \cdot n 個,所以兩者合起來花費 r(m+n)。只有當 r(m+n) < m\,n,也就是當 r < \dfrac{m\,n}{m+n} 時,才真的賺到。代入 m=n=1000:損益兩平的秩是 \dfrac{1{,}000{,}000}{2000}=500。選 r=50,你就只儲存 50 \cdot (1000+1000)=100{,}000 個數字,而不是 1{,}000{,}000 個——乾淨俐落地省下 10 倍,而且矩陣乘法也便宜了約 10 倍。
同樣的把戲也適用於卷積:一個昂貴的卷積可以拆成兩個便宜的(例如把一個 k\times k 卷積分解成一個 k\times 1 接著一個 1\times k,或把一個厚卷積拆成一個瘦的「瓶頸」投影再加一個擴張)。而這裡有個誠實的限制:分解只有在低秩就夠用時才有幫助。如果某一層真的需要很大的 r——比方在我們 1000\times1000 的例子裡 r=450,那就是 450\cdot2000=900{,}000 個參數對上原本的一百萬——你幾乎沒省到什麼,而且那點近似誤差可能根本不划算。許多現代架構(用 1\times1 卷積與瓶頸)在設計上本來就已經是低秩的,能再擠的空間就少了。在假設分解會划算之前,永遠先檢查一下奇異值的衰減情形。
import numpy as np W = np.random.randn(1000, 1000) # a big weight matrix: 1,000,000 params U, S, Vt = np.linalg.svd(W, full_matrices=False) r = 50 # keep only the top-50 directions sqrtS = np.sqrt(S[:r]) U_r = U[:, :r] * sqrtS # 1000 x 50 V_r = (sqrtS[:, None]) * Vt[:r, :] # 50 x 1000 W_approx = U_r @ V_r # close to W if singular values decay fast # stored params: 50*(1000 + 1000) = 100,000 vs 1,000,000 -> ~10x smaller
組合多種技巧並找回準確率
這些工具沒有一個是彼此的對手——它們是一個工具箱,而最強的成果來自把它們疊著用。一條典型而積極的流水線:先用知識蒸餾訓練一個本來就以小搏大的較小架構,再用結構化剪枝去修掉那些活下來卻沒在好好做事的通道,最後(下一篇)量化,把剩下的每一個數字都縮小。每一步都是對同一個目標的不同槓桿,而模型壓縮在實務上,就是針對你特定的預算去挑選並排序它們。
- 建立基準線:量測完整模型的準確率、大小與延遲,好讓之後每一個改動都拿真實數字來評斷,而不是憑感覺。
- 用師生損失把模型蒸餾進一個較小的架構,然後把學生微調到收斂。
- 在重要性分數的引導下,結構化地剪掉一小批通道(例如約 10%),然後微調,把你剛剛失去的準確率找回來。
- 反覆進行剪枝 → 微調,直到你撞上準確率的下限,或撞上你的延遲/大小目標。
- 量化(下一篇),然後微調或校準,以吸收因位元變少而產生的捨入誤差。
- 在出貨之前,於完整的測試集以及你已知的失敗情境上重新驗證。
該推到多遠?把它當成一個明確的預算。挑出真正會卡住你的限制——延遲目標、記憶體上限、App 下載大小——以及你不肯跌破的準確率下限。然後一路壓縮,直到撞上其中之一,過程中盯著準確率對大小的曲線。早期的壓縮幾乎是免費的:你刪掉的是真正的無用權重,所以曲線幾乎是平的。到某個點,它會急轉直下——那個膝點就是你該停手的地方。越過它,你就是拿大量準確率去換一點點大小,這通常不划算。
這就只剩下我們一直只是預告的那一種技巧:量化——用 8 位元整數(甚至更少)而非 32 位元浮點數來表示權重與激活值。它往往是其中最大的「免費午餐」:大小大約縮小 4 倍,而且和非結構化剪枝不同,它在為整數運算打造的硬體上能帶來真正的加速,常常幾乎不損失準確率。它重要到足以自成一篇——也就是下一篇。在那之後,我們會看看那些把這一切壓縮真正轉化成部署速度的執行環境與加速器(ONNX、TensorRT,以及整個硬體堆疊)。