SF-Mamba: Rethinking State Space Model for Vision

Die Arbeit stellt SF-Mamba vor, einen effizienten visuellen State Space Model, der durch einen Hilfs-Patch-Swap für bidirektionale Informationsflüsse und ein Batch-Folding-Verfahren für verbesserte GPU-Parallelisierung die Grenzen bestehender Mamba-Architekturen überwindet und dabei in verschiedenen visuellen Aufgaben sowie beim Durchsatz neue State-of-the-Art-Ergebnisse erzielt.

Das Wesentliche

Die große Herausforderung: Das Bild-Puzzle

Stell dir vor, du möchtest einem Computer beibringen, ein Foto zu verstehen. Früher haben wir dafür CNNs (wie einen sehr sorgfältigen Maler, der Bild für Bild streicht) oder Vision Transformer (ViT) (wie einen genialen, aber teuren Detektiv, der jedes Teil des Puzzles mit jedem anderen vergleicht) verwendet.

Der Detektiv (ViT) ist sehr gut, aber er hat ein Problem: Wenn das Puzzle riesig ist (ein hochauflösendes Foto), muss er so viele Vergleiche anstellen, dass er langsam wird und die Rechenleistung explodiert. Das ist wie wenn du in einer riesigen Bibliothek jedes Buch mit jedem anderen vergleichen müsstest, um eine Information zu finden – das dauert ewig.

Dann kam Mamba auf den Plan. Mamba ist wie ein effizienter Läufer, der durch die Bibliothek sprintet. Er liest die Bücher nacheinander (von links nach rechts) und merkt sich das Wichtigste. Das ist super schnell und spart Energie. Aber er hat einen Haken: Er kann nicht zurückblicken. Wenn er am Ende des Buches steht, weiß er nicht mehr, was am Anfang stand, es sei denn, er hat es sich perfekt gemerkt. Bei Bildern ist das aber ein Problem: Ein Pixel oben rechts ist oft wichtig für einen Pixel unten links. Ein reiner „Läufer" verpasst diese Zusammenhänge.

Bisherige Lösungen für dieses Problem waren wie ein chaotischer Tanz: Man ließ den Läufer mehrmals durch das Bild rennen (einmal vorwärts, einmal rückwärts, einmal diagonal). Das funktionierte gut für die Genauigkeit, war aber extrem langsam und energieaufwendig, weil man die Daten ständig neu sortieren musste.

Die Lösung: SF-Mamba (Der clevere Kurier)

Die Autoren von SF-Mamba haben sich gedacht: „Warum müssen wir den Läufer so oft hin und her jagen? Wir können ihn smarter machen." Sie haben zwei geniale Tricks erfunden:

1. Der „Geheim-Notizblock" (Auxiliary Patch Swapping)

Stell dir vor, unser Läufer (Mamba) läuft durch das Bild von links nach rechts. Er verpasst also alles, was hinter ihm liegt.

