이 글은 경량화 시리즈의 세 번째(#3)입니다. 데이터타입 → 양자화가 “숫자를 더 적은 비트로” 였다면, 이번엔 아예 연결을 잘라내는 프루닝입니다. 코드는 github.com/warpspaceinc/efficient-ml-practice의 pruning.ipynb 노트북에 있고, 아래 모든 plot은 그 코드를 Colab(A100)에서 직접 돌려 측정한 우리 숫자입니다.

1989년의 “뇌 손상”

프루닝을 이야기할 때 빼놓을 수 없는 논문이 하나 있습니다. 제목이 살벌합니다 — Optimal Brain Damage(최적 뇌 손상). 저자는 Yann LeCun, John Denker, Sara Solla, AT&T Bell Labs, 발표는 NeurIPS 1989입니다.

잠깐 멈춰서 생각해보면 놀랍습니다. 1989년. 지금으로부터 36년 전입니다. 월드와이드웹이 세상에 나오기도 전이고, LeCun이 CNN(LeNet)으로 유명해지기도 전이며, “딥러닝"이라는 말조차 없던 시절입니다. 그 LeCun은 지금 Meta의 수석 AI 과학자이자 튜링상 수상자입니다. “신경망에서 덜 중요한 연결을 잘라내면 더 작고, 더 잘 일반화되는 망을 얻는다” 는 프루닝의 핵심 아이디어가, 그렇게 오래전에 이미 정식화되어 있었습니다.

논문의 발상은 생물학적 은유와도 맞닿아 있습니다. 사람의 뇌도 시냅스를 무작정 늘리지 않습니다. 뉴런당 시냅스 수는 유아기에 약 2,500개 → 2~4세 무렵 15,000개까지 폭증했다가 → 성인이 되면 약 7,000개로 다시 줄어듭니다. 쓸모없는 연결을 쳐내는 이 “가지치기(pruning)“가 오히려 효율적인 뇌를 만든다는 것이죠.

이 글은 그 36년의 통찰이 오늘날 NVIDIA GPU의 2:4 sparsity로 어떻게 현실이 되는지를, 세 가지 질문으로 따라갑니다.

  1. 무엇을 자를 것인가 (criterion)
  2. 얼마나 자를 수 있는가 (재학습과 압축의 놀라움)
  3. 어떻게 잘라야 하드웨어가 빨라지는가 (granularity와 2:4)

프루닝이란 & 5단계 프레임워크

프루닝은 가중치를 0으로 만들어 연결을 삭제하는 것입니다. dense한 망을 sparse하게 만들어 저장·연산·전력을 아낍니다. 형식적으로는 “0이 아닌 가중치 수를 $N$개 이하로 제한하면서 손실을 최소화"하는 문제입니다.

$$\arg\min_{W_P} L(x; W_P) \quad \text{s.t.} \quad \lVert W_P\rVert_0 < N$$

Han et al.(NeurIPS 2015) 이후 프루닝은 보통 5단계로 정리됩니다: ① 학습 → ② granularity(어떻게) → ③ criterion(무엇을) → ④ ratio(얼마나) → ⑤ fine-tune(회복). 이 글은 그중 핵심인 criterion → ratio/회복 → granularity 순으로 봅니다.

동기는 양자화 편과 같습니다. 가중치는 대개 0 근처에 몰려 있고, 메모리 접근은 연산보다 수백 배 비쌉니다(Horowitz: DRAM 640pJ vs int ADD 0.1pJ). 그러니 “많은 가중치가 거의 0인데, 정말 다 필요할까?“라는 질문이 자연스럽습니다.

1. 무엇을 자를까 — Criterion

가장 먼저 답할 질문은 “어떤 가중치를 지울 것인가"입니다. 후보는 세 가지입니다.

  • Magnitude(크기): 중요도 $= |w|$. 작으면 지운다. 싸고 단순.
  • OBD saliency(2차/헤시안): OBD의 답입니다. 각 가중치를 지웠을 때 손실이 얼마나 늘어나는지를 테일러 전개로 근사합니다.
  • Random: 무작위. 비교용 하한선.

OBD의 유도는 이렇습니다. 손실 $L$을 가중치 변화 $\delta w$에 대해 2차까지 전개하면:

$$\delta L = \sum_i g_i\,\delta w_i + \frac{1}{2}\sum_i h_{ii}\,\delta w_i^2 + \frac{1}{2}\sum_{i\neq j} h_{ij}\,\delta w_i \delta w_j + O(\lVert\delta w\rVert^3)$$

여기서 세 가지 가정을 넣습니다 — 수렴 완료(1차 항 $g_i \approx 0$), 대각 근사(교차항 $h_{ij}$ 무시), 2차 근사(3차 이상 무시). 그러면 가중치 하나를 0으로 만들 때($\delta w_i = -w_i$)의 손실 증가, 즉 saliency가 깔끔하게 남습니다.

$$\text{saliency}_{w_i} = \frac{1}{2}\, h_{ii}\, w_i^2$$

크기가 큰 가중치($w_i^2$)일수록, 그리고 손실 곡률이 가파른 곳($h_{ii}$)일수록 중요하다는 뜻입니다. 문제는 헤시안 $h_{ii}$가 비싸다는 것 — 그래서 실무에서는 대각 헤시안을 empirical Fisher(기울기 제곱의 평균) 로 근사합니다. 우리 실험도 그렇게 했습니다.

직접 측정: 세 기준 비교

MNIST MLP(784→512→512→10)를 학습한 뒤, 세 기준으로 각각 sparsity를 올려가며 (재학습 없이) 정확도를 쟀습니다.

One-shot 프루닝: OBD·magnitude·random 정확도 vs sparsity

읽어야 할 핵심:

  • random은 50%만 넘어도 무너집니다. 무엇을 자르는지가 결정적이라는 증거입니다.
  • magnitude와 OBD는 거의 겹칩니다. 둘 다 90%까지 정확도를 지킵니다. 극단(98%)에서는 OBD의 2차 통찰이 조금 앞서지만(44.6% vs 29.3%), 그 앞까지는 magnitude가 오히려 근소하게 낫습니다.

아래에서 직접 만져보세요. 슬라이더로 sparsity를 올리면 잘린 시냅스가 오른쪽 신경망에서 사라지고, 연결이 하나도 안 남은 뉴런은 제거됩니다. 동시에 연산량(MAC·FLOPs)이 줄어드는 것도 보입니다 — 단, 이 절감은 하드웨어가 0을 실제로 건너뛸 때만 실현됩니다. 기준을 magnitude ↔ random으로 바꾸면 같은 sparsity(같은 MAC)라도 어떤 연결·뉴런이 먼저 사라지는지가 달라집니다.

이 결과가 바로 “왜 실무는 OBD 대신 magnitude를 쓰나” 에 대한 답입니다. OBD는 이론적으로 우아하지만 헤시안이 비쌉니다. 반면 magnitude는 공짜인데 이만큼 강력합니다. 그래서 36년이 지난 지금도 대부분의 프루닝은 magnitude 기반입니다. OBD가 연 문(“무엇을 자를지 원리적으로 고르자”)을, 정작 가장 단순한 기준이 실용적으로 통과한 셈입니다.

2. 얼마나 자를 수 있나 — 재학습과 압축의 놀라움

한 번에 많이 자르면 당연히 정확도가 떨어집니다. 하지만 자르고 → 재학습(fine-tune) 을 반복하면 이야기가 완전히 달라집니다. 재학습 학습률은 보통 원래의 1/10~1/100로 낮게 씁니다.

프루닝+재학습: 놀라울 만큼 높은 sparsity까지 정확도 유지

  • prune only(회색): 재학습 없이 자르면 95%에서 88%로 떨어집니다.
  • iterative prune+retrain(초록): 조금씩 자르며 매번 재학습하면, 가중치를 95% 잘라내도(=20배 적은 파라미터) dense와 사실상 같은 정확도(97.3% vs 97.6%) 를 유지합니다.

가중치의 20분의 1만 남겨도 성능이 그대로라는 게, 프루닝이 지금도 연구되는 이유입니다. 프루닝 후 가중치 분포를 보면 무슨 일이 일어났는지 한눈에 보입니다 — 정확히 0에 거대한 스파이크가 생깁니다.

90% 프루닝 후 가중치 분포: 0에 스파이크

이 “잘라내고 다시 학습” 아이디어를 극단으로 민 것이 Deep Compression(Han et al., 2016)입니다. iterative pruning으로 AlexNet의 프루닝 한계를 5배에서 9배까지 끌어올렸고, 양자화·허프만 코딩까지 얹어 최대 수십 배를 달성했습니다.

3. 어떻게 잘라야 빨라지나 — Granularity와 하드웨어

여기서 반드시 짚어야 할 함정이 있습니다. 프루닝으로 sparsity를 얻었다고 저절로 빨라지지 않습니다.

곰곰이 생각해보면, 가중치를 “자른다"는 건 사실 그 자리에 0을 넣는 것과 다를 게 없습니다. 그리고 0을 그냥 곱셈-누산기에 넣어 계산하면 — $0 \times a = 0$ — 연산 횟수는 하나도 줄지 않습니다. 위에서 정확도를 잰 실험도 실은 가중치를 0으로 만든 뒤 그대로 dense 커널로 돌린 것이라, 정확도는 측정됐지만 저장도 속도도 전혀 줄지 않았습니다(양자화 편의 “fake quantization"과 똑같은 상황입니다). 0을 곱하나 0.3을 곱하나 곱셈기 입장에선 같은 한 번의 곱셈입니다.

