기술 요약: 출력 오차 상한을 활용한 심층 상태-공간 모델 압축 방법

문제 정의

S5 와 Mamba 와 같은 심층 상태-공간 모델 (Deep SSMs) 은 어텐션 메커니즘 없이 장거리 의존성을 처리하는 데 있어 최첨단 성능을 입증해 왔습니다. 그러나 이러한 모델을 배포할 때는 성능을 유지하면서 매개변수 수를 줄이는 것이 종종 필요합니다. 기존 심층 SSM 을 위한 모델 차수 축소 (Model Order Reduction, MOR) 기술들은 일반적으로 균형 절단 (Balanced Truncation) 이나 $H_2$ 노름과 같은 국소 기준에 기반하여 각 선형 상태-공간 블록 (레이어 단위) 에 독립적으로 압축을 적용합니다.

저자들은 이러한 기존 접근법의 중요한 한계를 지적합니다: 개별 선형 블록의 근사 오차를 최소화하는 것이 전체 심층 SSM 의 전체 출력 오차 최소화를 보장하지는 않는다는 점입니다. 심층 SSM 은 비선형 상호작용 (특히 2 차 출력과 레이어 정규화) 을 가진 캐스케이드 레이어를 포함하기 때문에, 국소 오차는 예측 불가능하게 누적되거나 전파될 수 있습니다. 결과적으로, 레이어 단위 축소 전략과 심층 SSM 의 전역 출력 성능을 직접적으로 연결하는 이론적 프레임워크가 부재합니다.

방법론

1. 이론적 프레임워크: LQO 시스템 및 오차 상한

본 논문은 심층 SSM 을 선형 2 차 출력 (Linear Quadratic-Output, LQO) 시스템의 캐스케이드로 모델링합니다. LQO 시스템은 선형 상태-공간 업데이트를 followed by 2 차 출력 맵으로 정의되며, 선형 블록 다음에 2 차 활성화 함수가 따르는 S5 와 같은 아키텍처를 일반화합니다.

저자들은 전체 차수 심층 SSM 과 축소 차수 심층 SSM 사이의 $\ell_\infty$ 출력 오차에 대한 증명 가능한 상한을 유도합니다.

• 핵심 통찰: 전역 출력 오차는 각 레이어의 개별 LQO 시스템에 대한 $h_2$-노름 근사 오차의 가중 합으로 상한이 설정됩니다.
• 레이어 민감도: 유도된 상한은 각 레이어의 근사 오차가 전체 출력 오차에 기여하는 정도가 균일하지 않음을 보여줍니다. 구체적으로, 얕은 레이어(네트워크의 초기 부분) 의 오차는 깊은 레이어의 오차보다 최종 출력 오차에 훨씬 더 큰 영향을 미칩니다. 이는 레이어 정규화와 관련된 Lipschitz 상수의 곱셈적 효과와 후속 레이어를 통한 오차 전파 때문입니다.

2. 최적화 공식화

유도된 상한에 기반하여, 저자들은 모델 차수 축소 (MOR) 문제를 공식화합니다. 직접적인 출력 오차 (비연속적이고 폐쇄형 특성이 없음) 를 최소화하는 대신, 최적화는 유도된 상한의 최소화를 목표로 합니다.

• 목적: 각 레이어에서 원래 LQO 시스템과 축소된 LQO 시스템 간의 $h_2$-노름 오차의 가중 합을 최소화합니다.
• 설계 원칙: 이 상한을 최소화하기 위해, 방법은 얕은 레이어에는 더 큰 축소 상태 차원을 할당하고 깊은 레이어에는 더 작은 차원을 할당하도록 규정합니다.

3. 알고리즘: 경사 기반 MOR

본 논문은 MOR 문제를 해결하기 위한 경사 기반 최적화 알고리즘 (알고리즘 1) 을 개발합니다:

