Lottery ticket hypothesis 와 후속 연구 정리

Intro

AI는 대용량의 데이터와 큰 모델을 기반으로 많은 발전을 이루고 있다. 하지만 많은 데이터와 거대 모델을 학습시키는 것은 많은 비용이 들게 된다. 큰 모델은 저장 용량이 클 뿐 아니라 예측을 할 때 많은 연산을 요구하기 때문에 효율적이지 않다. 특히 스마트폰 같은 작은 기계에서 사용하기에 적합하지 않다. 그에 따라 모델의 크기를 줄여 메모리를 줄이며 연산량을 줄여보자 하는 시도가 많이 이루어졌고 대표적인 분야로 knowledge distilation, pruning 등이 있다. 이번 글에서는 pruning, 특히 ICLR 2019에서 best paper를 받은 lottery ticket hypothesis에 대한 해석과 단점을 극복한 후속 논문들을 살펴볼 것이다.

내용을 설명하기 전에 글을 읽기 위해 필요한 용어 몇 개를 먼저 소개한다.

* Global pruning & Local pruning

- Global pruning은 pruning시 모든(혹은 일부) layer를 한 번에 비교하여 pruning 할 parameter를 정하는 방법이다.

- Local pruning은 각 layer별로 따로 비교하여 pruning 할 parameter를 정하는 방법이다.

* One-shot pruning & Iterative pruning

- 전체 pruning 하고 싶은 비율이 80%라고 할 때, one-shot 방법은 한 번에 80%를 pruning 하는 것이고, iterative 방법은 여러 번에 걸쳐 pruning을 해서 80%까지 도달하는 것이다.

* Layer-collapse

- One-shot 방법이 iterative 방법에 비해 간단하고 cost가 적지만 보통 성능이 떨어진다. 주요한 요인으로 'layer collapse'가 있다. Layer collapse는 여러 layer를 한 번에 비교하여 pruning을 해버리면 특정 layer의 parameter가 모두 사라지는 경우가 생길 수 있다. 그런 경우는 보통 magnitude를 사용하여 pruning 하는 경우에 발생하는데, 각 layer의 size에 따라 parameter 분포가 달라지기 때문이다. 예를 들어, layer size가 크면 0 근처의 parameter가 많게 되고 size가 작으면 넓게 분포하여 0 근처의 값은 적게 된다. 그렇게 되면 size가 큰 layer부터 pruning 되는 현상이 발생할 수 있게 된다. 다만 iterative 하게 pruning을 하게 되면 처음엔 큰 layer부터 pruning 할 수 있겠지만 그다음 iteration에서는 size가 작아졌으므로 덜 pruning 되는 식으로 balance를 맞출 수 있게 된다.([2])

* Winning ticket

- Winning ticket은 lottery ticket 방법으로 얻어진 mask와 initial parameter $\theta_{0}$를 element-wise 하여 구한 model($m\odot\theta$)을 의미한다. 

 

Lottery Ticket Hypothesis (LTH)

Lottery ticket hypothesis는 지난 글에서 설명한 적이 있다. 다시 간단히 설명하면, Initialize -> Train -> Pruning-> Initialize -> Train의 과정으로 학습시키는 방법이다. 2번째 initialize 시 pruning 된 곳을 제외하고는 첫 번째 initialize와 같이 해주는 것이 특징이다.

figure 1. algorithm of LTH ([1])

이 방법은 기존의 방법 혹은, pruning 하지 않는 모델을 넘는 성능을 보여주어 주목받은 논문으로 왜 이 방법이 잘 되는지에 대한 후속 논문과 단점을 극복한 후속 논문들이 많이 나오게 되었다.

 

LTH의 특징과 그에 따른 후속 논문 소개

1. 훈련이 끝난 모델을 기준으로 pruning을 하고 다시 initialize를 하므로 훈련을 여러 번(최소 2번 이상) 반복해야 한다는 것이다.(LTH는 magnitude based pruning을 하므로 iterative pruning을 사용한다) 그러므로 computational cost가 너무 크다. 이는 치명적 한계이므로 후속 연구가 활발히 나오고 있다.

 

