기술 요약: AMP를 통한 효율적인 LLM: 어텐션 헤드 및 MLP 프루닝(Pruning)

문제 정의

대규모 언어 모델(LLM)은 인지적 작업에서 인간 수준의 성능을 달amat지만, 방대한 파라미터 수로 인해 높은 계산 비용과 느린 추론 속도라는 문제를 안고 있습니다. 이러한 요인들은 자원이 제한된 환경에서의 배포를 저해합니다. 프루닝(Pruning)은 모델 크기를 줄이는 유망한 전략이지만, 기존 방법들은 다음과 같은 상당한 트레이드오프(trade-offs)에 직면해 있습니다:

• 구조적 프루닝(Structured Pruning): 전체 구성 요소(예: 레이어, 헤드)를 제거하지만, 종종 비용이 많이 드는 미세 조정(fine-tuning)을 요구하거나 중요한 구성 요소가 제거될 경우 성능 저하를 초축합니다.
• 반구조적 프루닝(Semi-Structured Pruning): 속도 향상을 위해 특수 하드웨어(예: 2:4와 같은 특정 GPU 희소성 패턴)를 필요로 하며, 이는 소비자용 하드웨어에서의 배포를 제한합니다.
• 비구조적 프루닝(Unstructured Pruning): 불규칙한 희소 행렬을 생성하여, 특수 소프트웨어 최적화 없이는 일반 목적의 하드웨어에서 추론 가속화를 달성하지 못합니다.
• 일반적 한계: 기존의 많은 기술은 계산 비용이 많이 들거나, 유연성이 떨어지거나, 특수 하드웨어를 필요로 하는 손실 기반 메트릭(loss-based metrics), 테일러 전개(Taylor expansions) 또는 복잡한 휴리스틱에 의존합니다. 따라서 성능 손실을 최소화하면서 실질적인 추론 속도 향상을 이끌어낼 수 있는 하드웨어 불가지론적(hardware-agnostic)이고 효율적인 방법이 필요합니다.

방법론: AMP (Attention Heads and MLP Pruning)

저자들은 멀티 헤드 어텐션(MHA) 내의 덜 중요한 어텐션 헤드와 멀티 레이어 퍼셉트론(MLP) 내의 뉴런을 제거하는 구조적 프루닝 방법인 AMP를 제안합니다. 단순히 가중치 크기나 손실 변화에만 의존하는 방법과 달리, AMP는 각 구성 요소가 잔차 스트림(residual stream)에 기여하는 바를 평가하기 위해 입력 데이터를 가중치에 투영하여 구조적 중요도를 평가합니다.

핵심 메커니나즘

이 방법은 정보가 잔차 스트림을 통해 흐른다는 원리에 기반합니다. AMP는 각 구성 요소의 중요도 점수를 계산할 때, 레이어의 최종 출력에 대한 각 구성 요소의 활성화 기여도 크기를 측정합니다.

1. MHA 구성 요소:

   • MHA 출력은 개별 어텐션 헤드의 기여 합으로 분해됩니다 ($O = \sum h_n W_n$).
   • $n$번째 헤드의 중요도($I_n$)는 헤드의 출력($h_n$)과 그에 대응하는 출력 투영 행렬($W_n$)의 곱의 $\ell_1$-노름(norm)으로 정의됩니다:
     $$I_n = \|h_n W_n\|_1$$
   • 이 접근 방식은 단순한 가중치 노름(weight-norm) 프루닝과 달리, 가중치와 입력 사이의 비선형 상호작용을 고려합니다.

2. MLP 구성 요소:

   • 이 방법은 LLaMA와 같은 모델에서 사용되는 SwiGLU 메커니즘, 특히 Up 및 Gate 프로젝션을 대상으로 합니다.
   • SwiGLU는 요소별 곱셈(element-wise multiplication)을 사용하므로, 구조적 무결성을 유지하기 위해 뉴런을 두 프로젝션에서 쌍으로 프루닝해야 합니다.
   • $m$번째 뉴런 쌍의 중요도($I_m$)는 샘플 입력 토큰들에 대해 Gate 프로젝션과 Up 프로젝션 출력의 요소별 곱의 평균 $\ell_1$-노름으로 계산됩니다:
     $$I_m = \frac{1}{S} \sum_{s=1}^{S} | \text{SiLU}((x_s W_{\text{Gate}})_m) \cdot (x_s W_{\text{Up}})_m |$$

3. 프루닝 프로세스:

   • 평균 활성화 값을 추출하기 위해 작은 캘리브레이션 세트(예: Alpaca에서 무작위 추출한 50개 샘플)를 사용합니다.
   • 모든 레이어에 걸쳐 구성 요소(헤드 및 뉴런)의 중요도를 순위 매깁니다.
   • 목표 압축 비율 $c$를 달ach하기 위해, 하위 $c\%$의 어텐션 헤드와 하위 $c\%$의 MLP 뉴런을 모든 레이어에서 균등하게 제거합니다.
   • 사후 학습(Post-Training): 프루닝된 모델은 성능 회복을 위해 저계수 적응(Low-Rank Adaptation, LoRA)을 이용한 경량 미세 조정을 거칩니다.

