EfficientViM: Mamba de Visão Eficiente com Misturador de Estado Oculto baseado na Dualidade do Espaço de Estado

EfficientViM: Efficient Vision Mamba with Hidden State Mixer based State Space Duality

November 22, 2024

Autores: Sanghyeok Lee, Joonmyung Choi, Hyunwoo J. Kim

cs.AI

Resumo

Para a implementação de redes neurais em ambientes com recursos limitados, trabalhos anteriores construíram arquiteturas leves com convolução e atenção para capturar dependências locais e globais, respectivamente. Recentemente, o modelo de espaço de estados surgiu como uma interação eficaz de token global com seu favorável custo computacional linear no número de tokens. No entanto, espinhas dorsais eficientes de visão construídas com SSM foram menos exploradas. Neste artigo, apresentamos Efficient Vision Mamba (EfficientViM), uma arquitetura inovadora construída com base na dualidade de espaço de estados baseada em misturador de estado oculto (HSM-SSD) que captura eficientemente dependências globais com custo computacional reduzido. Na camada HSM-SSD, redesenhamos a camada SSD anterior para permitir a operação de mistura de canais dentro dos estados ocultos. Além disso, propomos a fusão de estados ocultos em múltiplos estágios para reforçar ainda mais o poder de representação dos estados ocultos, e fornecemos o design para aliviar o gargalo causado pelas operações limitadas pela memória. Como resultado, a família EfficientViM alcança uma nova relação velocidade-precisão de ponta no ImageNet-1k, oferecendo até 0,7% de melhoria de desempenho sobre o segundo melhor modelo SHViT com maior velocidade. Além disso, observamos melhorias significativas na taxa de transferência e precisão em comparação com trabalhos anteriores, ao dimensionar imagens ou empregar treinamento por destilação. O código está disponível em https://github.com/mlvlab/EfficientViM.

English

For the deployment of neural networks in resource-constrained environments, prior works have built lightweight architectures with convolution and attention for capturing local and global dependencies, respectively. Recently, the state space model has emerged as an effective global token interaction with its favorable linear computational cost in the number of tokens. Yet, efficient vision backbones built with SSM have been explored less. In this paper, we introduce Efficient Vision Mamba (EfficientViM), a novel architecture built on hidden state mixer-based state space duality (HSM-SSD) that efficiently captures global dependencies with further reduced computational cost. In the HSM-SSD layer, we redesign the previous SSD layer to enable the channel mixing operation within hidden states. Additionally, we propose multi-stage hidden state fusion to further reinforce the representation power of hidden states, and provide the design alleviating the bottleneck caused by the memory-bound operations. As a result, the EfficientViM family achieves a new state-of-the-art speed-accuracy trade-off on ImageNet-1k, offering up to a 0.7% performance improvement over the second-best model SHViT with faster speed. Further, we observe significant improvements in throughput and accuracy compared to prior works, when scaling images or employing distillation training. Code is available at https://github.com/mlvlab/EfficientViM.