1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 은 추론 및 생성 성능이 뛰어나지만, 서비스 비용이 매우 높다는 문제가 있습니다. 이는 다음과 같은 요인들 때문입니다:

• 높은 메모리 및 지연 시간: 모델 파라미터 수, 어텐션 계산, 그리고 시퀀스 길이에 비례하여 증가하는 KV 캐시 (Key-Value Cache) 가 상당한 메모리 사용량과 지연 시간을 유발합니다.
• 기존 구조적 가지치기 (Structured Pruning) 의 한계:
  1. 분리된 기준: FFN(Feed-Forward Network) 채널과 KV 헤드 그룹을 별도의 기준으로 가지치기하여, 동일한 배포 예산 (Budget) 을 경쟁하는 구조적 단위 간의 최적 할당을 놓칩니다.
  2. 고정된 레이어별 할당: 많은 기존 방법들이 레이어별 고정 비율이나 깊이 의존적 스케줄을 사용하여, 예산을 레이어와 구조 유형 간에 유연하게 학습하지 못합니다.
  3. 불연속적 선택과 미세 조정의 분리: 먼저 중요도 점수를 학습하고 나중에 예산을 적용하는 방식은 최종 희소 아키텍처와 학습 목표 간의 연결을 약화시키며, 추가적인 재구성 (Reconstruction) 이나 미세 조정을 필요로 하여 엔지니어링 오버헤드를 증가시킵니다.

2. 제안 방법론: GRASPrune

저자들은 GRASPrune을 제안하며, 이는 사전 학습 (Pretraining) 후 적용되는 구조적 가지치기 프레임워크입니다. 핵심 아이디어는 단일 전역 예산 (Single Global Budget) 하에서 FFN 채널과 KV 헤드 그룹을 동시에 (Jointly) 가지치기하는 것입니다.

주요 기술적 요소:

1. 전역 예산 기반 구조적 가지치기 (Global Budgeted Structured Pruning):

   • FFN 채널과 KV 헤드 그룹을 하나의 전역 선택 문제로 공식화합니다.
   • 각 구조 단위에 비용 (Cost) 을 부여합니다 (FFN 채널 단위 비용 = 1, KV 헤드 그룹 단위 비용 = $\alpha$). $\alpha$는 파라미터 수를 기반으로 계산되어 두 구조가 동일한 예산 내에서 경쟁하도록 합니다.
   • 목표는 주어진 예산 $B$ 내에서 언어 모델링 손실 ($L$) 을 최소화하는 이진 마스크 $z$를 찾는 것입니다.

2. 예산 준수 투사 (Projected STE) 학습:

   • 핵심 차별점: 기존 방법들이 중요도 점수를 제약 없이 학습한 후 나중에 예산을 적용하는 것과 달리, GRASPrune 은 최적화 루프 내부에서 매 단계 예산 제약이 활성화된 상태로 점수를 학습합니다.
   • 프로세스:
     1. 연속적인 게이트 점수 (Gate Scores) 를 학습합니다.
     2. 매 스텝마다 현재 점수를 예산 제약을 만족하는 하드 마스크 (Hard Mask) 로 투사 (Project) 합니다. (점수를 비용 순으로 정렬하여 예산 범위 내까지 선택)
     3. Straight-Through Estimator (STE) 를 사용하여, 순전파 (Forward) 에서는 하드 마스크를 적용하고 역전파 (Backward) 에서는 부드러운 점수 (Soft Probabilities) 를 통해 그래디언트를 전달합니다.
   • 이 방식은 게이트 점수가 최종 배포 가능한 마스크의 선택 규칙과 동일한 제약 하에서 학습되도록 보장합니다.

3. 유지된 단위에 대한 스케일 보정 (Scaling Calibration):

   • 가지치기 후 남은 단위들의 스케일 불일치를 완화하기 위해, 가지치기된 마스크가 고정된 후 스칼라 승수 (Scaling Factors) 를 학습합니다.
   • 이 보정은 가지치기된 가중치에 통합되어 추론 시 추가 파라미터 없이 더 작은 밀집 (Dense) 체크포인트를 생성합니다.

