GradPruner: Gradient-Guided Layer Pruning Enabling Efficient Fine-Tuning and Inference for LLMs

기술 요약: GradPruner

문제 정의

대규모 언어 모델(LLM)을 다운스트림 데이터에 맞춰 미세 조정(Fine-tuning)하는 것은 종종 계산 비용이 많이 들고 시간이 오래 걸립니다. 구조적 프루닝(Structured pruning) 방법들이 추론 효율성을 높이기 위해 존재하지만, 이들은 일반적으로 추가적인 학습, 지식 증류(Knowledge distillation) 또는 복잡한 구조 탐색을 필요로 하며, 이는 효율적인 미세 조정이라는 목표를 저해합니다. 또한, 기존의 프루닝 접근 방식은 파라미터 중요도를 평가하기 위해 캘리브레이션 데이터를 사용하는 경우가 많습니다. 그러나 LLM은 도메인 특화 데이터에서 성능이 최적화되지 않은 모습을 보이는 경우가 빈번하여, 편향된 중요도 평가를 초래할 수 있습니다.

저자들은 세 가지 핵심 과제를 식별했습니다:

1. 메모리나 학습 시간을 늘리지 않으면서 특정 다운스트림 데이터에 맞춤화된 파라미터 중요도를 측정하는 방법 개발.
2. 원래 모델의 구조를 최대한 보존하면서 파라미터 프루닝을 극대화하는 것.
3. 풀 파인 튜닝(FFT)과 LoRA와 같은 매개변수 효율적 미세 조정(PEFT) 방법을 모두 지원하는 것.

방법론: GradPruner

GradPruner는 학습 및 추론 효율성을 모두 향상시키기 위해 설계된 그래디언트 유도형 레이어 프루닝 프레임워크입니다. 이 방법은 세 가지 주요 단계로 작동합니다:

1. 파라미터 중요도 평가 (IGIA-Matrix)

미세 조정 초기 단계에서 손실(Loss)이 급격히 감소한다는 관찰(이는 다운스트림 태스크의 빠른 학습을 의미함)에서 영감을 받아, GradPreler는 **초기 그래디언트 정보 축적 행렬(IGIA-Matrix)**을 활용합니다.

• 프로세스: 전체 학습 대신, 모델은 짧은 LoRA 미세 조정 단계(전체 학습 단계의 약 1%)를 거칩니다.
• 그래디언트 시뮬레이션: 가중치 $W$와 LoRA 어댑터 $W_A, W_B$를 가진 선형 레이어에 대해, 이 방법은 LoRA 그래디언트의 곱인 $\nabla W_{sim} = \nabla W_B \cdot \nabla W_A$를 계산함으로써 $W$의 그래디언트를 시뮬레이션합니다.
• 축적: IGIA-Matrix ($F_W$)는 초기 $t$ 단계 동안의 시뮬레이션된 그래디언트 제곱의 평균을 구하여 계산됩니다:
  $$F_W = \frac{1}{t} \sum_{i=1}^{t} (\nabla W_{sim}^{(i)})^2$$
• 레이어 스코어링: 특정 블록 내의 모든 선형 레이어에 대한 IGIA-Matrix 값을 합산하여 해당 레이어의 중요도 점수를 도출합니다.

2. 레이어 프루닝 (Layer Pruning)

계산된 중요도 점수를 바탕으로, 가장 중요도가 낮은 레이어를 식별하고 제거합니다. 저자들은 레이어의 최대 30%까지 프루닝하더라도 정확도에 미치는 영향이 미미하지만, 이 지점을 넘어서는 공격적인 프루닝은 심각한 성능 저하를 일으킨다고 언급했습니다.

3. 레이어 병합 (Layer Merging)

정확도 손실 없이 30% 이상의 프루닝 비율을 달er하기 위해, GradPruner는 레이어 병합 전략을 도입합니다. 프루닝된 레이어를 완전히 버리는 대신, 이를 남은 레이어와 병합합니다.

