MambaMixer: Efficiënte Selectieve State Space Modellen met Dubbele Token- en Kanaalselectie

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

March 29, 2024

Auteurs: Ali Behrouz, Michele Santacatterina, Ramin Zabih

cs.AI

Samenvatting

Recente vooruitgang in deep learning heeft voornamelijk vertrouwd op Transformers vanwege hun data-afhankelijkheid en het vermogen om op grote schaal te leren. De aandachtmodule in deze architecturen vertoont echter een kwadratische tijd- en ruimtecomplexiteit ten opzichte van de invoergrootte, wat hun schaalbaarheid voor modellering van lange sequenties beperkt. Ondanks recente pogingen om efficiënte en effectieve architectuurkernen te ontwerpen voor multidimensionale data, zoals afbeeldingen en multivariate tijdreeksen, zijn bestaande modellen ofwel data-onafhankelijk, of slagen ze er niet in om communicatie tussen en binnen dimensies mogelijk te maken. Recentelijk hebben State Space Models (SSMs), en meer specifiek Selective State Space Models, met een efficiënte hardwarebewuste implementatie, veelbelovend potentieel getoond voor modellering van lange sequenties. Gemotiveerd door het succes van SSMs, presenteren we MambaMixer, een nieuwe architectuur met data-afhankelijke gewichten die een dubbele selectiemechanisme gebruikt over tokens en kanalen, genaamd Selective Token and Channel Mixer. MambaMixer verbindt selectieve mixers met behulp van een gewogen gemiddelde mechanisme, waardoor lagen direct toegang hebben tot vroege kenmerken. Als proof of concept ontwerpen we Vision MambaMixer (ViM2) en Time Series MambaMixer (TSM2) architecturen gebaseerd op het MambaMixer-blok en onderzoeken we hun prestaties in verschillende vision- en tijdreeksvoorspellingstaken. Onze resultaten benadrukken het belang van selectief mengen over zowel tokens als kanalen. In ImageNet-classificatie, objectdetectie en semantische segmentatietaken behaalt ViM2 competitieve prestaties met gevestigde vision-modellen en overtreft het SSM-gebaseerde vision-modellen. In tijdreeksvoorspelling behaalt TSM2 uitstekende prestaties vergeleken met state-of-the-art methoden, terwijl het een aanzienlijk verbeterde rekenkost demonstreert. Deze resultaten tonen aan dat hoewel Transformers, cross-channel aandacht en MLPs voldoende zijn voor goede prestaties in tijdreeksvoorspelling, geen van deze noodzakelijk is.

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.