* 후속 연구

 

- Efficient Lottery Ticket Finding: Less Data is More([3]): 모든 데이터를 다 사용하면 오래 걸리니 중요한 일부 데이터(core set)만 골라서 훈련하는 방법이다. 데이터에는 학습하기 어렵거나 pruning 때 쉽게 잊어버리는 데이터가 있고 그런 데이터만 잘 골라 훈련시키면 효율적인 훈련이 될 수 있다. 따라서 논문에선 학습이 잘 안 되는 데이터 혹은 full model과  subnetworks가 서로 다르게 inference 하는 데이터를 고르는 방법을 제시한다.

 

- Drawing Early-Bird Tickets: Towards More Efficient Training of Deep Networks([4]): Winning ticket이 학습 초반에 발견 가능하다는 사실을 발견한 논문이다. 이 사실을 이용하여 epoch 마다 mask를 만들고, 만들어지는 mask를 저장해두어 mask의 변화(distance)가 충분히 작아지면 종료하는 방법을 택한 것이다. 이렇게 하면 초반 epoch 까지만 학습시켜도 마지막까지 돌렸을 때 얻어지는 mask와 유사한 mask를 얻을 수 있다는 것이다.

 

- One ticket to win them all: generalizing lottery ticket initializations across datasets and optimizers([6]): Image domain에서 특정 dataset으로 만든 winning ticket으로 다른 데이터에 학습시켜도 잘 적용된다는 내용이다. Fashion-MNIST, Cifar10, Cifar100, ImageNet으로 실험하였다. 만약 domain이 비슷하다면 한번 찾아놓은 winning ticket으로 훈련해보는 것도 좋은 방법일 듯하다.

 

- SNIP: Single-shot Network Pruning based on Connection Sensitivity([7]): 이 논문은 학습 전 각 weight를 끊고(0으로 만들고) loss가 어떻게 변하는지 gradient를 얻어 전체 weight 중 작은 값을 갖는 것들을 pruning 하는 방법이다. 이러한 방법을 connection sensitivity라고 부른다. 학습을 하지 않아도 된다는 강력한 장점을 가지고 있지만, 각 weight를 고립시켜 끊어버리고 영향을 확인하는 방법이라 전체 network에서 중요한 connection도 끊을 수 있다는 단점이 있다.

 

- PICKING WINNING TICKETS BEFORE TRAINING BY PRESERVING GRADIENT FLOW([8]): SNIP은 weights의 interaction을 고려하지 않고 pruning 하기 때문에 network flow에 중요한 weight를 pruning 할 수 있다. 이런 현상은 gradient norm 감소로 드러난다(practical 하게 찾음). 그래서 이 논문에서는 Gradient Signal Preservation(GraSP)라는 방법을 제안한다. SNIP과 다르게 weights 간의 interaction을 고려하기 위해 gradient 뿐 아니라 hessian matrix를 사용하는 것이 특징이다(hessian-gradient product). GraSP는 pruning 후에도 gradient norm의 감소를 최소화하여 gradient flow를 유지하는 방법이다. SNIP과 마찬가지로 학습을 하지 않아도 된다는 장점이 있다.

 