이득을 보려면 0을 아예 저장하지도, 곱하지도 않고 건너뛰어야 합니다. 그런데 그게 어렵습니다. 지금까지의 프루닝은 0이 아무 데나 흩어지는 unstructured(비정형) 방식이라, 압축률은 최고지만 GPU 입장에서는 0의 위치가 불규칙해서 건너뛸 수가 없습니다. “다음 0이 어디 있나"를 매번 찾는 비용이 곱셈 한 번 아끼는 것보다 비싸서, 결국 dense로 돌리는 게 낫습니다.

그래서 프루닝에는 이 원칙이 따라붙습니다:

프루닝은 하드웨어가 그 sparsity 패턴을 지원해야만 실이득이 난다. 쓸 하드웨어가 무엇을 가속하는지 알고 잘라야 한다.

그동안 sparse 연산을 위한 전용 하드웨어(EIE, ISCA 2016 등)가 연구됐고, 오늘날 가장 널리 쓰이는 답이 NVIDIA의 2:4 structured sparsity입니다.

2:4 프루닝

규칙은 단순합니다. 연속된 4개 가중치마다 정확히 2개를 0으로 만듭니다(50% sparsity). 0의 위치가 “4개 중 2개"로 규칙적이기 때문에, NVIDIA Ampere(A100) 이후의 Sparse Tensor Core가 이 패턴을 하드웨어로 가속합니다. 저장도 효율적입니다 — 비영값 절반 + 각 값의 위치를 가리키는 2-bit 인덱스만 있으면 됩니다.

