譯者 | 朱先忠

審校 | 重樓

簡介

無論你是在準備面試，還是在工作中構建機器學習系統，模型壓縮都已成為一項必備技能。在大語言模型（LLM）時代，模型規模越來越大，如何壓縮這些模型以使其更高效、更小巧、更易于在輕量級機器上使用，這一挑戰從未如此嚴峻。

在本文中，我將介紹每位機器學習從業者都應該理解和掌握的四種基本壓縮技術。我將探討剪枝、量化、低秩分解和知識蒸餾，每種方法都各有優勢。我還將為每種方法添加一些精簡的PyTorch代碼示例。

模型剪枝

剪枝可能是最直觀的壓縮技術。其原理非常簡單：移除網絡的一些權重，可以是隨機移除，也可以是移除“不太重要”的權重。當然，在神經網絡中，“移除”權重指的是將權重設置為零。

模型修剪（圖片來自作者和ChatGPT，創作靈感來自于【引文3】）

結構化與非結構化修剪

讓我們從一個簡單的啟發式方法開始：刪除小于閾值的權重。

當然，這并不理想，因為我們需要找到一種方法來找到適合我們問題的閾值！更實際的方法是刪除某一層中幅度（范數）最小的指定比例的權重。在單層中實現剪枝有兩種常見的方法：

• 結構化修剪：刪除網絡的整個組件（例如，來自權重張量的隨機行，或卷積層中的隨機通道）。
• 非結構化剪枝：刪除單個權重，無論其位置和張量的結構如何。

我們也可以將上述兩種方法中的任意一種用于全局剪枝。這將移除多個層中選定比例的權重，并且根據每層參數的數量，移除率可能會有所不同。

PyTorch使這變得非常簡單（順便說一下，你可以在我的GitHub代碼倉庫中找到所有代碼片段）。

import torch.nn.utils.prune as prune

# 1. 隨機非結構化剪枝（隨機選取20%的權重）
prune.random_unstructured(model.layer, name="weight", amount=0.2) 

# 2.L1-范數非結構化剪枝（最小權重的20%）
prune.l1_unstructured(model.layer, name="weight", amount=0.2)

# 3. 全局非結構化剪枝（按L1范數計算，各層權重的40%）
prune.global_unstructured(
 [(model.layer1, "weight"), (model.layer2, "weight")],
 pruning_method=prune.L1Unstructured,
 amount=0.4
) 

# 4. 結構化修剪（刪除L2范數最低的30%行）
prune.ln_structured(model.layer, name="weight", amount=0.3, n=2, dim=0)

注意：如果你上過統計學課，你可能學過一些正則化方法，它們也會在訓練過程中隱式地使用L0或L1范數正則化來修剪一些權重。修剪與此不同，因為它是作為模型壓縮后的一項技術應用的。

剪枝為何有效？彩票假說

基于ChatGPT生成的圖像

我想用“彩票假說”來結束本節。它既是剪枝的一個應用，也對移除權重如何能夠改進模型進行了有趣的解釋。我建議你閱讀一下相關論文（引文7）以了解更多詳細信息。

論文作者采用了以下程序：

• 訓練完整模型，直至收斂
• 修剪最小幅度的權重（例如10%）
• 將剩余權重重置為其原始初始化值
• 重新訓練這個修剪后的網絡
• 重復該過程多次

重復30次之后，最終得到的參數只有原始參數的0.930（約4%）。令人驚訝的是，這個網絡的表現竟然和原始網絡一樣好。

這表明存在重要的參數冗余。換句話說，存在一個子網絡（“彩票”）實際上完成了大部分工作！

結論是：修剪是揭示這個子網絡的一種方法。

量化

修剪的重點是完全刪除參數，而量化則采用不同的方法：降低每個參數的精度。

請記住，計算機中的每個數字都是以位序列的形式存儲的。float32值使用32位（參見下圖），而8位整數（int8）僅使用8位。

