LEAP: Learnable End-to-End Adaptive Pruning of Large Language Models

기술 요약: LEAP: 대규모 언어 모델의 학습 가능한 엔드투엔드 적응형 가지치기

문제 정의
대규모 언어 모델 (LLM) 의 배포는 점점 더 메모리 및 연산 제약에 의해 병목 현상을 겪고 있습니다. 가중치 희소성 (weight sparsity) 이 해결책을 제시하지만, 그 환경은 구조화, 반구조화 (예: 2:4), 비구조화 (unstructured) 영역으로 나뉩니다. 최근 GPU 커널 (예: FlashLLM, SpInfer) 과 데이터 흐름 하드웨어의 발전으로 병목 현상이 비구조화 희소성을 수행하는 것에서 최소한의 정확도 손실로 이를 유도하는 것으로 이동했습니다.

비구조화 가지치기를 위한 현재 최첨단 방법들인 Wanda, SparseGPT, Thanos, ADMM 은 최적의 뇌 외과 의사 (Optimal Brain Surgeon, OBS) 원칙에서 파생된 레이어별 대리 모델 (layer-wise surrogates) 에 의존합니다. 이러한 방법들은 엔드투엔드 모델 손실이 아닌 국소 재구성 오차를 최소화합니다. 결과적으로 이들은 심층 네트워크에서 누적되는 국소 오차를 발생시켜, 특히 공격적인 희소성 (50~60%) 하에서 상당한 정확도 저하를 초래합니다.

반면, MaskLLM 과 PATCH 와 같은 엔드투엔드 학습 가능한 마스크 방법들은 언어 모델링 손실에 직접적으로 대항하여 마스크를 성공적으로 최적화하지만, 반구조화 패턴으로 제한됩니다. 그들의 파라미터화는 그룹 내 유효 패턴에 대한 범주형 분포 (categorical distribution) 에 의존합니다. 작은 그룹 (예: 2:4 희소성을 위한 $\binom{4}{2}=6$개의 패턴) 에는 실용적이지만, 이 접근 방식은 비구조화 희소성에서는 조합적으로 처리 불가능해집니다. 50% 희소성에서 너비 $d=4096$인 행의 경우, 유효 마스크의 수는 $\binom{4096}{2048} \approx 10^{1229}$로, 패턴 로짓 (logits) 의 저장 및 색인이 불가능합니다.

방법론: LEAP
저자들은 비구조화 마스크 학습을 실용적으로 만들기 위해 패턴에 대한 범주형 파라미터화를 가중치별 (per-weight) 접근법으로 대체하는 프레임워크인 LEAP(Learnable End-to-End Adaptive Pruning) 을 소개합니다.

• 가중치별 베르누이 완화 (Per-Weight Bernoulli Relaxation): 모든 유효 패턴에 로짓을 할당하는 대신, LEAP 은 각 개별 가중치에 학습 가능한 로짓을 할당합니다. 가중치 행렬 $W \in \mathbb{R}^{m \times n}$에 대해, 병렬 로짓 행렬 $P \in \mathbb{R}^{m \times n}$이 도입됩니다. 확률적 마스크 $M$은 Gumbel-시그모이드 완화를 통해 생성됩니다:
  $$M = \sigma\left(\frac{\alpha P + g}{\tau}\right)$$
  여기서 $g$는 Gumbel 노이즈, $\sigma$는 시그모이드 함수, $\alpha$는 스케일 인자, $\tau$는 온도 파라미터입니다. 이 공식은 유효 패턴의 수가 아닌 가중치 수 ($O(mn)$) 에 비례하여 선형적으로 확장됩니다.
• 최적화 전략:
  • 동결된 가중치 (Frozen Weights): 사전 학습된 가중치 $W$는 동결된 상태로 유지됩니다. 마스크 로짓 $P$만 훈련됩니다. 이는 사전 학습된 모델의 보정 (calibration) 을 유지하고 배포를 단순화합니다.
  • 초기화: $P$는 일회성 Wanda 마스크 (선택된 항목을 $+s$로, 나머지를 $-s$로 설정) 를 사용하여 초기화되어 합리적인 시작점을 제공하고 콜드 스타트 문제를 피합니다.
  • 스케줄: 스케일 스케줄 ($\alpha$) 은 시그모이드를 이진 값으로 밀어내기 위해 증가하고, 온도 스케줄 ($\tau$) 은 마스크를 선명하게 만들기 위해 감소합니다.
  • 손실 함수: 목적 함수는 보정 스트림의 언어 모델링 손실과 두 가지 정규화 항을 결합합니다:
    1. 전역 희소성 정규화 (Global Sparsity Regularizer): 모든 레이어에 걸쳐 전역적으로 목표 밀도를 강제하며, 개별 레이어가 엔드투엔드 중요도에 따라 밀도를 조정할 수 있게 합니다.
    2. 크기 인지 항 (Magnitude-Aware Term): 중요한 가중치가 제거되는 퇴화 최소값 (degenerate minima) 을 방지하기 위해 더 큰 크기의 가중치를 유지하도록 최적화를 편향시킵니다.
