Model Pruning

Cắt tỉa model loại bỏ các tham số ít ảnh hưởng đến đầu ra của mạng nơ-ron đã huấn luyện, từ từng trọng số riêng lẻ đến toàn bộ filter, head attention hoặc layer, nhằm giúp model nhỏ hơn và ít tốn chi phí vận hành hơn. Hầu hết mạng đều dư thừa tham số, vì vậy có thể loại bỏ phần lớn trọng số mà ít làm giảm độ chính xác, đặc biệt nếu model được tinh chỉnh sau đó. Một phương pháp cắt tỉa rồi huấn luyện lại năm 2015 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 mà không làm giảm độ chính xác trên ImageNet. Cắt tỉa là bước tiêu chuẩn để chuẩn bị model cho điện thoại, camera và phần cứng AI biên khác; hiện nay phương pháp này cũng được áp dụng cho mô hình ngôn ngữ lớn (LLM).

Cách cắt tỉa model hoạt động#

Cắt tỉa chấm điểm mức độ quan trọng của từng trọng số hoặc cấu trúc ứng viên, loại bỏ phần có điểm thấp nhất rồi tinh chỉnh các trọng số model còn lại để khôi phục độ chính xác. Tiêu chí đơn giản nhất và vẫn phổ biến nhất là độ lớn: các trọng số gần 0 hoặc các filter có chuẩn L1 hay L2 nhỏ sẽ làm thay đổi đầu ra ít nhất khi bị loại bỏ. Những tiêu chí khác dùng gradient hoặc activation để ước tính mức tăng của loss nếu thiếu từng thành phần. Loại bỏ cùng lúc một tỷ lệ lớn thường làm giảm độ chính xác, vì vậy quá trình cắt tỉa thường diễn ra theo nhiều vòng; mỗi vòng được tiếp nối bằng tinh chỉnh cho đến khi model đạt kích thước mục tiêu hoặc độ chính xác giảm quá ngưỡng cho phép:

Cắt tỉa cũng làm thay đổi cách các nhà nghiên cứu nhìn nhận kích thước mạng. Giả thuyết vé số cho rằng một mạng dày đặc chứa các mạng con nhỏ mà khi được huấn luyện riêng, bắt đầu từ trọng số khởi tạo ban đầu, vẫn đạt độ chính xác của mạng đầy đủ; các tác giả liên tục tìm thấy những mạng con như vậy với kích thước dưới 10-20% mạng gốc trên MNIST và CIFAR-10.

Cắt tỉa không cấu trúc và có cấu trúc#

Cắt tỉa không cấu trúc đặt các trọng số riêng lẻ ở bất kỳ vị trí nào trong tensor thành 0. Phương pháp này đạt độ thưa cao với mức giảm độ chính xác nhỏ, nhưng hình dạng tensor không đổi nên các kernel dense tiêu chuẩn vẫn thực hiện cùng lượng công việc; muốn tăng tốc phải có kernel thưa hoặc hỗ trợ phần cứng. GPU NVIDIA từ thế hệ Ampere có thể chạy một dạng như vậy — độ thưa 2:4, với hai số 0 trong mỗi bốn trọng số liên tiếp — nhanh gấp tối đa hai lần so với phép tính dense. Cắt tỉa có cấu trúc loại bỏ toàn bộ filter, channel, head hoặc layer, tạo ra model dense nhỏ hơn và chạy nhanh hơn trên mọi bộ xử lý, nhưng làm giảm độ chính xác nhiều hơn với cùng số lượng trọng số bị loại bỏ. Khi bỏ một filter, cũng cần bỏ channel đầu vào tương ứng của layer tiếp theo; các thư viện như Torch-Pruning tự động theo dõi những phụ thuộc này. Các tiện ích cắt tỉa của PyTorch minh họa sự khác nhau giữa hai phương pháp:

import torch
from torch.nn.utils import prune

# Two identical convolution layers, one for each pruning style
torch.manual_seed(0)
a = torch.nn.Conv2d(16, 32, kernel_size=3)
b = torch.nn.Conv2d(16, 32, kernel_size=3)

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

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

for name, layer in [("unstructured", a), ("structured", b)]:
    zeros = (layer.weight == 0).float().mean().item()
    empty = (layer.weight.flatten(1) == 0).all(dim=1).sum().item()
    print(f"{name}: {zeros:.0%} of weights zero, {empty} of 32 filters empty")

Cả hai layer đều có 50% trọng số bằng 0. Chỉ layer được cắt tỉa có cấu trúc mới có các filter rỗng — 16 trong số 32 filter — và xóa chúng sẽ tạo ra một layer dense có số filter giảm một nửa, nhờ đó giảm một nửa lượng tính toán trên phần cứng thông thường. Phần cứng được thiết kế cho tính thưa cũng hưởng lợi từ phương pháp không cấu trúc: tích hợp CVflow cho Ultralytics YOLO của Ambarella, do Ambarella duy trì trong fork riêng của package, cắt tỉa trọng số dần trong quá trình huấn luyện YOLO và duy trì độ thưa đó qua bước xuất ONNX để biên dịch cho chip AI biên của hãng.

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

• Lượng tử hóa model: Giảm độ chính xác biểu diễn của mọi trọng số, chẳng hạn từ số thực 32-bit xuống số nguyên 8-bit, nhưng vẫn giữ lại toàn bộ trọng số. Cắt tỉa và lượng tử hóa thường được kết hợp: cắt tỉa giảm số phép toán, còn lượng tử hóa giúp mỗi phép toán ít tốn kém hơn.
• Chưng cất tri thức: Huấn luyện model student nhỏ hơn bắt chước model teacher lớn hơn; phương pháp này thường được áp dụng sau cắt tỉa có cấu trúc để khôi phục độ chính xác bị mất khi loại bỏ toàn bộ filter hoặc layer.
• Dropout: Tạm thời đặt ngẫu nhiên một số activation thành 0 trong quá trình huấn luyện như một hình thức regularization; khi suy luận, model đầy đủ được sử dụng.

Cắt tỉa mô hình ngôn ngữ lớn#

Huấn luyện lại LLM sau mỗi vòng cắt tỉa quá tốn kém, vì vậy các phương pháp một lần cắt tỉa model trong một lượt bằng tập hiệu chỉnh nhỏ. SparseGPT cắt tỉa các model OPT 175 tỷ tham số và BLOOM 176 tỷ tham số xuống độ thưa không cấu trúc 60% trong chưa đầy 4,5 giờ, với mức tăng perplexity không đáng kể; phương pháp này cũng hỗ trợ dạng 2:4 được GPU tăng tốc. Cắt tỉa có cấu trúc kết hợp với chưng cất tri thức tạo ra các model dense nhỏ hơn: Minitron của NVIDIA tạo ra các model 8B và 4B từ model pretrained 15B bằng cách cắt tỉa channel embedding, neuron MLP và head attention, sau đó huấn luyện lại với số lượng token ít hơn tối đa 40 lần so với huấn luyện từ đầu.