brunch

라이킷 1 댓글

You can make anything
by writing

C.S.Lewis

계정을 잊어버리셨나요?

by Qscar Feb 28. 2023

[AI]Optimization Technique 02

Pruning

2. Pruning

What is Pruning?

Pruning Strategy - Image by Tivadar Danka[1]

가지치기로 표현될 수 있는 Pruning 전략은 덜 중요한 가중치 연결을 끊어내고(Drop Connect), 만약 모든 연결이 끊어진 뉴런(혹은 노드)이 있다면 해당 뉴런을 비활성화하는 것이다. 만약 딥러닝에 대해 경험이 있는 사람이라면 알겠지만 이는 모델 학습 과정에서 사용되는 DropConnect 혹은 Dropout의 기능을 추론 단계에서 적용한 것이다.

일반적으로 학습 과정에서 사용된 Drop 전략들은 실제 모델에선 확률값으로 반영되어 적용되기에 실질적으로 모델의 경량화(최적화)에는 도움이 되지 않는다. 만약 학습 과정에서 사용되는 Drop 전략들이 그대로 모델을 최적화시킬 수 있다면 별도의 Pruning 전략은 필요치 않다.

오히려 학습 과정 중에 사용되는 Drop 전략들이 그 이후에도 온전할 수 있다면 Pruning보다 효율적이며, 효과적일 것으로 예상된다. 그 이유는 후술한다.

Tensorflow 공식 블로그 글을 살펴보면, MNIST 모델의 경우 Pruning을 통해 기존의 12MB 사이즈의 모델을 2MB로 무려 1/6로 줄였으며, 여기에 지난 글에서 포스팅했던 동적 양자화를 가미함으로써 0.5MB로 1/24로 줄였다고 한다.

Sparsity results across different models and tasks - Image by Tensorflow Blog

심지어 위 그림을 살펴보면 각기 다른 영역에 Pruning을 적용해본 결과 Pruning에 따른 다소의 성능 하락이 관찰되었지만, 그 이상으로 모델 압축이 가능해 성능을 희생하더라도 메모리 측면이 더욱 중요한 도메인 혹은 서비스라면 충분히 적용할 가치가 있음을 시사한다.

다만 이전 포스팅에서 소개한 Quantizatino 혹은 다음 시간에 소개할 Distilation과 달리 Pruning은 일반적으로 사용되는 weight만을 저장하는 방식이 아닌 모델 전체를 저장해야 하고, 이로 인한 추가적인 비효율성, 비호환성을 감안해야 하기에 신중한 사용이 요구된다.

How do Pruning?

Pruning은 크게 가중치 가지치기(Drop Connect)와 뉴런 가지치기(Drop out)으로 구성되며, 지정한 수준의 희소성을 지닐 때까지 비교적 중요한 가중치만을 남기는 것으로써, 어떤 기준으로 가지치기를 진행할 것인지에 따라 구분된다. 어떤 가중치를 삭제해야하고(성능을 유지할 수 있도록), 이러한 삭제해도 되는 가중치를 어떻게 계산할 것인지(빠르고 효율적인 계산을 위해)에 대한 전략이 중요한 셈.

모든 가지치기 방법은 중요도 순으로 최상위 k%의 가중치만 남도록 하는 것으로써, 남겨진 가중치에는 1을 제거되는 가중치에는 0을 부여하는 하나의 희소행렬을 구성하는 것이다. 이렇게 생성되는 행렬을 마스크 행렬이라고 하며, 마스크 행렬을 적용해 모델 학습을 진행해봄으로써 성능이 본래 모델에 수렴하도록 하는 것이다.

절대값 가지치기

첫 번째 Pruning 방법은 절대값 가지치기(Magnitude Pruning)이다.

가중치 절대값 크기에 따라 스코어를 계산하고, 마스크 행렬을 생성한다. 일반적으로 가장 중요한 가중치가 무엇인지 학습하도록 모델을 훈련하고, 덜 중요한 가중치를 제거하는 과정을 반복함으로써 원하는 비율인 k%의 가중치만 남도록 가지치기 하는 것이다.

Sparsity ramp-up function with a schedule of pruning - Image by Tensorflow Blog

이 과정은 계산량이 많고, 가지치기의 매 단계마다 모델을 훈련해 성능을 수렴시켜야 한다는 단점을 가지고 있다. 이에 따라 희소성 초기값을 점진적으로 증가시키며 진행하는 것이 일반적이며 위 그림과 같이 중복된 가중치가 많은 초기 단계에선 빠른 속도로 학습되지만 점차 속도가 느려지게 된다.

이동 가지치기

두 번째 가지치기는 이동 가지치기(movement pruning)로, 가지치기의 단점이라 할 수 있는 전이학습 시 성능하락을 개선하기 위해 고안되었다. 일반적으로 가지치기는 해당 도메인에 대해 최적화하는 마지막 단계로서 사용되고, 가지치기가 된 모델은 미세튜닝, 혹은 더이상 다른 도메인이나 태스크로의 전이학습이 온전히 이뤄지지 않는다.

Difference between Magnitude pruning and Movement pruning - official paper[4]

너무 완벽한 맞춤정장은 비슷한 체형의 다른 사람에게 주기 어려운 것과 같은 맥락이다. 때문에 이동 가지치기에서는 가중치와 가중치 점수를 함께 학습해 경사하강법을 통해 최적화한다. 이렇게 학습된 점수에 근거해 마스크 행렬을 생성하고, 원점에서 가장 크게 이동하는 가중치를 중요하다고 판단한다.

이는 위 그림과 같이 사전훈련에서의 가중치가 큰 영향을 미쳤던 절대값 가지치기와 달리 이동 가지치기의 경우 사전훈련의 영향을 줄여 전 영역에서 고르게 가지치기가 가능하다. 이를 가중치 분포를 통해 살펴보면 아래 그림과 같다.

Distribution of remaining weights - official paper[4]

Code Example 코드

일반적으로 가지치기 전략은 모델을 재학습하며 진행되야 하기에 모델의 사이즈가 클수록 필요성이 높아지지만, 그만큼 비용도 함께 높아진다는 문제가 있다. 또한 이 과정에서 기존에 사용하던 딥러닝 라이브러리 상의 코드와 제대로 호환되지 않거나 예상치 못한 문제를 발생시키기도 한다.

때문에 가지치기는 도메인 최적화를 위한 마지막 단계에 사용하길 권장하며, 지나치게 큰 모델을 경량화시키는 목적이라기 보다는 메모리 이슈가 존재하는 Single Board Computer 등에 적용할 때와 같은 상황에 적용해볼 수 있을 것으로 기대되며 예시코드를 실습해볼 수 있는 url과 함께 이번 포스팅을 마친다.

1) Pytorch Pruning Tutorial

2) Tensorflow Pruning Tutorial

3) Huggingface Pruning Library Official Github(nn_pruning)