← 소식 목록

목표 구조 없이 LLM을 줄이는 법: 유전 알고리즘을 품은 가지치기, ElasticPruner

연구 목적

대규모 언어 모델(LLM)은 파라미터가 수십억 개라 메모리와 추론 지연이 큽니다. 구조적 가지치기(structured pruning)는 레이어, 어텐션 헤드, 채널 같은 덩어리를 통째로 지워 작은 밀집(dense) 모델을 만듭니다. 그래서 일반 GPU에서도 바로 빨라집니다.

문제는 "얼마나, 어디를 지울지" 정하는 일입니다. 기존 방법에는 두 가지 한계가 있습니다.

  • 목표 구조에 의존한다. 레이어 수, 히든 차원, FFN 중간 차원, 헤드 수를 미리 정해야 합니다. 보통 같은 크기로 처음부터 학습한 모델(예: Pythia)의 구조를 빌려 옵니다. 하지만 원하는 크기의 참조 모델이 없을 수도 있고, 빌린 구조가 원본 모델에 최적이라는 보장도 없습니다.
  • 로컬 미니멈에 빠진다. 제약 항은 모델 크기만 봅니다. 그래서 경사하강법은 예산은 맞지만 좋지 않은 구조 조합에 일찍 머물 수 있습니다.

핵심 아이디어

ElasticPruner는 전체 파라미터 수 하나만 제약으로 둡니다. 나머지 구조 배분은 학습이 스스로 정합니다.

  1. 미분 가능한 마스크 학습. 레이어, 어텐션 헤드, 히든 차원, FFN 중간 차원에 마스크를 두고 hard-concrete 분포로 완화해 end-to-end로 학습합니다. 마스크로 계산한 미분 가능한 파라미터 수를 라그랑지안 제약 항에 넣어 목표 크기에 맞춥니다.
  2. 유전 알고리즘(GA)으로 두 번째 목표 추가. 학습 도중 주기적으로 GA를 돌립니다. 각 개체는 구조 크기를 비트 벡터로 표현합니다. 적합도는 원본 모델과의 logit KL-divergence입니다. 즉 "원본과 출력이 가장 비슷한 구조"를 찾습니다. 찾은 결과는 마스크에 logit shift로 주입됩니다. 선택된 위치는 올리고 나머지는 내립니다.
  3. 가우시안 감쇠. shift 세기는 warm-up 시점에 가장 크고 이후 줄어듭니다. 초중반에는 GA가 큰 그림을 바꾸고, 후반에는 경사하강법이 세부를 다듬습니다.
  4. 동적 GA 간격. 연속된 GA 결과가 비슷하면(Sim > 0.9) 간격을 1.5배로 늘리고, 많이 다르면 0.5배로 줄입니다. 초기 간격은 400 step입니다.

실험 설정

  • 원본 모델: LLaMA-2-7B. 약 1.3B와 2.7B로 줄였습니다.
  • 학습 데이터: Sheared LLaMA 저자들이 공개한 샘플링된 RedPajama. GA의 calibration 데이터도 여기서 무작위로 뽑았습니다.
  • 비교 대상: Sheared LLaMA, LLM-Pruner. 원 논문 수치를 인용하지 않고 공개 구현으로 전 과정을 직접 재현했습니다. 복원 파인튜닝은 모든 방법을 같은 데이터, 0.4B 토큰으로 통일했습니다.
  • 평가: lm-evaluation-harness로 11개 벤치마크(MMLU, BoolQ, HellaSwag, WinoGrande, ARC-E, ARC-C, SciQ, LogiQA, PIQA, Natural Questions, LAMBADA). 평가 시드 3개 평균입니다.

결과

아래 표는 논문의 주요 결과표에서 일부 벤치마크와 평균만 옮긴 것입니다. mid-GA는 기본 설정(동적 간격, 가우시안 스케줄)이고, last-only는 GA를 마지막에만 쓴 변형입니다.

MethodHS (10)WGARC-ELAMBADAAverage
LLaMA2-7B78.669.376.428.864.6
Sheared LLaMA 1.3B28.0549.4934.2627.0135.84±0.05
LLM-Pruner 1.4B32.1550.7547.1434.7639.24±0.01
Ours 1.3B (GA last-only)30.9950.6741.6739.1439.15±0.08
Ours 1.3B (mid-GA)31.0650.7544.0238.6439.47±0.08
Sheared LLaMA 2.7B32.0649.3342.3039.6739.94±0.06
Ours 2.7B (GA last-only)39.1253.2054.2153.3344.73±0.05
Ours 2.7B (mid-GA)39.3456.1252.5754.2045.63±0.04
  • 1.3B에서 mid-GA는 평균 39.47로, Sheared LLaMA 1.3B(35.84)보다 3.63점 높았습니다. LLM-Pruner 1.4B(39.24)보다도 조금 높았습니다.
  • 2.7B에서 mid-GA는 평균 45.63으로, Sheared LLaMA 2.7B(39.94)보다 5.69점 높았습니다.
  • 두 규모 모두 ElasticPruner는 같은 크기의 Sheared LLaMA보다 11개 벤치마크에서 통계적으로 유의하게 앞섰습니다(one-sided Wilcoxon signed-rank test, p<0.05).

찾아낸 구조도 다릅니다. 예산만 주었는데도 Sheared LLaMA가 빌려 쓴 구조보다 "얕고 넓은" 구조가 나왔습니다.

MethodLhIH
Sheared LLaMA 1.3B242048550416
Ours 1.3B (last-only)20281653329
Ours 1.3B (mid-GA)183200315620
Sheared LLaMA 2.7B322560691220
Ours 2.7B (last-only)223840592023
Ours 2.7B (mid-GA)243584547625

특히 1.3B에서 경사하강법만 쓰면(last-only) 헤드 수가 9개까지 줄었지만, GA를 중간에 넣으면 20개로 회복되었습니다. 논문은 이를 GA가 마스크를 로컬 미니멈에서 꺼내 준다는 근거로 봅니다. 또 2.7B에서 같은 마스크에 Sheared LLaMA 구조를 적용하면 평균 42.06, GA로 찾은 구조를 적용하면 44.73이었습니다.

동적 간격은 2.7B에서 GA를 4회만 실행해, 항상 상한 8회를 쓰는 고정 간격보다 탐색 비용을 절반으로 줄였습니다. 1.3B에서는 둘 다 8회였고 성능 차이는 0.29점이었습니다.

한계와 다음 단계

논문이 직접 밝힌 한계는 세 가지입니다.

  • 실험은 LLaMA-2-7B 한 모델에서만 했습니다. 다른 LLM에서도 통하는지는 아직 검증하지 않았습니다.
  • 복원 파인튜닝이 0.4B 토큰으로 적습니다. 더 오래 학습해도 성능 격차가 유지되는지는 확인하지 못했습니다.
  • 동적 GA 스케줄링은 각 모델 규모에서 고정 간격과 성능·비용을 동시에 통제해 비교하지 못했습니다. 체계적인 검증은 향후 과제로 남겼습니다.