• 희소화 (Sparsification): 프로닝된 레이어는 먼저 IGIA-Matrix를 사용하여 희소화되며, 상위 $p\%$의 파라미터(크기 기준)만을 유지하고 나머지는 0으로 만듭니다.
• 부호 기반 병합 (Sign-Based Merging): 충돌하는 파라미터 부호(양수와 음수 값이 서로 상쇄되는 경우)로 인한 간섭을 방지하기 위해, 이 방법은 동일한 부호를 가진 요소만 병합합니다. 구체적으로, 프루닝된 레이어의 파라미터는 부호가 일치하는 경우에만 대응하는 잔존 레이어의 파라미터와 병합됩니다. 부호가 다르면 프루닝된 파라미터는 폐기됩니다.

주요 기여

1. GradPruner 프레임워크: 광범위한 캘리브레이션이나 전체 재학습 없이 파라미터 중요도를 평가하기 위해 LoRA 미세 조정의 초기 그래디언트를 사용하여 IGIA-Matrix를 계산하는 방법을 제안했습니다.
2. 부호 인지형 레이어 병합 (Sign-Aware Layer Merging): 프루닝된 레이어를 희소화하고 부호 일관성에 따라 잔존 레이어와 결합하는 새로운 병합 기술을 통해, 정확도를 유지하면서 더 높은 프루닝 비율(최대 40%)을 달성할 수 있게 합니다.
3. 실증적 검증: 의료, 금융 및 일반 벤치마크 전반에서 GradPruner가 40%의 파라미터 감소를 달성하면서도 다운스트림 태스크 정확도 저하를 0.99%로 억제했음을 입증하는 광범위한 실험을 수행했습니다.

실험 결과

저자들은 두 가지 LLM(Llama3.1-8B 및 Mistral-7B)을 8개의 다운스트림 데이터셋(PubMedQA, MedMCQA, BillSum, FinGPT, HellaSwag, WinoGrande, ARC, PIQA 포함)에 대해 평가했습니다.

• 정확도: GradPruner는 밀집(Dense) 모델 대비 정확도 저하가 0.99%에 불과하면서 평균 40%의 파라미터 감소를 달성했습니다. 몇몇 사례에서는 프루닝된 Llama3.1-8B가 더 작은 크기의 프루닝되지 않은 Llama3.2-3B보다 우수한 성능을 보였습니다.
• 효율성: 밀집 모델과 비교했을 때, GradPruner는 학습 시간을 약 36%, 학습 메모리를 약 34% 줄였습니다. 추론 시간과 메모리는 각각 약 39%, 39% 감소했습니다.
• 비교: GradPruner는 정확도와 효율성 지표 모두에서 다른 구조적 프루닝 베이스라인(LLMPruner, LaCo, MINITRON, SAT, APT)을 지속적으로 능가했습니다.
• 절제 연구 (Ablation Studies):
  • 레이어 병합 단계를 제거하면 30% 이상의 레이어를 프루닝할 때 상당한 정확도 저하가 발생했습니다.
  • 부호 기반 병합을 가중 평균(Weighted averaging)으로 대체하면 정확도가 낮아졌으며, 이는 부호 일관성의 중요성을 확인시켜 줍니다.
  • 이 방법은 학습 데이터를 원래 크기의 10%로 줄여도 견고함을 유지했으나, 극도로 적은 데이터 양(1%)에서는 성능이 약간 변동되었습니다.

의의 및 주장

본 논문은 GradPruner가 미세 조정 효율성과 모델 성능 사이의 트레이드오프 문제를 해결한다고 주장합니다. 미세 조정 초기 단계의 빠른 학습 역학을 활용함으로써, 전체 학습이나 복잡한 증류 과정의 오버헤드 없이 핵심 파라미터를 식별합니다. 저자들은 자신들의 부호 기반 병합 전략이 공격적인 프루닝과 관련된 간섭를 효과적으로 완화하여, 모델의 특화된 다운스트림 태스크 능력을 보존하면서도 실질적인 파라미터 감소(40%)를 가능하게 한다고 단언합니다. 이 연구는 레이어 수준의 프루닝과 지능적인 병합의 결 조합이 효율적인 도메인 적응형 LLM을 배포하기 위한 실행 가능한 경로임을 시사합니다.