Pruning

Pruning là một kỹ thuật tối ưu hóa model loại bỏ trọng số, neuron, filter, attention head hoặc toàn bộ layer đóng góp ít vào đầu ra của mạng nơ-ron đã huấn luyện, tạo ra model nhỏ hơn cần ít bộ nhớ, năng lực tính toán và năng lượng hơn. Trong deep learning, thuật ngữ này đồng nghĩa với pruning model. Hầu hết mạng đều dư thừa tham số đáng kể: Han và cộng sự giảm AlexNet từ 61 triệu xuống 6,7 triệu tham số và VGG-16 từ 138 triệu xuống 10,3 triệu trên ImageNet mà không mất độ chính xác, bằng cách huấn luyện, loại bỏ các kết nối có độ lớn nhỏ rồi huấn luyện lại phần còn lại. Pruning là bước phổ biến trước khi triển khai AI biên, nơi ngân sách bộ nhớ, điện năng và độ trễ suy luận rất hạn chế.

Cách Pruning hoạt động#

Mỗi phương pháp pruning quyết định cần loại bỏ tham số nào và vào thời điểm nào. Điểm độ quan trọng thường dựa trên độ lớn của trọng số, vì trọng số gần bằng 0 ít làm thay đổi đầu ra nhất, trong khi điểm dựa trên gradient ước tính trực tiếp hơn tác động của từng tham số lên hàm mất mát. Loại bỏ một tỷ lệ lớn cùng lúc thường làm giảm độ chính xác nhiều hơn pruning lặp, vốn loại bỏ một tỷ lệ nhỏ, fine-tune các trọng số còn lại để khôi phục hiệu năng rồi lặp lại cho đến khi đạt mục tiêu về kích thước hoặc tốc độ. Giả thuyết vé số bổ sung rằng mạng dày đặc chứa các mạng con thưa mà khi được khởi tạo lại bằng giá trị ban đầu, có thể tự huấn luyện để đạt độ chính xác tương đương mạng đầy đủ.

Pruning không cấu trúc, có cấu trúc và bán cấu trúc#

Pruning không cấu trúc đặt các trọng số riêng lẻ ở bất kỳ vị trí nào trong tensor về 0. Phương pháp này đạt độ thưa cao nhất ở một mức độ chính xác nhất định, nhưng tạo ra ma trận thưa không đều mà các kernel GPU và CPU dày đặc vẫn xử lý với toàn bộ chi phí, trừ khi runtime thưa bỏ qua các giá trị 0. Pruning có cấu trúc loại bỏ toàn bộ filter, channel, attention head hoặc layer khỏi mạng nơ-ron tích chập hoặc transformer, vì vậy mỗi layer sau pruning trở thành một layer dày đặc nhỏ hơn, chạy nhanh hơn trên mọi phần cứng, đổi lại mức giảm độ chính xác trên mỗi tham số bị loại bỏ lớn hơn. Tiện ích pruning của PyTorch áp dụng cả hai kiểu dưới dạng mask trên một layer hiện có:

import copy

from torch import nn
from torch.nn.utils import prune

# Two copies of one convolution layer with 32 filters of shape 16x3x3
unstructured = nn.Conv2d(16, 32, kernel_size=3)
structured = copy.deepcopy(unstructured)

# Unstructured: zero the 50% of individual weights with the smallest magnitude
prune.l1_unstructured(unstructured, name="weight", amount=0.5)

# Structured: zero the 50% of whole filters (dim=0) with the smallest L2 norm
prune.ln_structured(structured, name="weight", amount=0.5, n=2, dim=0)

for name, layer in [("Unstructured", unstructured), ("Structured", structured)]:
    zeros = (layer.weight == 0).float().mean()
    empty_filters = (layer.weight.flatten(1).abs().sum(dim=1) == 0).sum()
    print(f"{name}: {zeros:.0%} of weights zeroed, {empty_filters}/32 filters entirely empty")

Cả hai layer đều mất 50% trọng số, nhưng chỉ layer có cấu trúc mới còn 16 trong số 32 filter hoàn toàn rỗng; các filter này có thể được loại bỏ vật lý để tạo thành một phép tích chập dày đặc gồm 16 filter.

Pruning bán cấu trúc nằm giữa hai phương pháp. Mẫu độ thưa 2:4 của NVIDIA đặt hai trong mỗi bốn trọng số liên tiếp thành 0; Sparse Tensor Cores trên Ampere và các GPU đời sau bỏ qua các trọng số này để thực hiện cùng phép nhân ma trận trong một nửa thời gian. NVIDIA đo được mức tăng hiệu năng đầu-cuối lên đến 20% với ResNeXt-101 thưa trên TensorRT trên A100. Trên phần cứng xử lý dày đặc, các số 0 rải rác tự chúng không giúp tăng tốc: pruning Ultralytics YOLOv5x xuống độ thưa phi cấu trúc toàn cục 30% bằng tiện ích prune(), nhưng không fine-tune, làm mAP_50-95 trên COCO giảm từ 0.507 xuống 0.489 và thời gian inference trên V100 vẫn ở mức 5,2 ms mỗi ảnh.

Ứng dụng thực tế#

Pruning có ý nghĩa nhất khi model cần chạy trên phần cứng có tài nguyên hạn chế. Các mạng phát hiện đối tượng được pruning theo channel phù hợp với ngân sách bộ nhớ và điện năng của drone, camera thông minh và điện thoại, cho phép suy luận thời gian thực mà không cần kết nối đám mây, đồng thời giữ hình ảnh trên thiết bị để bảo vệ quyền riêng tư dữ liệu. Trong hoạt động kiểm tra sản xuất, cùng phương pháp nén này cho phép bộ phát hiện lỗi chạy trên một module biên giá rẻ bên cạnh dây chuyền sản xuất.

Mô hình ngôn ngữ lớn cũng được pruning. SparseGPT pruning OPT-175B và BLOOM-176B để đạt độ thưa không cấu trúc 60% chỉ trong một lần, không cần huấn luyện lại và chỉ làm perplexity tăng không đáng kể. Công trình Minitron của NVIDIA đã pruning Llama 3.1 8B thành model 4B bằng cách loại bỏ layer hoặc chiều ẩn, sau đó khôi phục chất lượng bằng chưng cất tri thức từ model gốc.

Pruning và các kỹ thuật liên quan#

Pruning thường được kết hợp với các phương pháp nén khác, và thuật ngữ này cũng có một nghĩa cũ hơn ngoài lĩnh vực mạng nơ-ron.

• Lượng tử hóa model: Lượng tử hóa giữ nguyên mọi trọng số và giảm độ chính xác số học, chẳng hạn từ số thực 32 bit xuống số nguyên 8 bit. Model đã pruning thường được lượng tử hóa sau đó để tiết kiệm thêm tài nguyên.
• Chưng cất tri thức: Chưng cất huấn luyện một student nhỏ gọn riêng biệt để mô phỏng teacher lớn hơn; đây là cách phổ biến để khôi phục độ chính xác mà model bị pruning nhiều đã mất.
• Layer Dropout: Dropout đặt ngẫu nhiên các activation về 0 trong một bước huấn luyện để regularization và giữ nguyên số lượng tham số, trong khi pruning loại bỏ vĩnh viễn các tham số.
• Pruning cây quyết định: Trong machine learning cổ điển, pruning cắt các nhánh khỏi cây quyết định đã khớp để giảm overfitting, như phương pháp pruning theo độ phức tạp chi phí của scikit-learn, loại bỏ các cây con yếu nhất khi tham số ccp_alpha tăng.