Streamlining Redundant Layers to Compress Large Language Models

ICLR 2025 Spotlight

Code

👀 요약 👀
LLM-Streamline
1. Layer pruning
    * cosine 유사도로 제거할 연속된 레이어 선택
   (다른 메트릭은 벡터의 크기를 고려하기 때문에 제외, ppl은 data dependency높아서 제외)
2. Layer replacement
     * 대체하는 레이어 아키텍처: FNN, Transformer block(원본모델이랑 동일한 구조)
     * finetuning:  대체 맨앞레이어의 input과 대체 맨뒤레이어의 아웃풋 히든벡터로 학습
     * finetuning 할 때 loss: MSE loss


* 아예 제거하고 싶은 레이어 개수를 정해두고, 그 간격끼리의 cosine sim을 구한다. 
예를들면 7개를 제거할 목적이라면, 0번째와 6번째 hidden vector의 cos sim, 1-7의 cos sim .... 이런식으로 구하고, 가장 유사도가 높은 애를 구해서 뭉텅이로 날려버리는 방법. (공식 코드 참고)




✨ 볼만한 부분✨
.* .. 따라서 벡터 크기의 영향을 받지 않는 코사인 유사도를 선택함 .
(흐음..코사인유사도는 높지만, magnitude 유사도가 작으면 이건 어떻게 고려할건디?? - '벡터의 유사함'에 대한 좀 더 명확한 이유가 있으면 좋을듯)
* 반면 LoRA를 사용하여 프루닝 후 모델을 학습시키는 과정은, 제거된 레이어의 기능을 남은 레이어에 재분배하는 과정
으로 볼 수 있다.
* - 모든 프루닝 방식이 GSM8K 벤치마크를 잘 못잡고 있음!!
* calibratin/train dataset을 SlimPajama만 사용함
* -> FFN 레이어가 transformer layer보다 수렴이 빠르다

* 벤치마크 많이 씀
* Stability라는 새로운 메트릭 제안

--
오픈리뷰 리뷰어들 리뷰가 매우 좋다... 왤까 궁금하다.
리뷰까지 읽어봤는데.. 그냥 세세한 비교실험들이 많아서 그런 것 같음 ..
더 다양한 (종류/사이즈) 모델에 대한 실험 결과 비교, LoRA와 비교 실험 -> 리뷰지적받고 추가됨

 

 

Abstract로 흐름 파악하기

가장 덜 중요한 레이어를 식별하고 제거하는 LLM의 layer pruning에 대한 연구로, LLM-Streamline을 제안한다.

 

LLM-Streamline은 두 단계로 나뉜다.

1. layer pruning: 가장 덜 중요한 연속된 레이어들을 제거하는 방법

2. layer replacement: lightweight network를 학습하고 pruned layer를 대체하는 방식. 성능 손실을 완화하기 위함이다.

 

추가적으로, Stability라는 새로운 metric을 제안한다.

이 메트릭은 model compression 테스크에서 accuracy만 사용하는 한계를 극복하기 위함이다.

 

다른 sota pruning method를 성능과 학습효율성 측면에서 아웃퍼폼한 결과를 보였다. 

 

---

1. Introduction

LLM의 사이즈가 커지면서, 하드웨어의 요구가 상당히 심해지고, 따라서 real-world scenario에 적용하기에 제약이 되고 있다. 이런 제약을 없애기 위해, model compression을 통해 높은 성능을 유지하면서 compact한 모델을 만들고자 하는 연구들이 쏟아지고 있다.

model compression => {kd, quantization, pruning}

Knowledge distillation achieves compression by transferring the capabilities of a larger teacher model to a smaller student model. Quantization compresses the model by quantizing the weights to lower precision. Alternatively, pruning compresses the model by eliminating unimportant parameters and modules.

 

이번 연구에서는 popular pruning method에 포커싱한다. 이전 프루닝 연구에서의 프루닝 단위는 dense matrices (SliceGPT), attention heads, filters, parameters 등이 있다. 이러한 방법들이 효과적이긴 하지만, 모델 구조의 불규칙성(structural irregularity)을 초래하는 경우가 많아, 프루닝된 모델을 저장하거나 배포하기에 inflexible하다는 한계가 있다.

반면 layer pruning method는 단순히 LLM의 depth를 줄이는 방법이다. nn.ModuleList 와 같은 데이터구조 안에 저장되어있는 레이어를 단순하게 제거하면 되는 아주 간단한 방법이다. 따라서 효율적인 layer-wise pruning 방법을 탐구하는 것은 중요하다.

 