• 경사 유도: 저자들은 축소된 LQO 시스템의 매개변수 (행렬 $\hat{A}, \hat{B}, \hat{C}, \hat{M}$) 에 대한 $h_2$-노름 오차에 대한 폐쇄형 경사를 유도합니다. 이러한 경사는 대각 상태 행렬을 가진 이산 시간 복소수 시스템에 맞게 조정됩니다.
• 안정성 보장: 알고리즘은 Armijo 조건을 사용하는 백트래킹 라인 서치를 사용합니다. 특히, 표준 MOR 방법에서 발생할 수 있는 불안정 모델 생성을 방지하기 위해 모든 반복 단계에서 축소된 상태 행렬 $\hat{A}$에 대해 **슈르 안정성 (Schur stability, 모든 고유값이 단위 원판 내부)**을 명시적으로 강제합니다.
• 초기화: 이 방법은 기존 시간 제한 MOR 기술 (예: 시간 제한 균형 절단 또는 시간 제한 $H_2$ 축소) 로 초기화할 수 있지만, 특정 심층 SSM制약을 충족시키기 위해 제안된 경사 하강법을 사용하여 이를 정제합니다.

주요 기여

1. 이론적 상한: 본 논문은 구성 LQO 블록의 $h_2$-노름 오차 측면에서 심층 SSM 의 출력 오차에 대한 상한을 최초로 이론적으로 유도합니다. 이는 레이어 단위 축소와 전역 성능 간의 직접적인 연결을 확립합니다.
2. 레이어 단위 민감도 분석: 이 분석은 깊은 레이어의 오차를 줄이는 것보다 얕은 레이어의 근사 오차를 줄이는 것이 전역 출력 오차 상한을 최소화하는 데 더 효과임을 증명합니다.
3. 최적화 알고리즘: 유도된 상한을 최소화하면서 축소된 모델의 안정성을 보장하고 심층 SSM 의 대각 구조를 유지하는 새로운 경사 기반 MOR 알고리즘을 제안합니다.
4. 재학습 없는 압축: 이 방법은 재학습 (one-shot 압축) 없이 사전 학습된 가중치에서 정확한 압축 모델을 구축하도록 설계되었습니다.

실험 결과

저자들은 Long Range Arena (LRA) 벤치마크의 IMDb 감성 분석 작업과 WikiText-103 언어 모델링 작업에서 이 방법을 평가했습니다.

• IMDb 작업:
  • 207,490 개의 매개변수를 가진 사전 학습된 심층 SSM 은 86.66% 의 테스트 정확도를 달성했습니다.
  • 제안된 방법을 사용하여 특정 레이어별 차수 할당 ($r_{list} = [96, 48, 36, 12]$) 을 적용한 결과, 모델은 83,650 개의 매개변수로 압축되어 (약 60% 감소) 테스트 정확도 **86.69%**를 유지했습니다.
  • 결과는 동일한 총 매개변수 수에 대해 균일한 할당 (예: $[16] \times 4$) 보다 얕은 레이어에 더 큰 차원을 할당하는 것 (예: $[96, 48, 36, 12]$) 이 훨씬 더 나은 정확도를 제공함을 확인했습니다.
  • 제안된 방법은 비교 가능한 상태 차수에서 기존 재학습 없는 압축 베이스라인 (예: [12] 의 방법들) 보다 우수한 성능을 발휘했습니다.
• WikiText-103 작업:
  • 유도된 상한을 최소화하는 것이 퍼플렉시티 (PPL) 개선과 상관관계가 있다는 유사한 경향이 관찰되었습니다.
  • 비균일 차수 할당 (얕은 레이어에서 더 큰 차수) 을 가진 모델들은 원래 사전 학습된 모델과 가까운 PPL 을 달성한 반면, 균일한 할당은 더 높은 PPL 을 초래했습니다.

중요성 및 주장

본 논문은 심층 SSM 을 위한 이론적으로 근거를 둔 재학습 없는 압축 전략을 제공하는 데 그 주요 의미가 있다고 주장합니다. 국소 블록 근사에서 전역 오차 상한으로 초점을 이동시킴으로써, 이 방법은 모델을 어디에서 압축할지 결정하는 체계적인 방법을 제공합니다.

저자들은 계산 제약으로 인해 재학습이 비실용적인 시나리오에서 저비용, 고성능 배포를 가능하게 한다는 점을 강조합니다. 이 방법은 "얕은 레이어 우선" 축소 전략이 단순한 휴리스틱이 아니라 유도된 출력 오차 상한에 의해 지지됨을 검증하여, 전역적 고려 없이 레이어 단위로 적용될 때 전체 시스템 성능을 유지하지 못하는 기존 MOR 기술에 대한 실용적인 대안을 제공합니다.