MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

기술 요약: MaskPro – LLM 을 위한 엄격한 (N:M) 희소성에 대한 선형 공간 확률적 학습

문제 정의
대규모 언어 모델 (LLM) 의 급속한 확장으로 추론 효율성에 심각한 병목 현상이 발생하여 효과적인 압축 기술이 필요해졌습니다. 반구조적 희소성, 즉 매 M 개의 연속된 가중치 중 N 개의 요소를 유지하는 (N:M) 패턴은 가속화 및 메모리 감소를 위한 하드웨어 친화적인 해결책을 제공합니다. 그러나 최적의 (N:M) 마스크를 식별하는 것은 formidable 한 조합 최적화 문제입니다. 기존 접근 방식은 고유한 한계를 가진 두 가지 범주로 나뉩니다:

1. 규칙 기반 방법: 이러한 방법은 수동으로 설계된 지표 (예: 계층별 재구성 오차 최소화) 를 사용하는 탐욕적이고 계층별 검색을 수행합니다. 계산 비용은 저렴하지만 내재된 편향과 엔드 투 엔드 작업 성능 간의 상당한 격차로 고통받습니다.
2. 학습 기반 방법 (예: MaskLLM): 이러한 방법은 생성 손실에 대한 경사 하강을 통해 마스크를 직접 최적화합니다. 효과적이지만 prohibitive 한 훈련 비용을 초래합니다. 구체적으로 MaskLLM 은 M 개의 가중치 그룹 내의 모든 가능한 마스크에 대한 로짓을 저장해야 하므로 메모리 복잡도가 $O(\binom{M}{N} \frac{d}{M})$가 됩니다. M 이 크면 이는 기하급수적으로 확장되어 대규모 모델에는 실현 불가능합니다. 또한, 광활한 조합 공간은 정책 경사에서의 높은 분산을 유발하여 훈련 불안정성을 초래합니다.

방법론: MaskPro
저자들은 (N:M) 희소성을 효율적이고 안정적으로 학습하도록 설계된 새로운 선형 공간 확률적 프레임워크인 MaskPro를 제안합니다.

1. 선형 공간 확률적 완화:
   M 개의 가중치 그룹에 대한 모든 $\binom{M}{N}$개의 가능한 마스크에 대한 범주형 분포를 정의하는 대신, MaskPro 는 문제를 재구성합니다. 각 그룹 내의 M 개 개별 요소에 대한 범주형 분포를 설정합니다. 그런 다음 (N:M) 마스크는 이 분포에서 N-way 비복원 샘플링을 수행하여 생성됩니다.

   • 이론적 근거: 저자들은 임의의 (N:M) 마스크가 N 개의 서로 다른 기저 벡터의 확률적 합으로 표현될 수 있음을 증명합니다.
   • 메모리 효율성: 이 재구성은 학습 가능한 매개변수 (로짓) 의 수를 $O(\binom{M}{N} \frac{d}{M})$에서 선형인 $O(d)$로 줄입니다. 여기서 $d$는 총 가중치 수입니다. 이는 이전 학습 기반 방법과 관련된 메모리 폭발을 제거합니다.

2. 정제된 정책 경사 추정기 (PGE):
   로짓을 최적화하기 위해 MaskPro 는 정책 경사를 활용합니다. 저자들은 미니배치 차이로 인한 높은 분산 (데이터 샘플링으로 인한 손실 변동이 마스크 샘플링으로 인한 손실 변동을 종종 압도함) 으로 인해 LLM 에서 원시 손실 값을 사용하는 표준 정책 경사가 불안정하다는 점을 확인했습니다.

   • 손실 잔차: 원시 손실 $f(m \odot w, \xi)$를 사용하는 대신, MaskPro 는 초기 기준 마스크 $m_0$에 대한 손실 잔차를 기반으로 로짓을 업데이트합니다: $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$. 이는 마스크 변경의 영향을 분리합니다.
   • 이동 평균 추적기: 훈련을 더욱 안정화하고 수치적 폭발을 방지하기 위해 기준선을 부드럽게 하는 이동 평균 추적기 $\delta$가 도입됩니다. 업데이트 규칙은 다음과 같습니다:
     $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
   • 이론적 보장: 이 논문은 이러한 정제된 추정기가 실제 정책 경사의 불편성 (unbiased) 추정기로 남아 있음을 보여주는 증명을 제공합니다. 또한, 샘플링된 마스크의 손실이 기준선 손실의 두 배 미만인 조건 하에서 (실제에서 쉽게 충족됨), 정제된 추정기의 분산은 순식간 정책 경사보다 엄격하게 낮습니다.