layer purning은 LLM에서 덜 중요한 레이어를 찾고 없애는 방법이다. 구체적으로, 각 레이어는 Hidden states를 변환하는 역할을 한다고 볼 수 있으며, 따라서 특정 레이어의 Input/output hidden state의 유사도가 높다면, 레이어의 영향이 작다고 할 수 있다.

 

관련연구

without further training - SLEB, ShortGPT

with finetuning - Shortened llama, LaCo, Gromov at al.

그러나 레이어를 직접 제거하는 경우 성능 저하가 더 크게 발생할 수 있다.

또한 LoRA(Hu et al., 2021)와 같은 parameter-efficient fine-tuning 기법을 사용해 프루닝된 LLM을 학습할 수 있지만,

원래의 비연속적인(non-contiguous) 레이어들이 성능 저하를 보완하도록 모델을 미세조정하는 과정은 쉽지 않다(2.3 참조)

 

본 연구에서는 LLM-Streamline이라는 프루닝 방법을 제안한다.

적은 학습데이터로, 높은 성능, 학습효율성을 동시에 달성할 수 있음.

   1. 레이어 프루닝 2. 레이어 replacement 

lightweight network는 다양한 구조를 가질 수 있다. (FFN, SwiGLU, Transformer 등..)

 

추가적으로, 모델 compression method를 평가하는 metricd의 한계를 발견하였다. 

구체적으로, multiple-choice classification를 포함하는 NLU 과제에서, 압축된 모델은 원래 모델이 불확실했던 샘플에 대해 우연히 정답을 맞추는 경우가 발생할 수 있다. 이러한 현상은 성능을 과대평가하게 만드는 문제를 초래한다.

이를 해결하기 위해, 새로운 metric인 stability을 제안한다. 이 metric은 프루닝 전후 예측 일관성(consistency)을 측정하며, 특히 원래 모델의 prediction confidence를 함께 고려한다.

2. LLM-Streamline

 

2.1. Layer Redundancy in LLMs

그냥 LLM 레이어 식

- 각 레이어의 input/output hidden vector를 cosine similarity로 유사도 계산

- layer importace를 측정하기 위한 데이터는 pre-training data에서 랜덤샘플링 해옴.

그냥 레이어마다 hidden vector의 cosine sim 구하는 것 식으로 쓴 거

 

- 모델 사이즈와 종류에 따른 영향을 완화하고자, 4개 모델을 사용함. 

 -> 모든 모델에서, 연속된 레이어들의 input/output유사도가 높게 나타남. 

 

Discussion I: Why not use other similarity to measure the importance of layers? 

내적(dot product)와 유클리드 거리도 사용되지만, 이들은 벡터의 크기를 추가적으로 고려함.

에 따르면, pre-norm 구조를 사용하는 transformer의 hidden states는, 레이어의 깊이가 증가함에 따라 점점 커지는 경향을 보인다.

이로 인해 후반 레이어에서는 dot product similarity가 높아지고,

초반 레이어에서는 duclidean distance가 작아지는 편향이 발생한다.

 

따라서 벡터 크기의 영향을 받지 않는 코사인 유사도를 선택함 ..

(흐음..코사인유사도는 높지만, magnitude 유사도가 작으면 이건 어떻게 고려할건디?? - '벡터의 유사함'에 대한 좀 더 명확한 이유가 있으면 좋을듯)

 

 

Discussion II: Why not use perplexity as the metric to measure the importance of layers?

ppl을 사용하는 이전 연구에서는,각 레이어를 하나씩 제거하면서, pre-training 데이터에서 모델의 퍼플렉서티 변화를 측정하고, 퍼플렉서티 변화가 가장 작은 레이어를 제거하는 방식으로 진행한다.

그러나  ppl 지표는 매우 data-sensitive하다고 판단, 즉, 서로 다른 사전학습 데이터를 사용할 경우 제거되는 레이어가 달라지며,

결과적으로 프루닝에 사용된 데이터에서는 퍼플렉서티가 낮더라도, 다른 데이터셋에서는 성능이 저하되는 문제가 발생한다.

반면 cosine similarity는 매우 안정적이며, 항상 동일한 레이어가 선택되는 consistency를 보인다 .

Appendix A

2.2. Layer Pruning (step1)

 

 

 

