[논문 리뷰] On the Parameterization and Initialization of Diagonal State Space Models

Albert Gu, Ankit Gupta외 2명|arXiv (Cornell University)|2022. 06. 23.

EEG and Brain-Computer Interfaces인용 수 73

한 줄 요약

본 논문은 대각 상태 공간 모델(SSM)을 매개변수화하고 초기화하는 방법을 분석하고, 간단한 대각 SSM(S4D)이 S4의 성능에 도달할 수 있으며 이미지, 오디오, 의학 시계열 작업에서 강력한 성과를 보이고 Long Range Arena에서 85%에 달한다는 점을 보여준다. 이 논문은 이론적 통찰과 대각 SSM과 DPLR SSM을 비교하는 경험적 연구를 제공한다.

ABSTRACT

State space models (SSM) have recently been shown to be very effective as a deep learning layer as a promising alternative to sequence models such as RNNs, CNNs, or Transformers. The first version to show this potential was the S4 model, which is particularly effective on tasks involving long-range dependencies by using a prescribed state matrix called the HiPPO matrix. While this has an interpretable mathematical mechanism for modeling long dependencies, it introduces a custom representation and algorithm that can be difficult to implement. On the other hand, a recent variant of S4 called DSS showed that restricting the state matrix to be fully diagonal can still preserve the performance of the original model when using a specific initialization based on approximating S4's matrix. This work seeks to systematically understand how to parameterize and initialize such diagonal state space models. While it follows from classical results that almost all SSMs have an equivalent diagonal form, we show that the initialization is critical for performance. We explain why DSS works mathematically, by showing that the diagonal restriction of S4's matrix surprisingly recovers the same kernel in the limit of infinite state dimension. We also systematically describe various design choices in parameterizing and computing diagonal SSMs, and perform a controlled empirical study ablating the effects of these choices. Our final model S4D is a simple diagonal version of S4 whose kernel computation requires just 2 lines of code and performs comparably to S4 in almost all settings, with state-of-the-art results for image, audio, and medical time-series domains, and averaging 85\% on the Long Range Arena benchmark.

연구 동기 및 목표

• 대각 제한이 초기화 및 커널 구조를 통해 S4 성능을 보존할 수 있는 이유를 이해한다.
• 대각 SSM의 매개변수화 및 계산 선택을 체계적으로 분류한다.
• 대각 SSM이 무한 상태 한계에서 S4 역학을 재현할 수 있음을 보이고 실용적인 초기화 스킴을 제공한다.
• 이미지, 오디오 및 의학 시계열 작업에서 대각 SSM 변형들을 실험적으로 평가한다.
• 대각 SSM을 위한 간단하고 구현 가능한 커널 계산을 제시하여 DPLR 접근법에 필적하는 성능을 보인다.

제안 방법

• 대각 SSM들(A 대각, B, C)과 그 커널 K(t)=C e^{tA} B를 정의하고 분석한다.
• 커널 계산이 Vandermonde 기반의 행렬 연산으로 축소되어 구현에 따라 O(N+L) 또는 O(NL) 시간 복잡도를 가능하게 한다.
• S4, DSS, S4D 간의 매개변수화 선택(이산화, B와 C 다루기, 고유값 제약)을 비교한다.
• HiPPO 기반 A 행렬의 대각 근사체가 극한에서 S4와 같은 커널을 회복함을 증명한다(N→∞, 정리 3).
• 간단한 A 초기화를 가진 S4D 변형(S4D-Inv, S4D-Lin)을 제안하고 그 경험적 성능을 분석한다.
• 초기화, 이산화 및 B의 훈련에 대한 비실험적 분석을 통해 성능에 대한 핵심 효과를 고립한다.

실험 결과

연구 질문

• RQ1 적절한 매개화와 초기화로 대각 SSM이 원래 S4의 성능에 필적할 수 있는가?
• RQ2 대각 SSM의 매개화 및 계산에서 필수적인 설계 선택은 무엇이며 성능에 어떤 영향을 미치는가?
• RQ3 대각 HiPPO 기반 초기화가 큰 상태 수에서 S4 역학과 어떤 관계가 있는가?
• RQ4 간단하고 구현 가능한 대각 SSM(S4D)이 이미지, 오디오 및 시계열 도메인에서 경쟁력 있는 결과를 얻는가?
• RQ5 대각 SSM 매개변수화에서 B와 같은 구성 요소를 훈련하는지 고정하는지가 성능에 어떤 영향을 미치는가?

주요 결과

• S4D, 대각 SSM은 대부분의 설정에서 S4와 비교할 만한 성능을 보이며 이미지, 오디오 및 의학 시계열 벤치마크에서 강력한 성과를 달성한다.
• HiPPO 기반 행렬의 대각 근사는 무한 상태 한계에서 S4와 같은 커널을 회복한다(정리 3).
• 대각 SSM의 커널 계산은 Vandermonde 유사 곱을 통해 간단하며 몇 줄의 코드로 구현할 수 있다.
• 이산화 선택과 B의 훈련이 A의 초기화에 비해 영향이 제한적임을 보여주는 비측정 실험으로 S4D의 단순성을 뒷받침한다.
• Inv 및 Lin 초기화를 가진 S4D 변형은 해석 가능한 기저(감쇠된 푸리에 유사)를 제공하며 Long Range Arena에서 거의 최상위 성능에 근접한다(85%).
• DSS와 비교했을 때 S4D는 소프트맥스 정규화를 피하므로 커널 계산이 더 단순하고 성능이 안정적이다.