float32數字如何用32位表示的示例（圖片來自作者和ChatGPT，創作靈感來自引文2）

大多數深度學習模型都使用32位浮點數（FP32）進行訓練。量化會將這些高精度值轉換為低精度格式，例如16位浮點數（FP16）、8位整數（INT8），甚至4位表示。

這里的節省是顯而易見的：INT8所需的內存比FP32少75%。但是，我們如何在不破壞模型性能的情況下實際執行這種轉換呢？

量化背后的數學

要將浮點數轉換為整數表示，我們需要將連續的數值范圍映射到一組離散的整數。對于INT8量化，我們將其映射到256個可能的值（從-128到127）。

假設我們的權重在-1.0和1.0之間標準化（在深度學習中很常見）：

然后，量化值由下式給出：

這里，zero_point=0因為我們希望0映射到0。然后，我們可以將這個值四舍五入到最接近的整數，以獲得-127到128之間的整數。

而且，你猜對了：為了將整數恢復為浮點數，我們可以使用逆運算：

注意：實際上，縮放因子是根據我們量化的范圍值確定的。

如何應用量化？

量化可以應用于不同的階段，并采用不同的策略。以下是一些值得了解的技巧：（下文中的“激活”一詞指的是每一層的輸出值）

• 訓練后量化（PTQ）：A.靜態量化：離線量化權重和激活（訓練之后和推理之前）。
  B.動態量化：離線量化權重，但在推理過程中動態激活。這與離線量化不同，因為縮放因子是根據推理過程中迄今為止看到的值確定的。
• 量化感知訓練（QAT）：通過對值進行舍入來模擬訓練過程中的量化，但計算仍然使用浮點數進行。這使得模型學習到對量化更具魯棒性的權重，這些權重將在訓練后應用。其底層思想是添加一些“假”操作：x -> dequantize(quantize(x))，這個新值接近x，但仍有助于模型容忍8位舍入和削波噪聲。

import torch.quantization as tq

# 1. 訓練后靜態量化（權重+離線激活）
model.eval()
model.qconfig = tq.get_default_qconfig('fbgemm') # 分配靜態量化配置
tq.prepare(model, inplace=True)
# 我們需要使用校準數據集來確定值的范圍
with torch.no_grad():
 for data, _ in calibration_data:
 model(data)
tq.convert(model, inplace=True) # 轉換為全int8模型

# 2.訓練后動態量化（權重離線，激活實時）
dynamic_model = tq.quantize_dynamic(
 model,
 {torch.nn.Linear, torch.nn.LSTM}, # layers to quantize
 dtype=torch.qint8
)

# 3. 量化感知訓練（QAT）
model.train()
model.qconfig = tq.get_default_qat_qconfig('fbgemm') # 設置QAT配置
tq.prepare_qat(model, inplace=True) #插入假量子模塊
# [here, train or fine?tune the model as usual]
qat_model = tq.convert(model.eval(), inplace=False) # 在QAT之后轉換為真正的int8

量化非常靈活！你可以對模型的不同部分應用不同的精度級別。例如，你可以將大多數線性層量化為8位，以實現最大速度和內存節省，同時將關鍵組件（例如注意力頭或批量規范層）保留為16位或全精度。

低秩分解

現在我們來談談低秩分解——一種隨著LLM的興起而流行的方法。

【關鍵觀察點】神經網絡中許多權重矩陣的有效秩遠低于其維度所暗示的秩。簡而言之，這意味著參數中存在大量冗余。

注意：如果你曾經使用過主成分分析（PCA）進行降維，那么你已經遇到過一種低秩近似的形式。主成分分析（PCA）將大矩陣分解為較小、低秩因子的乘積，從而盡可能多地保留信息。

低秩分解背后的線性代數

取權重矩陣W。每個實數矩陣都可以用奇異值分解（SVD）來表示：

其中Σ是一個奇異值非增階的對角矩陣。正系數的數量實際上對應于矩陣W的秩。

秩為r的矩陣的SVD可視化（圖片來自作者和ChatGPT，創作靈感來自引文5）

