[논문 리뷰] Simplifying and Understanding State Space Models with Diagonal Linear RNNs

Ankit Gupta, Harsh Mehta외 1명|arXiv (Cornell University)|2022. 12. 01.

Topic Modeling인용 수 4

한 줄 요약

이 논문은 상태공간모델의 간소화된 연속시간 공식화인 대각선 선형 RNN(DLR)을 제안한다. 이는 이산화 과정이 필요 없게 하여 더 깔끔하고 해석 가능한 아키텍처를 가능하게 한다. 단순함에도 불구하고, DLR는 장거리 시퀀스 작업, 즉 Long Range Arena와 원시 음성 분류에서 이전의 SSM 및 어텐션 모델을 능가하거나 그 성능을 유지하면서도 O(L) 추론 복잡도를 유지한다.

ABSTRACT

Sequence models based on linear state spaces (SSMs) have recently emerged as a promising choice of architecture for modeling long range dependencies across various modalities. However, they invariably rely on discretization of a continuous state space, which complicates their presentation and understanding. In this work, we dispose of the discretization step, and propose a model based on vanilla Diagonal Linear RNNs ($\mathrm{DLR}$). We empirically show that, despite being conceptually much simpler, $\mathrm{DLR}$ is as performant as previously-proposed SSMs on a variety of tasks and benchmarks including Long Range Arena and raw speech classification. Moreover, we characterize the expressivity of SSMs (including $\mathrm{DLR}$) and attention-based models via a suite of $13$ synthetic sequence-to-sequence tasks involving interactions over tens of thousands of tokens, ranging from simple operations, such as shifting an input sequence, to detecting co-dependent visual features over long spatial ranges in flattened images. We find that while SSMs report near-perfect performance on tasks that can be modeled via $ extit{few}$ convolutional kernels, they struggle on tasks requiring $ extit{many}$ such kernels and especially when the desired sequence manipulation is $ extit{context-dependent}$. Despite these limitations, $\mathrm{DLR}$ reaches high performance on two higher-order reasoning tasks $\mathrm{ListOpsSubTrees}$ and $\mathrm{PathfinderSegmentation} ext{-}\mathrm{256}$ with input lengths $8K$ and $65K$ respectively, and gives encouraging performance on $\mathrm{PathfinderSegmentation} ext{-}\mathrm{512}$ with input length $262K$ for which attention is not a viable choice.

연구 동기 및 목표

• 이산화 단계를 제거함으로써 상태공간모델의 이해와 구현을 복잡하게 만드는 요소를 제거함으로써 상태공간모델을 단순화한다.
• 장거리 시퀀스 모델링을 위한 기초로 사용할 수 있는 순수한 대각선 선형 RNN(DLR) 기반의 직접적이고 해석 가능한 공식화를 제안한다.
• 합성 및 실제 장시계열 작업에서 DLR의 표현력과 성능을 실험적으로 평가하고, SSM 및 어텐션 기반 모델과 비교한다.
• 현재 SSM이 문맥 의존적 및 고핵심 작업을 처리하는 데 가지는 한계를 조사하고, 어텐션 기반 메커니즘과의 표현력 상대적 특성을 규명한다.
• 13개의 합성 시퀀스-투-시퀀스 작업을 사용하여 장거리 모델링 능력을 평가할 수 있는 청결하고 접근 가능한 벤치마킹 프레임워크를 제공한다.

제안 방법

• 이산 상태공간 이산화 과정을 피하는 대로, 대각선 선형 RNN의 직접적 파arameterization인 DLR를 제안한다.
• 유닛 원주에 가까운 고유값을 가진 복소수 전이 행렬을 사용하여 수백만 단계에 걸쳐 장거리 신호 전파를 가능하게 한다.
• 크기가 N인 DLR가 길이 N인 임의의 컨볼루션 커널을 표현할 수 있도록 하는 주기적 초기화 기법을 도입한다.
• 간단한 시퀀스 이동부터 편평화된 이미지에서의 복잡한 공간적 추론에 이르기까지, 13개의 합성 작업 세트에서 DLR를 평가한다.
• Long Range Arena, 음성 분류, 인과적 언어 모델링에서 S4, S4D, 트랜스포머, LocalAttention과의 비교를 통해 DLR를 벤치마크한다.
• 스케일링 및 초장거리 시퀀스에서의 성능 테스트를 위해 입력 길이를 최대 262K 토큰까지 늘였다.

실험 결과

연구 질문

• RQ1 이산화 과정이 복잡한 연속시간 이산화 없이도 직접적인 대각선 선형 RNN(DLR) 공식화가 이전의 이산화 SSM의 성능을 따라하거나 능가할 수 있는가?
• RQ2 장거리 의존성과 복잡한 상호작용을 요구하는 작업에서 DLR의 표현력은 SSM 및 어텐션 기반 모델과 비교해 어떻게 되는가?
• RQ3 현재 SSM은 문맥 의존적 시퀀스 조작이나 많은 컨볼루션 커널이 필요한 작업을 처리하는 데 어떤 한계를 가지는가?
• RQ4 DLR는 원시 음성 분류 및 인과적 언어 모델링과 같은 실제 장시계열 작업에 경쟁적인 성능과 효율성으로 일반화될 수 있는가?
• RQ5 DLR의 파arameterization에서 주기성이 임의의 컨볼루션 커널 표현을 지원하는 원리적인 초기화를 가능하게 하는가?

주요 결과

• DLR는 Path-X 및 ListOps를 포함한 Long Range Arena의 모든 작업에서 S4 및 S4D와 동등하거나 뛰어난 성능을 기록하며, 더 단순하고 해석 가능한 공식화를 제공한다.
• 262K 입력 길이를 가진 PathfinderSegmentation-512 작업에서, DLR는 어텐션 모델이 이 크기에서 계산적으로 불가능한 상황에서도 강력한 성능을 기록한다.
• ListOps-SubTrees 및 PathfinderSegmentation-256에서 DLR는 LocalAttention을 능가하며, 특히 더 긴 시퀀스 길이에서 높은 처리량을 유지한다.
• 10가지 음성 명령어 데이터셋에서 DLR는 S4 및 S4D와 동등하거나 뛰어난 성능을 기록하며, O(L) 디코딩 복잡도로 97.1%의 정확도를 달성한다.
• PG-19에서의 인과적 언어 모델링에서 DLR는 1.2B 파라미터와 1.0× 처리량으로 트랜스포머와 유사한 테스트 교차엔트로피 손실(2.44–2.88)을 기록한다.
• Transformer가 이러한 아키텍처 선택에 매우 민감한 데 비해, DLR는 레이어 정규화의 위치에 더 강인하다.