SiMBA: 비전 및 다변수 시계열을 위한 단순화된 Mamba 기반 아키텍처

SiMBA: Simplified Mamba-Based Architecture for Vision and Multivariate Time series

March 22, 2024

저자: Badri N. Patro, Vijay S. Agneeswaran

cs.AI

초록

트랜스포머는 시퀀스 혼합을 위해 어텐션 네트워크를, 채널 혼합을 위해 MLP를 널리 채택하며 다양한 분야에서의 혁신을 이끄는 데 중요한 역할을 해왔습니다. 그러나 최근 연구에서는 어텐션 네트워크의 낮은 귀납적 편향과 입력 시퀀스 길이에 대한 2차 복잡성과 같은 문제점들이 지적되고 있습니다. 이러한 문제를 해결하고 더 긴 시퀀스 길이를 처리하기 위해 S4와 같은 상태 공간 모델(SSM)들(Hippo, Global Convolutions, liquid S4, LRU, Mega, Mamba 등)이 등장했습니다. Mamba는 최신 SSM이지만 컴퓨터 비전 데이터셋을 위한 대규모 네트워크로 확장할 때 안정성 문제가 있습니다. 우리는 채널 모델링을 위해 특정 고유값 계산을 통해 Einstein FFT(EinFFT)를 도입하고 시퀀스 모델링을 위해 Mamba 블록을 사용하는 새로운 아키텍처인 SiMBA를 제안합니다. 이미지 및 시계열 벤치마크에 대한 광범위한 성능 연구를 통해 SiMBA가 기존 SSM들을 능가하며 최신 트랜스포머와의 성능 격차를 줄이는 것을 입증했습니다. 특히, SiMBA는 ImageNet 및 Stanford Car, Flower와 같은 전이 학습 벤치마크뿐만 아니라 7개의 시계열 벤치마크 데이터셋에서도 새로운 최신 SSM으로 자리매김했습니다. 프로젝트 페이지는 ~https://github.com/badripatro/Simba에서 확인할 수 있습니다.

English

Transformers have widely adopted attention networks for sequence mixing and MLPs for channel mixing, playing a pivotal role in achieving breakthroughs across domains. However, recent literature highlights issues with attention networks, including low inductive bias and quadratic complexity concerning input sequence length. State Space Models (SSMs) like S4 and others (Hippo, Global Convolutions, liquid S4, LRU, Mega, and Mamba), have emerged to address the above issues to help handle longer sequence lengths. Mamba, while being the state-of-the-art SSM, has a stability issue when scaled to large networks for computer vision datasets. We propose SiMBA, a new architecture that introduces Einstein FFT (EinFFT) for channel modeling by specific eigenvalue computations and uses the Mamba block for sequence modeling. Extensive performance studies across image and time-series benchmarks demonstrate that SiMBA outperforms existing SSMs, bridging the performance gap with state-of-the-art transformers. Notably, SiMBA establishes itself as the new state-of-the-art SSM on ImageNet and transfer learning benchmarks such as Stanford Car and Flower as well as task learning benchmarks as well as seven time series benchmark datasets. The project page is available on this website ~https://github.com/badripatro/Simba.