[논문 리뷰] How to Train Your HiPPO: State Space Models with Generalized Orthogonal Basis Projections

Albert Gu, Isys Johnson외 3명|arXiv (Cornell University)|2022. 06. 24.

Blind Source Separation Techniques인용 수 11

한 줄 요약

이 논문은 상태공간모델(SSM)에 대한 일반화된 직교기저 투사 프레임워크를 제안하며, S4 모델이 장거리 의존성을 어떻게 잘 포착하는지에 대한 수학적으로 엄밀한 해석을 제공한다. 히포 매트릭스를 지수적으로 왜곡된 레지오르 다항식으로의 분해로 재해석함으로써, 저자들은 S4-FouT와 같은 새로운 SSM 변종—푸리에 기저를 사용한 모델—을 유도한다. 이는 Long Range Arena에서 86%, Path-X에서 96%의 성능을 달성하며, 시간스케일 매개변수의 명확한 초기화 규칙을 제공한다.

ABSTRACT

Linear time-invariant state space models (SSM) are a classical model from engineering and statistics, that have recently been shown to be very promising in machine learning through the Structured State Space sequence model (S4). A core component of S4 involves initializing the SSM state matrix to a particular matrix called a HiPPO matrix, which was empirically important for S4's ability to handle long sequences. However, the specific matrix that S4 uses was actually derived in previous work for a particular time-varying dynamical system, and the use of this matrix as a time-invariant SSM had no known mathematical interpretation. Consequently, the theoretical mechanism by which S4 models long-range dependencies actually remains unexplained. We derive a more general and intuitive formulation of the HiPPO framework, which provides a simple mathematical interpretation of S4 as a decomposition onto exponentially-warped Legendre polynomials, explaining its ability to capture long dependencies. Our generalization introduces a theoretically rich class of SSMs that also lets us derive more intuitive S4 variants for other bases such as the Fourier basis, and explains other aspects of training S4, such as how to initialize the important timescale parameter. These insights improve S4's performance to 86% on the Long Range Arena benchmark, with 96% on the most difficult Path-X task.

연구 동기 및 목표

• S4 모델이 장기간 시퀀스를 모델링하는 데 성공한 배경에 대한 이론적 이해 부족을 해결하기 위해, 특히 히포 매트릭스의 역할을 규명하는 것.
• 원래 S4 수식의 범위를 초월한 직교 기저 함수를 사용하여 SSM을 구성하는 일반적이고 수학적으로 일관된 프레임워크를 제공하는 것.
• S4에서 시간스케일 매개변수 Δ의 역할과 적절한 초기화 방법을 명확히 하는 것. 이는 모델이 포착하는 의존성의 길이를 제어한다.
• 푸리에 기반 변환과 국소 컨볼루션을 자연스럽게 포함하는 새로운 SSM 변종—예를 들어 S4-FouT—를 유도하는 것.
• 함수 재구성 및 기억 작업에서 이론적 프레임워크를 경험적으로 검증하여 이론과 성능 간의 일치를 보여주는 것.

제안 방법

• 저자들은 SSM을 가중 측도 하에서 직교 기저 함수로의 투사로 재정의함으로써 히포 프레임워크를 일반화한다. 특히 지수 감쇠 측도를 사용한다.
• 그들은 원래 S4 상태 행렬 A가 지수적으로 왜곡된 레지오르 다항식으로의 직교 투사에 해당함을 증명하며, 이는 장기적 컨텍스트 모델링에서의 효과성을 설명한다.
• 이 프레임워크를 통해 다른 기저 함수를 선택함으로써 새로운 SSM을 도출할 수 있으며, 예를 들어 잘라낸 푸리에 기저 함수를 사용해 S4-FouT 변종을 유도한다.
• 시간스케일 매개변수 Δ는 커널의 효과적인 길이를 제어하는 것으로 해석되며, 원하는 컨텍스트 창 크기에 기반한 명확한 초기화 규칙을 제공한다.
• 이 방법은 스케일링, 이동, 유니터리 변환에 대해 닫힘 성질을 갖는 일반화된 구조적 상태공간모델(TOSSMs) 클래스를 도입하여, 정규직교성과 안정성을 유지한다.
• 이론적 결과는 패드 근사와 연속 분수에 기반하며, SSM 커널이 유리 함수를 통해 지수 함수를 근사함을 보여준다.

실험 결과

연구 질문

• RQ1 S4의 히포 매트릭스를 경험적 초기화가 아닌 수학적으로 엄밀한 방식으로 어떻게 해석할 수 있는가?
• RQ2 S4가 장거리 의존성을 어떻게 포착하는가에 대한 이론적 메커니즘은 무엇이며, 다른 기저 함수로 일반화할 수 있는가?
• RQ3 S4에서 시간스케일 매개변수 Δ의 정확한 역할과 적절한 초기화 방법은 무엇인가?
• RQ4 푸리에 변환 또는 국소 컨볼루션과 같은 일반적인 신호 처리 연산을 자연스럽게 포함하는 새로운 SSM 변종을 도출할 수 있는가?
• RQ5 기저 함수의 정렬에 대한 이론적 예측이 시퀀스 모델링 작업에서의 경험적 성능과 얼마나 잘 일치하는가?

주요 결과

• 원래 S4 모델은 지수 감쇠 측도 하에서 지수적으로 왜곡된 레지오르 다항식으로의 투사로 이론적으로 해석되며, 이는 장기 컨텍스트 성능이 뛰어난 이유를 설명한다.
• 잘라낸 푸리에 기저 함수를 기반으로 한 새로운 S4-FouT 변종은 Path-X 작업에서 96%의 정확도로 최신 기술 수준을 달성하며, 기억 및 연속 신호 작업에서 다른 SSM들을 능가한다.
• 일반화된 프레임워크를 통해 시간스케일 매개변수 Δ의 원칙적인 초기화가 가능해지며, 이는 SSM이 모델링하는 의존성의 효과적인 길이를 직접 제어한다.
• 기저 함수 정렬에 대한 이론적 예측은 경험적 성능과 정확히 일치하며, S4-FouT는 스파ike 함수 인코딩과 국소 패턴 탐지가 필요한 작업에서 뛰어난 성능을 보인다.
• 가장 도전적인 Long Range Arena 벤치마크에서 원래 S4 모델은 평균 86%의 정확도를 기록하며, 장거리 모델링에 있어 그 견고함을 입증한다.
• 프레임워크의 스케일링, 이동, 유니터리 변환에 대한 닫힘 성질은 안정적이고 분산을 유지하는 SSM 파rameterization을 보장하며, 신뢰할 수 있는 학습을 가능하게 한다.