本文將深入研究修剪、量化、蒸餾等輕量化機(jī)器學(xué)習(xí)的五種核心技術(shù)，從而使你的神經(jīng)網(wǎng)絡(luò)更高效、更易于部署。

簡介

無論你是在準(zhǔn)備面試，還是在工作中構(gòu)建機(jī)器學(xué)習(xí)系統(tǒng)，模型壓縮都已成為一項必備技能。在大語言模型（LLM）時代，模型規(guī)模越來越大，如何壓縮這些模型以使其更高效、更小巧、更易于在輕量級機(jī)器上使用，這一挑戰(zhàn)從未如此嚴(yán)峻。

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

模型剪枝

剪枝可能是最直觀的壓縮技術(shù)。其原理非常簡單：移除網(wǎng)絡(luò)的一些權(quán)重，可以是隨機(jī)移除，也可以是移除“不太重要”的權(quán)重。當(dāng)然，在神經(jīng)網(wǎng)絡(luò)中，“移除”權(quán)重指的是將權(quán)重設(shè)置為零。

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

結(jié)構(gòu)化與非結(jié)構(gòu)化修剪

讓我們從一個簡單的啟發(fā)式方法開始：刪除小于閾值的權(quán)重。

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

• 結(jié)構(gòu)化修剪：刪除網(wǎng)絡(luò)的整個組件（例如，來自權(quán)重張量的隨機(jī)行，或卷積層中的隨機(jī)通道）。?
• 非結(jié)構(gòu)化剪枝：刪除單個權(quán)重，無論其位置和張量的結(jié)構(gòu)如何。?

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

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

import torch.nn.utils.prune as prune

# 1. 隨機(jī)非結(jié)構(gòu)化剪枝（隨機(jī)選取20%的權(quán)重）
prune.random_unstructured(model.layer, name="weight", amount=0.2) 

# 2.L1-范數(shù)非結(jié)構(gòu)化剪枝（最小權(quán)重的20%）
prune.l1_unstructured(model.layer, name="weight", amount=0.2)

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

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

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

剪枝為何有效？彩票假說

基于ChatGPT生成的圖像

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

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

• 訓(xùn)練完整模型，直至收斂?
• 修剪最小幅度的權(quán)重（例如10%）?
• 將剩余權(quán)重重置為其原始初始化值?
• 重新訓(xùn)練這個修剪后的網(wǎng)絡(luò)?
• 重復(fù)該過程多次?

重復(fù)30次之后，最終得到的參數(shù)只有原始參數(shù)的0.930（約4%）。令人驚訝的是，這個網(wǎng)絡(luò)的表現(xiàn)竟然和原始網(wǎng)絡(luò)一樣好。

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

結(jié)論是：修剪是揭示這個子網(wǎng)絡(luò)的一種方法。

量化

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

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

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

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

這里的節(jié)省是顯而易見的：INT8所需的內(nèi)存比FP32少75%。但是，我們?nèi)绾卧诓黄茐哪Ｐ托阅艿那闆r下實際執(zhí)行這種轉(zhuǎn)換呢？

量化背后的數(shù)學(xué)

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

假設(shè)我們的權(quán)重在-1.0和1.0之間標(biāo)準(zhǔn)化（在深度學(xué)習(xí)中很常見）：

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

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

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

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

如何應(yīng)用量化？

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

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

import torch.quantization as tq

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

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

# 3. 量化感知訓(xùn)練（QAT）
model.train()
model.qconfig = tq.get_default_qat_qconfig('fbgemm') # 設(shè)置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之后轉(zhuǎn)換為真正的int8

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

低秩分解

現(xiàn)在我們來談?wù)劦椭确纸狻环N隨著LLM的興起而流行的方法。

【關(guān)鍵觀察點】神經(jīng)網(wǎng)絡(luò)中許多權(quán)重矩陣的有效秩遠(yuǎn)低于其維度所暗示的秩。簡而言之，這意味著參數(shù)中存在大量冗余。

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

