MambaMixer: Effiziente selektive Zustandsraummodelle mit doppelter Token- und Kanalauswahl

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

March 29, 2024

Autoren: Ali Behrouz, Michele Santacatterina, Ramin Zabih

cs.AI

Zusammenfassung

In den letzten Fortschritten im Bereich des Deep Learning wurde hauptsächlich auf Transformer zurückgegriffen, aufgrund ihrer Datenabhängigkeit und ihrer Fähigkeit, im großen Maßstab zu lernen. Das Aufmerksamkeitsmodul in diesen Architekturen zeigt jedoch eine quadratische Zeit- und Speichernutzung in Bezug auf die Eingabegröße, was ihre Skalierbarkeit für das Modellieren von langen Sequenzen einschränkt. Trotz jüngster Versuche, eine effiziente und effektive Architekturgrundlage für multidimensionale Daten wie Bilder und multivariate Zeitreihen zu entwerfen, sind bestehende Modelle entweder datenunabhängig oder versagen darin, eine Kommunikation zwischen und innerhalb der Dimensionen zu ermöglichen. In letzter Zeit haben State Space Models (SSMs) und speziell Selektive State Space Models mit einer effizienten, hardwarebewussten Implementierung vielversprechendes Potenzial für das Modellieren langer Sequenzen gezeigt. Inspiriert vom Erfolg der SSMs präsentieren wir MambaMixer, eine neue Architektur mit datenabhängigen Gewichten, die einen dualen Auswahlmechanismus über Tokens und Kanäle verwendet, genannt Selektiver Token- und Kanalmixer. MambaMixer verbindet selektive Mixer mittels eines gewichteten Mittelungsmechanismus, der es Schichten ermöglicht, direkten Zugriff auf frühe Merkmale zu haben. Als Machbarkeitsnachweis entwerfen wir Vision MambaMixer (ViM2) und Time Series MambaMixer (TSM2) Architekturen auf Basis des MambaMixer-Blocks und untersuchen ihre Leistung in verschiedenen Vision- und Zeitreihenprognoseaufgaben. Unsere Ergebnisse unterstreichen die Bedeutung des selektiven Mischens über Tokens und Kanäle hinweg. Bei der ImageNet-Klassifikation, der Objekterkennung und der semantischen Segmentierung erzielt ViM2 eine wettbewerbsfähige Leistung im Vergleich zu etablierten Vision-Modellen und übertrifft SSM-basierte Vision-Modelle. Bei der Zeitreihenprognose erreicht TSM2 eine herausragende Leistung im Vergleich zu modernsten Methoden und zeigt dabei eine signifikant verbesserte Rechenleistung. Diese Ergebnisse zeigen, dass während Transformer, Cross-Channel-Aufmerksamkeit und MLPs für eine gute Leistung bei der Zeitreihenprognose ausreichen, keines davon notwendig ist.

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.