• Der alte Weg: Man ließ ihn viermal durch das Bild rennen, um alles zu sehen. (Langsam!)
• Der SF-Mamba-Weg: Wir geben dem Läufer zwei kleine, unsichtbare Notizblöcke mit.
  • Einen am Anfang des Laufs (den „Kopf-Block").
  • Einen am Ende des Laufs (den „Schwanz-Block").
  • Während er läuft, sammelt der Schwanz-Block alle Informationen aus dem Bild.
  • Der Trick: Sobald er fertig ist, tauschen wir die Blöcke! Der „Schwanz-Block" (der jetzt alles weiß) wird zum neuen „Kopf-Block" für den nächsten Durchgang.
  • Das Ergebnis: Der Läufer läuft nur einmal, aber durch den Tausch bekommt er quasi „Rückblick"-Informationen, ohne den ganzen Weg nochmal laufen zu müssen. Es ist, als würde ein Kurier, der eine Nachricht von A nach B bringt, am Zielort einen Zettel mit der gesamten Route hinterlassen, den der nächste Kurier sofort mitnimmt.

2. Der „Zusammenklapp-Trick" (Batch Folding)

Stell dir vor, du hast eine Fabrik (deine Grafikkarte/GPU), die Pakete verarbeitet. Normalerweise verarbeitet die Maschine ein Paket nach dem anderen. Wenn die Pakete (die Bild-Ausschnitte) aber sehr klein sind, steht die Maschine oft nur herum und wartet, weil sie für jedes kleine Paket einen ganzen Arbeitsschritt einrichten muss. Das ist ineffizient.

• Der SF-Mamba-Weg: Wir nehmen viele kleine Pakete und falten sie zu einem langen, dicken Paket zusammen, als wären sie ein einziges großes Paket.
• Das Problem dabei: Wenn wir sie zusammenfalten, könnte die Maschine denken, dass die Informationen aus Paket A in Paket B „lecken" (vermischt werden).
• Die Lösung: Wir setzen an den Grenzen der kleinen Pakete eine unsichtbare Mauer (ein „Reset"). Sobald die Maschine ein kleines Paket fertig bearbeitet hat, setzt sie den Gedächtnis-Speicher sofort zurück, bevor sie mit dem nächsten kleinen Paket im langen Bündel weitermacht.
• Das Ergebnis: Die Maschine läuft jetzt im vollen Sprint, weil sie einen langen Strom von Daten verarbeitet, vergisst aber nichts von den einzelnen Paketen. Es ist wie ein Zug, der viele kleine Waggons hat: Der Lokführer fährt schnell durch, aber an jeder Waggon-Grenze wird der Schalter für den nächsten Waggon neu gesetzt, damit die Passagiere nicht durcheinandergeraten.

Warum ist das so toll?

Durch diese zwei Tricks (den Notizblock-Tausch und das Zusammenklappen der Daten) erreicht SF-Mamba das Beste aus beiden Welten:

1. Geschwindigkeit: Es ist viel schneller als die bisherigen Mamba-Modelle, besonders bei kleinen Bildern oder vielen Bildern gleichzeitig.
2. Genauigkeit: Es versteht Bilder fast so gut wie die schweren, langsamen Transformer-Modelle, weil es durch den „Notizblock-Tausch" den ganzen Bildzusammenhang versteht.

Fazit:
Stell dir SF-Mamba wie einen extrem effizienten Logistik-Manager vor. Anstatt dass alle Mitarbeiter (die Daten) wild durcheinanderrennen, um sich abzustimmen, gibt es einen klaren Plan: Jeder läuft seinen Weg, tauscht aber kurzzeitig seine Notizen mit dem Chef aus, damit alle Bescheid wissen. Und gleichzeitig werden die kleinen Aufträge so gebündelt, dass die Maschinen nie stillstehen. Das macht KI schneller, günstiger und schlauer.

Technische Zusammenfassung

1. Problemstellung

State Space Models (SSMs), insbesondere Mamba, haben sich als vielversprechende Alternative zu Vision Transformers (ViTs) etabliert, da sie eine lineare Komplexität bieten und die quadratische Komplexität der Self-Attention überwinden. Dennoch leiden bestehende visuelle Mamba-Architekturen unter zwei wesentlichen Nachteilen:

• Eingeschränkte nicht-kausale Interaktionen: Mamba verarbeitet Daten typischerweise in einer rekursiven, unidirektionalen (links-nach-rechts) Scan-Reihenfolge. Dies verhindert, dass frühere Bild-Patches auf Informationen späterer Patches zugreifen können. Um dies zu kompensieren, verwenden viele Ansätze Multi-Scan-Strategien (z. B. bidirektionale oder Kreuz-Scans), bei denen die Daten mehrfach neu sortiert und aus verschiedenen Richtungen gescannt werden. Dies führt jedoch zu erheblichem Overhead durch Datenmanipulation und reduziert die Inferenzgeschwindigkeit erheblich.
• Ineffizienz bei kurzen Token-Längen: Die parallele Scan-Implementierung von Mamba auf GPUs (basierend auf Warp-Scans) ist für kurze Sequenzen (typisch bei visuellen Aufgaben mit niedriger bis mittlerer Auflösung) ineffizient. Da moderne GPUs oft mindestens 32 Threads pro Sequenz benötigen, führt eine kurze Sequenzlänge zu einer schlechten Auslastung der Hardware-Ressourcen, was Mamba in diesen Szenarien langsamer macht als Attention-basierte Modelle.

2. Methodik: SF-Mamba

Die Autoren schlagen SF-Mamba vor, eine neue visuelle Mamba-Architektur, die zwei Kerninnovationen einführt, um diese Probleme zu lösen:

A. Auxiliary Patch Swapping (Hilfliche Patch-Swapping)

Um die Notwendigkeit langsamer Multi-Scan-Strategien zu umgehen und dennoch bidirektionale Informationsflüsse zu ermöglichen, wird ein unidirektionaler Scan beibehalten, ergänzt durch einen leichten Mechanismus:

• Mechanismus: Zu Beginn jedes Mamba-Blocks werden zwei lernbare Hilfstokens (ein Kopf-Token $x_{head}$ und ein Schwanz-Token $x_{tail}$) an die Patch-Sequenz angehängt.
• Ablauf: Der kausale selektive Scan aggregiert den globalen Kontext der gesamten Sequenz in den Schwanz-Token ($y_{tail}$). Anschließend wird ein parameterfreier Swap-Operation durchgeführt, bei dem der Schwanz-Token des aktuellen Layers mit dem Kopf-Token des nächsten Layers getauscht wird.
• Effekt: Dies ermöglicht es, dass die Patches des nächsten Layers auf den globalen Kontext (Zukunftsinformationen aus der aktuellen Sicht) zugreifen können, ohne die Sequenz neu sortieren oder parallele Pfade zu benötigen. Der Overhead ist minimal (nur zwei Tokens), aber der Informationsfluss ist bidirektional.

B. Batch Folding mit Periodischem State Reset

Um die GPU-Ineffizienz bei kurzen Sequenzen zu beheben, wird eine Strategie zur besseren Parallelisierung entwickelt:

• Batch Folding: Die Batch-Dimension wird in die Sequenz-Dimension umgeformt. Statt $B$ kurze Sequenzen der Länge $T$ zu verarbeiten, werden $B_2$ Sequenzen zu einer längeren Sequenz der Länge $B_2 \cdot T$ zusammengefasst. Dies erhöht die effektive Sequenzlänge und ermöglicht eine effizientere Nutzung der 32 GPU-Threads pro Warp.
• Periodic State Reset: Da das Zusammenführen von Batches die Unabhängigkeit der Sequenzen verletzen würde (Informationen würden zwischen verschiedenen Bildern „lecken"), wird ein Trick angewendet: Alle $T$ Schritte (am Ende der ursprünglichen Sequenzlänge) wird die Matrix $A_t$ im SSM auf Null gesetzt. Dies bricht die Rekursion und setzt den versteckten Zustand ($h_t$) zurück, ohne die Eingabe- oder Ausgabe-Projektionen ($B_t, C_t$) zu beeinflussen.
• Ergebnis: Die Berechnung bleibt mathematisch äquivalent zur Verarbeitung einzelner Sequenzen, nutzt aber die GPU-Parallelität wie bei langen Sequenzen.

Die Architektur selbst folgt einem Mamba-Transformer-Hybrid-Design (ähnlich wie MambaVision), wobei CNNs für die Patch-Embedding und Mamba- sowie Attention-Blöcke in den späteren Stufen verwendet werden.

3. Hauptbeiträge

1. Effiziente Unidirektionale Scan-Strategie: Einführung des „Auxiliary Patch Swapping", das bidirektionale Informationsflüsse mit vernachlässigbarem Overhead ermöglicht und teure Multi-Scan-Strategien ersetzt.
2. GPU-Parallelisierung für visuelle Aufgaben: Entwicklung von „Batch Folding" mit periodischem State Reset, das die Inferenzgeschwindigkeit für kurze Sequenzen (niedrige Auflösung) drastisch steigert, indem es die GPU-Auslastung optimiert.
3. Umfassende Validierung: Die Methode wurde erfolgreich auf Bildklassifizierung, Objekterkennung und semantischer/Instanz-Segmentierung getestet und übertrifft den Stand der Technik (SOTA) in Bezug auf Genauigkeit und Durchsatz.

4. Ergebnisse

Die Experimente wurden auf ImageNet-1K (Klassifizierung), ADE20K (Segmentierung) und MS COCO (Detektion/Segmentierung) durchgeführt.

• Klassifizierung (ImageNet-1K): SF-Mamba erreicht eine überlegene Trade-off zwischen Genauigkeit und Durchsatz. Beispielsweise erreicht das SF-Mamba-T (Tiny) Modell bei 82,5 % Top-1-Genauigkeit einen Durchsatz von 7.600 Bildern/s, verglichen mit 6.662 Bildern/s bei MambaVision-T (bei ähnlicher Genauigkeit). Es schlägt auch reine Transformer- und CNN-basierte Modelle wie Swin, ConvNeXt und FasterViT.
• Segmentierung & Detektion: Auf ADE20K und COCO zeigt SF-Mamba konsistent bessere Ergebnisse als Baselines. Besonders bei hohen Auflösungen profitiert die Architektur von der linearen Komplexität von Mamba im Vergleich zur quadratischen Komplexität von Attention.
• Geschwindigkeitsanalyse: Die Analyse zeigt, dass Multi-Scan-Strategien (wie bei VMamba oder Vim) bis zu 40 % der Rechenzeit für das Neuordnen von Tokens und parallele Pfade verschwenden. SF-Mamba eliminiert diesen Overhead fast vollständig.
• Batch-Folding-Effekt: Durch das Batch-Folding konnte eine Beschleunigung des SSM-Kerns um 110 % bis 180 % erzielt werden, insbesondere bei kurzen Sequenzen.

5. Bedeutung und Fazit

SF-Mamba adressiert kritische Engpässe in der Anwendung von State Space Models für die Bildverarbeitung.

• Es beweist, dass unidirektionale Scans in Kombination mit intelligenten Informations-Routing-Mechanismen (Swapping) leistungsfähiger sein können als komplexe Multi-Scan-Ansätze.
• Es löst das Problem der GPU-Ineffizienz bei kurzen Sequenzen, was Mamba für eine breite Palette visueller Aufgaben (insbesondere solche mit niedriger bis mittlerer Auflösung) praktikabel macht.
• Die Arbeit legt einen neuen Standard für effiziente visuelle Backbones, die sowohl hohe Genauigkeit als auch hohen Durchsatz bieten, und zeigt, dass Mamba-basierte Modelle ViTs in vielen Szenarien nicht nur ersetzen, sondern übertreffen können.

Die Autoren kündigen an, den Quellcode nach der Veröffentlichung freizugeben, was die Reproduzierbarkeit und weitere Forschung in diesem Bereich fördern wird.