[논문 리뷰] Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning

Mengzhou Xia, Tianyu Gao외 2명|arXiv (Cornell University)|2023. 10. 10.

Topic Modeling인용 수 20

한 줄 요약

논문은 LLM-shearing을 소개합니다. 이는 사전 학습된 더 큰 모델을 대상으로 한 Targeted structured pruning과 이어지는 사전 학습(continued pre-training), 데이터 도메인 간 불균형을 해소하기 위한 동적 배치 로딩(dynamic batch loading)을 통해 더 작고 경쟁력 있는 언어 모델을 만들어내는 방법입니다.

ABSTRACT

The popularity of LLaMA (Touvron et al., 2023a;b) and other recently emerged moderate-sized large language models (LLMs) highlights the potential of building smaller yet powerful LLMs. Regardless, the cost of training such models from scratch on trillions of tokens remains high. In this work, we study structured pruning as an effective means to develop smaller LLMs from pre-trained, larger models. Our approach employs two key techniques: (1) targeted structured pruning, which prunes a larger model to a specified target shape by removing layers, heads, and intermediate and hidden dimensions in an end-to-end manner, and (2) dynamic batch loading, which dynamically updates the composition of sampled data in each training batch based on varying losses across different domains. We demonstrate the efficacy of our approach by presenting the Sheared-LLaMA series, pruning the LLaMA2-7B model down to 1.3B and 2.7B parameters. Sheared-LLaMA models outperform state-of-the-art open-source models of equivalent sizes, such as Pythia, INCITE, OpenLLaMA and the concurrent TinyLlama models, on a wide range of downstream and instruction tuning evaluations, while requiring only 3% of compute compared to training such models from scratch. This work provides compelling evidence that leveraging existing LLMs with structured pruning is a far more cost-effective approach for building competitive small-scale LLMs

연구 동기 및 목표

• 기존 대규모 사전 학습 모델에서 강력하고 중간 규모의 LLM을 비용 효율적으로 생성하려는 동기를 제공한다.
• 가지치기 방법을 개발하여 검증된 사전 학습 구성과 일치하는 대상 아키텍처를 얻는다.
• 계속된 사전 학습 중 데이터 도메인 간의 지식 유지 불균형을 다룬다.
• 도메인별 손실 감소율을 기반으로 데이터를 도메인에 따라 할당하는 동적 배치 로딩을 도입한다.

제안 방법

• 소스 모델을 특정 대상 아키텍처로 모양을 바꾸는 대상화된 구조화 가지치기를 제안한다(레이어, 히든 차원, 헤드, 중간 차원에서 가지치기 마스크를 사용).
• 하드-콘크리트 분포와 최소-최대 목적을 통한 제약 최적화를 통해 가지치기 마스크를 학습하여 대상 형태를 만족시킨다.
• 가지치기된 모델에 대해 계속된 사전 학습을 수행하여 성능을 회복하거나 초과한다.
• 도메인별 손실 감소율에 따라 훈련 중 도메인 데이터 비율을 조정하는 동적 배치 로딩을 도입한다.
• 도로의 두 단계 프로세스: 가지치기 후 계속된 사전 학습으로 도메인 인식 데이터 전략에 guided된다.

실험 결과

연구 질문

• RQ1 대규모 사전 학습 LLM을 지정된 대상 아키텍처로 가지치면 더 적은 계산으로도 경쟁력 있는 작은 LLM을 얻을 수 있는가?
• RQ2 가지치기 후 계속된 사전 학습이 가지치기만으로 얻는 성능 회복보다 더 나은가?
• RQ3 동적 배치 로딩이 도메인별 손실 감소를 균형 있게 조정하여 전체 하위다운 스트림 성능을 개선하는가?

주요 결과

• Sheared-LLaMA 모델(1.3B 및 2.7B)은 유사 규모의 최신 오픈 소스 모델보다 11개의 다운스트림 태스크 및 지시문 튜닝 벤치마크에서 더 높은 성능을 발휘한다.
• LLaMA2-7B에서 1.3B/2.7B로의 가지치기는 가지치기 및 계속된 사전 학습에 약 50B 토큰만 사용했고, 처음부터 학습하는 데 필요한 계산의 약 3% 수준으로 경쟁력 있는 결과를 달성했다.
• 동적 배치 로딩은 도메인 간 손실 감소를 정렬하고 더 어려운 도메인에 대한 데이터 사용량을 늘려 하위 스트림 성능을 향상시킨다.
• 일부 비교에서 가지치기된 모델은 단순히 비슷한 크기의 기존 모델에서 시작하는 것보다 계속된 사전 학습을 위한 초기화가 더 우수하다.
• 대상화된 구조화 가지치기는 동등한 희소성에서 비균일 가지치기 방식보다 더 높은 추론 처리량을 제공한다.