Batch-wise Adaptive Pruning: Periodic Neuron Activation-Aware Weight Pruning for Language Reasoning Model

기술 요약: 대규모 추론 모델을 위한 배치 단위 적응형 프루닝 (Batch-wise Adaptive Pruning)

문제 정의
대규모 추론 모델(LRM)은 긴 사고 사슬(chain-of-thought) 생성을 통해 복잡한 작업에서 높은 성능을 달-성하지만, 상당한 계산 비용을 발생시킵니다. 프로덕션 환경에서는 높은 처리량을 위해 배치 추론(여러 샘플을 동시에 처리하는 것)이 필수적입니다. 그러나 TEAL 및 CATS와 같은 기존의 훈련 불필요(training-free) 적응형 프루닝 방식은 이 환경에서 심각하게 실패합니다.

핵심 실패 메커니즘은 두 가지입니다:

1. 분포 변화(Distribution Shift): 기존 방식들은 전체 배치에 대해 단일 프루닝 마스크를 사용합니다. 이들은 개별 샘플의 활성화 값(예: 평균)을 집계하고, 단일 샘플 분포에서 오프라인으로 보정된 임계값(threshold)을 적용합니다. 이러한 집계는 활성화 분포를 변화시켜, 임계값이 더 이상 데이터와 일치하지 않게 만들며, 결과적으로 실현된 희소성 비율(sparsity ratio)의 드리프트와 추론 정확도의 붕괴를 초래합니다.
2. 정적 대 동적 트레이드오프(Static vs. Dynamic Trade-off): 정적 프루닝 방식(고정된 패턴)은 긴 추론 과정 중에 진화하는 활성화 패턴을 추적할 수 없습니다. 반대로, 매 토큰 단계마다 새로운 마스크를 선택하는 것은 계산 비용이 너무 커서 속도 향상 효과를 상쇄해 버립니다.

방법론
저자들은 배치 추론을 위해 특별히 설계된 훈련 불필요 방식인 **배치 단위 적응형 프루닝(Batch-wise Adaptive Pruning)**을 제안합니다. 이 접근 방식은 두 가지 주요 구성 요소에 의존합니다:

1. 주기적 Top-k 선택:

   • 임계값 기반 방식(분포 변화에 민감함) 대신, 이 방법은 집계된 중요도 점수에 대한 top-k 선택을 사용합니다. Top-k 선택은 점수의 상대적 순위에만 의존하므로, 배치 집계로 인한 분포 변화에 강건합니다.
   • 매 토큰 단계마다 마스크를 업데이트하는 계산 오버헤드를 피하기 위해, 프루닝 마스크는 디코딩 단계마다가 아니라 주기적으로($T_{trans}$ 단계마다) 업데이트됩니다.
   • 배치 집계: $B$개의 샘플로 구성된 배치의 경우, 개별 샘플의 중요도 점수는 **원소별 최댓값(element-wise maximum)**을 사용하여 하나의 벡터로 집계됩니다. 이러한 보수적인 접근 방식은 배치 내 다른 샘플에서 활성화가 낮더라도, 적어도 하나의 샘플에 중요한 뉴런이 높은 점수를 유지하도록 하여 프루닝되는 것을 방지합니다.

2. 활성화 메모리(Activation Memory):

   • LLM 활성화 패턴 관찰 결과, 중요한 뉴런은 긴 추론 생성 과정 동안 **주기적으로 재발화(re-firing)**하는 경향이 있습니다(예: 약 22개 토큰마다 재발화하며, 재발화 비율은 약 76%).
   • 이를 포착하기 위해, 이 방법은 업데이트 단계 전반에 걸쳐 중요도 점수를 축적하는 활성화 메모리를 유지합니다. 메모리는 원소별 최댓값($m_n \leftarrow \max(m_{n-1}, s_n)$)을 통해 업데이트됩니다.
   • 이를 통해 이전 단계에서 중요하다고 식별된 뉴런이 일시적으로 활성화가 떨어지더라도 계속 유지 후보로 남게 되어, 핵심 구성 요소를 잃지 않으면서도 진화하는 패턴에 적응할 수 있습니다.

알고리즘 흐름
이 방법은 자기회귀 생성(autoregressive generation) 중에 세 단계로 작동합니다:

1. 입력 단계(Input Phase): 초기 중요도 점수를 계산하기 위해 초기 프롬프랜트를 처리합니다.
2. 밀집 탐색 단계(Dense Exploration Phase): 추론 방향을 설정하고 문맥 특화된 활성화 통계를 수집하기 위해 정해진 수의 밀집 단계($T_{init}$)를 수행합니다. 이때 메모리가 업데이트되고 첫 번째 프루닝 마스크가 생성됩니다.
3. 주기적 적응형 프루닝 단계(Periodic Adaptive Pruning Phase): 희소 계산($T_p$ 단계)과 밀집 탐색($T_E$ 단계)을 교대로 수행합니다. 각 밀집 탐색 단계가 끝날 때마다 활성화 메모리가 업데이트되며, 축적된 메모리를 기반으로 새로운 top-k 마스크가 생성됩니다.

주요 기여

1. 배치 붕괴의 경험적 확인: 본 논문은 기존의 훈련 불료 프루닝 방식(정적 및 적응형 모두)이 배치 추론 시 추론 작업에서 심각하게 성능이 저하된다는 것을 최초로 경험적으로 입증했습니다. 구체적으로, TEAL과 같은 적응형 방식은 배치 크기 4에서 밀집 모델 대비 평균 정확도가 58.4~67.7 퍼센트 포인트 하락합니다.
2. 강건한 배치 프러닝 프레임워크: 저자들은 주기적 top-k 선택을 도입하고, 반복되는 뉴런의 중요성을 추적하기 위해 활성화 메모리를 결합함으로써 배치 환경에서도 정확도를 유지하는 방법을 소개합니다.
3. 보정 불필요 운영(Calibration-Free Operation): 외부 코퍼스에 대한 오프라인 보정이 필요한 방식(예: Wanda, TEAL)과 달리, 이 방법은 런타임 활성화에만 의존하므로 보정이 필요 없습니다.

실험 결과
실험은 DeepSeek-R1-Distill-Qwen-7B 및 DeepSeek-R1-Distill-Llama-8B 모델을 대상으로 5가지 추론 벤치마크(GSM8K, MATH500, MINERVA, AMC23, GPQA-DIAMOND)에서 수행되었습니다.

• 정확도: 50% 목표 희소성 및 배치 크기 4에서, 제안된 방법은 Qwen-7B 모델에서 이전 SOTA(TEAL)보다 평균 정확도 측면에서 39.7 퍼센트 포인트(54.0% vs. 14.3%) 앞섰으며, Llama-8B 모델에서는 24.1 포인트를 앞섰습니다.
• 안정성: 배치 크기가 증가함에 따라 성능이 붕괴되는 TEAL과 달리, 제안된 방법은 모든 배치 크기(1, 4, 8, 16)에서 일관된 성능을 유지합니다.
• 속도 향상: Qwen-7B 모델에서 실제 희소성 50% 기준 밀집 추론 대비 1.40배의 속도 향상을 달s성했습니다. 이 속도 향상은 배치 설정에서 더 두드러지는데, 이는 제안된 방법이 가중치 행렬의 차원을 구조적으로 줄여 계산량을 직접 낮추는 반면, 0인 활성화를 건너뛰는 방식은 워크로드가 메모리 제한적(memory-bound)에서 계산 제한적(compute-bound)으로 변함에 따라 효과가 떨어지기 때문입니다.
• 높은 희소성: 제안된 방법은 높은 희소성 수준에서도 강력한 성능을 유지합니다(예: 50% 목표 희소성에서 71.0% 정확도). 반면 베이스라인 모델들은 40% 희소성을 넘어서면 급격히 성능이 저하됩니다.

의의 및 주장
본 논문은 제안된 방법이 대규모 추론 모델 배포의 핵심 병목 구간인 '기존 적응형 프러닝 기술과 배치 추론 간의 부조화'를 해결한다고 주장합니다. 배치 집계에 내재된 분포 변화 문제를 해결하고 뉴런의 재발화 특성을 활용함으로써, 성능 붕괴 없이 효율적이고 높은 처리량의 추론을 가능하게 합니다. 저자들은 이러한 성과가 단순히 집계 연산자(max vs. mean)의 선택 때문이 아니라, 주기적 top-k 선택과 활성화 메모리에 의해 주도되었다는 점을 강조합니다.