4. 효율성:

   • 백본 (Backbone) 가중치는 고정된 채로, 소량의 레이블 없는 보정 데이터 (Calibration Data) 만을 사용하여 경량 게이트와 스케일 파라미터만 학습합니다.
   • LLaMA-2-7B 기준, 4 에포크 학습 및 보정 단계가 단일 A100 GPU 에서 약 6 분 만에 완료됩니다.

3. 주요 기여 (Key Contributions)

1. 단일 전역 가지치기 공식화: 서로 다른 단위 비용을 가진 FFN 채널과 KV 헤드 그룹을 하나의 전역 선택 문제로 통합하여, 단일 예산 제약 하에서 최적의 마스크를 학습합니다.
2. 예산 준수 투사 STE 학습: 학습 과정에서 매 단계 예산 실현 가능성 (Feasibility) 을 강제하는 프로젝션 기법을 도입하여, 중요도 점수가 실제 배포 제약 하에서 최적화되도록 합니다.
3. 보정을 통한 더 작은 밀집 체크포인트: 가지치기된 구조에 대한 경량 스케일 보정을 수행하여, 추론 시 추가 오버헤드 없이 성능을 회복하고 더 작은 모델을 생성합니다.

4. 실험 결과 (Results)

실험은 LLaMA-2-7B 를 중심으로 진행되었으며, LLaMA-3.1, Qwen 등 다른 모델에서도 검증되었습니다.

• 성능 (Perplexity 및 Zero-shot 정확도):

  • 50% 파라미터 유지 (0.5 Ratio): WikiText-2 에서 12.18 perplexity를 기록하며, 기존 SOTA 방법들 (LLM-Pruner, SliceGPT, FLAP) 보다 월등히 우수한 성능을 보였습니다. (예: LLM-Pruner 는 312.96, SliceGPT 는 26.10)
  • 40% 파라미터 유지 (0.4 Ratio): 고압축 영역에서도 16.65 perplexity 를 유지하며, 기존 방법들이 급격히 성능이 저하되는 상황에서도 안정적인 성능을 유지했습니다.
  • 다운스트림 태스크: ARC, HellaSwag, PIQA 등 5 개 벤치마크의 평균 Zero-shot 정확도에서도 모든 가지치기 비율에서 가장 높은 점수를 기록했습니다.

• 메모리 및 처리량 (Throughput):

  • KV 캐시 감소: KV 헤드 그룹 가지치기로 인해 컨텍스트 길이에 따른 메모리 증가율이 크게 감소했습니다.
  • 처리량 향상: 가지치기된 밀집 체크포인트는 배치 크기 (Batch Size) 가 커질수록 더 큰 처리량 (Tokens/s) 향상을 보여주었습니다.

• 학습 효율성:

  • 단일 A100 GPU 에서 약 6 분 만에 가지치기 마스크를 학습할 수 있어, 기존 재학습 (Fine-tuning) 기반 방법들에 비해 매우 효율적입니다.

5. 의의 및 결론 (Significance)

• 효율적인 LLM 압축: GRASPrune 은 복잡한 재학습이나 지식 증류 없이도, 전역적인 예산 할당과 경량 게이트 학습을 통해 LLM 의 구조적 압축을 가능하게 합니다.
• 배포 최적화: 단순히 파라미터 수만 줄이는 것이 아니라, KV 캐시 메모리 감소와 처리량 향상을 동시에 달성하여 실제 서비스 환경에서의 배포 비용을 절감합니다.
• 일반화 능력: LLaMA 계열뿐만 아니라 Qwen 등 다양한 최신 모델 아키텍처에서도 일관된 성능 향상을 보여주어, 구조적 가지치기 프레임워크로서의 확장성을 입증했습니다.

결론적으로, GRASPrune 은 전역적 예산 제약 하에서 구조적 단위를 통합적으로 선택하는 메커니즘을 통해, 고압축 영역에서도 언어 모델링 능력을 유지하면서 추론 비용과 지연 시간을 획기적으로 줄일 수 있는 실용적인 솔루션을 제시합니다.