EfficientViM: Effiziente Vision Mamba mit Hidden State Mixer basierend auf dem Dualismus des Zustandsraums

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

November 22, 2024

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

cs.AI

Zusammenfassung

Für den Einsatz von neuronalen Netzwerken in ressourcenbeschränkten Umgebungen haben frühere Arbeiten leichte Architekturen mit Faltungs- und Aufmerksamkeitsmechanismen entwickelt, um lokale und globale Abhängigkeiten zu erfassen. In letzter Zeit hat sich das Zustandsraummodell als effektive globale Token-Interaktion mit günstigen linearen Rechenkosten in der Anzahl der Token herausgestellt. Dennoch wurden effiziente Vision-Backbones, die mit SSM erstellt wurden, weniger erforscht. In diesem Artikel stellen wir Efficient Vision Mamba (EfficientViM) vor, eine neuartige Architektur, die auf dem Hidden-State-Mixer-basierten Zustandsraum-Dualität (HSM-SSD) aufbaut, um globale Abhängigkeiten effizient mit weiter reduzierten Rechenkosten zu erfassen. In der HSM-SSD-Schicht gestalten wir die vorherige SSD-Schicht neu, um den Kanalmischbetrieb innerhalb der Hidden States zu ermöglichen. Darüber hinaus schlagen wir eine mehrstufige Verschmelzung der Hidden States vor, um die Repräsentationskraft der Hidden States weiter zu verstärken, und bieten das Design zur Linderung des Engpasses durch die speichergebundenen Operationen. Als Ergebnis erzielt die EfficientViM-Familie einen neuen Stand der Technik im Speed-Accuracy-Trade-off auf ImageNet-1k und bietet eine Leistungssteigerung von bis zu 0,7% gegenüber dem zweitbesten Modell SHViT bei schnellerer Geschwindigkeit. Darüber hinaus beobachten wir signifikante Verbesserungen in Durchsatz und Genauigkeit im Vergleich zu früheren Arbeiten beim Skalieren von Bildern oder beim Einsatz von Distillationstraining. Der Code ist verfügbar unter 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.