[논문 리뷰] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

Tri Dao, Albert G. Gu|arXiv (Cornell University)|2024. 05. 31.

Neural Networks and Applications인용 수 66

한 줄 요약

본 논문은 Structured state-space models (SSMs)와 어텐션 사이의 형식적 연결을 수립하고, Structured State Space Duality (SSD)를 도입하여 효율적인 SSM 기반 아키텍처(Mamba-2)를 도출한다. 이는 언어 모델링에서 Transformer와 경쟁하면서 확장성 이점을 제공한다.

ABSTRACT

While Transformers have been the main architecture behind deep learning's success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely related, and develop a rich framework of theoretical connections between SSMs and variants of attention, connected through various decompositions of a well-studied class of structured semiseparable matrices. Our state space duality (SSD) framework allows us to design a new architecture (Mamba-2) whose core layer is an a refinement of Mamba's selective SSM that is 2-8X faster, while continuing to be competitive with Transformers on language modeling.

연구 동기 및 목표

• SSMs와 Transformers를 연결하는 동기를 부여하여 긴 시퀀스에서 확장성과 효율성을 향상시키려 한다.
• 구조화된 행렬 및 어텐션 변형과 SSMs를 연결하는 이론적 프레임워크(SSD)를 개발한다.
• 효율적 계산을 위한 선형 재귀와 2차 듀얼 형식을 통합하는 새로운 알고리즘(SSD)을 제시한다.
• SSD와 구조화된 어텐션을 활용하여 처리량을 향상시키는 Mamba-2 아키텍처 디자인을 제안한다.

제안 방법

• SSMs를 행렬 변환으로 형식화하고 이를 semiseparable 행렬(SSS 표현)과 연결한다.
• Structured State Space Duality(SSD)를 도입하여 선형(재귀)과 2차(어텐션 유사) 형식을 통합한다.
• 구조화된 행렬을 사용하여 선형 어텐션을 Structured Masked Attention(SMA)로 일반화한다.
• 하드웨어 효율성을 위한 semiseparable 행렬의 블록 분해를 기반으로 SSD 알고리즘을 개발한다.
• 그룹화된 값 어텐션과 병렬 데이터 프로젝션을 갖춘 Mamba-2 아키텍처를 제안한다.
• 계산적 트레이드오프를 분석하고 하드웨어를 고려한 구현 가이드를 제공한다.

실험 결과

연구 질문

• RQ1 SSMs를 semiseparable 행렬로 표현하여 어텐션과의 알고리즘적 이중성을 어떻게 드러낼 수 있는가?
• RQ2 통합된 SSD 프레임워크가 Transformer에 버금가는 SSM의 효율적인 학습 및 추론을 가능하게 할 수 있는가?
• RQ3 현대 하드웨어에 맞게 SSM을 최적화하는 어떠한 수정(예: Mamba-2)이 언어 모델링에서 성능을 보존하면서 성능을 향상시킬 수 있는가?
• RQ4 구조화된 마스킹된 어텐션 변형은 선형 어텐션 프레임워크와 어떻게 연관되며 이를 일반화하는가?

주요 결과

• SSMs는 적절한 구조를 갖춘 semiseparable 행렬과 동등하며, 이를 통해 O(TN) 계산이 가능하다.
• SSD는 선형 재귀와 어텐션 유사 2차 형식 사이의 다리 역할을 하여 효율적인 하이브리드를 가능하게 한다.
• 새로운 SSD 알고리즘은 Mamba에서 선택적 스캔을 능가하고 더 큰 재귀 상태를 가능하게 한다.
• 병렬 데이터 프로젝션과 SMA 변형을 갖춘 Mamba-2 아키텍처가 표준 다운스트림 과제에서 Mamba 및 오픈 소스 Transformer를 상회하거나 일치한다.
• 언어 모델링, 효율성 지표, 그리고 어려운 기억(recall) 작업에 대한 경험적 검증은 경쟁력 있는 성능을 뒷받침한다.