• 훈련 효율성: 이 방법은 작은 일반 텍스트 보정 스트림 (SlimPajama) 에서 약 2,000 회 반복 훈련을 수행합니다.

주요 기여

1. 조합적 장애물 식별: 이 논문은 MaskLLM/PATCH 에서 사용하는 패턴에 대한 범주형 파라미터화가 유효 패턴의 기하급수적 증가로 인해 비구조화 희소성으로 확장될 수 없음을 공식적으로 식별합니다.
2. 실용적 재형성: 저자들은 LLM 규모에서 비구조화 마스크를 위한 엔드투엔드 미분 가능 학습을 가능하게 하는 Gumbel-시그모이드를 통한 가중치별 베르누이 재형성을 제안합니다.
3. LEAP 프레임워크: 동결된 가중치에서 작동하며 소규모 보정 데이터셋과 약 2,000 단계의 훈련만 요구하는 경량 가지치기 프레임워크입니다.
4. 실증적 검증: 0.5B 에서 8B 파라미터에 이르는 다섯 가지 LLM 계열 (Qwen-2.5, Gemma-3, LLaMA-3.2, LLaMA-3.1) 에 대해 50% 및 60% 희소성으로 광범위한 평가가 수행되었습니다.

결과
LEAP 은 레이어별 베이스라인 (Wanda, SparseGPT, Thanos, ADMM) 과 반구조화 학습 가능한 베이스라인인 MaskLLM 에 대해 평가되었습니다.

• 정확도 향상: 열 개의 (모델, 희소성) 설정에서 LEAP 은 최상의 레이어별 베이스라인인 ADMM 대비 6 개 작업 평균 제로샷 정확도를 +2.59 포인트개선했습니다.
• 구체적 개선: 가장 큰 개선은 60% 희소성에서 LLaMA-3.2 1B 에서 +5.40 포인트 (50.39 대 44.99) 였습니다.
• 퍼플렉시티: LEAP 은 베이스라인보다 일관되게 낮은 퍼플렉시티 (PPL) 를 달성했습니다. 예를 들어, 50% 희소성에서 LLaMA-3.1 8B 의 경우 LEAP 은 7.66 PPL 을 달성한 반면 ADMM 은 9.12 였습니다.
• 희소성 할당: 전역 예산을 가진 엔드투엔드 최적화 하의 실험은 LEAP 이 반구조화 방법에서 보이는 비균일 변동과 대조적으로 거의 균일한 블록당 밀도 할당으로 수렴함을 보여주었습니다.
• 연산 비용: 일회성 방법 (H100 에서 몇 분) 보다 비싸지만 (H100 에서 몇 시간), LEAP 은 MaskLLM (수천 개의 A100 시간) 보다 훨씬 저렴합니다 (예: LLaMA-3.1 8B 의 경우 276 GPU 시간 대 MaskLLM 의 수천 시간).

의의 및 주장
이 논문은 LEAP 이 비구조화 희소성에서 엔드투엔드 마스크 학습을 방해하는 근본적인 장애물을 해결한다고 주장합니다. 마스크 학습을 가중치 업데이트와 분리함으로써, 사전 학습된 가중치의 보정을 희생하지 않고 파운데이션 모델을 압축할 수 있는 실용적인 경로를 제공합니다.

저자들은 LEAP 이 기존 가속 커널 (예: FlashLLM, SpInfer, MACKO) 과 직접 호환되는 50%~60% 비구조화 영역의 마스크를 생성한다고 강조합니다. 의의는 새로운 추론 커널을 만드는 데 있는 것이 아니라, 레이어별 방법으로는 이전까지 달성할 수 없었던 고품질 가지치기 모델에서 이러한 커널의 배포를 가능하게 하는 데 있습니다. 이 연구는 엔드투엔드 학습을 통해 유도된 비구조화 희소성이 상용 하드웨어 배포를 위해 압축과 정확도 사이의 유리한 트레이드오프를 제공할 수 있음을 시사합니다.

한계 및 범위

• 동결된 가중치: 현재 범위는 가중치를 동결된 상태로 유지합니다. 가중치와 마스크의 공동 최적화는 자연스러운 확장으로 언급되지만, 이 작업에서는 구현되지 않았습니다.
• 균일 할당: 균일한 블록당 밀도에 대한 관찰은 연구된 모델과 희소성 수준에 대한 영역별 발견으로 제시되며, 비균일 할당에 대한 일반적 규칙은 아닙니다.
• 안전성: 저자들은 압축된 모델이 롱테일 입력에서 다른 실패 모드를 보일 수 있음을 경고하며, 안전이 중요한 배포의 경우 재평가를 권장합니다.