2.3. Layer Replacement (step2)

 

Discussion: Layer Replacement of Fine-Tuning Pruned LLMs?

우선, resource overhead 관점에서, layer replacement가 다른 방법에 비해 hardware 리소스 제약이 적다. PEFT 방법들은 모델의 모든 weight, activation value, PEFT모듈의 optimizer 상태 등을 gpu에 올려야 한다. 반면 layer replacement 방법은, 첫 번째 단계에서는 모델 가중치와 순전파(forward) 연산 오버헤드만 저장하고, 두 번째 단계에서는 lightweight network 가중치, activation 값, 옵티마이저 상태만 저장하면 된다.

두번째로, MSE 손실 함수를 사용해 제거된 레이어의 지식을 경량 네트워크에 증류(distill)한다.
반면 LoRA를 사용하여 프루닝 후 모델을 학습시키는 과정은, 제거된 레이어의 기능을 남은 레이어에 재분배하는 과정으로 볼 수 있다.
따라서 프루닝된 레이어를 경량 네트워크로 대체하는 것이, 남은 레이어에 기능을 재분배하는 것보다 학습 난이도가 낮을 수 있다.

3. Metrics for Evaluating Pruned Models

 

3.1. Shortcoming of Accuracy Metric

TP와 TN의 표준편차(std) 가 FN과 FP보다 현저히 높게 나타났다.
→ 이는 모델이 FN과 FP 샘플에 대해 상대적으로 불확실성(uncertainty) 이 크다는 것을 의미

3.2. Stability Metric

 

정확도(accuracy)와 달리, stability는 모델의 답변에 대한 신뢰도(confidence)와 프루닝 전후 모델의 일관성(consistency)에 초점을 맞춘다.

따라서 stability는 프루닝 후 모델이 원래 모델과 최대한 유사하게 유지되는가라는 모델 압축의 본래 목표에 더 부합하는 지표이다.

 

4. Experiments

4.1. Setup

Model

Llma2-7B, 13B

 

pruning ratio : 25%

lightweight network: 1) FFN (랜덤init) 2) Transformer layer(프루닝된 첫번째 레이어의 파라미터를 따름)

calibration dataset : SlimPajama

데이터를 무작위 샘플링하여, 최종적으로 30,000개의 데이터로 구성된 데이터셋

 500개의 샘플을 무작위로 선택하여 LLM에 입력하고, Fig. 2를 생성했으며, 이 500개의 데이터 샘플은 레이어 프루닝(layer pruning)에 사용되었다. 나머지 30,000개의 데이터는 경량 네트워크(lightweight network) 학습에 사용

 * SlimPajama: CommonCrawl 약 52.2 %, C4 약 26.7 %, GitHub 약 5.2 %, Books 약 4.2 %, arXiv 약 4.6 %, Wikipedia 약 3.8 %, StackExchange 약 3.3 %.

 

 

4.2. Benchmark

12개 NLU task.

CMNLI, HellaSwag, PIQA, CHID, WSC, CommonSencseQA, BoolQ, MMLU, CMMLU, Race-High/Middle/ C3

추가 3개. (OpenCompass 프레임워크 사용)

XSum, GSM8K, StrategyQA

4.3. Baseline

LLM-Pruner
SliceGPT

LaCo

4.4 Main Results

-> benchmark들에 대한 accuracy 결과

 

-> benchmark들에 대한 Stability (얘네들이 제안한 metric) 결과

 

- 모든 프루닝 방식이 GSM8K 벤치마크를 잘 못잡고 있음!!

 

 

OPT-1.3B, OPT-2.7B, OPT-6.7B, Baichuan-7B, Baichuan-13B, Baichuan2-7B, Baichuan2-13B(Yang et al., 2023), Llama3.1-8B, Llama3.1-70B(Dubey et al., 2024), Mixtral-8x7B-v0.1(Jiang et al., 2024)에서도 실험을 수행 (Appendix E)

 

Appendix E

-> purning ratio=50%

 

 

 

4.5. Impact of Different Lightweight Networks

Why FFN achieves the best result, Transformer layer still has performance potential. 

- 다양한 lightweight network 구조에 대한 실험

1) FNN 2)SwiGLU 기반 FNN 3)Transformer layer

 

+ Transformer layer를 초기화하는 방법

3-1) 랜덤 3-2) 첫번째 프루닝된 레이어 상속 3-3)마지막 프루닝된 레이어 상속 3-4)프루닝레이어들 평균

 

