Content-Aware Mamba for Learned Image Compression

Probleemstelling

Recente ontwikkelingen in Learned Image Compression (LIC) maken gebruik van Mamba-achtige State-Space Modellen (SSM's) vanwege hun vermogen om een globaal receptief veld te bieden met lineaire complexiteit. Echter, de standaard Mamba-architectuur heeft twee fundamentele beperkingen wanneer deze wordt toegepast op beeldcompressie:

1. Rigiditeit en Content-Agnosticisme: Standaard Mamba gebruikt een vooraf gedefinieerde, vaste raster-scan (of multi-directionele scans) die onafhankelijk is van de inhoud van de afbeelding. Dit is suboptimaal voor compressie, omdat effectieve redundantiemodellering interacties vereist tussen tokens die semantisch vergelijkbaar zijn, maar ruimtelijk ver van elkaar verwijderd. De vaste scanpaden groeperen ongerelateerde tokens en splitsen gerelateerde tokens, wat de modellering van lange-afstandsafhankelijkheden verzwakt.
2. Strict Causaliteit: Mamba is een causaal sequentieel model; een token kan alleen informatie van voorgaande tokens in de scanvolgorde benutten. Afbeeldingen zijn echter niet-causaal (informatie aan de rechterkant is even belangrijk als aan de linkerkant). Bestaande oplossingen, zoals multi-directionele scans, verdubbelen of verviervoudigen de computationele complexiteit en loss het fundamentele probleem van content-agnostische paden niet op.

Methodologie: Content-Aware Mamba (CAM)

De auteurs introduceren Content-Aware Mamba (CAM), een dynamische SSM die specifiek is ontworpen om deze beperkingen te overwinnen door de verwerking aan te passen aan de inhoud van de afbeelding. De kern van de methode bestaat uit twee innovatieve mechanismen:

1. Content-Adaptive Token Permutatie (CTP)

In plaats van tokens in een vaste rastervolgorde te verwerken, reorganiseert CAM de tokensequentie op basis van inhoudelijke gelijkenis.

• Codebook-based Clustering: Er wordt een gedeelde, leerbare codebook (centroïden) gebruikt om tokens te clusteren op basis van cosine-afstand. Dit is efficiënter en stabieler dan online K-means per sample.
• Permutatie: Tokens die tot dezelfde cluster behoren (en dus vergelijkbare visuele kenmerken hebben), worden gegroepeerd en achtereenvolgens verwerkt in de 1D-sequentie.
• Effect: Hierdoor worden tokens die ruimtelijk ver van elkaar liggen maar semantisch gerelateerd zijn, dicht bij elkaar gebracht in de scanvolgorde. Dit maximaliseert de interactie tussen relevante tokens binnen de lineaire complexiteit van Mamba.

2. Global-Prior Prompting (GPP)

Om de strikte causaliteit van Mamba te doorbreken zonder de kosten van multi-directionele scans, wordt een prompt-mechanisme geïntroduceerd.

• Redundantie-bewuste Dictionary: Een leerbare dictionary wordt gegenereerd uit de clustercentroïden. Elke entry in deze dictionary fungeert als een "prompt" vector die de globale statistieken van een specifieke semantische cluster encodeert.
• Conditionering: Voor een gegeven afbeelding wordt een sample-specifiek prompt gegenereerd op basis van de cluster-toewijzingen van de tokens. Dit prompt wordt geïnjecteerd in de output-projectiematrix ($C$) van de State-Space vergelijking.
• Effect: Hierdoor krijgt elke stap in de sequentiemodeling toegang tot globale priors van de hele afbeelding. Dit vermindert de afhankelijkheid van alleen voorgaande tokens en maakt het model "niet-causaal" in de zin dat het globale context begrijpt, zonder de computationele last te verhogen.

Architectuur: CMIC

Het voorgestelde model, CMIC (Content-Aware Mamba-based Image Compression), volgt een standaard VAE-architectuur (Variational Autoencoder) met:

• Non-lineaire Transformatie: Bestaande uit zes stadia met een combinatie van window-attention (voor lokale afhankelijkheden) en CAM-blokken (voor globale, content-afhankelijke afhankelijkheden).
• Entropy Model: Een verbeterde versie van het SCCTX-model (Space-Channel Context) voor het modelleren van de verdeling van de gekwantiseerde latent features.
• Training: Het model wordt getraind op de Flickr2W dataset met een rate-distortion verliesfunctie.

Belangrijkste Resultaten

CMIC behaalt state-of-the-art (SOTA) prestaties op drie standaard datasets (Kodak, Tecnick, CLIC) en overtreft zowel traditionele codecs als recente LIC-modellen.

• Rate-Distortion Prestaties:

  • CMIC verbetert de BD-rate (bitrate besparing ten opzichte van VTM-21.0) met 15,91% (Kodak), 21,34% (Tecnick) en 17,58% (CLIC).
  • Het presteert significant beter dan andere Mamba-gebaseerde modellen zoals MambaVC en MambaIC.
  • Het overtreft Transformer-gebaseerde modellen (zoals FTIC) en hybride CNN-Transformer modellen (zoals TCM-L) in zowel PSNR als MS-SSIM.

• Efficiëntie en Complexiteit:

  • Ondanks de geavanceerde mechanismen behoudt CMIC een lineaire complexiteit.
  • Het is aanzienlijk efficiënter dan MambaIC: 56% minder parameters, 57% minder FLOPs en 78% minder GPU-geheugengebruik.
  • De inferentie-tijd is laag (0,405s voor een 2K-afbeelding), wat slechts een minimale toename is ten opzichte van basismodellen.

• Visualisaties:

  • Effective Receptive Field (ERF): Visualisaties tonen aan dat CMIC een groter en content-adaptief receptief veld heeft. In tegenstelling tot andere modellen die isotrope, content-agnostische patronen tonen, richt CMIC zich op semantisch gerelateerde gebieden (bijv. veren, wolken, randen), ongeacht hun ruimtelijke afstand.

Betekenis en Conclusie

Dit paper is significant omdat het een fundamentele beperking van Mamba-architecturen in visuele taken oplost: de mismatch tussen de lineaire, causale scan en de 2D, niet-causale aard van beelden.

1. Paradigmaverschuiving: Het bewijst dat het mogelijk is om Mamba's lineaire complexiteit te behouden terwijl men toch effectief globale, content-afhankelijke redundantiemodellering bereikt, zonder de noodzaak voor kostbare multi-directionele scans.
2. Content-Awareness: Door de scanvolgorde dynamisch aan te passen aan de inhoud (via clustering) en globale priors in te voegen, creëert het model een veel efficiëntere representatie voor compressie.
3. Praktische Toepasbaarheid: De combinatie van SOTA compressieprestaties met lage latentie en geheugenefficiëntie maakt CMIC een sterke kandidaat voor praktische implementaties in beeldcompressie-applicaties.

Kortom, CMIC introduceert een nieuwe standaard voor Learned Image Compression door de statische aard van Mamba te transformeren in een dynamisch, inhoudsbewust systeem.