[AI System] Model Sparsification

"Model Sparsification"

[Model Pruning]

· Model에서 덜 중요한 부분을 제거 (Weight Pruning : 개별 Weight 제거 / Neuron Pruning : Neuron이나 채널 단위로 제거)

▣ Weight Pruning(= Unstructured Pruning)
· 개별 Weight 제거

△ 덜 중요한 Weight를 세밀하게 골라 제거할 수 있어서 정확도 유지
▼ 살아남은 Weight들이 불규칙하게 흩어지기 때문에 일반 GPU나 CPU가 이를 효율적으로 활용하기 어려움
▼ Weight 값뿐만 아니라 Index 정보도 저장해야 함

▣ Neuron Pruning(= Structured Pruning)
· Neuron이나 Channel 단위로 제거

△ 더 작은 Dense Network로 표현 가능하여 표준 Library를 그대로 활용해 연산 속도 향상 가능
△ 별도의 희소 Format 필요 없이 일반 Matrix Multiplication으로 바로 실행 가능
▼ 구조 자체가 크게 바뀌므로 정확도 손실
▼ Fine-Tuning이 필요할 수 있음

· Unstructured Pruning : Weight Matrix 안의 일부 값만 0으로 변경 → 남은 값들의 위치가 불규칙

 △ 정확도 유지
 ▼ Sparse Computation Library 필요 (∵ 남은 값들의 위치가 불규칙)

· Structured Pruning : 행 또는 열 전체를 제거 → 결과 Model을 더 작은 Dense Matrix로 표현 가능

 △ 기존 연산 Library를 그대로 사용 가능
 ▼ 정확도 손실

[Which Weights should we Prune? - Unstructured Pruning]

▣ Magnitude-Based Pruning
· 절댓값이 큰 Weight는 Model 출력에 큰 영향을 줄 가능성이 높으므로 남기고, 절댓값이 작은 Weight는 영향이 작다고 보고 제거
· 33.3% 제거 : |wi| 값이 가장 작은 Weight 2개가 0으로 변경
· 50% 제거 : |wi| 값이 가장 작은 Weight 3개가 0으로 변경
    → Pruning Ratio : 사용자가 선택

∴ Pruning Ratio가 높을수록 더 많은 Weight가 제거되어 Model이 가벼워지지만, 성능 손실 위험도 증가

▣ Gradient-Based Pruning
· Weight를 제거했을 때 Loss가 거의 변하지 않는다면, 해당 Weight는 덜 중요하므로 제거 가능
· 33.3% 제거 : |wi · gi| 값이 가장 작은 Weight 2개가 0으로 변경
· 50% 제거 : |wi · gi| 값이 가장 작은 Weight 3개가 0으로 변경
    → Magnitude-Based Pruning 결과와 다를 수 있음 (∵ Weight가 작더라도 Gradient가 크다면, 해당 연결 중요)

> Taylor Series
· 어떤 함수 f(x)를 점 a 근처에서 근사 가능

· w + u(= x) : 우리가 가고자 하는 새로운 Weight 상태
· w(= a) : 현재 우리가 알고 있는 Weight 상태
· u(= x - a) : Weight를 제거하면서 발생하는 변화량 (- wiei)

→ |gi · wi|가 작으면 해당 Weight를 제거해도 Loss가 많이 변하지 않으므로 Pruning 가능

▣ Layer-Wise Pruning
· 전역 Pruning : 전체 Model에서 n%를 전역적으로 제거
→ 특정 Layer의 Weight가 거의 전부 제거될 수 있음

· Layer-Wise Pruning : 각 Layer에서 n%씩 제거
→ 특정 Layer가 과도하게 제거되는 것을 막고, 각 Layer가 일정 수준의 표현력 유지 가능

Q) Dropout, Sparse Backpropagation, Pruning 모두 "Sparsity"를 고려하는데, 차이점은?
· Dropout : Overfitting 방지가 목적이며, 학습 시 무작위로 Neuron을 끔
· Sparse Backpropagation : 학습 과정에서의 연산량을 줄이는 것이 목적
· Pruning : 추론 단계에서 사용할 가볍고 작은 최종 Model을 만드는 것이 목적
→ 3가지 방법을 함께 사용하면 학습 효율과 Model 경량화 모두 달성 가능

▣ Hessian-Based Pruning
· g : Gradient (Loss를 Weight에 대해 1번 미분)
· H : Hessian (Loss를 Weight에 대해 2번 미분)

· Model이 충분히 잘 학습되었다고 가정 → 잘 학습된 지점에서는 Loss가 거의 최소점 근처 → Gradient가 0에 가까움
∴ Taylor 근사식에서 1차항 무시 가능