=> FFN이 가장 우수한 성능. 한편, Transformer 레이어에서는 프루닝된 첫 번째 레이어를 상속한 경우가 가장 좋은 결과를 나타냈다. 반대로 LaCo에서 영감을 받은 Layer-Avg는 가중치 평균화가 프루닝된 첫 번째 레이어만큼 효과적이지 않음을 보여준다

(LaCo는 평균이 아니라 차이를 더해주는 건데...)

 

-> FFN 레이어가 수렴이 빠르다

 

4.6. Impact of Different Pruning Ratios

The performance of the pruned model is linearly correlated with the number of parameters at modest pruning ratios.

파라미터 수와 선형적으로 성능이 저하되는 모습을 보였고, 이는 LLM-Strimeline 방법으로 프루닝된 모델의 성능이 동일한 파라미터 수를 가진 사전 학습 모델과 비교할만하다 (견준다)는 것을 시사한다.

(에엥 라마 3b 랑 왜 비교안함? 논리 깨져서 그런듯 ㅋㅋ 이래도돼?)

 

4.7. Comparison of Layer Replacement and LoRA

Layer Replacement outperforms LoRA in both performance and GPU memory consumption

- layer replacement는 LoRA와 학습 목적이 다르므로, 추자거으로 1epoch LM loss로 학습을 진행한 결과임.

- layer replacement는 30,000개의 데이터, LoRA는 300,000개의 데이터로 학습함.

- LoRA의 rank는 비슷하게 맞추기 위해 128로 설정

LoRA보다 항상 우수함. 즉 훨씬 적은 GPU 메모리와 학습 데이터를 요구하는 방법임.

 

AppendixE.8.

[layer replacement 학습 데이터 개수에 관하여]

SlimPajama-6B 전체로 post training을 했을 때, 성능이 약간 오르긴 하지만 computational time이 100배 증가한 것에 비하면 그저 그렇다.

 

5. Related Work

LLM-Streamline과 동시에 진행된 레이어 프루닝 관련 연구에는 LaCo (Yang et al., 2024), ShortGPT (Men et al., 2024), UIDL (Gromov et al., 2024), SLEB (Song et al., 2024), Shortened Llama (Kim et al., 2024) 등이 있다.

• LaCo (Yang et al., 2024) 는 연속된 여러 레이어를 하나의 그룹으로 묶고, 그들의 파라미터를 평균 내어 레이어를 압축(compress) 한다. (아니 laco 평균 아니잔아!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!)
• ShortGPT (Men et al., 2024) 는 코사인 유사도(cosine similarity) 와 동일한 BI 점수를 사용하여 레이어의 중요도를 평가하고, 덜 중요한 레이어를 제거한다.
• UIDL (Gromov et al., 2024) 역시 코사인 유사도에 해당하는 각 거리(angular distance) 를 이용해 덜 중요한 레이어를 제거하며, 성능 향상을 위해 QLoRA 를 함께 사용한다.
• SLEB (Song et al., 2024) 는 퍼플렉서티(perplexity) 를 통해 레이어의 중요도를 계산하고, 중요하지 않은 레이어를 제거한다.
• Shortened Llama (Kim et al., 2024) 는 다양한 레이어 선택 기준(metric) 을 탐색하고, 프루닝 이후 연속 사전학습(continual pre-training) 과 LoRA 적용의 효과를 분석한다.

전통적인 레이어 프루닝 기법과 달리, LLM-Streamline은 프루닝된 레이어를 단순히 제거하거나 프루닝된 모델을 재학습(retrain)하는 대신, 가벼운 대체 모델(lightweight model)을 학습시켜 그 레이어를 대체한다.

 

 

6. Conclusion

본 논문에서는 LLM-Streamline이라는 LLM용 레이어 프루닝 및 대체(layer pruning-and-replacement) 알고리즘을 제안한다.
또한 기존의 정확도(accuracy) 지표의 한계를 지적하고, 모델 압축 성능을 평가하기 위한 새로운 지표인 stability를 제안한다.
광범위한 실험 결과, 가벼운 네트워크(lightweight network) 를 활용한 본 레이어 대체 방식은 기존의 SOTA 프루닝 방법들을 능가하며, 동시대의 다른 레이어 프루닝 기법들보다 효율성과 성능 모두에서 우수한 결과를 보였다.