Progressive Split Mamba: Effective State Space Modelling for Image Restoration

Het artikel introduceert Progressive Split-Mamba (PS-Mamba), een hiërarchisch framework dat de beperkingen van bestaande State Space Modellen voor beeldherstel overwint door topologie-bewuste partitie en kruis-schaal shortcuts te combineren voor zowel lokale structuurbehoud als efficiënte globale coherentie.

De kern

Stel je voor dat je een oude, beschadigde foto probeert te repareren. Je wilt de kleine details (zoals een rimpel in een gezicht of een haartje in een bos) scherp houden, maar je moet ook het grote plaatje (zoals de vorm van een gebouw of de horizon) correct laten kloppen.

Dit is precies wat PS-Mamba doet, maar dan voor computers. Hier is een uitleg in gewone taal, met een paar creatieve vergelijkingen.

Het Probleem: De "Reis" van de Pixels

Om een foto te verbeteren, moeten computers de pixels (de kleine puntjes van de foto) in een rijtje zetten en ze één voor één bekijken. Dit heet een "sequentie".

1. De oude manier (CNN's): Dit is alsof je door een klein raampje kijkt. Je ziet heel goed wat er direct voor je neus gebeurt, maar je ziet niet wat er verderop in de straat gebeurt. Je mist het grote plaatje.
2. De nieuwe manier (Transformers): Dit is alsof je een gigantische kaart hebt en overal tegelijk naar kunt kijken. Je ziet het grote plaatje perfect, maar je bent zo druk met het hele landschap dat je de kleine details (zoals een rimpel) soms vergeet of vaag maakt.
3. De "Mamba" manier (de basis): Dit is een slimme nieuwe methode die snel is en goed kan kijken over lange afstanden. Maar er zit een addertje onder het gras:
   • Het probleem: Omdat de computer de 2D-foto platlegt in één lange rij (van linksboven naar rechtsonder), raken buur-pixels elkaar kwijt. Een pixel die direct naast een andere zat, staat nu misschien helemaal aan het einde van de rij. Het is alsof je een familiefoto maakt, maar iedereen in de rij zet, zodat de broer en zus die naast elkaar zaten, nu aan tegenovergestelde kanten van de zaal staan. De computer vergeet dan dat ze familie zijn.
   • Het vervagingsprobleem: Omdat de computer zo lang moet "rekenen" om van het begin van de rij naar het einde te komen, vergeten de informatie uit het begin (de "oude" pixels) langzaam hun kracht. Het is alsof je een boodschap fluistert door een lange rij mensen; aan het einde is de boodschap nauwelijks nog te horen.

De Oplossing: PS-Mamba (De Slimme Verdelingsstrategie)

De auteurs van dit papier hebben een oplossing bedacht die we PS-Mamba noemen. Ze gebruiken twee slimme trucs om de foto te repareren zonder de structuur te verstoren.

1. De "Puzzelstukjes"-methode (Progressive Split)

In plaats van de hele foto in één lange, chaotische rij te zetten, snijden ze de foto eerst in stukjes.

• De analogie: Stel je voor dat je een grote puzzel moet maken. In plaats van alle 1000 stukjes in één grote hoop te gooien, leg je ze eerst in vierkante blokken van 10x10.
• Hoe het werkt: De computer kijkt eerst naar een klein blokje (bijvoorbeeld een halve foto, dan een kwart, dan een achtste). Omdat de stukjes klein zijn, blijven de buren (de pixels die naast elkaar zitten) ook echt naast elkaar zitten in de rij. De computer ziet de lokale details (zoals een oog of een blad) heel scherp, omdat ze niet door de hele foto hoeven te "reizen" om elkaar te vinden.

2. De "Express-lijn" (Symmetrische Shortcuts)

Om het probleem van de "vervaging" (waarbij informatie aan het einde van de rij zwak wordt) op te lossen, bouwen ze een snelle express-lijn.

• De analogie: Stel je voor dat je een bericht moet doorgeven in een lange trein. Normaal gaat het van de ene naar de andere passagier, en naarmate je verder komt, is het bericht vager. PS-Mamba bouwt echter een tunnel of een expresslift die direct van het begin van de trein naar het einde loopt.
• Hoe het werkt: De computer stuurt de "grote lijnen" van de foto (zoals de vorm van een gebouw) direct door via deze tunnels, zonder dat ze de lange, vermoeiende weg hoeven te lopen. Zo blijft de informatie helder en sterk, zelfs aan het andere einde van de foto.

Waarom is dit geweldig?

• Het beste van twee werelden: PS-Mamba houdt de kleine details scherp (door de kleine puzzelstukjes) én het grote plaatje helder (door de express-tunnels).
• Snel en efficiënt: Het is niet traag. Het werkt net zo snel als de snelle "Mamba"-methode, maar dan veel slimmer.
• Resultaat: Als je dit toepast op foto's die wazig zijn, ruis hebben of beschadigd zijn door JPEG-compressie (zoals bij oude internetfoto's), krijg je een resultaat dat eruitziet alsof de foto net van de camera is gekomen.

