MambaMixer : Modèles efficaces d'espaces d'états sélectifs avec sélection duale de tokens et de canaux

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

March 29, 2024

Auteurs: Ali Behrouz, Michele Santacatterina, Ramin Zabih

cs.AI

Résumé

Les récents progrès en apprentissage profond ont principalement reposé sur les Transformers en raison de leur dépendance aux données et de leur capacité à apprendre à grande échelle. Cependant, le module d'attention dans ces architectures présente une complexité quadratique en temps et en espace par rapport à la taille de l'entrée, limitant ainsi leur évolutivité pour la modélisation de longues séquences. Malgré les tentatives récentes de concevoir des architectures efficaces et performantes pour les données multidimensionnelles, telles que les images et les séries temporelles multivariées, les modèles existants sont soit indépendants des données, soit incapables de permettre une communication inter et intra-dimensionnelle. Récemment, les modèles à espace d'états (SSMs), et plus spécifiquement les modèles à espace d'états sélectifs, avec une implémentation optimisée pour le matériel, ont montré un potentiel prometteur pour la modélisation de longues séquences. Motivés par le succès des SSMs, nous présentons MambaMixer, une nouvelle architecture avec des poids dépendants des données qui utilise un mécanisme de sélection dual à travers les tokens et les canaux, appelé Sélecteur de Tokens et de Canaux. MambaMixer connecte des mélangeurs sélectifs en utilisant un mécanisme de moyenne pondérée, permettant aux couches d'accéder directement aux caractéristiques initiales. À titre de preuve de concept, nous concevons les architectures Vision MambaMixer (ViM2) et Time Series MambaMixer (TSM2) basées sur le bloc MambaMixer et explorons leurs performances dans diverses tâches de vision et de prévision de séries temporelles. Nos résultats soulignent l'importance du mélange sélectif à la fois des tokens et des canaux. Dans les tâches de classification ImageNet, de détection d'objets et de segmentation sémantique, ViM2 atteint des performances compétitives par rapport aux modèles de vision bien établis et surpasse les modèles de vision basés sur les SSMs. Dans la prévision de séries temporelles, TSM2 obtient des performances exceptionnelles par rapport aux méthodes de pointe tout en démontrant une amélioration significative du coût computationnel. Ces résultats montrent que bien que les Transformers, l'attention inter-canaux et les MLPs soient suffisants pour de bonnes performances en prévision de séries temporelles, aucun n'est nécessaire.

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.