EfficientViM: Visión Eficiente Mamba con Mezclador de Estado Oculto basado en la Dualidad del Espacio de Estados.

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

Resumen

Para la implementación de redes neuronales en entornos con recursos limitados, trabajos previos han construido arquitecturas ligeras con convolución y atención para capturar dependencias locales y globales, respectivamente. Recientemente, el modelo de espacio de estados ha surgido como una interacción de tokens globales efectiva con su favorable costo computacional lineal en el número de tokens. Sin embargo, se ha explorado menos el uso de espinazos visuales eficientes construidos con SSM. En este documento, presentamos Vision Mamba Eficiente (EfficientViM), una arquitectura novedosa construida sobre la dualidad de espacio de estados basada en mezclador de estados ocultos (HSM-SSD) que captura eficientemente dependencias globales con un costo computacional reducido. En la capa HSM-SSD, rediseñamos la capa SSD anterior para habilitar la operación de mezcla de canales dentro de los estados ocultos. Además, proponemos una fusión de estados ocultos en múltiples etapas para reforzar aún más el poder de representación de los estados ocultos, y proporcionamos el diseño que alivia el cuello de botella causado por las operaciones limitadas por la memoria. Como resultado, la familia EfficientViM logra un nuevo equilibrio entre velocidad y precisión de vanguardia en ImageNet-1k, ofreciendo hasta un 0.7% de mejora en el rendimiento sobre el segundo mejor modelo SHViT con una velocidad más rápida. Además, observamos mejoras significativas en el rendimiento y la precisión en comparación con trabajos anteriores al escalar imágenes o al emplear entrenamiento por destilación. El código está disponible en 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.