[논문 리뷰] Shortened LLaMA: Depth Pruning for Large Language Models with Comparison of Retraining Methods

Bo-Kyeong Kim, Geonmin Kim외 5명|arXiv (Cornell University)|2024. 02. 05.

Natural Language Processing Techniques인용 수 8

한 줄 요약

LoRA 재훈련과 함께 Transformer 블록 제거를 통한 LLM의 깊이 가지치기는 제로샷 태스크에서 너비 가지치기에 필적할 수 있으며, 작은 배치 메모리 제약 조건에서 더 빠른 추론을 제공한다.

ABSTRACT

Structured pruning of modern large language models (LLMs) has emerged as a way of decreasing their high computational needs. Width pruning reduces the size of projection weight matrices (e.g., by removing attention heads) while maintaining the number of layers. Depth pruning, in contrast, removes entire layers or blocks, while keeping the size of the remaining weights unchanged. Most current research focuses on either width-only or a blend of width and depth pruning, with little comparative analysis between the two units (width vs. depth) concerning their impact on LLM inference efficiency. In this work, we show that simple depth pruning can effectively compress LLMs while achieving comparable or superior performance to recent width pruning studies. Our pruning method boosts inference speeds, especially under memory-constrained conditions that require limited batch sizes for running LLMs, where width pruning is ineffective. In retraining pruned models for quality recovery, continued pretraining on a large corpus markedly outperforms LoRA-based tuning, particularly at severe pruning ratios. We hope this work can help build compact yet capable LLMs. Code and models can be found at: https://github.com/Nota-NetsPresso/shortened-llm

연구 동기 및 목표

• 메모리 제약이 있는 소배치 설정에서 LLM 추론 효율성을 개선할 동기를 제시한다.
• 나머지 가중치 형태를 고정한 채 전체 Transformer 블록을 제거하는 간단한 깊이 가지치기 접근법을 제안한다.
• 공개 LLM(LLaMA-7B 및 Vicuna-7B/13B)에 대해 깊이 가지치기를 Wanda-sp, FLAP, LLM-Pruner 등의 너비 가지치기 기준선과 비교 평가한다.
• 깊이 가지치기와 LoRA 재훈훈을 결합하면 생성 속도를 높이면서 제로샷 태스크 성능이 경쟁력 있음을 보여준다.

제안 방법

• Transformer 블록을 가지치기 가능한 단위로 간주하여 추론 지연 시간을 줄인다.
• Mag, Taylor, PPL 기반 기준으로 블록 중요도를 평가하고, Taylor+ 및 PPL을 가지치기 결정에 선택한다.
• 가장 덜 중요한 블록을 제거하고 목표 매개변수 수를 충족시키기 위해 원샷 가지치기를 수행한다(처음 4개 블록과 마지막 2개 블록은 유지).
• 보정 데이터셋에서 LoRA(저랭크 적응)로 가지치기된 모델을 효율적으로 재훈련하여 성능의 빠른 회복을 가능하게 한다.
• 제로샷 태스크에서 깊이 가지치기를 너비 가지치기 기준선과 비교하고, 소배치 조건에서 지연 시간, 처리량, 메모리 사용량을 측정한다.

실험 결과

연구 질문

• RQ1 트랜스포머 블록의 간단한 깊이 가지치기가 대형 언어 모델의 제로샷 성능에서 너비 가지치기에 필적하거나 초과할 수 있는가?
• RQ2 메모리 제약이 있는 소배치 시나리오에서 깊이 가지치기된 모델이 자기회귀 생성에서 실제 속도 향상을 제공하는가?
• RQ3 블록 수준의 어떤 중요도 기준과 가지치기 세밀도가 정확도와 효율성 간의 최적의 트레이드오프를 제공하는가?
• RQ4 LoRA 재훈련으로 한 번의 원샷 깊이 가지치기가 반복 가지치기 접근법과 실제로 경쟁력이 있는가?

주요 결과

• 깊이 가지치기는 소배치 조건에서 원래 모델보다 더 빠른 생성 속도를 제공합니다.
• 깊이 가지치기는 LoRA 재훈련 시 Wanda-sp, FLAP, LLM-Pruner과 같은 너비 가지치기 기준선과 경쟁력 있는 제로샷 태스크 성능을 달성합니다.
• Taylor+는 상식 추론 정확도를 향상시키고, PPL은 생성 품질을 향상시킵니다.
• 깊이 가지치기된 모델은 GPU 메모리 요구를 줄이고 하드웨어 제약 설정에서 비가지치기 모델보다 더 큰 배치 크기나 더 긴 출력을 처리할 수 있습니다.
• LoRA 재훈련을 통한 원샷 가지치기는 반복 가지치기의 성능에 근접하여 효율적인 배포를 가능하게 합니다.
• 더 큰 규모에서 전체 Transformer 블록 가지치기가 개별 MHA/FFN 모듈 가지치기보다 일반적으로 더 나은 결과를 내며, 작은 모델의 경우에도 블록 가지치기가 여전히 유리합니다.