LaCo: Large Language Model Pruning via Layer Collapse

 

EMNLP 2024

Code

👀 요약 👀


✨ Point ✨

 

Abstract로 흐름 파악하기

 

quantization, knowledge distillation, model pruning 같은 방법들은 많은 이슈들에 의해 제약이 있다. (hardware support, 방대한 학습, 모델 내부 구조 변화)

 

-> 간결한(concise) layer-wise pruner인 Layer Collapse(LaCo)를 제안한다. 

이 방법은 모델의 후반 layer를 앞쪽 layer에 합치는(collapse) 방법이다. -> 모델 구조를 유지하면서 사이즈를 줄일 수 있음

 

25-30%의 pruning ratio에서도 80%를 웃도는 성능을 유지한다. (현존하는 SOTA 모델보다 아웃퍼포밍함.) ㄹㅇ..?

추가적으로 post-training 실험 진행, layer-wise similarity, various pruning ratio에 대해 논의한다.

 

---

1. Introduction

트랜스포머 기반 LLM은 다양한 테스크에서 상당한 능력을 보이고 있으나, 모델의 크기가 커지면서 computational resource의 필요도도 높아지고 있다.

추론 속도를 향상시키고, 학습 cost를 줄이고, 작은 모델을 만드는 방법들 : quantization, knowledge distillation, model pruning

그러나 이 방법들에 단점이 존재함.

- quantization : 특정한 하드웨어가 필요하다. 때로는 모델 성능에 영향을 미친다 (이렇게 당연하고 다른 거에도 해당하는 말을 써도 되는 거임?)

- kd : 작은 모델을 재학습해야 한다.

- non-structured pruning: 모델이 sparse해지며 성능 저하 유발됨, 특정한 하드웨어가 필요하다.

- strucrured pruning: 모델의 구조가 바뀌거나 모델의 portability가 감소한다.

 

위와 같은 이슈들을 고려하여, 새로운 방법을 제안한다.

- 이미 학습된 LLM에서 몇 layer를 pruning한다.

- 한 레이어의 파라미터를 다른 여러 레이어로 대체한다. (substitute the parameters of one layer for multiple layers) (합친다는 뜻)

 

특정 레이어의 파라미터의 차이(differentials)와 이후레이어들을 mergeg해도 모델 성능에 큰 영향을 주지 않는 것을 발견. 

Reserving-Differences-whileSeeking-Common (RDSC) Layer Merge 이라고 부르기로 함.

 

 

In this paper :

- 30-50% 를 제거해도 추가학습없이 성능 유지함. 다양한 benchmarks 테스트를 통해 SOTA모델보다 뛰어남을 보여줌

- LLM의 내부 구조를 유지함. 시스템 구현을 변경하지 않고도 기존 애플리케이션에 원활하게 통합될 수 있음

- 압축된 모델이 효율적으로 파라미터를 받아왔으며(interit), 최소한의 학습만으로 원본 모델수준으로 복구가 가능한지 확인하기 위해 post-training을 진행.

 

 

2. Method

2.1. Reserving-Differences-wile-Seeking-Common Layer Merge

l : LLM의 레이어
$𝜃_l$ : l번째 레이어의 모든 파라미터들
$𝜃^*_l$ : 합쳐진 최종 레이어

여기서 $theta_(l+k) - theta_l$은 layer-wise 파라미터의 차이(difference)를 의미한다. 

실제로 계산할 때는 self-attention(SAN)과 MLP 레이어를 각각 처리함.

그리고 반영이 된 m개의 레이어는 없애버린다.

 

이후의 pruning 과정에서 계속 RDSC Layer Merge가 포함되며, 이는 특정 레이어로의 연속적인 layer collapse로 볼 수 있다. 이러한 이유로 ‘Layer Collapse’라는 이름을 붙임

 

2.2. Layer Collapse

최상위 레이어(topmost layer, ==후반 레이어)부터 인접한 레이어를 dynamically하게 merge한다.

few-shot calibration sample을 사용하여 원래 모델과의 성능 손실을 최소화한다. 

(1) Preparation

M : LLM
C : merge할 레이어 수
[L, H] : merge할 레이어 범위
I(i) : merge 연산 사이의 최소 간격
D : few-shot calibration data
T : 원본 모델과 merge된 모델의 유사도 threshold

 

(2) Pruning (lline 1-17)

 

l : layer pointer (H-C) : 즉 후반 레이어부터 내려오면서 계산한다.

K :C -1(합칠 레이어 개수-1) 과 M^* - l(전체 레이어 수에서 l을 뺀 값) 중 작은 값

 

RDSC Layer Merge (line 4-5)

-레이어 l 바로 다음의 K개 레이어를 레이어 l에 merge한 후, 중복된 K개의 레이어를 제거

 

 

Calculate similarity (line6)

