기술 요약: AIRE-Prune

문제 정의

심층 상태 공간 모델(Deep State Space Models, SSMs)은 입력 이력을 내부 상태로 압축함으로써 순차적 데이터를 모델링하는 데 탁와한 효능을 입증해 왔다. 그러나 상태 차원 $n$은 계산 및 메모리 비용을 결정하는 주요 요인으로 남아 있다. 현대적 아키텍처들은 효율성과 안정성을 보장하기 위해 종종 대각(diagonal) 또는 대각화 가능한 파라미터화(예: S5와 같은 multi-SISO 또는 MIMO 스택)를 활용하지만, 이러한 모델들은 빈번하게 과잉 매개변수화(over-parameterization) 문제를 겪는다. 기존의 사후 학습(post-training) 최적화 메커니즘은 $n$을 줄이는 데 한계가 있다. 비구조적 가중치 프루닝(unstructured weight pruning)과 같은 기존 방식은 상태 차원을 직접적으로 줄이지 못한다. 최근의 LAST (Gwak et al., 2025)와 같은 다른 구조적 방법들은 최악의 경우 주파수 영역 이득($H_\infty$)에 의존하여 상태를 점수화한다. 이 방식은 효과적이긴 하지만, 최악의 경우를 측정하는 기준은 지나치게 보수적일 수 있다. 이는 작업이 거의 자극하지 않을 수도 있는 피크 증폭(peak amplification)을 강조함으로써, 잠재적으로 최적의 압축 비율을 얻는 것을 방해할 수 있기 때문이다.

방법론: AIRE-Prune

저자들은 대각(또는 대각화 가능한) SSM을 위해 설계된 구조적, 사후 학습 기반, 레이어 적응형 프루닝 방법인 AIRE-Prune(Asymptotic Impulse-Response Energy for State PRUN(E)ing)을 소개한다. 최악의 이득에 집중하는 이전 방법들과 달리, AIRE-Prune은 점근적 임펄스 응답 에너지(asymptotic impulse-response energy), 즉 무한한 시간 지평 동안 상태가 기여하는 총 출력 에너지라는 '전형적인 경우(typical-case)'의 기준에 프루닝을 정렬한다.

핵심 메커니즘

1. 에너지 기반 랭킹:
   $(\Lambda, B, C)$로 정의되는 안정적인 이산 시간 대각 SSM 레이어(여기서 $\Lambda = \text{diag}(\lambda_1, \dots, \lambda_n)$ 이고 $|\lambda_i| < 1$)에 대해, 이 방법은 각 상태(모드) $i$에 대해 생성되는 단위 임펄스 입력에 의한 총 제곱 출력 에너지를 계산한다.
   $i$번째 모드에 대한 폐쇄형(closed-form) 점수는 다음과 같이 유도된다:
   $$\text{EnergyScore}_{\text{local}}(x_i) = E_i = \frac{\|C_{:,i}\|_2^2 \|B_{i,:}\|_2^2}{1 - |\lambda_i|^2}$$
   이 점수는 세 가지 요소를 통합한다:

   • 가제어성(Controllability): $\|B_{i,:}\|_2^2$로 표현됨 (흥분의 용이성).
   • 가관측성(Observability): $\|C_{:,i}\|_2^2$로 표현됨 (측정의 용이성).
   • 감쇠/메모리(Damping/Memory): 폴(pole) $\lambda_i$로 표현됨; 단위 원에 가까운 폴(느린 감쇠)을 가진 상태는 장기적인 에너지에 더 많이 기여한다.

2. 교차 레이어 정규화(Cross-Layer Normalization):
   레이어마다 에너지 스케일이 크게 다를 수 있으므로(예: 인코더/디코더 이득에 의해), 저자들은 전역적 비교를 가능하게 하는 접두사 정규화 점수화(prefix-normalized scoring) 전략을 제안한다.

   • 각 레이어 $\ell$ 내에서, 모드들은 로컬 에너지 $E^{(\ell)}_i$에 따라 내림차순으로 정렬된다.
   • 정렬된 에너지에 대해 접두사 합(prefix sum) $S^{(\ell)}_{(i)}$가 계산된다.
   • 전역 점수는 "위험률(hazard-rate)" 비율로 정의된다:
     $$\text{AIREPrune}(x^{(\ell)}_{(i)}) = \frac{E^{(\ell)}_{(i)}}{S^{(\ell)}_{(i)} + \epsilon}$$
     이 정규화는 레이어 내에서 점수가 단조 감소하도록 보장하여, 전체 모델에 걸쳐 단일 전역 임계값 $\tau$를 적용할 수 있게 한다.

3. 프루닝 실행:
   목표 프루닝 비율 또는 임계값이 주어지면, 알고리즘은 각 레이어에서 정규화된 점수가 $\tau$를 초과하는 가장 긴 접두사(longest prefix)의 상태들을 유지한다. 이를 통해 일부 레이어는 대폭 프루닝되거나 완전히 제거될 수 있는 반면, 다른 레이어는 거의 온전하게 유지되는 레이어 적응형 프루닝이 가능하다.

