SiMBA: Arquitetura Simplificada Baseada em Mamba para Visão e Séries Temporais Multivariadas

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

March 22, 2024

Autores: Badri N. Patro, Vijay S. Agneeswaran

cs.AI

Resumo

Os Transformers adotaram amplamente redes de atenção para a mistura de sequências e MLPs (Multi-Layer Perceptrons) para a mistura de canais, desempenhando um papel fundamental na conquista de avanços em diversos domínios. No entanto, a literatura recente destaca problemas com as redes de atenção, incluindo baixo viés indutivo e complexidade quadrática em relação ao comprimento da sequência de entrada. Modelos de Espaço de Estados (SSMs, na sigla em inglês), como S4 e outros (Hippo, Convoluções Globais, liquid S4, LRU, Mega e Mamba), surgiram para abordar esses problemas e ajudar a lidar com sequências mais longas. O Mamba, embora seja o SSM mais avançado, apresenta um problema de estabilidade quando dimensionado para grandes redes em conjuntos de dados de visão computacional. Propomos o SiMBA, uma nova arquitetura que introduz a Einstein FFT (EinFFT) para a modelagem de canais por meio de cálculos específicos de autovalores e utiliza o bloco Mamba para a modelagem de sequências. Estudos extensivos de desempenho em benchmarks de imagens e séries temporais demonstram que o SiMBA supera os SSMs existentes, reduzindo a lacuna de desempenho em relação aos transformers mais avançados. Notavelmente, o SiMBA se estabelece como o novo SSM mais avançado no ImageNet e em benchmarks de aprendizado por transferência, como Stanford Car e Flower, além de benchmarks de aprendizado de tarefas e sete conjuntos de dados de séries temporais. A página do projeto está disponível no site ~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.