· i번째 Weight 제거 시 Loss 변화 (Hii : Hessian 행렬의 대각 성분)
∴ Hessian-Based Pruning에서는 이 값이 작은 Weight를 먼저 제거
    (∵ 값이 작을수록 해당 Weight를 지웠을 때 Model의 손실 변화가 적음)

· Hessian 행렬은 Parameter 수가 N개일 때 N2의 원소 필요 → Model이 크면 계산과 저장이 매우 어려움
· Hii(Hessian의 대각 성분)만 사용하므로 필요한 값이 N개로 감소

▣ Different Pruning Criteria
· Magnitude-Based Pruning : |wi|
· Gradient-Based Pruning : |wi · gi|
· Hessian-Based Pruning : w2Hii

· Activation-Based Pruning : Activation이 작은 Neuron을 제거
→ 어떤 Neuron이 데이터가 들어와도 거의 활성화되지 않는다면, 해당 Neuron은 Model 출력에 크게 기여하지 않을 가능성이 높음

[When should we Prune? - Unstructured Pruning]

▣ Pruning at Initialization(Sparse Training)
· 학습을 시작하기 전, 초기 Model에서 중요도가 낮아 보이는 Weight를 제거

△ GPU Memory와 학습 비용 절약
▼ 어떤 Weight가 중요할지 미리 알기 어렵기 때문에 성능이 매우 불안정

▣ Pruning after Training Ends(One-Shot Pruning)
· 전체 학습이 완전히 끝난 Model을 한 번에 제거
· 전체 학습 → Pruning → Fine-Tuning

△ 이미 학습된 Pretrained Model을 활용
▼ 한 번에 많은 Weight를 제거하면 성능이 크게 저하

▣ Iterative Pruning
· Full Model을 학습한 후, 일부 Weight를 조금 제거하고 Fine-Tuning하는 과정을 여러 번 반복

△ 성능 저하를 최소화하면서도 높은 압축률을 달성 가능
▼ 원래 학습 시간보다 2~10배 이상의 시간 소요

▣ 성능 비교
· Iterative Pruning의 성능이 가장 좋음
· Pruning만 하는 것보다 Pruning 후 Fine-Tuning을 하는 것이 더 좋은 성능을 보임

[Implementing Unstructured Pruning in Practice - Unstructured Pruning]

▣ COO(Coordinate List) Format
· 희소 행렬에서 0이 아닌 값의 Row Index, Column Index, Value를 각각 별도의 List로 저장
· Row Index : K개 / Column Index : K개 / Value : K개 → 총 3K개를 저장

∴ 전체 원소 수가 N일 때 3K < N이면 저장 공간 절약 가능

> COO Format을 이용한 Sparse Matrix-Vector Multiplication
· 모든 행렬 요소를 곱하지 않고, COO Format에 저장된 값들(Non-Zero)만 계산
· 0인 Weight는 아예 곱하지 않아 계산량 감소

※ CSR Format : CUDA Sparse Library인 cuSPARSE에서 기본적으로 많이 쓰이는 Sparse Row 기반 형식

** COO/CSR 같은 Sparse Format을 사용하지 않으면 Unstructured Pruning의 장점을 제대로 얻기 어려움
∵ 0인 값들까지 메모리에 저장하고 계산 → 메모리 절약 X & 계산량 감소 X

[Which Weights should we Prune? - Structured Pruning]

▣ Neuron Importance Score
· 각 Neuron과 연결된 Weight들의 중요도 점수를 모아서 Neuron Importance Score를 계산
· Importance Score가 작은 Neuron은 Model 출력에 덜 중요하다고 보고 제거

[When should we Prune? - Unstructured Pruning]

· Structured Pruning에서도 Pruning 시점은 Unstructured Pruning과 비슷

▣ Pruning at Initialization(Sparse Training)
▣ Pruning after Training Ends(One-Shot Pruning)
▣ Iterative Pruning(Repeated Pruning and Fine-Tuning)
→ 구현 측면에서는 Structured Pruning이 더 쉬움
∵ Neuron이나 Channel을 통째로 제거하면 결과 Model이 여전히 Dense Matrix 구조를 가짐
    → COO/CSR 같은 Sparse Matrix Format 필요 없이, 기존 행렬 곱셈 연산 그대로 사용 가능

[Reference]

· Model Sparsification (AI System : Week 10) - Dongjun Han

댓글

이 블로그의 인기 게시물

[Mini-NPU RTL] NN Reference Model

[Mini-NPU RTL] TPU (Study Paper)