MambaMixer: Modelli Efficienti di Spazio degli Stati Selettivi con Selezione Duale di Token e Canali

MambaMixer: Efficient Selective State Space Models with Dual Token and Channel Selection

March 29, 2024

Autori: Ali Behrouz, Michele Santacatterina, Ramin Zabih

cs.AI

Abstract

I recenti progressi nel deep learning si sono principalmente basati sui Transformer grazie alla loro dipendenza dai dati e alla capacità di apprendere su larga scala. Tuttavia, il modulo di attenzione in queste architetture presenta una complessità quadratica in termini di tempo e spazio rispetto alla dimensione dell'input, limitando la loro scalabilità per la modellazione di sequenze lunghe. Nonostante i recenti tentativi di progettare architetture efficienti ed efficaci per dati multidimensionali, come immagini e serie temporali multivariate, i modelli esistenti sono o indipendenti dai dati o non consentono la comunicazione inter- e intra-dimensione. Recentemente, i modelli a spazio di stato (SSM), e più specificamente i Selective State Space Models, con un'implementazione efficiente e consapevole dell'hardware, hanno mostrato un potenziale promettente per la modellazione di sequenze lunghe. Motivati dal successo degli SSM, presentiamo MambaMixer, una nuova architettura con pesi dipendenti dai dati che utilizza un meccanismo di selezione duale tra token e canali, chiamato Selective Token and Channel Mixer. MambaMixer collega i mixer selettivi utilizzando un meccanismo di media ponderata, consentendo ai livelli di accedere direttamente alle feature iniziali. Come prova di concetto, progettiamo le architetture Vision MambaMixer (ViM2) e Time Series MambaMixer (TSM2) basate sul blocco MambaMixer e ne esploriamo le prestazioni in vari compiti di visione e previsione di serie temporali. I nostri risultati sottolineano l'importanza della miscelazione selettiva sia tra token che tra canali. Nelle attività di classificazione su ImageNet, rilevazione di oggetti e segmentazione semantica, ViM2 raggiunge prestazioni competitive rispetto ai modelli di visione consolidati e supera i modelli di visione basati su SSM. Nella previsione di serie temporali, TSM2 ottiene prestazioni eccezionali rispetto ai metodi all'avanguardia, dimostrando al contempo un significativo miglioramento del costo computazionale. Questi risultati mostrano che, sebbene i Transformer, l'attenzione cross-canale e le MLP siano sufficienti per ottenere buone prestazioni nella previsione di serie temporali, nessuno di essi è necessario.

English

Recent advances in deep learning have mainly relied on Transformers due to their data dependency and ability to learn at scale. The attention module in these architectures, however, exhibits quadratic time and space in input size, limiting their scalability for long-sequence modeling. Despite recent attempts to design efficient and effective architecture backbone for multi-dimensional data, such as images and multivariate time series, existing models are either data independent, or fail to allow inter- and intra-dimension communication. Recently, State Space Models (SSMs), and more specifically Selective State Space Models, with efficient hardware-aware implementation, have shown promising potential for long sequence modeling. Motivated by the success of SSMs, we present MambaMixer, a new architecture with data-dependent weights that uses a dual selection mechanism across tokens and channels, called Selective Token and Channel Mixer. MambaMixer connects selective mixers using a weighted averaging mechanism, allowing layers to have direct access to early features. As a proof of concept, we design Vision MambaMixer (ViM2) and Time Series MambaMixer (TSM2) architectures based on the MambaMixer block and explore their performance in various vision and time series forecasting tasks. Our results underline the importance of selective mixing across both tokens and channels. In ImageNet classification, object detection, and semantic segmentation tasks, ViM2 achieves competitive performance with well-established vision models and outperforms SSM-based vision models. In time series forecasting, TSM2 achieves outstanding performance compared to state-of-the-art methods while demonstrating significantly improved computational cost. These results show that while Transformers, cross-channel attention, and MLPs are sufficient for good performance in time series forecasting, neither is necessary.