[논문 리뷰] Efficiently Modeling Long Sequences with Structured State Spaces

Albert Gu, Karan Goel외 1명|arXiv (Cornell University)|2021. 10. 31.

Topic Modeling참고 문헌 41인용 수 481

한 줄 요약

S4는 매우 긴 시퀀스의 효율적 처리를 가능하게 하는 구조화된 상태 공간 시퀀스 모델을 도입하고, Path-X를 해결하며 Transformers에 비해 생성 속도를 크게 높이고, 장거리 의존성 벤치마크에서 최첨단 성능을 달성합니다.

ABSTRACT

A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies. Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of $10000$ or more steps. A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) \( x'(t) = Ax(t) + Bu(t), y(t) = Cx(t) + Du(t) \), and showed that for appropriate choices of the state matrix \( A \), this system could handle long-range dependencies mathematically and empirically. However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution. We propose the Structured State Space sequence model (S4) based on a new parameterization for the SSM, and show that it can be computed much more efficiently than prior approaches while preserving their theoretical strengths. Our technique involves conditioning \( A \) with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of a Cauchy kernel. S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91\% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation $60 imes$ faster (iii) SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X task of length 16k that all prior work fails on, while being as efficient as all competitors.

연구 동기 및 목표

• 다양한 모달리티와 작업 전반에 걸친 장거리 의존성을 처리해야 하는 필요성을 제시한다.
• 매우 긴 시퀀스까지 확장 가능한 실용적이고 효율적인 SSM 기반 시퀀스 모델을 제안한다.
• 정규 플러스 저랭크(NPLR) 매개변수화가 빠른 계산과 안정적인 학습을 가능하게 함을 보여준다.
• 이미지, 텍스트, 음성 벤치마크에 걸친 S4의 성능과 Transformer와의 경쟁력을 보여준다.

제안 방법

• 상태공간 행렬 A를 안정적인 대각화가 가능하도록 normal plus low-rank(NPLR) 형태로 재매개변수화한다.
• 대각 형태로 켤 수 있도록 변환하고 저랭크 보정을 위해 Woodbury 항등식을 적용하여 이산 SSM 커널을 효율적으로 계산한다.
• SSM 합성곱 커널을 Cauchy 커널로 표현하고, 유리근에서 샘플링된 잘린 생성 함수로 평가한 뒤 역 FFT를 수행한다.
• 긴 범위 의존성을 다루기 위해 HiPPO 기반의 연속 시간 기억 이론을 활용한다.
• 특성들 간에 매개변수를 공유하는 아키텍처(H 독립 복사본)를 제공하고, 다중 특성 입력에 대해 깊이 방향의 브로드캐스팅과 유사한 접근을 사용한다.

실험 결과

연구 질문

• RQ1 S4 매개변수화를 가진 SSM이 표준 벤치마크에서 Transformer 성능에 부합하거나 이를 능가하면서 매우 긴 시퀀스(L 최대 16k 이상)도 효율적으로 모델링할 수 있는가?
• RQ2 언어 및 이미지 모델링에서 어텐션 제로 혹은 저어텐션 모델이 Transformer에 얼마나 근접하면서도 더 빠른 생성을 제공할 수 있는가?
• RQ3 SSM 기반 모델이 최소한의 아키텍처 변경으로 이미지, 텍스트, 음성 같은 다양한 도메인에 일반화될 수 있는가?
• RQ4 순환 및 합성곱 표현에 대해 NPLR S4 매개변수화가 제공하는 이론적·계산적 보장(복잡도, 안정성)은 무엇인가?

주요 결과

• S4는 데이터 증가나 보조 손실 없이 순차 CIFAR-10에서 91% 정확도를 달성하여 더 큰 2-D ResNet과 동등한 성능을 보인다.
• S4는 이미지 및 언어 모델링 작업에서 Transformer와의 격차를 크게 줄이고 생성 속도는 약 60배 빨라진다.
• Long Range Arena 작업에서 최첨단을 달성하며, Path-X(길이 16k) 해법에서 88% 정확도를 달성했다(이전 연구는 무작위 추측).
• 길이 16000 시퀀스의 음성 분류에서 테스트 오차를 1.7%로 절반으로 줄여 전문 음성 CNN보다 더 나은 성능을 보이고, 기존 베이스라인보다 우수했다.
• WikiText-103 언어 모델링에서 S4는 Transformer 베이스라인보다 0.8 퍼플렉시티 이내로, 어텐션 없이도 경쟁력을 보여준다.
• S4는 빠른 자기회귀 생성, 다중 도메인 적용성(이미지, 텍스트, 음성), 재훈련 없이 샘플링 속도 변화에 대한 강건성을 보여준다.