低秩分解背后的線性代數(shù)

取權(quán)重矩陣W。每個實數(shù)矩陣都可以用奇異值分解（SVD）來表示：

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

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

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

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

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

在PyTorch中

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

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. 構(gòu)造因子矩陣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：低秩近似的應(yīng)用

LoRA微調(diào)：W是固定的，A和B經(jīng)過訓(xùn)練（來源：引文1）

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

這個想法很簡單：在微調(diào)過程中，LoRA不會修改原始模型權(quán)重W，而是凍結(jié)它們并學(xué)習(xí)可訓(xùn)練的低秩更新：

其中，A和B是低秩矩陣。這使得僅使用一小部分參數(shù)就可以實現(xiàn)特定任務(wù)的自適應(yīng)。

甚至更好：QLoRA通過將量化與低秩自適應(yīng)相結(jié)合，進(jìn)一步實現(xiàn)了這一點！

再次強(qiáng)調(diào)，這是一種非常靈活的技術(shù)，可以應(yīng)用于各個階段。通常，LoRA僅應(yīng)用于特定的層（例如，注意力層的權(quán)重）。

知識蒸餾

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

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

蒸餾損失

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

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

實踐中，訓(xùn)練學(xué)生模型非常簡單。我們將常規(guī)損失（基于硬標(biāo)簽的標(biāo)準(zhǔn)交叉熵?fù)p失）與“蒸餾”損失（基于教師的軟目標(biāo)）結(jié)合起來：

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

至于其他方法，可以并且鼓勵根據(jù)用例調(diào)整此框架：例如，還可以比較學(xué)生和教師模型之間網(wǎng)絡(luò)中的中間層的logit和激活，而不僅僅是比較最終輸出。

實踐中的知識蒸餾

與之前的技術(shù)類似，有兩種選擇：

• 離線蒸餾：預(yù)先訓(xùn)練好的教師模型是固定的，并訓(xùn)練一個單獨的學(xué)生模型來模仿它。這兩個模型完全獨立，并且在蒸餾過程中教師的權(quán)重保持不變。?
• 在線蒸餾：兩個模型同時訓(xùn)練，知識轉(zhuǎn)移發(fā)生在聯(lián)合訓(xùn)練過程中。?

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

import torch.nn.functional as F

def distillation_loss_fn(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
 # 帶有硬標(biāo)簽的標(biāo)準(zhǔn)交叉熵?fù)p失
 student_loss = F.cross_entropy(student_logits, labels)

 # 軟目標(biāo)的蒸餾損失(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需要將對數(shù)概率作為第一個參數(shù)的輸入！
 distill_loss = F.kl_div(
 soft_student_log_probs,
 soft_teacher_probs.detach(), #不為教師計算梯度
 reductinotallow='batchmean'
 ) * (temperature ** 2) # 可選，縮放因子

 #根據(jù)公式計算損失
 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()

結(jié)論

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

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

除了向你介紹主要方法之外，我希望本文還能激發(fā)你進(jìn)行實驗并開發(fā)自己的創(chuàng)造性解決方案！

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

參考文獻(xiàn)

【1】Hu, E.等人，2021?！洞笮驼Z言模型的低秩自適應(yīng)》（??Low-rank Adaptation of Large Language Models??）。arXiv preprint arXiv：2106.09685。?

【2】Lightning AI?！妒褂没旌暇燃夹g(shù)加速大型語言模型。》（??Accelerating Large Language Models with Mixed Precision Techniques??）。Lightning AI博客。?

【3】TensorFlow博客?！禩ensorFlow模型優(yōu)化工具包中的剪枝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）?！恫势奔僭O(shè)：尋找稀疏、可訓(xùn)練的神經(jīng)網(wǎng)絡(luò)》（??The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks??）。arXiv preprint arXiv：1803.03635。?

譯者介紹

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

原文標(biāo)題：??Model Compression: Make Your Machine Learning Models Lighter and Faster??，作者：Maxime Wolf?