- Winning the Lottery Ahead of Time: Efficient Early Network Pruning([10]): 기존 LTH에 대한 연구는 unstructured pruning에 집중되어 연구가 많이 되었으나, unstructured pruning은 weight를 0으로 만드는 것뿐이라 실질적인 메모리 개선이나 훈련 속도 증가, inference 속도 증가 등을 기대하긴 어려웠다. 그래서 이 논문은 기존 논문들보다 structured pruning을 통해 좋은 구조를 찾고 성능까지 얻어내는 모델을 얻는 것을 목표로 한다. 우선 pruning은 학습을 아예 하지 않는 방법을 선택하진 않고, NTK가 안정될 때까지(거의 변하지 않을 때까지) 학습한 후, pruning 하는 방법을 택한다. NTK는 network의 width가 매우 클 때 파라미터가 거의 변하지 않을 경우 변하지 않게 된다. 하지만 width가 작더라도 훈련이 충분히 된 상태에서 파라미터가 거의 변하지 않게 될 때도 NTK는 constant에 수렴한다. 이를 이용한 방법으로 structured pruning도 안정적으로 수행하여 좋은 성능을 보여준다. (이 방법은 init으로 rewind하진 않고 pruning 한 곳에서 추가 학습하는 방법을 취한다. [5]랑 관련되어 noise가 적은 곳에서 시작한다고 생각할 수 있을 듯하다.)

 

- Comparing Rewinding and Fine-tuning in Neural Network Pruning([14]): (이 논문은 이전처럼 시간과 직접 관련된 내용은 아니다.) LTH이전 pruning에서는 Initialize -> Train -> Pruning -> Finetuning 방법으로 pruning 했으나 이 방법보단 다시 rewind 하여 훈련하는 LTH의 성능이 더 좋음이 여러 논문에서 나타났다. 이 논문에서도 weight rewind가 성능이 좋다는 것을 보여주었다. 추가로 모델 대신 learning rate을 rewind 하고 이전처럼 pruning 후 finetuning 하는 방법도 비슷하게 좋은 성능을 보여준다고 주장하였다.

 

2. 작은 크기의 모델은 winning ticket을 찾기 수월하나, 크기가 큰 모델은 작은 learning rate를 사용하거나 warm up scheduling을 사용해야 찾을 수 있었다. (LTH는 큰 neural network 안에는 많은 sub network가 있고 그 안에는 원래의 모델보다 좋거나 비슷한 것(winning ticket)이 있다는 것을 가정한다. 모델이 클수록 sub network가 기하급수로 늘어나고 그중 winning ticket을 찾는 것은 더 어렵다고 해석할 수 있다.)

 

* 후속 연구

 

- Linear Mode Connectivity and the Lottery Ticket Hypothesis([5]): LTH의 저자가 작성한 논문으로, lottery ticket을 찾고 rewind 후 재학습할 때, 학습 초기의 SGD noise가 LTH의 성공을 좌우한다는 내용이다. 즉, iterative magnitude pruning(IMP)는 SGD noise가 적을 때 학습이 잘 된다는 이야기다. 다행히 SGD noise는 학습 초기에 안정화되어 이 시기가 지난 곳으로 rewind 하는 방법을 선택할 수 있다. 다만 이 논문의 한계로 그 지점을 찾을 수 있는 방법을 제시하지 않았고 그냥 여러 번의 실험을 통해 그런 곳이 존재한다는 사실만 보여준다. 이 연구는 기존에 warm up scheduling이 왜 잘 동작했는지 알려준다. 초기에 매우 작은 learning rate에서 시작하여 점점 커지는 이 방법은 학습 초기 큰 SGD noise에 강건한 최적화를 가능하게 하므로 IMP가 잘 작동하게 해 준 것이다. ([9] 참고)

 

- Gradient Flow in Sparse Neural Networks and How Lottery Tickets Win([11]): 초기 SGD noise의 제어가 winning ticket을 찾는 것에 중요한 역할을 한다는 것을 보았다. 그런데 SGD noise가 큰 것이 무엇이 문제가 되는 것일까? [5] 논문에도 나왔지만, 초기 SGD noise가 큰 구간 이후부터 학습시킨 모델은 linear 하게 파라미터를 interpolation 했을 때 loss가 증가하지 않는 것으로 나왔다. 이를 linear mode connectivity라고 한다. [11] 논문에서도 이를 보여주는 데, 같은 initialize에서 다른 random seed를 사용하여 여러 번 학습했을 때 winning ticket을 잘 찾는 경우는 항상 거의 비슷한 (loss landscape) 영역으로 최적화가 된다는 것이다. 단순히 같은 곳으로 최적화가 된다는 것이 중요한 것은 아니고, winning ticket solution 영역으로 잘 찾아갔다고 생각하면 된다. 이는 지난 글에서 설명한 바 있다.

 

 

