[논문 리뷰] Simplified State Space Layers for Sequence Modeling

Jimmy T. H. Smith, Andrew Warrington외 1명|arXiv (Cornell University)|2022. 08. 09.

Parallel Computing and Optimization Techniques인용 수 77

한 줄 요약

이 논문은 S5 계층을 소개하는데, 다중 입력 다중 출력 상태-공간 계층으로 S4와 효율성이 동일하고, 시간 영역에서의 완전 순환 방식으로 작동하여 장기 시퀀스 작업에서 최첨단 성능을 달성합니다.

ABSTRACT

Models using structured state space sequence (S4) layers have achieved state-of-the-art performance on long-range sequence modeling tasks. An S4 layer combines linear state space models (SSMs), the HiPPO framework, and deep learning to achieve high performance. We build on the design of the S4 layer and introduce a new state space layer, the S5 layer. Whereas an S4 layer uses many independent single-input, single-output SSMs, the S5 layer uses one multi-input, multi-output SSM. We establish a connection between S5 and S4, and use this to develop the initialization and parameterization used by the S5 model. The result is a state space layer that can leverage efficient and widely implemented parallel scans, allowing S5 to match the computational efficiency of S4, while also achieving state-of-the-art performance on several long-range sequence modeling tasks. S5 averages 87.4% on the long range arena benchmark, and 98.5% on the most difficult Path-X task.

연구 동기 및 목표

• 더 간단한 순환 시간 영역 계층으로 효율적인 장기 시퀀스 모델링을 동기화하고 가능하게 만든다.
• S4의 단일 입력 단일 출력 SSM 뱅크를 하나의 MIMO SSM으로 대체하여 효율성과 확장성을 개선한다.
• Long-range 의존성을 포착하기 위해 대각화된 HiPPO에서 영감을 받은 초기화를 활용한다.
• 일반적인 장기 벤치마크에서 경쟁력 있거나 우수한 성능을 보여주면서 불규칙한 샘플링과 가변 관찰 간격을 지원한다.]
• method_words_placeholder
• Replace the S4 bank of independent SISO SSMs with one diagonalized MIMO SSM of latent size P.
• Diagonalize the state matrix A to enable efficient parallel scans in the time domain.
• Initialize with diagonalized HiPPO-inspired matrices (HiPPO-N) and relate S5 dynamics to S4 under guiding assumptions.
• Discretize continuous-time SSMs using zero-order hold with learnable timescales Δ.
• Use a nonlinear activation and a mixing step to produce layer outputs from SSM preactivations.
• Show that with Latent size P = O(H), S5 matches S4 in runtime and memory complexity.

제안 방법

• S4 뱅크의 독립적인 SISO SSM들을 하나의 잠재 크기 P의 대각화된 MIMO SSM으로 대체한다.
• 시간 도메인에서 효율적인 병렬 탐색이 가능하도록 상태 행렬 A를 대각화한다.
• 대각화된 HiPPO에서 영감을 받은 행렬(HiPPO-N)로 초기화하고, 가이드 가정 하에서 S5의 다이나믹스를 S4와 연관시킨다.
• 0차 보정(Zero-Order Hold)으로 학습 가능한 시간 규모 Δ를 사용하여 연속 시간 SSM을 이산화한다.
• 비선형 활성화와 혼합 단계를 사용하여 SSM 프리액티베이션으로부터 층 출력을 생성한다.
• 잠재 크기 P가 O(H)일 때 S5가 런타임 및 메모리 복잡도에서 S4와 일치함을 보인다.

실험 결과

연구 질문

• RQ1 대각화된 다이내믹스를 갖는 단일 MIMO SSM이 S4의 SISO SSM 뱅크와 같은 기능적 능력을 달성할 수 있는가?
• RQ2 HiPPO에서 영감을 받은 초기화(diagonalized HiPPO-N)가 MIMO 설정에서 S5의 효과적인 시작점을 제공하는가?
• RQ3 S5 계층이 비정형 샘플링 및 시간 가변 매개변수를 처리하면서 시퀀스 길이에 대해 선형 복잡도를 유지할 수 있는가?
• RQ4 S4와 다른 베이스라인에 비해 장기 시퀀스 벤치마크에서 S5의 성능은 어떠한가?
• RQ5 핵심 가정(예: 동일한 A, 동일 Δ, P ≈ H)을 완화하는 것이 S5 설계에 어떤 실질적인 함의를 가지는가?

주요 결과

• S5는 시퀀스 길이에 선형 복잡도를 가지면서 S4의 효율성과 일치하는 최첨단 성능을 장기 시퀀스 작업에서 달성한다.
• Long Range Arena 벤치마크에서 S5는 평균 점수 87.4%, 도전적인 Path-X 작업에서 98.5%를 달성한다.
• S5는 음성 분류에서 여러 베이스라인을 능가하고 여러 작업에서 S4 변형들과 일치하거나 그 이상을 보인다.
• 대각화된 HiPPO-N 초기화는 S4에서 사용된 HiPPO-LegS 초기화에 비해 경쟁력 있는 성능을 제공하며 안정적인 대각화 가능한 다이나믹스를 가능하게 한다.
• S5는 시간 영역 병렬 스캔을 통해 불규칙 샘플링과 시간 가변 SSM을 지원한다. 이는 합성 곱(convolution) 기반의 S4 커널과 달리 가능하다.