MambaMixer: Modelos de Espacios de Estados Selectivos Eficientes con Selección Dual de Tokens y Canales

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

March 29, 2024

Autores: Ali Behrouz, Michele Santacatterina, Ramin Zabih

cs.AI

Resumen

Los avances recientes en aprendizaje profundo han dependido principalmente de los Transformers debido a su dependencia de los datos y su capacidad para aprender a gran escala. Sin embargo, el módulo de atención en estas arquitecturas exhibe un tiempo y espacio cuadráticos en función del tamaño de la entrada, lo que limita su escalabilidad para el modelado de secuencias largas. A pesar de los intentos recientes de diseñar arquitecturas eficientes y efectivas para datos multidimensionales, como imágenes y series temporales multivariadas, los modelos existentes son independientes de los datos o no permiten la comunicación inter e intra-dimensional. Recientemente, los Modelos de Espacio de Estados (SSMs), y más específicamente los Modelos de Espacio de Estados Selectivos, con una implementación eficiente orientada al hardware, han mostrado un potencial prometedor para el modelado de secuencias largas. Motivados por el éxito de los SSMs, presentamos MambaMixer, una nueva arquitectura con pesos dependientes de los datos que utiliza un mecanismo de selección dual a través de tokens y canales, denominado Mezclador Selectivo de Tokens y Canales. MambaMixer conecta mezcladores selectivos mediante un mecanismo de promediado ponderado, permitiendo que las capas tengan acceso directo a características tempranas. Como prueba de concepto, diseñamos las arquitecturas Vision MambaMixer (ViM2) y Time Series MambaMixer (TSM2) basadas en el bloque MambaMixer y exploramos su rendimiento en diversas tareas de visión y predicción de series temporales. Nuestros resultados subrayan la importancia de la mezcla selectiva tanto en tokens como en canales. En tareas de clasificación en ImageNet, detección de objetos y segmentación semántica, ViM2 logra un rendimiento competitivo con modelos de visión bien establecidos y supera a los modelos de visión basados en SSMs. En la predicción de series temporales, TSM2 alcanza un rendimiento sobresaliente en comparación con los métodos más avanzados, demostrando además una mejora significativa en el costo computacional. Estos resultados muestran que, aunque los Transformers, la atención entre canales y las MLPs son suficientes para un buen rendimiento en la predicción de series temporales, ninguno de ellos es necesario.

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.