稀疏性概述¶
稀疏性是一種從神經網路中移除引數的技術，旨在降低其記憶體開銷或延遲。透過仔細選擇元素的修剪方式，可以在不顯著犧牲模型質量（準確率/f1 分數）的情況下，大幅減少記憶體開銷和延遲。
目標¶
我們認為，當前稀疏性研究人員/使用者面臨的主要問題是碎片化。研究人員理應致力於展示端到端的結果，但這通常意味著需要花費大量時間來研究如何與 PyTorch 整合以及解決諸如以下實現問題：
應該何時應用掩碼？
我應該何時/如何儲存壓縮表示？
我想要就地更新還是非就地更新掩碼？
如何呼叫稀疏矩陣乘法而不是密集矩陣乘法？
我們認為，上述問題可以由 torchao 一次性解決，讓研究人員專注於真正重要的事情——提升稀疏核效能或改進修剪演算法的準確性。
更具體地說，我們希望為稀疏核（張量子類化）和修剪演算法（torch.ao.pruning.Sparsifier）提供教程和 API，供使用者擴充套件。我們旨在提供模組化的構建塊，這些構建塊不僅可以用於加速推理，還可以用於加速訓練，並且可以與 torchao 的量化工作流程良好地結合。
從頭開始訓練帶有硬體加速的稀疏模型，並實現最小的準確率損失。
使用自定義修剪演算法恢復修剪模型的準確率損失。
在支援稀疏性的硬體上加速掩碼/修剪模型，以實現效能提升。
設計¶
稀疏性，就像量化一樣，是一種準確率/效能的權衡，我們不僅關心加速效果，也關心架構最佳化技術帶來的準確率下降。
在量化中，理論效能增益通常取決於我們量化到的資料型別——從 float32 量化到 float16 可帶來理論上的 2 倍加速。對於修剪/稀疏性，類似的變數是稀疏度水平/稀疏模式。對於半結構化稀疏性，稀疏度水平固定為 50%，因此我們期望理論上實現 2 倍的提升。對於塊稀疏矩陣和非結構化稀疏性，加速效果是可變的，取決於張量的稀疏度水平。
稀疏性和量化之間的一個關鍵區別在於準確率下降的決定因素：通常，量化的準確率下降由選擇的 scale 和 zero_point 決定。然而，在修剪中，準確率下降由掩碼決定。稀疏性和量化密切相關，並共享諸如量化/稀疏感知訓練之類的準確率緩解技術。
透過仔細選擇指定的元素並重新訓練網路，修剪可以實現微不足道的準確率下降，在某些情況下甚至能略微提升準確率。這是一個活躍的研究領域，尚未達成一致共識。我們期望使用者心中有一個目標稀疏模式，並按照該模式進行修剪。
給定一個目標稀疏模式，修剪模型可以被視為兩個獨立的子問題：
準確率 - 我如何找到一組滿足目標稀疏模式的稀疏權重，從而最大程度地減少模型的準確率下降？
效能 - 我如何加速稀疏權重的推理並減少記憶體開銷？
我們的工作流程設計包含兩個獨立回答這些問題的部分：
一個面向使用者的 Python 前端 API，用於查詢任意稀疏模式的稀疏權重。
一個後端稀疏核/操作集合，用於減少記憶體/延遲。
這兩部分之間的交接點是以密集格式儲存的稀疏權重，其中缺失元素的位置填充 0。這是一個自然的交接點，因為使用這個張量進行稀疏矩陣乘法和密集矩陣乘法在數值上是等效的。這使得我們可以為使用者提供後端清晰的契約，對於給定的稀疏模式：
如果您可以將密集矩陣轉換為 2:4 稀疏格式，我們可以在沒有數值損失的情況下將矩陣乘法加速高達 1.7 倍。
這也允許現有以密集格式儲存稀疏權重的使用者利用我們快速的稀疏核。我們預計許多使用者會提出自己的自定義前端掩碼解決方案或使用其他第三方解決方案，因為這是一個活躍的研究領域。
下面，我們提供一個使用我們的 PyTorch API 加速帶有 2:4 稀疏性 + bf16 的模型的示例。
import torch
from torch.sparse import to_sparse_semi_structured, SparseSemiStructuredTensor
from torch.ao.pruning import WeightNormSparsifier
# bfloat16 CUDA model
model = model.half().cuda()
# Accuracy: Finding a sparse subnetwork
sparse_config = []
for name, mod in model.named_modules():
if isinstance(mod, torch.nn.Linear):
sparse_config.append({"tensor_fqn": f"{name}.weight"})
sparsifier = WeightNormSparsifier(sparsity_level=1.0,
sparse_block_shape=(1,4),
zeros_per_block=2)
# attach FakeSparsity
sparsifier.prepare(model, sparse_config)
sparsifier.step()
sparsifier.squash_mask()
# now we have dense model with sparse weights
# Performance: Accelerated sparse inference
for name, mod in model.named_modules():
if isinstance(mod, torch.nn.Linear):
mod.weight = torch.nn.Parameter(to_sparse_semi_structured(mod.weight))
從根本上講，流程是透過操作 torch.Tensors 來工作的。在前端，我們在 sparse_config 字典中透過它們的完全限定名指定張量。前端設計遵循量化 API，帶有 prepare 函式，該函式將 FakeSparsity 引數化附加到配置中指定的張量上。
FakeSparsity 是一種引數化，它模擬非結構化稀疏性，其中每個元素都有一個掩碼。因此，我們可以用它來模擬我們想要的任何稀疏模式。
然後，使用者將使用自己的自定義程式碼訓練準備好的模型，並在必要時呼叫 .step() 更新掩碼。一旦找到合適的掩碼，他們會呼叫 squash_mask() 將掩碼融合到權重中，從而在正確位置建立帶有 0 的密集張量。
然後，使用者可以透過使用量化流程進行量化塊稀疏 CPU 推理，或者透過在指定的權重張量上呼叫 to_sparse_semi_structured 來轉換模型以進行加速稀疏推理。
背景¶
本節提供了一些關於神經網路修剪/稀疏性的背景資訊，並定義了一些常見的修剪/稀疏性術語。在學術界/工業界，修剪（pruning） 和 稀疏性（sparsity） 經常互換使用，指代同一件事。這可能會引起混淆，特別是因為稀疏性是一個具有多重含義的術語，可以指代許多其他事物，例如稀疏張量表示。
請注意，本節側重於修剪（pruning），而非稀疏訓練（sparse training）。區別在於，在修剪中，我們從預訓練的密集模型開始；而在稀疏訓練中，我們從頭開始訓練一個稀疏模型。
為了避免混淆，我們通常儘量使用稀疏性來指代張量。請注意，稀疏張量可以指代包含許多零值的密集張量，或使用稀疏表示儲存的張量。我們將整個流程描述為修剪，將由此產生的模型描述為修剪模型（pruned model）。
粗略地說，實現效能更好的修剪模型的流程如下：
修剪背後的總體思想是，我們可以對訓練好的神經網路的某些權重進行掩碼處理，並恢復由此帶來的準確率損失。由此產生的修剪模型可以在利用這種稀疏性進行加速推理的最佳化核上執行。
直接將修剪後的引數置零並不會影響模型的延遲/記憶體開銷。這是因為密集張量本身仍然包含被修剪的元素（那些 0 元素），並在矩陣乘法期間仍會使用這些元素進行計算。為了實現效能提升，我們需要將密集核替換為稀疏核。
粗略地說，這些稀疏表示允許我們跳過涉及修剪元素的計算，以加快矩陣乘法。為此，這些最佳化的稀疏核處理的是以更高效格式儲存的稀疏矩陣。某些稀疏張量佈局與特定後端緊密耦合，例如 NVIDIA 2:4，而另一些則更通用，受多個後端支援（CSC 受 FBGEMM 和 QNNPACK 支援）。
| 名稱 | 描述 | 稀疏矩陣的儲存方式 |
| COO (sparse_coo) | 用於儲存稀疏矩陣的 COOrdinate 格式。矩陣儲存為非稀疏資料向量和這些元素在密集矩陣中的索引位置的組合。 | 稀疏矩陣 = {Index: 座標位置張量, Data: 與索引位置對應的值張量 } |
| BSR (sparse_bsr) | 用於儲存稀疏矩陣的塊稀疏行格式。矩陣儲存為資料塊和這些塊在密集矩陣中的索引位置。與 COO 非常相似，區別在於單個數據由塊組成，而非標量。 | 稀疏矩陣 = {Index: 座標位置張量（對於矩陣是二維的）, Data: 與索引位置對應的塊張量 } 其中塊是與稀疏模式對應的矩陣。 |
| CSR (sparse_csr) / CSC (sparse_csc) | 用於儲存稀疏矩陣的壓縮稀疏行/列格式。稀疏矩陣儲存為列/行上的資料塊以及這些行/列在密集矩陣中的索引。這是儲存塊稀疏矩陣最緊湊的格式。 | 稀疏矩陣 = {Index: 列索引的 1D 張量, IndexPtr: 指定行（從第 0 行開始）的列開始和結束索引的 1D 張量, Data: 與 Index 位置對應的塊張量。} |
| NVIDIA 2:4 壓縮表示 | 適用於 2:4 半結構化稀疏性的自定義 NVIDIA 壓縮儲存格式。我們將稀疏矩陣儲存為一個壓縮的密集矩陣（½ 原大小），其中包含未修剪的元素和一個位掩碼索引。當我們將稀疏矩陣與另一個密集矩陣相乘時，我們使用掩碼來索引密集矩陣並與我們的壓縮密集矩陣相乘。 | 稀疏矩陣 = {Bitmask: 修剪元素的 2bit 索引 壓縮密集矩陣: 包含所有未修剪的元素，大小是原始密集矩陣的一半} |
表 4.1: 常見稀疏張量佈局概述。
雖然修剪的總體思想相當簡單，但在成功修剪模型之前，使用者必須弄清楚許多細節。
這些可以大致分解如下：
修剪配置 - 我應該修剪哪些層？應該修剪到什麼稀疏度水平？
修剪標準 - 我應該如何決定移除哪些引數？
修剪策略 - 移除引數後，我如何恢復任何準確率下降？
稀疏模式 - 修剪模型時，我應該嘗試使用特定的稀疏模式嗎？不同的硬體後端支援不同稀疏模式的加速推理。
修剪配置¶
並非神經網路中的所有層都一樣。有些層可能比其他層對修剪更敏感。使用者必須決定修剪哪些層以及每層的稀疏度水平，即該權重張量中 0 的百分比。修剪配置會影響修剪模型的準確率和加速效果。
確定給定模型的最佳修剪配置和稀疏度水平是一個開放問題，尚無通用解決方案。這部分是因為最優修剪配置取決於後續的修剪標準和策略，並且決定如何修剪模型以及如何恢復損失的準確率有無數種方法。
一種確定修剪哪些層以及修剪程度的常用方法是進行敏感性分析，即在不同稀疏度水平下修剪模型中的每一層，並觀察隨後的準確率下降（不進行重新訓練）。這為使用者提供了每層的稀疏度-準確率曲線，使用者可以將其用作確定最佳修剪配置的代理。
修剪標準¶
使用者必須決定從神經網路中移除引數的標準。就像確定最佳修剪配置一樣，確定最佳修剪標準也是一個開放的研究問題，並且取決於上述其他因素。
最常見的修剪標準是使用權重幅值。其思想是，低幅值的權重對模型輸出的貢獻小於高幅值的權重。如果我們要移除引數，我們可以移除絕對值最小的權重。
然而，即使使用權重幅值這樣簡單的修剪標準，使用者也必須考慮其他因素：
區域性範圍 vs 全域性範圍
區域性範圍意味著稀疏性掩碼僅根據層的統計資訊計算。
優點：掩碼計算簡單
缺點：準確率與稀疏度的權衡可能不是最優。
全域性範圍意味著稀疏性統計資訊不受單個層的限制，如果需要，可以跨越多個層。
優點：無需逐層閾值。張量統計資訊在層之間共享，並使用跨層歸一化來實現。
缺點：計算掩碼時複雜度增加。
用於掩碼計算的張量
權重：僅使用權重張量來計算掩碼。對於推理而言，這種方法最簡單，因為權重張量是恆定的。
梯度：根據權重和梯度範數計算重要性。常見於基於預訓練的方法。目前 CTR_mobile_feed 使用基於梯度的修剪演算法。
啟用值：在一些研究論文中，與相關權重一起應用的啟用值的範數被用來計算重要性得分。
就地或非就地掩碼更新
就地更新透過執行 W = W (Mask) 來更新稀疏張量。權重張量更新後，稀疏值被置零，無法恢復。
優點：只需儲存一份稀疏張量（+ 掩碼）
缺點：一旦掩碼應用於權重，該權重即被置零，所有歷史資訊丟失。這些權重無法“重新生長”（regrow）。
非就地更新不直接修改張量，而是執行以下操作：W’ = W (Mask) 和 dW’ = dW (Mask)
優點：原始張量得以保留（被掩碼的元素不會透過反向傳播更新）。如果掩碼改變，權重可以“重新生長”。這對 PAT 是必需的。
缺點：除了未被掩碼的權重 (W)，還需要計算被掩碼的權重 (W’)，並在前向/後向計算期間駐留在記憶體中。
| 名稱 | 描述 | 注意事項 |
| 幅值 / 顯著性 | 移除範數最低的引數（常用 L1 範數） | 已證明與 2:4 半結構化稀疏性配合良好。透過在一次性幅值修剪後重複訓練迴圈，能夠達到與原始模型相同的準確率。 |
| Movement Pruning | 這些方法旨在利用梯度資訊來決定移除哪些引數。其思想是移除在微調過程中變化不大的引數。 | 常用於預訓練模型。 |
| 低秩分解 | 這些方法旨在將 Wx 替換為 SQx，其中 S 和 Q 是低秩矩陣。 | 通常這些方法使用某種層級重建，其中不是透過訓練模型來恢復損失的準確率，而是尋求匹配層級的統計資訊（找到使得 L2(SQx, Wx) 最小化的 SQx）。 |
| 隨機 | 隨機移除引數 |
表 4.2: 一些常見修剪標準描述。
修剪策略¶
這是一個通用術語，描述了使用者試圖從修剪模型中恢復任何準確率下降的方法。修剪模型後，通常會看到模型準確率下降，因此使用者通常會重新訓練修剪後的模型以彌補這一點。修剪策略還決定了在模型訓練期間何時以及多久進行一次修剪。
修剪策略和修剪標準之間的界限並不明確，尤其是在修剪感知訓練方法中，它們在訓練期間更新掩碼。我們有時使用術語修剪演算法（pruning algorithm）來指代這兩者的組合。這兩個因素，連同修剪配置，最終決定了修剪模型的最終準確率。
| 修剪策略 | 描述 | 注意事項 |
| 零樣本 (Zero-shot) | 修剪一次，不重新訓練模型 | 這些方法依賴於更復雜的修剪標準。 在文獻中有時也稱為“一次性（one-shot）”，但我們將“一次性（one-shot）”用於指修剪一次並重新訓練一次的情況。 |
| 一次性 (One-shot) | 修剪一次，重新訓練模型一次 | NVIDIA 已表明，一次性 2:4 半結構化稀疏修剪在各種常見視覺/NLP 模型上具有良好的泛化能力。 \ \ 重新訓練策略就是簡單地再次重複訓練過程。 |
| 迭代式 (Iterative) | 修剪模型，重新訓練，重複 | 我們可以迭代地增加稀疏度水平，或者迭代地修剪模型中的不同層。 |
| 修剪感知訓練 (Pruning Aware Training) | 掩碼在訓練過程中學習 | CTR_feed 用於其當前修剪演算法。 |
| NAS / 多掩碼 (Multimask) | 訓練期間使用多個掩碼。這可以被認為是神經架構搜尋的一種形式。 | PySpeech 使用 (FastNAS) |
| 層級重建 (Layer-wise reconstruction) | 我們不是使用損失函式進行重新訓練，而是透過使用類似於知識蒸餾的雙模型方法，嘗試從每一層恢復儘可能多的資訊。 | 參見 https://arxiv.org/pdf/2204.09656.pdf |
表 4.3: 一些常見修剪策略描述。
稀疏模式¶
稀疏模式描述了修剪後的引數在模型/張量中的排列方式。
回想一下，通常需要使用最佳化的稀疏核才能實現效能提升。根據權重張量的格式和稀疏度水平，稀疏矩陣乘法可能比其對應的密集乘法更快。如果張量不夠稀疏，也可能更慢。
在最一般的層面上，修剪是非結構化的——每個引數都有自己的掩碼。這提供了最大的靈活性，但需要非常高的稀疏度（>98%）才能提供效能優勢。為了在較低稀疏度水平下提供加速推理，硬體後端增加了對特殊稀疏模式的支援。
我們尋求修剪模型，使得權重張量呈現與我們的推理後端相同的稀疏模式。如果在保持稀疏模式的同時能夠恢復損失的準確率，我們就可以在稀疏硬體上執行此模型進行加速推理，而不會犧牲準確率。我們也可以在目標後端上執行修剪到不同稀疏模式的模型，但這會帶來一些額外的準確率損失。
特定的後端硬體及其對應的稀疏模式，以及修剪配置，最終決定了我們觀察到的效能加速效果。如果我們使用不同的修剪標準修剪模型，只要它遵循相同的稀疏模式和稀疏度水平，其效能特徵將是相同的。例如，如果我們決定移除最高幅值的權重而不是最低幅值的權重，我們預計這不會改變修剪模型的效能特徵。
| 稀疏模式 | 掩碼視覺化
（50% 稀疏度水平） |
||||||||||||||||||||||||||||||||
| 非結構化稀疏性 |
|
||||||||||||||||||||||||||||||||
| 2:4 半結構化 |
|
||||||||||||||||||||||||||||||||
| 塊稀疏性 |
|
||||||||||||||||||||||||||||||||
| 結構化稀疏性 |
|
表 4.4: 一些常見稀疏模式描述。
有關我們支援的 API 和基準測試的更多資訊，請參閱稀疏性 README。