為了用秩k<r的矩陣近似W，我們可以選擇sigma的k個最大元素，以及相應的U和V的前k列和前k行：

看看新矩陣如何分解為A與B的乘積，現在參數總數是m*k+k*n=k*(m+n)而不是m*n！這是一個巨大的進步，尤其是當k遠小于m和時n。

實際上，它相當于用兩個連續的線性層x→Wx替換線性層x→A(Bx)。

在PyTorch中

我們可以在訓練前應用低秩分解（將每個線性層參數化為兩個較小的矩陣——這并非真正的壓縮方法，而是一種設計選擇），也可以在訓練后應用（對權重矩陣應用截斷奇異值分解）。第二種方法是迄今為止最常見的方法，如下所示。

import torch

# 1.提取重量并選擇秩
W = model.layer.weight.data # (m, n)
k = 64 # 期望的秩

# 2. 近似低秩SVD
U, S, V = torch.svd_lowrank(W, q=k) # U: (m, k), S: (k, k), V: (n, k)

# 3. 構造因子矩陣A和B
A = U * S.sqrt() # [m, k]
B = V.t() * S.sqrt().unsqueeze(1) # [k, n]

# 4. 替換為兩個線性層，并插入矩陣A和B
orig = model.layer
model.layer = torch.nn.Sequential(
 torch.nn.Linear(orig.in_features, k, bias=False),
 torch.nn.Linear(k, orig.out_features, bias=False),
)
model.layer[0].weight.data.copy_(B)
model.layer[1].weight.data.copy_(A)

LoRA：低秩近似的應用

LoRA微調：W是固定的，A和B經過訓練（來源：引文1）

我認為有必要提一下LoRA：如果你一直關注LLM微調的發展，你可能聽說過LoRA（低秩自適應）。雖然LoRA嚴格來說不是一種壓縮技術，但它因能夠有效地適應大型語言模型并使微調非常高效而變得非常流行。

這個想法很簡單：在微調過程中，LoRA不會修改原始模型權重W，而是凍結它們并學習可訓練的低秩更新：

其中，A和B是低秩矩陣。這使得僅使用一小部分參數就可以實現特定任務的自適應。

甚至更好：QLoRA通過將量化與低秩自適應相結合，進一步實現了這一點！

再次強調，這是一種非常靈活的技術，可以應用于各個階段。通常，LoRA僅應用于特定的層（例如，注意力層的權重）。

知識蒸餾

知識蒸餾過程（圖片來自作者和ChatGPT，創作靈感來自引文4）

知識蒸餾與我們迄今為止所見的方法截然不同。它不是修改現有模型的參數，而是將“知識”從一個龐大而復雜的模型（“老師”）遷移到一個規模更小、更高效的模型（“學生”）。其目標是訓練學生模型模仿老師的行為并復制其表現，這通常比從頭開始解決原始問題更容易。

蒸餾損失

我們來解釋一下分類問題中的一些概念：

• 教師模型通常是一個大型、復雜的模型，可以在當前任務中取得高性能。
• 學生模型是第二個較小的模型，具有不同的架構，但針對相同的任務進行定制。
• 軟目標：這些是教師模型的預測（概率，而不是標簽！）。學生模型將使用它們來模仿教師的行為。請注意，我們使用原始預測而不是標簽，因為它們也包含有關預測置信度的信息。
• 溫度：除了教師模型的預測之外，我們還在softmax函數中使用了一個系數T（稱為溫度），以便從軟目標中提取更多信息。增加T可以柔化分布，并幫助學生模型更加重視錯誤的預測。

實踐中，訓練學生模型非常簡單。我們將常規損失（基于硬標簽的標準交叉熵損失）與“蒸餾”損失（基于教師的軟目標）結合起來：

蒸餾損失只不過是教師分布和學生分布之間的KL散度（你可以將其視為兩個分布之間距離的度量）。

