MambaMixer: Эффективные модели пространства состояний с выбором двойных токенов и каналов.

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

March 29, 2024

Авторы: Ali Behrouz, Michele Santacatterina, Ramin Zabih

cs.AI

Аннотация

Недавние достижения в области глубокого обучения в основном опирались на трансформеры из-за их зависимости от данных и способности обучаться в масштабе. Однако модуль внимания в этих архитектурах демонстрирует квадратичную сложность по времени и пространству в зависимости от размера ввода, что ограничивает их масштабируемость для моделирования длинных последовательностей. Несмотря на недавние попытки разработать эффективную и эффективную архитектурную основу для многомерных данных, таких как изображения и многомерные временные ряды, существующие модели либо независимы от данных, либо не позволяют обеспечить взаимодействие между измерениями и внутри измерений. Недавно модели пространства состояний (SSM), а точнее, выборочные модели пространства состояний, с эффективной аппаратно-ориентированной реализацией, показали многообещающий потенциал для моделирования длинных последовательностей. Вдохновленные успехом SSM, мы представляем MambaMixer, новую архитектуру с весами, зависящими от данных, которая использует двойной механизм выбора между токенами и каналами, называемый выборочным микшером токенов и каналов. MambaMixer соединяет выборочные микшеры с помощью механизма взвешенного усреднения, позволяя слоям иметь прямой доступ к начальным признакам. В качестве доказательства концепции мы разрабатываем архитектуры Vision MambaMixer (ViM2) и Time Series MambaMixer (TSM2) на основе блока MambaMixer и исследуем их производительность в различных задачах видения и прогнозирования временных рядов. Наши результаты подчеркивают важность выборочного смешивания как между токенами, так и между каналами. В классификации ImageNet, обнаружении объектов и семантической сегментации ViM2 достигает конкурентоспособной производительности с хорошо установленными моделями видения и превосходит модели видения на основе SSM. В прогнозировании временных рядов TSM2 достигает выдающейся производительности по сравнению с передовыми методами, продемонстрировав значительно улучшенные вычислительные затраты. Эти результаты показывают, что хотя трансформеры, внимание между каналами и многослойные перцептроны достаточны для хорошей производительности в прогнозировании временных рядов, ни одно из них не является обязательным.

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.