magnitude로 4열마다 작은 2개를 자르는 마스크는 이렇게 만듭니다.

def two_four_mask(weight):
    """4열 단위로 |w|가 큰 2개만 남기는 2:4 마스크."""
    w = weight.detach().abs(); R, C = w.shape
    mask = torch.zeros_like(weight)
    for j in range(0, C - C % 4, 4):
        idx = w[:, j:j+4].topk(2, dim=1).indices   # 각 4개 중 큰 2개
        mask[:, j:j+4].scatter_(1, idx, 1.0)
    return mask

우리 MLP에 2:4를 적용한 결과: dense 97.6% → 2:4(재학습 없이) 96.8% → 2:4 + 재학습 98.2%. 절반을 규칙적으로 잘라도 재학습이면 완전히 회복됩니다.

파이썬에서 2:4 활용하기

PyTorch는 2:4 패턴을 압축 저장하고 Sparse Tensor Core로 곱하는 API를 제공합니다(fp16, Ampere+ 필요).

from torch.sparse import to_sparse_semi_structured

mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).bool()
W = (torch.rand(64, 64) * mask).half().cuda()
Ws = to_sparse_semi_structured(W)          # 비영값 + 2-bit 인덱스로 압축
x  = torch.rand(64, 64).half().cuda()

torch.mm(Ws, x)                            # Sparse Tensor Core 사용

정말 빨라지나 — 직접 벤치마크

A100에서 dense vs 2:4 sparse 행렬곱 속도를 크기별로 쟀습니다. 결과가 정직하고 교훈적입니다.

A100에서 2:4 sparse vs dense fp16 행렬곱 속도

  • 작은 행렬(N < 4096)에서는 2:4가 오히려 느립니다. 압축·인덱싱 오버헤드가 이득보다 큽니다.
  • 큰 행렬에서만 이득이 납니다. N=8192에서 약 1.72배 빨라집니다.

이것이 “하드웨어와 워크로드를 알아야 한다"의 실체입니다. 같은 2:4 프루닝도 Ampere 이상 GPU에서, 충분히 큰 행렬에서만 빨라집니다. Turing(T4)이나 작은 행렬에서는 이득이 없거나 오히려 손해입니다. 프루닝은 알고리즘만의 문제가 아니라, 처음부터 끝까지 하드웨어와 함께 설계해야 하는 문제입니다.

4. 프루닝 × 양자화는 직교한다 — Deep Compression

여기까지 프루닝을 봤고, 앞선 양자화 편에서는 비트를 줄였습니다. 핵심은 이 둘이 서로 직교(orthogonal) 한다는 것입니다. 프루닝은 “어떤 가중치를 지울지“를, 양자화는 “남은 가중치를 몇 비트로 표현할지“를 정합니다. 축이 다르니 함께 쓰면 이득이 곱해집니다.

이를 가장 잘 보여준 것이 Deep Compression(Han et al., 2016)의 3단계 파이프라인입니다.

  1. 프루닝 + 재학습 — 약한 연결을 지우고 다시 학습해 회복.
  2. 양자화 + 재학습 — 살아남은 가중치를 K-Means 코드북 몇 개로 묶고(코드북 재학습), 인덱스만 저장.
  3. 허프만 코딩 — 인덱스 분포가 한쪽으로 쏠려 있으니 무손실로 한 번 더 압축.