Samenvattend

Stel je voor dat je een oude, versleten tapijt wilt repareren.

• De oude methoden keken alleen naar één draadje per keer en wisten niet hoe het tapijt eruitzag.
• De andere methoden keken naar het hele tapijt, maar vergeten hoe de draden precies met elkaar verweven waren.
• PS-Mamba is als een meester-restaurator die het tapijt eerst in kleine, logische stukken snijdt om de patronen te zien, en tegelijkertijd een magische draad gebruikt die direct de vorm van het hele tapijt vasthoudt. Het resultaat is een tapijt dat er weer perfect uitziet, met alle details en de juiste vorm.

De wetenschappers hebben dit getest op super-resolutie (kleine foto's groot maken), ruisverwijdering en het verbeteren van JPEG-foto's, en in alle gevallen was hun methode beter dan de huidige beste methoden.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Progressive Split Mamba: Effective State Space Modelling for Image Restoration" in het Nederlands.

Titel: Progressive Split Mamba: Effectieve State Space Modelling voor Beeldherstel

1. Het Probleem

Beeldherstel (image restoration), zoals super-resolutie, ruisreductie en JPEG-artefactreductie, vereist een delicate balans tussen het behoud van fijne lokale structuren en het handhaven van lange-afstands ruimtelijke coherentie. Bestaande methoden hebben beperkingen:

• CNN's: Hebben een beperkt receptief veld en kunnen lange-afstandsrelaties moeilijk modelleren.
• Transformers: Hebben een kwadratische complexiteit voor globale aandacht en verwaarlozen vaak lokale nabuurschappen door hun window-based of globale mechanismen.
• State Space Models (SSM's / Mamba): Bieden een lineaire tijdscomplexiteit voor lange-afstandsmodelling, maar wanneer ze naïef op 2D-beelden worden toegepast, treden er twee fundamentele problemen op:
  1. Locality Distortion (Lokale Vervorming): Het platdrukken van 2D-kenmerkkaarten naar 1D-sequenties (rasterisatie) verstoort de ruimtelijke topologie. Naburige pixels in het beeld kunnen ver uit elkaar komen in de sequentie, wat de precisie bij het herstellen van lokale structuren ondermijnt.
  2. Long-Range Decay (Verval op Lange Afstand): De recurrente dynamiek van SSM's is ontworpen voor stabiliteit, wat leidt tot een exponentieel verval van informatie naarmate de sequentie langer wordt. Dit verzwakt globale consistentie en zorgt ervoor dat informatie van verre posities verdwijnt voordat deze de relevante plekken bereikt.

2. Methodologie: Progressive Split-Mamba (PS-Mamba)

De auteurs stellen PS-Mamba voor, een hiërarchisch, topologie-bewust framework dat de voordelen van Mamba behoudt terwijl het de bovengenoemde beperkingen oplost. De kernarchitectuur bestaat uit de volgende componenten:

• Progressieve Splitsing (Progressive Split):
  In plaats van het hele beeld in één lange sequentie te plakken, deelt PS-Mamba het beeld op in geometrisch consistente gebieden (patches) via een hiërarchie: halveren, kwadranten en octanten.

  • Dit behoudt de 4-verbonden nabuurschap van pixels binnen elke patch.
  • Het verkort de recurrente diepte per patch aanzienlijk, wat het verval van informatie beperkt en lokale details beter behoudt.
  • De complexiteit blijft lineair ($O(HW)$), omdat de som van de lengtes van de gesplitste sequenties gelijk blijft aan de totale beeldgrootte.

• Convolutionele Voorverwerking:
  Voordat de data de Mamba-module ingaat, wordt er een lichtgewicht convolutie-tak toegepast. Dit zorgt voor lokale continuïteit en versterkt structurele patronen (randen), waardoor de Mamba-module niet hoeft te worstelen met ruwe, ongeordende activaties.

• Symmetrische Cross-Scale Shortcut Paden:
  Om het probleem van het verval op lange afstand te compenseren, introduceert het model symmetrische shortcuts tussen gespiegelde niveaus in de hiërarchie. Deze paden transporteren direct laagfrequente globale context over het netwerk, waardoor de exponentiële afname van informatie in de lineaire SSM wordt tegengegaan en de gradient flow wordt gestabiliseerd.

• Fusie en Aandacht:
  De features van de convolutie-tak en de Mamba-tak worden gefuseerd via een content-adaptieve gating-mechanisme. Dit wordt gevolgd door een dubbele aandacht-verfijning (Channel Attention en Spatial Attention) om discriminatieve structuren te benadrukken en globale consistentie te versterken.

3. Belangrijkste Bijdragen

1. Topologie-bewuste State Space Module: Een nieuwe module die kenmerkkaarten verwerkt via een progressieve hiërarchie (halven, kwadranten, octanten) zonder de 2D-ruimtelijke structuur te verstoren. Dit elimineert raster-geïnduceerde fouten.
2. Symmetrische Cross-Scale Skip Links: Een mechanisme dat lange Mamba-ketens omzeilt om globale structurele informatie direct te leveren, wat het verval in lineaire SSM's effectief tegengaat.
3. Geïntegreerde Architectuur: Een unificatie van state-space modelling, lokale convolutie en dubbele aandacht, wat resulteert in een architectuur die zowel fijne details als globale context perfect kan modelleren zonder token-shifting of complexe herschikkingen.

4. Resultaten

PS-Mamba is uitgebreid getest op diverse beeldherstel-taken en presteert consequent beter dan de state-of-the-art (SOTA) methoden, waaronder recente Mamba-varianten (MambaIR, MambaIRv2) en Transformer-gebaseerde modellen (SwinIR, HAT, DAT).

• Super-Resolutie (SR):
  • Lightweight: PS-Mamba Light overtreft bestaande lichte modellen (zoals MambaIRv2-light) op alle datasets (Set5, Set14, BSDS100, Urban100, Manga109) met een duidelijke marge, terwijl het minder parameters en MACs (Multi-Adds) vereist.
  • Classic: De grote variant (PS-Mamba-Large) bereikt nieuwe SOTA-resultaten, zelfs met 34% minder parameters dan MambaIRv2-Large, maar met betere PSNR/SSIM scores.
• Ruisreductie (Denoising):
  • Op de taak van kleurbeelddenoising (Gaussian noise $\sigma=15$) behaalt het model de beste resultaten op alle geteste datasets (CBSD68, Kodak24, McMaster, Urban100), zelfs vergeleken met U-Net gebaseerde modellen zoals Restormer.
• JPEG Artefact Reductie:
  • PS-Mamba toont superioriteit bij het herstellen van beelden met compressie-artefacten, met name bij het reconstrueren van structurele details onder verschillende kwaliteitsfactoren.
• Efficiëntie:
  • Het model behoudt de lineaire complexiteit van Mamba, maar verbetert de stabiliteit en lokale precisie aanzienlijk zonder de rekentijd exponentieel te laten toenemen.

5. Betekenis en Conclusie

De paper introduceert een fundamentele verbetering in het toepassen van State Space Models op computer vision-taken. PS-Mamba lost het inherente dilemma op tussen lokale precisie en globale efficiëntie dat bestaande Mamba-implementaties voor 2D-beelden parten speelt.

De belangrijkste implicatie is dat het geometrisch behoud van lokale nabuurschap vóór de sequentiemodeling cruciaal is voor hoogwaardig beeldherstel. Door dit te combineren met mechanismen die het verval van lange-afstandsinformatie tegengaan, biedt PS-Mamba een robuust, efficiënt en schaalbaar framework. Het bewijst dat State Space Models een krachtig alternatief kunnen zijn voor zowel CNN's als Transformers, mits ze correct worden aangepast aan de 2D-natuur van beelddata. Dit opent de deur voor verdere toepassingen in videoherstel en andere domeinen waar zowel lokale details als globale coherentie essentieel zijn.