From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

기술 요약: 레이어에서 서브모듈로 (SUBFIT)

문제 정의
사후 학습 압축(Post-training compression)은 아키텍처 구성 요소를 제거함으로써 대규모 언어 모델(LLM)의 배포 비용을 줄이는 것을 목표로 합니다. 전체 레이어나 연속된 블록을 삭제하는 방식은 효과적이지만, 잔차 스트림(residual stream)을 방해하여 급격한 품질 저하를 초래하곤 합니다. 최근의 "교체 기반(replacement-based)" 접근 방식은 제거된 구성 요소를 가볍게 피팅된 모듈로 대체함으로써 이 문제를 해결합니다. 그러나 기존 방법들은 두 가지 제한적인 설계 제약 조건을 공유합니다:

1. 전체 레이어 단위(Full-layer granularity): 트랜스포머 레이어 전체를 하나의 원자적 제거 단위로 취급합니다.
2. 연속적 선택(Contiguous selection): 연속된 깊이 범위 내에서 구성 요소를 제거합니다.

저자들은 이러한 제약이 지나치게 제한적이라고 주장합니다. 사전 학습된 트랜스포머의 중복성은 연속적인 영역에 국한되지 않으며, Attention과 FeedForward(FFN) 출력 사이에도 균등하게 분포되어 있지 않습니다. 결과적으로, 서로 다른 서브모듈 유형을 근사하기 위해서는 서로 다른 전략이 필요하며, 제거 가능한 구성 요소가 반드시 연속적인 깊이 범위 내에 클러스터링될 필요는 없습니다.

방법론: SUBFIT
본 논문은 개별 Attention 및 FFN 서브모듈 단위로 작동하는 사후 학습 압축 방법인 SUBFIT(Submodule-level Fitted residual replacement)를 소개합니다. 과정은 크게 두 단계로 나뉩니다:

1. 비연속적 선택 (Non-Contiguous Selection):

   • SUBFIT는 모델 깊이 전반에 걸쳐 희소성 예산(sparsity budget) $s$에 따라 목표로 하는 수의 Attention 및 FFN 서브모듈을 독립적으로 선택합니다.
   • 선택은 **비연속적(non-contiguous)**입니다. 모듈은 연속적인 블록을 형성하지 않고 점수에 따라 독립적으로 선택됩니다.
   • 스코어링 전략:
     • Attention: 서브모듈이 잔차 스트림을 얼마나 회전시키는지와 그 상대적 크기를 측정하는 캘리브레이션 활성화값 기반의 "영향력 점수(Impact score)"를 통해 모듈의 순위를 매깁니다.
     • FFN: 베이스 모델의 경우, 제거 효과와 대리 오차(surrogate error)를 결고합한 교체 인식 점수를 사용합니다. 인스트럭션 튜닝된 모델의 경우, 제대로 캘리브레이션되지 않은 잔차 크기로 인한 노이즈를 피하기 위해 코사인 유사도 점수만을 사용합니다.
   • 순서: Attention 서브모듈이 먼저 선택되고 교체됩니다. 그 후, Attention 교체로 인해 유도된 잔차 스트림의 변화를 반영하기 위해, 부분적으로 압축된 모델 위에서 FFN 서브모듈의 점수를 매기고 선택합니다.

2. 피팅된 잔차 교체 (Fitted Residual Replacement):

   • 선택된 각 서브모듈은 내부 연산이 아닌 원래 서브모듈의 잔차 기여($y = F_\ell(x)$)를 근사하는 가벼운 파라미터 대리 모델 $\hat{F}_\ell$로 교체됩니다.
   • 대리 구조 (Surrogate Structure): 대리 모델은 아핀 보정(element-wise gate, bias, centering)과 저계수 연산자 $U_\ell = V_\ell^\top W_\ell$로 정의된 구조적 상호작용 항을 사용합니다.
   • 비대칭 파라미터화 (Asymmetric Parameterization):
     • Attention: 각 제거된 Attention 서브모듈은 깊이에 따른 이질적인 역할을 반영하여 고유의 독립적인 저계수 기저($V_\ell$)와 파라미터를 가집니다.
     • FFN: 모든 선택된 FFN 레이어에 걸쳐 공유된 입력 기저 $V$를 사용합니다. 오직 출력 인자($W_\ell$)와 아핀 항만이 레이어별로 유지됩니다. 이러한 교차 레이어 파라미터 공유는 배포 비용을 크게 줄입니다.
   • 피팅: 파라미터는 캘리브레이션 통계량(ridge-regularized least squares 및 eigen-decomposition)을 통한 폐쇄형(closed-form) 피팅을 통해 분석적으로 복구되며, 그래디언트 기반 최적화나 전체 모델 미세 조정이 필요하지 않습니다.