주요 기여

1. 선형 공간 프레임워크: MaskPro 는 M 개 요소에 대한 범주형 분포에서 N-way 비복원 샘플링으로 선택 과정을 모델링함으로써 (N:M) 희소성 학습에 필요한 메모리 요구량을 기하급수적 $O(\binom{M}{N} \frac{d}{M})$에서 선형 $O(d)$로 줄입니다.
2. 향상된 정책 경사: 저자들은 손실 잔차와 이동 평균 추적기를 사용하여 분산이 감소된 업데이트 메커니즘을 도입하여, 대규모 조합 공간에서 정책 경사의 높은 분산으로 인한 훈련 불안정성을 해결합니다.
3. 이론적 분석: 추정기의 불편성과 제안된 업데이트 규칙의 분산 감소 특성에 대한 포괄적인 증명이 제공됩니다.

실험 결과
다양한 LLM(Gemma-7B, Vicuna-7B, LLaMA-2-7B, DeepSeek-7B) 에서 다양한 하위 작업 (예: HellaSwag, RACE, PIQA) 과 희소성 패턴 ((2:4), (4:8), (8:16)) 에 걸쳐 광범위한 실험이 수행되었습니다.

• 성능: MaskPro 는 최신 규칙 기반 방법 (SparseGPT, Wanda, Pruner-Zero) 보다 일관되게 우월한 성능을 발휘하며, 메모리 집약적인 MaskLLM 과 비교 가능한 성능을 달성합니다. 종종 정확도에서 평균 2% 이상의 개선을 보입니다.
• 메모리 확장성: MaskLLM 은 7B 모델의 (4:8) 희소성에서 8x A100 GPU 에서 Out-of-Memory(OOM) 오류를 발생시키는 반면, MaskPro 는 희소성 패턴에 관계없이 안정적인 메모리 사용량 (~36-37 GB) 을 유지합니다.
• 데이터 효율성: MaskPro 는 데이터 양에 대한 놀라운 견고성을 보여줍니다. MaskLLM 이 수렴하기 위해 수십만 개의 샘플이 필요한 반면, MaskPro 는 단 1 개의 데이터 샘플로만 훈련되어도 안정적인 성능을 달성합니다.
• 훈련 효율성: 이 방법은 MaskLLM 에 비해 훨씬 적은 훈련 샘플과 훈련 시간이 필요하면서도 규칙 기반 방법과 비교 가능한 메모리 사용량을 유지합니다.

의의 및 주장
이 논문은 MaskPro 가 학습 기반 반구조적 희소성의 메모리 및 안정성 병목 현상을 해결함으로써 LLM 의 실용적이고 확장 가능한 로컬 배포를 위한 실현 가능한 경로를 제공한다고 주장합니다. 저자들은 그들의 접근 방식이 학습 기반 방법의 높은 성능과 규칙 기반 방법의 효율성 사이의 격차를 해소한다고 강조합니다. 메모리 발자국을 선형 공간으로 줄이고 정제된 정책 경사를 통해 훈련 과정을 안정화함으로써 MaskPro 는 이러한 작업과 관련하여 이전에 prohibitive 했던 계산 오버헤드 없이 대규모 모델에서 최적의 (N:M) 마스크를 학습할 수 있게 합니다. 이 방법은 데이터 부족과 모델 규모에 강건한 확률적 마스크 학습을 위한 일반화 가능한 프레임워크로 제시됩니다.