- calibration data (D)를 사용하여, 원본 모델과 압축모델 각각의 마지막 레이어의 Output hidden state를 구한다.

- 구한 두 개의 값의 similarity score (s)를 구한다.

 

 

Merge Evaluation and Adjustment (line 7-15)

- similarity score 가 threshold(T)를 넘으면 merge한다.

그리고 포인터 l은 간격 설정값(i)만큼 내려간다.

- 진행하다보면 포인터 l이 레이어 개수보다 적어질 수 있기 때문에 l을 M^* - C로 재설정한다. (line11)

 

2.3. Complexity Analysis

complexity는 모델의 inference 속도에 달려있다.

최악의 경우, L = 0, H=전체레이어수 , 모든 반복에서 s < T 이면 모든 레이어를 순회하게 된다.

-> O(H × ||D||)

e.g., Llama2-13B (40layers) 와 calibration data 10개 사용하면, 최대 inference 횟수는 400번이기 떄문에

single GPU 환경에서 몇분내로 완료할 수 있다.

 

 

3. Experiments

3.1. Models

Llama2-7B, 13B

Baichuan2-7B, 13B (중국어, 영어)

 

3.2. Benchmarks

평가 툴: OpenCompass evaluation framework

 

- Reasoning: CMNLI, HellaSwag, PIQA

- Language: CHID, WSC

- Knowledge: CommonSenseQA, BoolQ

- Examination: MMLU, CMMLU

- Understanding: Race-Higt/Midddle, XSum, C3

 

제로샷이거나 few샷 (추가 학습 x)

 

Evaluation

- perplexity(PPL), generation(GEN) for CHID, XSum, WSC

(평가 툴인 OpenCompass에 따라 점수가 변환되어, 높은 점수가 좋은 성능을 의미한다)

 

3.3. Baselines

SOTA structured pruning 방법들을 선택함.

- LLM-Pruner, SliceGPT (여기 두 모델은 SparseGPT를 능가한 methode들임)

 

3.4. Settings

Hyperparameter Setting

Appendix A.

 

Calibration data

Llama2 : English Widipedia 에서 랜덤 10개 

Baichuan2 : eng/cn wikipedia에서 각각 랜덤 5개

- Eng: English Widipedia 에서 랜덤 10개 

- Cn: Chinese Widipedia

 

GPU

8 Nvidia A100 80GB GPU를 사용

 

3.5. Main Results

 

- 다른 baseline 모델들과 비교했을 때, LaCo가 pruning 비율이 더 높음에도 불구하고 약간 더 좋은 성능을 보이고 있다.

 

- Reasoning 능력은 약간 떨어지지만, 그래도 비슷한 수준이다.

- 전체적으로 LoCo의 성능이 우수하다. 원본 모델의 80%정도의 성능 유지를 하고 있음. (반면 다른 baseline은 70%도 넘지 못함)

 

Apendix D.5

- 주목할 점은, GEN 모드로 평가한 세 가지 벤치마크(CHID, XSUM, WSCG)에서 LaCo로 pruning한 LLM은 비교적 안정적인 성능을 유지하는 반면, 기존 방식으로 pruning한 모델들은 성능이 저하되어 일부 결과는 0.00까지 떨어졌다는 것

- 기존 방식으로 pruning한 모델은 의미 없는 반복 출력을 생성하는 경향 (Table.23)

- Llama2-70B에서도 outperform한 결과

 

결론적으로 LaCo는 우수한 pruner이며, 

 모델의 내부 구조를 변경하지 않고, 파라미터의 차이와 추가에만 의존하기 때문에 간결하고 효율적인 pruning 방법이다.

 

3.6. Comparison of Perplexity

- Llama2-7B

27% sparsity

500 sentences selected from Wikipedia (length of 512 tokens)

 

 

 

 

3.7. Pruning Time

llama2-7B / 27% sparsity / A100 GPU

모델 로딩, 데이터 로딩, 모델 저장 시간은 제외하고 주요 pruning 과정만 측정

- LaCo는 더 낮은 시간 복잡도와 더 빠른 pruning 속도

 

3.8. Memory Usage and Inference Speed

llama2-13B / English Wiki dataset / bf16 / A100 GPU

-> consume less memory / achieve faster inference speed !

- baseline들은 dense 모델보다 추론속도가 느려졌다 (오호..) 반면 LaCo는 그런 문제 없음.

 

 

 

4. Further Analysis

4.1. Post-training and Re-pruning

4.1.1. Post-training

purning으로 인한 불가피한 성능 손실 때문에, LaCo 모델이 원래 모델의 파라미터를 잘 보존하고 잇으며, post-training으로 회복이 가능한지 테스트하였다.

llama2-7b / Baichuan2-7b

- LLaMA-Factory framework를 사용해서 post-training을 진행.

 

- 학습 과정에서 빠르게 수렴하며, 약 250 스텝 이후 손실이 급격히 감소한 뒤 안정화.