주요 기여

• 새로운 중요도 메트릭: AMP는 전통적인 가중치 전용 또는 손실 기반 휴리스틱과 달리, 구성 요소의 중요도를 평가하기 위해 입력 활성화를 가중치에 투영하는 데이터 기반 메트릭을 도입했습니다.
• 하드웨어 불가지론적 속도 향상: 이 방법은 특수 하드웨어나 희소 텐서 코어를 필요로 하지 않고도 실질적인 추론 속도 향상(최대 1.25$\times$)을 달성하여, 소비자용 GPU에서도 사용이 가능합니다.
• 효율성: 프루닝 결정 과정은 소비자용 GPU에서 단 몇 분 만에 완료되며, 이후의 미세 조정은 전체 재학습에 비해 매우 적은 자원을 소모합니다.
• 일관성 검증: 저자들은 "일관성 체크(coherence check)" 실험을 제안하고 실행하여, AMP가 식별한 가장 중요한 구성 요소를 제거했을 때 치명적인 성능 저하가 발생함을 보여줌으로써 메트릭의 능력을 입증했습니다.
• 유연성: 이 방법은 다양한 LLM 제품군(LLaMA, Phi)에 적용 가능하며, 고정된 희소성 패턴(예: 2:4)에 얽매이지 않고 유연한 압축 비율을 허용합니다.

실험 결과

저자들은 LLaMA 7B, LLaMA-2 7B, Phi-1.5, Phi-2 모델에 대해 제로샷 벤치마크(WinoGrande, HellaSwag, ARC-e/c, PIQA) 및 퍼플렉시티(WikiText2)를 사용하여 AMP를 평가했습니다.

• 성능:
  • LLaMA 7B: 20% 프루닝 비율에서 AMP는 평균 정확도 **64.52%**를 달성하여, LLM-Pruner(59.71%) 및 Shortened LLaMA(63.46%)와 같은 최신 기술(SOTA)을 능가했습니다.
  • LLaMA-2 7B: 30% 프루닝 비율에서 AMP는 **61.02%**의 평균 정확도를 기록하여, SliceGPT(51.50%) 및 LLM Surgeon(59.03%)을 앞질렀습니다.
  • 개선 사항: AMP는 기존의 구조적 프루닝 기술보다 평균 태스크 정확도에서 최대 1.49 퍼센트 포인트 더 높은 성능을 보였습니다.
• 추론 속도:
  • AMP는 LLaMA 7B에서 1.25$\times$, LLaMA-2 7B에서 1.19$\times$의 속도 향상을 달성했습니다.
  • 이는 Shortened LLaMA(1.13$\times$)보다 우수하며, 특수 하드웨어 없이도 Wanda(1.24$\times$)와 같은 반구조적 방법과 대등한 성능을 보여줍니다.
• 퍼플렉시티: 미세 조정된 프루닝 모델은 프루닝되지 않은 버전보다 WikiText2에서 현저히 낮은 퍼플렉시티를 보였으며, 높은 프루닝 비율에서 그 격차가 더 커졌습니다.
• 절제 연구(Ablation Studies):
  • 일관성 체크: 프루닝 순서를 뒤집어(가장 중요한 구성 요소를 제거) 실행했을 때 정확도가 25 퍼센트 포인트 이상 급락하여, 메트릭의 타당성을 확인했습니다.
  • 구성 요소 격리: MHA 헤드만 프루닝하거나 MLP 뉴런만 프루닝하는 것은 두 가지를 동시에 프루닝하는 것보다 성능이 훨씬 낮았으며, 이는 균형 잡힌 접근 방식의 필요성을 나타냅니다.

의의 및 주장

본 논문은 AMP가 다음과 같은 특징을 가진 모델 압축의 핵심적인 트레이드오프를 해결한다고 주장합니다:

1. 효과적임: 높은 압축률을 유지하면서도 정확도 측면에서 현재의 최신 구조적 프루닝 방법들을 능가합니다.
2. 효율적임: 프루닝 결정(몇 분 소요)과 성능 회복(경량 LoRA 미세 조정) 모두에 최소한의 계산 자원을 필요로 합니다.
3. 실용적임: 표준 하드웨어에서 실질적인 추론 속도 향상을 가능하게 하여, 자원이 제한된 환경에서의 배포를 용이하게 합니다.
4. 지속 가능함: 추론 지연 시간과 파라미터 수 감소를 통해 운영 비용을 약 17.39%, 탄소 배출량을 약 19.97% 줄임으로써 Green AI 원칙에 부합합니다.

저자들은 자신들의 방법이 비필수적인 구조를 성공적으로 식별하고 제거함으로써, 특수 하드웨어 인프라 없이도 고도로 최적화된 LLM을 배포할 수 있게 한다고 강조합니다. 또한 프루닝 후 성능을 복구하는 데 있어 사후 학습(미세 조정)의 결정적인 역할을 강조합니다.