Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning

논문 개요

대규모 언어 모델(LLM)의 사전학습은 막대한 컴퓨팅 자원과 시간을 요구한다. 이를 해결하기 위한 한 가지 접근법은 이미 잘 학습된 대형 모델로부터 소형 모델을 효율적으로 파생시키는 것이다. Xia et al.(2024)이 ICLR 2024에서 발표한 Sheared LLaMA는 이 아이디어를 구조적 프루닝(structured pruning)의 형태로 구체화하였다.

핵심 아이디어는 간단하다. LLaMA-2-7B처럼 강력하게 사전학습된 모델을 출발점으로 삼고, 중요도가 낮은 레이어·어텐션 헤드·FFN 중간 차원을 체계적으로 제거하여 목표 크기(1.3B 또는 2.7B)의 서브네트워크를 추출한다. 이후 소량의 토큰으로 계속학습(continued pre-training)을 수행하여 성능을 회복한다. 이 과정 전체가 동일 크기 모델을 처음부터 학습하는 비용의 수 퍼센트 수준에서 완결된다.

핵심 기여

목표 아키텍처 기반 구조적 프루닝: 단순히 임계값 기반으로 가중치를 제거하는 것이 아니라, 원하는 목표 아키텍처(레이어 수, 헤드 수, 차원)를 명시하고 이에 맞는 서브네트워크를 선택하는 방식을 사용한다.
Dynamic Batch Loading (DBL): 계속학습 단계에서 각 도메인 데이터의 손실 기반으로 배치 샘플링 비율을 동적으로 조정하여, 고정 비율 대비 더 균형 잡힌 성능을 달성한다.
학습 비용 절감 실증: LLaMA-2-1.3B 수준의 모델을 처음부터 학습하는 데 필요한 토큰의 약 3%만으로도 경쟁력 있는 성능을 얻을 수 있음을 보였다.

방법론 상세

1단계: 구조적 프루닝

목표 아키텍처 $T = (L_{T}, H_{T}, D_{T}, D_{T}^{FFN})$ 를 먼저 정의한다. 여기서 $L_{T}$ 는 레이어 수, $H_{T}$ 는 어텐션 헤드 수, $D_{T}$ 는 히든 차원, $D_{T}^{FFN}$ 는 FFN 중간 차원이다.

각 구조적 단위(레이어, 헤드, 차원 슬라이스)에 대해 중요도 점수(importance score)를 계산한다. 중요도는 마스크 기반 기울기 방법으로 추정된다:

$s_{i} = \frac{\partial L}{\partial m _{i}}$

여기서 $m_{i}$ 는 $i$ 번째 구조적 단위에 대한 이진 마스크, $L$ 은 소규모 캘리브레이션 데이터셋에서 계산된 손실이다.

이 점수를 기반으로 각 차원에서 상위 $k$ 개의 단위를 유지한다. 즉, 레이어 단위에서는 중요도 높은 $L_{T}$ 개의 레이어를, 어텐션 헤드 단위에서는 $H_{T}$ 개의 헤드를, FFN 뉴런 단위에서는 $D_{T}^{FFN}$ 개의 뉴런을 선택한다.

중요한 점은 이 프루닝이 **구조적(structured)**이라는 것이다. 즉, 개별 가중치가 아닌 전체 레이어나 헤드 단위로 제거되므로, 추가적인 희소 행렬 연산 라이브러리 없이도 표준 밀집 행렬(dense matrix) 연산으로 추론할 수 있다.

2단계: Dynamic Batch Loading (DBL)

프루닝 이후 서브네트워크는 대형 모델의 가중치를 이어받아 계속학습을 수행한다. 이때 훈련 데이터는 여러 도메인(웹 텍스트, 코드, 학술 문서 등)으로 구성된다.