주요 기여

• 폐쇄형 점근적 에너지 점수: 본 논문은 무한 지평 출력 에너지를 기반으로 학습된 모든 모드에 대한 수학적으로 유도된 폐쇄형 점수를 제공한다. 이는 비용이 많이 드는 수치 시뮬레이션이나 재학습을 피할 수 있게 한다.
• 모드 절단(Modal Truncation)의 확장: 이 연구는 단일 선형 시스템에 주로 적용되는 고전적인 모드 절단을 심층 비선형 SSM 스택으로 확장하며, 효율성에 필수적인 대각 구조를 보존한다.
• 전형적인 경우 vs 최악의 경우: 중요도 지표를 최악의 경우 $H_\infty$ 이득에서 총 임펄스 응답 에너지로 전환함으로써, 이 방법은 희귀한 피크 증폭보다는 시스템의 집합적 행동을 목표로 한다.
• 아키텍처 불가지론(Architecture Agnosticism): 이 방법은 multi-SISO 구성 및 MIMO 스택(예: S5 스타일)뿐만 아니라 S4D 및 Mamba와 같은 다른 SSM 계열에도 적용 가능하다.

실험 결과

저자들은 S5, S4D, Mamba 백본을 사용하여 Long Range Arena (LRA) 벤치마크와 Speech Commands 데이터셋에서 AIRE-Prune을 평가했다.

• 압축 효율: 다양한 벤치마크 전반에 걸쳐, AIRE-Prune은 재학습 없이 평균 **60.8%**의 상태 프루닝을 달성하면서도 평균 정확도 하락은 단 **0.29%**에 불과했다.
• 베이스라인과의 비교:
  • LRA 제품군에서, AIRE-Prune은 Random, Uniform Magnitude, Global Magnitude, LAMP, Uniform $H_\infty$, Global $H_\infty$, 그리고 LAST를 포함한 베이스라인들을 유의미하게 능가했다.
  • $\le 1$ 퍼센티지 포인트(pp)의 정확도 손실이라는 엄격한 제약 조건 하에서, AIRE-Prune은 가장 강력한 베이스라인인 LAST보다 훨씬 높은 프루닝 비율을 허용했다. 예를 들어, Path-X 작업에서 AI-RE-Prune은 LAST가 30%를 프루닝할 때 안전하게 70%를 프루닝했다.
  • Speech Commands 작업에서, 무시할 만한 정확도 손실(전체 모델 96.43% 대비 96.40%)과 함께 45%의 상태가 프루닝되었다.
• 추론 효율성: 프루닝은 시스템 수준의 이득으로 직결되었다. NVIDIA H100 GPU에서, 이 방법은 작업 및 프루닝 비율에 따라 1.23$\times$에서 2.86$\times$의 추론 속도 향상과 **19%에서 65%**의 파라미터 감소를 실현했다.
• 일반화: 이 방법은 S4D 및 Mamba 아키텍처에서도 견고한 성능을 보여주었으며, S4D의 Text 작업에서 최대 90%, Mamba의 Image 작업에서 80%의 프루닝을 최소한의 정확도 저하와 함께 달성했다.

의의 및 주장

본 논문은 AIRE-Prune이 훈련된 SSM의 상태 공간에 존재하는 상당한 중복성을 드러내며, 이를 재학습 없이 제거할 수 있다고 주장한다. 프루닝을 최악의 이득이 아닌 점근적 응답 에너지에 정렬함으로써, 이 방법은 작업 성능에 대한 비용을 거의 들이지 않고 공격적인 압축을 달성한다.

저자들은 AIRE-Prune에서 관찰되는 "계단형(step-like)" 정확도 하락 프로파일(높은 프루닝 임계값에 도달할 때까지 정확도가 안정적으로 유지되다가 급격히 떨어지는 현상)이 중요한 상태와 중요하지 않은 상태 사이의 강력한 구분을 나타낸다고 강조한다. 이는 부드러운 저하를 보이는 베이스라인들과 대조적이며, 베이스라인들의 점수화 방식이 중요하거나 중요하지 않은 상태를 혼합하고 있음을 시사한다.

실질적으로, 이 방법은 다음을 가능하게 한다:

1. 전역 예산 책정(Global Budgeting): 레이어별 튜닝 없이 단일 목표 예산을 설정할 수 있다.
2. 블록 단위 감소(Block-Level Reduction): 상태가 임계값 미만으로 떨어질 때 전체 레이어를 프루닝할 수 있어, 미세한 희소성을 하드웨어 친화적인 블록 구조적 감소로 변환할 수 있다.
3. 학습이 필요 없는 최적화(Training-Free Optimization): 추가적인 미세 조정이 필요 없는 사후 학습 메커니즘으로서, 배포된 모델에 즉시 적용 가능하다.

결론적으로, AIRE-Prune은 딥 SSM의 모델 차수 축소를 위한 원칙적이고 에너지 기반의 접근 방식을 제공하며, 압축 비율과 정확도 보존 측면 모두에서 기존의 최첨단 프루닝 전략을 능가한다.