- 5B 크기의 pruned llama2-7B와 Baichuan2-7B 모델의 최종 convergence loss는 각각 1.6과 2.0으로, 이는 Llama2-7B(1.75)와 Baichuan2-7B(1.9)의 기술 보고서에 기재된 값과 상당히 유사

- Nvidia A100 80GB GPU 4개 사용, 학습 시간은 각각 약 28시간과 35시간

(참고로, 5B(50억) 파라미터 규모의 LLM을 처음부터 학습하려면 수백 개의 A100 GPU를 몇개월 동안 사용해야 함.)

 

---

[ 평가 ]

Appendix E

- llama2-7b의 경우, post-training을 진행했을 때 성능이 더 오름

   -> 일관된 점수 향상은, LaCo 방법으로 pruning된 모델이 원래 모델의 파라미터를 효과적으로 계승하고, low-cost post-training을 통해 성능을 회복할 수 있음을 보여줌

 

- 반면 baichuan2-7b의 경우, 일부 벤치마크에서는 향상하고 일부에서는 하락함

   -> 사전 학습 데이터가 다양한 출처를 포함하고 있어, 우리의 post-training 데이터와 데이터 분포가 다르기 때문에 post-training의 효과가 제한되었다고 추측

4.1.2. Re-pruning

post-training으로 성능을 회복시켰기 때문에, 여기서 더 purning하여 50%의 압축률이 가능할까?에 대한 실험을 진행.

- llama2-7b / 17 layers (55%)

Appendix E

- 원본 7b의 70%의 성능을 보존하는 결과가 나타남.

- 더 나은 data와 더 많은 data를 사용하면 더 좋은 결과가 나타날 것임.

 

4.2. Layer-wise Similarity

[ 가중치 유사도 분석 ]

가장 큰 L2 값이 200을 넘지 않는다. 즉 인접한 레이어끼리 매우 유사하다.

MLP matrix 사이즈(11008*4096) 와 SAN q,k,v 사이즈 (4096*4096)를 고려했을때, 인접한 레이어간의 값 변화는 작다는 것을 알 수 있다.

(L2 distance가 200이면, 각 원소당 평균적인 차이가 0.05정도인 것임)

 

[ 블럭 아웃풋 유사도 분석 (5-a) ]

3-28 layer에서 cosine 유사도가 거의 1에 가깝다.

 

[ 블럭 merge (5-b) ]

RDSC Layer Merge가 여러 층을 하나로 대체할 수 있음을 검증하기 위해 실험

- 10층부터 19층 사이의 연속된 4개 층을 하나로 병합

- 병합된 층의 출력과 원래 마지막 층의 출력 간 코사인 유사도를 평가 (뭔말이야이게)

4096차원 벡터에 대한 최저 코사인 유사도는 0.996 이상으로 나타나, RDSC Layer Merge가 표현을 잘 보존함을 확인하였다.

 

 

 

4.3. Varying Pruning Ratio

- llama2-7b / llama2-13b

- 10%, 25%, 50%

 

- pruning 비율이 커질수록 성능도 하락한다.

- 그러나 10-25%에서는 비슷한 성능을 유지하고 있기 떄문에, 이 range에서는 LaCo가 안정적으로 작동함.

- 50%의 ratio에서도 원본의 70% 성능을 유지하고 있다. 

 

5. Related Work

- Model Quantization

- Knowledge Distillation

- Model Pruning

 

6. Conclusion

이 논문에서는 Layer Collapse(LaCo)라는 간결한 층별 구조화 pruning 방법을 제안한다. LaCo는 모델의 뒤쪽 층들을 앞쪽 층에 병합하여 빠르게 모델 크기를 줄인다. LaCo는 특수 하드웨어 지원이 필요 없으며 모델의 고유 구조를 보존한다. 실험 결과 LaCo는 현재의 SOTA structured pruning 방법들보다 현저히 뛰어난 성능을 보였으며, 기존 LLM에서 잠재적인 파라미터 중복성을 드러냈다. 또한, 다양한 LaCo 설정에 대한 제거(ablation) 연구를 수행하였다. pruned 모델에 대해 post-training을 진행하여 LaCo가 원본 모델의 파라미터를 효과적으로 계승함을 확인했다. 아울러 층별 유사성 관점에서 동기를 논의하고, 서로 다른 pruning 비율에서 LaCo pruning 모델의 성능을 탐구하였다.

 

 

Limitations

1. layer-wise 이기 때문에, pruning 비율을 자유롭게 설정할 수 없다.

2. 𝛕(원본 모델과 merge된 모델의 유사도 threshold)와 같은 하이퍼파라미터를 조정해야 한다.

3. 기존 연구들(baselines)과 마찬가지로 이론적인 증명이 부족하다. (our method lacks a complete theoretical proof)