[25] Efficiently Modeling Long Sequences with Structured State Spaces (S4)

[dhkim0225]

paper code

openreview 에 S3 라고 제출되어 있는데, 코드도 그렇고, arxiv 에서도 그렇고, S4 가 밀고 있는 name으로 보인다.
#51 LSSL 은 state space model (SSM) 을 simulating 하느라 느리고, memory 도 많이 먹는다.

이를 해결하기 위해 S4 를 제안한다.
LSSL 은 state matrix A 를 잘 선택하는 것이 Long Range Dependency (LRD) 에 좋은 영향이 있다고 한다. S4 는 A matrix 에 low-rank correction (stably diagonalized)을 적용하고, cauchy kernel 을 이용해서 SSM computation을 낮춘다. Path-X 문제를 풀어냈고, LRA 에서 SOTA 를 찍는다.

Motivation: Diagonalization

LSSL 에서 A를 구하기 위해서 L개의 successive multiplication 을 수행해야 하는데 이게 bottleneck이다. A 를 canonical form 으로 바꿔버리면, 훨씬 빨라질 수 있다. 예를 들어, A 가 diagonal matrix라면 tractable 하기도 하다. 근데 RNN 관점으로 LSSL 을 보면, Vandermonde product 가 있는데, 이걸 diagonalize 하는 건 굉장히 어려운 문제라고 한다.

S4

1. klylov subspace 를 그냥 계산하는 대신, truncated generation function을 이용해서 spectrum을 구하고, inverse FFT 를 이용해서 구했다.
2. low-rank term 을 woodbury identity 를 통해 correct 한다.
3. cauchy kernel을 적용한다.

~무슨 소린지 잘 모르겠다~ 결론만 보자면, HiPPO matrices A 들은 모두 diagonalize 될 수 있다. 저자들은 Normal Plus Low-Rank (NPLR) 이라고 부른다.

음... 헷갈리니까 코드를 봤는데,,,, 잘 이해가 안간다. 어쨌든 요는, HiPPO matrices (A matrix in SSM) 들을 diagonalize 하는 데 성공했고, 이를 이용해서 굉장히 빠른 속도, 성능을 낼 수 있다는 것이다.

마지막으로 Informer 와 비교한 그림이다.

아래는 속도들이다. 확실히 빠르다.

Results

speech classification

wikitext-103

LRA

Pixel-level classification