至于其他方法，可以并且鼓勵根據用例調整此框架：例如，還可以比較學生和教師模型之間網絡中的中間層的logit和激活，而不僅僅是比較最終輸出。

實踐中的知識蒸餾

與之前的技術類似，有兩種選擇：

• 離線蒸餾：預先訓練好的教師模型是固定的，并訓練一個單獨的學生模型來模仿它。這兩個模型完全獨立，并且在蒸餾過程中教師的權重保持不變。
• 在線蒸餾：兩個模型同時訓練，知識轉移發生在聯合訓練過程中。

下面是應用離線蒸餾的簡單方法（本文的最后一個代碼塊）：

import torch.nn.functional as F

def distillation_loss_fn(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
 # 帶有硬標簽的標準交叉熵損失
 student_loss = F.cross_entropy(student_logits, labels)

 # 軟目標的蒸餾損失(KL散度）
 soft_teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
 soft_student_log_probs = F.log_softmax(student_logits / temperature, dim=-1)

  # kl_div需要將對數概率作為第一個參數的輸入！
 distill_loss = F.kl_div(
 soft_student_log_probs,
 soft_teacher_probs.detach(), #不為教師計算梯度
 reductinotallow='batchmean'
 ) * (temperature ** 2) # 可選，縮放因子

 #根據公式計算損失
 total_loss = alpha * student_loss + (1 - alpha) * distill_loss
 return total_loss

teacher_model.eval()
student_model.train()
with torch.no_grad():
 teacher_logits = teacher_model(inputs)
  student_logits = student_model(inputs)
  loss = distillation_loss_fn(student_logits, teacher_logits, labels, temperature=T, alpha=alpha)
  loss.backward()
  optimizer.step()

結論

感謝你閱讀本文！在LLM時代，由于參數數量高達數十億甚至數萬億，模型壓縮已成為一個基本概念，幾乎在每種情況下都至關重要，可以提高模型的效率和易部署性。

但正如我們所見，模型壓縮不僅僅是為了減小模型大小，而是為了做出深思熟慮的設計決策。無論是選擇在線還是離線方法，壓縮整個網絡，還是針對特定的層或通道，每種選擇都會顯著影響性能和可用性。現在，大多數模型都結合了其中幾種技術（例如，查看這個模型）。

除了向你介紹主要方法之外，我希望本文還能激發你進行實驗并開發自己的創造性解決方案！

不要忘記查看我的GitHub存儲庫，你可以在其中找到所有代碼片段以及本文討論的四種壓縮方法的并排比較。

參考文獻

【1】Hu, E.等人，2021。《大型語言模型的低秩自適應》（Low-rank Adaptation of Large Language Models）。arXiv preprint arXiv：2106.09685。

【2】Lightning AI。《使用混合精度技術加速大型語言模型。》（Accelerating Large Language Models with Mixed Precision Techniques）。Lightning AI博客。

【3】TensorFlow博客。《TensorFlow模型優化工具包中的剪枝API》（Pruning API in TensorFlow Model Optimization Toolkit）。TensorFlow博客，2019年5月。

【4】Toward AI。《知識蒸餾的簡單介紹》（A Gentle Introduction to Knowledge Distillation）。Towards AI，2022年8月。

【5】Ju, A。《ML算法：奇異值分解（SVD）》（ML Algorithm: Singular Value Decomposition (SVD)）。LinkedIn Pulse。

【6】Algorithmic Simplicity。《這就是大型語言模型能夠理解世界的原因》（THIS is why large language models can understand the world）。YouTube，2023年4月。

【7】Frankle, J.與Carbin, M。（2019）。《彩票假設：尋找稀疏、可訓練的神經網絡》（The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks）。arXiv preprint arXiv：1803.03635。

譯者介紹

朱先忠，51CTO社區編輯，51CTO專家博客、講師，濰坊一所高校計算機教師，自由編程界老兵一枚。

原文標題：Model Compression: Make Your Machine Learning Models Lighter and Faster，作者：Maxime Wolf