Model Pruning

모델 프루닝은 학습된 신경망의 출력에 거의 기여하지 않는 파라미터를 개별 가중치부터 필터 전체, 어텐션 헤드 또는 레이어까지 제거하여 모델을 더 작게 만들고 실행 비용을 낮춥니다. 대부분의 네트워크는 파라미터가 과도하게 많으므로 가중치의 상당 부분을 제거해도 정확도 손실은 거의 없으며, 이후 모델을 미세 조정하면 특히 그렇습니다. 2015년 프루닝 및 재학습 방법은 ImageNet 정확도 손실 없이 AlexNet의 파라미터 수를 6,100만 개에서 670만 개로, VGG-16의 파라미터 수를 1억 3,800만 개에서 1,030만 개로 줄였습니다. 프루닝은 휴대폰, 카메라 및 기타 엣지 AI 하드웨어용 모델을 준비할 때 표준적으로 거치는 단계이며, 이제는 대규모 언어 모델(LLMs)에도 적용됩니다.

모델 가지치기 작동 방식#

가지치기는 후보 가중치나 구조의 중요도를 점수화하고, 점수가 가장 낮은 비율을 제거한 뒤 남은 모델 가중치를 미세 조정해 정확도를 회복합니다. 가장 간단하고 여전히 가장 일반적인 점수 기준은 크기입니다. 0에 가까운 가중치나 L1 또는 L2 노름이 작은 필터는 제거해도 출력 변화가 가장 적습니다. 다른 기준은 그래디언트나 활성화를 사용해 각 부분을 제거할 때 손실이 얼마나 커질지 추정합니다. 한 번에 많은 부분을 제거하면 정확도가 떨어지는 경향이 있으므로, 가지치기는 보통 모델이 목표 크기에 도달하거나 정확도가 허용 범위 아래로 떨어질 때까지 라운드별로 수행하고 각 라운드 이후 미세 조정합니다:

가지치기는 네트워크 크기에 대한 연구자들의 관점도 바꾸었습니다. 복권 티켓 가설은 밀집 네트워크에 포함된 소형 서브네트워크를 원래 초기 가중치에서 독립적으로 학습하면 전체 네트워크와 같은 정확도를 달성한다고 제안합니다. 가설 제안자들은 MNIST와 CIFAR-10에서 원래 크기의 10~20% 미만인 서브네트워크를 일관되게 발견했습니다.

비구조적 가지치기와 구조적 가지치기#

비구조적 가지치기는 텐서의 어느 위치에 있는 개별 가중치든 0으로 설정합니다. 정확도 손실을 거의 일으키지 않으면서 높은 희소성에 도달할 수 있지만 텐서의 형태는 유지되므로, 표준 밀집 커널은 동일한 연산을 수행합니다. 따라서 속도를 높이려면 희소 커널이나 하드웨어 지원이 필요합니다. Ampere 세대 이후의 NVIDIA GPU는 연속된 가중치 네 개 중 두 개를 0으로 만드는 2:4 희소성 패턴을 실행할 수 있으며, 밀집 연산보다 최대 두 배의 속도를 냅니다. 구조적 가지치기는 필터, 채널, 헤드 또는 레이어 전체를 제거해 어떤 프로세서에서도 더 빠르게 실행되는 소형 밀집 모델을 만듭니다. 다만 같은 수의 가중치를 제거할 때 정확도 비용이 더 큽니다. 필터를 제거하면 다음 레이어에서 이에 대응하는 입력 채널도 제거됩니다. Torch-Pruning 같은 라이브러리는 이러한 종속성을 자동으로 추적합니다. PyTorch의 가지치기 유틸리티는 두 방식의 차이를 보여줍니다:

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")

두 레이어 모두 가중치의 50%가 0이 됩니다. 구조적 레이어에만 비어 있는 필터 16개가 있으며, 필터 32개 중 이를 삭제하면 절반의 필터로 구성된 밀집 레이어가 되어 일반 하드웨어에서 연산량이 절반으로 감소합니다. 희소성에 최적화된 하드웨어는 비구조적 가지치기의 이점도 활용할 수 있습니다. Ambarella의 CVflow 통합은 Ultralytics YOLO용으로 Ambarella가 자체 패키지 포크에서 유지 관리합니다. 이 통합은 YOLO 학습 중 가중치를 점진적으로 가지치기하고 ONNX 내보내기 이후에도 희소성을 유지해 엣지 AI 칩용으로 컴파일합니다.

가지치기와 관련 기법 비교#

• 모델 양자화: 예를 들어 모든 가중치의 정밀도를 32비트 부동소수점에서 8비트 정수로 낮추되 가중치 자체는 모두 유지합니다. 가지치기는 연산 수를 줄이고 양자화는 연산 하나의 비용을 낮추므로 두 방식을 함께 사용하는 경우가 많습니다.
• 지식 증류: 작은 학생 모델이 큰 교사 모델을 모방하도록 학습합니다. 필터나 레이어 전체를 제거하면서 발생한 정확도 손실을 회복하기 위해 구조적 가지치기 후에 적용하는 경우가 많습니다.
• 드롭아웃: 정규화 방법의 하나로 학습 중 무작위 활성화를 일시적으로 0으로 만들며, 추론 시에는 전체 네트워크를 사용합니다.

대규모 언어 모델 가지치기#

가지치기 라운드마다 LLM을 재학습하는 데는 비용이 너무 많이 들므로, 원샷 방식은 소규모 보정 데이터셋을 사용해 한 번에 가지치기를 수행합니다. SparseGPT는 1,750억 파라미터 OPT 모델과 1,760억 파라미터 BLOOM 모델을 4.5시간 이내에 비구조적 희소성 60%로 가지치기했으며, 퍼플렉서티는 미미하게 증가했습니다. 또한 GPU가 가속하는 2:4 패턴도 지원합니다. 지식 증류와 결합한 구조적 가지치기는 더 작은 밀집 모델을 만듭니다. NVIDIA의 Minitron은 사전 학습된 150억 파라미터 모델에서 임베딩 채널, MLP 뉴런, 어텐션 헤드를 가지치기해 80억 및 40억 파라미터 모델을 만들었으며, 처음부터 학습하는 것보다 최대 40배 적은 토큰으로 재학습했습니다.