고정된 도메인 혼합 비율을 사용하면 특정 도메인에 과적합되거나 과소학습되는 문제가 생긴다. DBL은 이를 해결하기 위해 각 도메인 $d$ 에 대한 현재 손실 $L_{d}$ 를 주기적으로 측정하고, 다음 배치에서의 샘플링 비율 $p_{d}$ 를 손실에 비례하게 조정한다:

$p_{d} \propto L_{d} - L_{d}^{*}$

여기서 $L_{d}^{*}$ 는 해당 도메인의 목표 손실(대형 모델의 성능에 기반). 현재 손실이 목표보다 높은 도메인에 더 많은 배치를 할당함으로써, 모든 도메인에 걸쳐 균형 있게 학습이 진행된다.

목표 아키텍처 선택

논문에서는 두 가지 목표 크기를 실험하였다:

모델	레이어 수	어텐션 헤드	히든 차원	FFN 차원
Sheared-LLaMA-1.3B	24	16	2048	5504
Sheared-LLaMA-2.7B	32	32	2560	6912

두 크기 모두 기존 LLaMA-2-7B(32레이어, 32헤드, 4096 히든)에서 구조적으로 축소된 것이다.

실험 결과

다운스트림 과제 성능

Sheared-LLaMA-2.7B는 동일 크기로 처음부터 학습된 모델(OpenLLaMA-3B, LLaMA-2-3B 등)과 비교하여 대부분의 벤치마크에서 동등하거나 우수한 성능을 보였다. 특히 다음 과제들에서 두드러진 결과를 보였다:

ARC-Easy/Challenge: 상식 추론 과제
HellaSwag: 문장 완성 과제
MMLU: 지식 집약적 다지선다 과제
BoolQ: 이진 질의응답 과제

학습 효율성

방법	학습 토큰 수	추정 계산량(GPU-시간)
처음부터 학습 (1.3B)	~1T	기준
Sheared-LLaMA-1.3B	~50B	~5%

약 50B 토큰(처음부터 학습 대비 약 5%)만으로도 유사한 수준의 성능에 도달할 수 있음을 보였다. 이는 강력한 사전학습 모델에서 지식 증류가 얼마나 효율적으로 이루어질 수 있는지를 보여준다.

Ablation 연구

DBL의 효과: DBL 없이 고정 비율로 학습하면 특정 도메인(특히 코드)에서 성능이 저하된다.
프루닝 전략 비교: 레이어만 제거하거나 헤드만 제거하는 방식보다 모든 차원을 함께 제거하는 통합 접근법이 더 우수하다.
계속학습 토큰 수: 50B 토큰 이후에는 성능 향상이 둔화되며, 더 많은 토큰이 반드시 더 좋은 성능을 보장하지 않는다.

의의 및 한계

의의

자원 효율적 LLM 개발: 고성능 대형 모델을 보유한 조직은 이 방법을 통해 다양한 크기의 효율적인 소형 모델을 빠르게 파생시킬 수 있다.
구조적 프루닝의 실용화: 비구조적 프루닝과 달리 특수 하드웨어 없이 표준 추론 환경에서 즉시 활용 가능하다.
DBL의 일반성: DBL은 Sheared LLaMA에 국한되지 않고, 다양한 도메인 혼합 학습 시나리오에 적용 가능한 범용 기법이다.

한계

원본 모델 의존성: Sheared LLaMA의 성능은 출발점이 된 LLaMA-2-7B의 질에 크게 의존한다. 더 약한 기반 모델에서 시작하면 결과도 제한적일 수 있다.
아키텍처 제약: 목표 아키텍처가 반드시 원본 모델의 서브구조여야 한다. 완전히 다른 아키텍처로의 변환은 불가능하다.
긴 문맥 일반화: 프루닝 이후 계속학습에 사용된 문맥 길이가 제한적이어서, 매우 긴 문맥에서의 성능 저하 가능성이 있다.
지식 손실: 대형 모델의 전체 용량으로 학습된 지식 중 일부는 소형 아키텍처의 용량 한계로 인해 필연적으로 손실된다.

도형준 지식 볼트

탐색기