주요 기여

1. 제약 조건 식별: 저자들은 전체 레이어 단위의 입도와 연속적 선택이 기존의 교체 기반 압축에서 제한 요소임을 식별했습니다.
2. SUBFIT 프레임워크: 비연속적 선택과 함께 서브모듈 수준(Attention 및 FFN)에서 조인트 선택 및 교체를 수행하는 새로운 방법을 제시합니다.
3. 비대칭 설계: Attention 바이패스는 독립적인 기저를 사용하는 반면, FFN 바이패스는 공통의 입력 기저를 공유하여 파라미터 오버헤드를 최소화하는 특정 아키텍처를 제안합니다.
4. 포괄적 평가: 10개의 LLM(베이스 5개, 인스트럭션 튜닝 5개)과 5개의 희소성 수준(12.5% ~ 37.5%)에 걸친 광범리한 테스트를 통해 4개의 베이스라인(Streamline, ReplaceMe 포함) 대비 개선된 트레이드오프를 입증했습니다.

실험 결과

• 퍼플렉시티 (베이스 모델): SUBFIT는 10개 모델-희소성 설정 중 9개에서 가장 낮은 퍼플렉시티 저하를 달로했습니다. 강력한 베이스라인과의 격차는 공격적인 압축 상황(예: 25% 희소성에서 SUBFIT는 2.42배의 퍼플렉시티 저하를 보이는 반면, 최상위 베이스라인은 4.34배)에서 크게 벌어집니다.
• 다운스트림 정확도 (인스트럭션 튜닝 모델): 25% 희소성에서 SUBFIT는 평균적으로 덴스(dense) 모델 다운스트림 정확도의 **84.6%**를 유지하며, 가장 강력한 베이스라인(81.6%)을 능가했습니다. 이는 25% 희소성에서 80% 이상의 정확도를 유지하고 37.5%에서 73% 이상을 유지하는 유일한 방법입니다.
• 안정성: SUBFIT는 높은 희소성에서 불안정성을 보이는(예: 표준 편차가 몇 자릿수 증가하는) 베이스라인들에 비해 훨씬 낮은 퍼플렉시티 저하 분산을 보여줍니다.
• 효율성: 이 방법은 특수 커널이나 커스텀 하드웨어 없이도 측정 가능한 추론 속도 향상(첫 토큰 생성 시간 기준 1.18배~1.40배)과 KV-캐시 절감 효과를 제공합니다.
• 절제 연구 (Ablation Studies):
  • 잔차 보상: FFN을 교체하는 것이 안정성에 결정적인 요인이며, Attention과 FFN을 모두 교체할 때 최상의 결과를 얻습니다.
  • 공유 기저: 공유된 FFN 기저는 레이어별 기저와 유사한 성능을 내면서도 파라미터 수는 현저히 적습니다.
  • 랭크 민감도: FFN 품질은 랭크 4096까지 개선되는 반면, Attention은 랭크 256에서 포화되는데, 이는 비대칭 랭크 설정의 타당성을 입증합니다.

의의 및 주장
본 논문은 입도(granularity)가 사후 학습 압축에서 과소 탐구된 축이라고 주장합니다. 레이어 수준에서 서브모듈 수준의 교체로 이동함으로써, SUBFIT는 동일한 파라미터 예산 하에서 연속적인 전체 레이어 교체보다 덴스 모델의 동작을 더 많이 복구합니다. 결과는 비연속적이고 서브모듈 특화된 복구가 특히 공격적인 압축 상황에서 더 강력한 종합적 품질-효율성 트레이드오프를 제공함을 시사합니다. 이 방법은 저렴하게 배포 가능하도록 설계되었으며(폐쇄형 피팅), 기존 추론 스택과 호환됩니다. 저자들은 이 접근 방식이 다른 제거 단위(예: 어텐션 헤드 또는 개별 뉴런)로 일반화될 수 있음을 완곡하게 제안하지만, 본 연구에서 해당 단위들에 대한 즉각적인 일반화를 주장하지는 않습니다.