3. LTH는 실제로 많은 경우에 좋은 성능을 보이는 것으로 보인다. 거기다 unstructured pruning의 경우까지 훈련 속도를 증가시킨다. 왜 winning ticket은 더 빠르게 converge 하고 좋은 성능을 낼 수 있을까?

각 선은 남아있는 parameter의 비율에 따른 학습 진행 결과이다.([1])

 

* 후속 연구

 

- Deconstructing Lottery Tickets: Zeros, Signs, and the Supermask([12]): LTH가 빠르게 학습이 될 수 있는 이유는, 학습이 완료된 모델의 parameter 중 magnitude가 작은 부분을 아예 0으로 만든 후, 나머지 부분만 원래의 initialize로 돌려놓고 재학습하는 것이 가장 주요한 원인이다. [12] 논문에서는 manitude가 작은 parameter를 0으로 만드는 것이 중요함을 실험으로 보여준다. 또한 random initialize + mask (winning ticket)가 random initialize만 한 모델보다 (둘 다 intialize임에도) 높은 성능을 보여준다. 그 이유는 mask 자체가 훈련이 끝난 optimal 한 모델을 참고하여 작은 파라미터를 0으로 무시하도록 만든 것이니 winning ticket은 random initialize보다 optimal에 좀 더 가까운 모델이라고 할 수 있기 때문이다. 결론적으로 winning ticket은 optimal에 더 가까운 포인트에서 훈련했으므로 같은 iteration 만큼 학습시켜도 그냥 random initi보다 더 많이 학습한 효과가 날 것이고 그에 따라 더 빨리 converge 하고 좋은 성능을 내는 것이라 생각된다([13]을 참고한 뇌피셜이다).

 

Reference

[1] https://arxiv.org/abs/1803.03635

The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks

 

[2] https://arxiv.org/abs/2006.05467?fbclid=IwAR3nVFrrgTmNhgv9TKJ7pezEgS-KIVyOFnw-3y_crZ3xsb8teGYZ6z9G9DY 

Pruning neural networks without any data by iteratively conserving synaptic flow

[3] https://arxiv.org/abs/2106.03225

Efficient Lottery Ticket Finding: Less Data is More

[4] https://arxiv.org/abs/1909.11957

Drawing Early-Bird Tickets: Towards More Efficient Training of Deep Networks

[5] https://arxiv.org/abs/1912.05671

Linear Mode Connectivity and the Lottery Ticket Hypothesis

[6] https://arxiv.org/abs/1906.02773

One ticket to win them all: generalizing lottery ticket initializations across datasets and optimizers

[7] https://arxiv.org/abs/1810.02340

SNIP: Single-shot Network Pruning based on Connection Sensitivity

[8] https://arxiv.org/abs/2002.07376

Picking Winning Tickets Before Training by Preserving Gradient Flow

[9] https://arxiv.org/abs/1903.01611

Stabilizing the Lottery Ticket Hypothesis

[10] https://arxiv.org/abs/2206.10451

Winning the Lottery Ahead of Time: Efficient Early Network Pruning

[11] https://arxiv.org/abs/2010.03533

Gradient Flow in Sparse Neural Networks and How Lottery Tickets Win

[12] https://arxiv.org/abs/1905.01067

Deconstructing Lottery Tickets: Zeros, Signs, and the Supermask

[13] https://arxiv.org/abs/1705.08741

Train longer, generalize better: closing the generalization gap in large batch training of neural networks

[14] https://arxiv.org/abs/2003.02389

Comparing Rewinding and Fine-tuning in Neural Network Pruning