우리 MLP로 직접 돌려봤습니다. 각 단계에서 fc2 가중치 분포가 어떻게 변하는지 보면, 프루닝과 양자화가 분포의 서로 다른 부분을 건드린다는 게 한눈에 보입니다. (0이 된 가중치는 숨기고 살아남은 것만 그렸습니다.)

가중치 분포 진화: 정규 → 프루닝(중간 제거) → 재학습(2-모드) → 양자화(이산 레벨)

  • ① dense — 익숙한 종 모양 정규분포.
  • ② 프루닝(80%) — 0 근처의 작은 가중치가 통째로 사라집니다. 가운데가 날아가고 양쪽 꼬리만 남습니다.
  • ③ 재학습 — 살아남은 가중치가 데이터에 다시 맞춰지며 뚜렷한 2-모드(bimodal) 분포가 됩니다.
  • ④ 양자화(4-bit) — 그 연속 분포가 16개의 이산 레벨로 접힙니다 (여기가 [양자화 편]의 K-Means가 하는 일).

프루닝은 가로축 가운데(작은 값)를 비우고, 양자화는 남은 값들을 세로선 몇 개로 이산화합니다. 손대는 곳이 다르니 겹치지 않고 쌓입니다. 그 결과 모델 크기는 단계마다 곱으로 줄어듭니다.

모델 크기 축소: dense → 프루닝(4.3×) → +양자화(17.8×) → +허프만(18.6×)

  • dense 2,612 KB → 프루닝 604 KB (4.3×) → +양자화 147 KB (17.8×) → +허프만 141 KB (18.6×)
  • 그동안 정확도는 오히려 97.5% → 98.3% 로 유지·개선됐습니다(재학습 덕분).

즉 프루닝만으로 4.3×, 양자화를 얹어 17.8×, 허프만까지 18.6×. 각 기법의 이득이 더해지는 게 아니라 곱해집니다. 원 논문은 이 파이프라인으로 AlexNet을 35×, VGG-16을 49× 줄였습니다(정확도 손실 없이). 데이터타입 → 양자화 → 프루닝, 이 시리즈가 다룬 세 축은 결국 하나의 모델에 동시에 얹을 수 있는 직교한 도구들입니다.

직접 돌려보기

위 모든 plot은 아래 노트북을 Colab에서 실행해 나온 것입니다. 런타임 → 모두 실행이면 됩니다(2:4 속도 측정은 A100 등 Ampere GPU 선택).

정리

  • 무엇을: random은 금방 무너진다. magnitude(|w|)는 싸고 강력해 OBD의 2차 saliency와 맞먹는다 — 그래서 실무 표준이 됐다.
  • 얼마나: prune→retrain 반복이면 95%(20배)를 잘라도 정확도가 유지된다.
  • 어떻게: unstructured는 압축률만 높다. 2:4 structured라야 Ampere Sparse Tensor Core가 실제로 가속한다 — 단, 큰 행렬에서만.
  • 함께: 프루닝과 양자화는 직교한다. Deep Compression처럼 쌓으면 이득이 곱해진다(우리 실측 18.6×, 정확도 유지).

Optimal Brain Damage가 1989년에 심은 통찰은 지금도 유효합니다. 다만 그것이 “실이득"이 되려면, 36년 뒤의 하드웨어가 그 sparsity를 받아줘야 했습니다. 프루닝의 역사는 곧 알고리즘과 하드웨어가 서로를 기다려온 역사이기도 합니다.

다음 편에서는 sparsity를 층마다 자동으로 정하는 방법(sensitivity analysis · AMC · NetAdapt)을 다뤄보겠습니다.

References

  • Y. LeCun, J. S. Denker, S. A. Solla. Optimal Brain Damage. Advances in Neural Information Processing Systems (NeurIPS), 1989.
  • S. Han, J. Pool, J. Tran, W. J. Dally. Learning both Weights and Connections for Efficient Neural Networks. NeurIPS, 2015.
  • S. Han, H. Mao, W. J. Dally. Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. ICLR, 2016.
  • A. Mishra et al. Accelerating Sparse Deep Neural Networks (2:4 Structured Sparsity). arXiv:2104.08378, 2021.
  • S. Han et al. EIE: Efficient Inference Engine on Compressed Deep Neural Network. ISCA, 2016.
  • M. Horowitz. Computing’s Energy Problem (and what we can do about it). ISSCC, 2014.
  • 강의 참고: MIT 6.5940 TinyML and Efficient Deep Learning Computing (Song Han), Lec 3–4 Pruning & Sparsity.