Trajectory-aware Shifted State Space Models for Online Video Super-Resolution

Dit artikel presenteert TS-Mamba, een nieuwe methode voor online videosuperresolutie die gebruikmaakt van trajectbewuste verschuifde State Space Models om lange-termijn temporele afhankelijkheden efficiënt te modelleren en zo state-of-the-art prestaties te bereiken met aanzienlijk minder rekenkracht.

De kern

Hoe een slimme "tijdsreiziger" je video's scherper maakt: De TS-Mamba

Stel je voor dat je een oude, wazige video bekijkt. Elke frame is onscherp, alsof je door een vieze bril kijkt. De taak van Video Super-Resolution (VSR) is om deze wazige beelden om te toveren naar haarscherpe, kristalheldere video's.

Maar er is een probleem: in het echte leven (zoals bij videogames of live-streams) kun je niet wachten tot de hele video is opgeladen. Je moet het nu doen, frame voor frame, terwijl de video nog draait. Dit heet Online Video Super-Resolution.

De meeste bestaande methoden zijn als een slome fotograaf die alleen naar de foto direct voor de huidige foto kijkt om te raden hoe de volgende eruit moet zien. Ze vergeten de rest van de film. Dat werkt goed voor simpele dingen, maar als er snel beweging is, wordt het beeld rommelig.

Hier komt TS-Mamba in het spel. Het is een nieuwe, slimme manier om video's scherper te maken, ontwikkeld door onderzoekers die hun werk presenteerden op het prestigieuze ICLR-congres. Laten we kijken hoe het werkt, zonder de moeilijke wiskunde.

1. De "Tijdsreiziger" (Trajecten)

Stel je voor dat je een bal gooit. Een simpele camera kijkt alleen naar waar de bal nu is en probeert te raden waar hij vandaan komt.
TS-Mamba doet iets anders: het tekent een spoor (een traject) van de bal. Het kijkt niet alleen naar de vorige foto, maar volgt de beweging van objecten door de tijd heen, alsof het een lijntje trekt door de hele film.

• De Analogie: Stel je voor dat je een groep mensen ziet dansen. De oude methoden kijken alleen naar de persoon links van jou. TS-Mamba kijkt naar iedereen in de kamer die een vergelijkbare dansbeweging maakt, zelfs als ze drie seconden geleden al in beeld waren. Het selecteert de "meest vergelijkbare dansers" uit het verleden om te helpen bij het tekenen van de huidige danser.

2. De "Slimme Scanner" (Mamba & Hilbert)

Om al deze informatie te verwerken, gebruikt TS-Mamba een nieuw type AI-model genaamd Mamba. Mamba is als een super-efficiënte scanner die een heel groot boek (de video) in één keer kan lezen, zonder dat het duizelig wordt.

Maar er is een valkuil: als je een 3D-ruimte (zoals een video) omzet in een lange lijst (1D) om te lezen, kun je de verbinding tussen de "buurman" en "buurvrouw" kwijtraken. Het is alsof je een puzzel uit elkaar haalt en de stukjes in een rij legt; soms liggen stukjes die dicht bij elkaar horen, nu ver uit elkaar.

• De Oplossing (Shifted Blocks): De onderzoekers hebben een trucje bedacht. Ze gebruiken een speciale manier van scannen (de Hilbert-scan) en voegen daar een verschuiving (shift) aan toe.
• De Analogie: Stel je voor dat je een tapijt afstoft. Als je alleen recht vooruit loopt, mis je de hoeken. TS-Mamba loopt niet alleen recht, maar maakt ook kleine zijwaartse sprongetjes en draait om. Zo zorgt het ervoor dat elk stukje van het tapijt (elk detail in de video) perfect wordt schoongemaakt en dat de patronen (de continuïteit) niet worden verbroken.

3. De "Slimme Leraar" (Verliesfunctie)

Tijdens het trainen van de AI moet de computer leren welke sporen (trajecten) het belangrijkst zijn.
De onderzoekers hebben een speciale straf-systeem (verliesfunctie) bedacht. Als de AI een verkeerd spoor trekt (bijvoorbeeld: "die auto bewoog naar links" terwijl hij naar rechts ging), krijgt hij een "traagheidspunt". Dit dwingt de AI om heel precies te kijken naar hoe objecten zich echt bewegen, zodat hij de juiste stukjes uit het verleden selecteert om de huidige foto te verbeteren.

Waarom is dit geweldig?

1. Snelheid: Het is niet alleen scherp, maar ook razendsnel. Het is ontworpen om op je telefoon of laptop te draaien zonder dat je computer in brand vliegt.
2. Efficiëntie: Het gebruikt ongeveer 22% minder rekenkracht dan de beste bestaande methoden, terwijl het resultaat vaak nog beter is.
3. Lange termijn: In plaats van alleen naar de vorige seconde te kijken, kan het kijken naar bewegingen die al een stukje verder in het verleden begonnen zijn. Dit helpt bij het reconstrueren van snelle bewegingen die anders wazig zouden worden.

Samenvattend

TS-Mamba is als een meester-restaurator die niet alleen naar het huidige schilderij kijkt, maar de hele geschiedenis van de penseelstreken volgt. Door slimme "verschuivingen" in zijn manier van kijken en een speciale focus op bewegingssporen, kan hij wazige video's omtoveren tot haarscherpe beelden, terwijl hij minder energie verbruikt dan zijn concurrenten.

Het is een grote stap voorwaarts voor live-video, videobellen en het streamen van content, waar snelheid en kwaliteit hand in hand moeten gaan.

Technische samenvatting

Probleemstelling

Online video super-resolution (VSR) is een cruciale techniek voor real-time toepassingen zoals videoconferenties en live-uitzendingen. Het doel is om een huidige hoogresolutie (HR) frame te reconstrueren op basis van een lage-resolutie (LR) input en eerdere frames.
Bestaande methoden hebben echter twee belangrijke beperkingen:

1. Beperkte temporale modellering: De meeste online VSR-methoden gebruiken slechts één voorgaand frame voor tijdelijke uitlijning. Dit beperkt het vermogen om langetermijn-tijdelijke afhankelijkheden te modelleren, wat essentieel is voor hoge reconstructiekwaliteit.
2. Complexiteit vs. Performance: Methoden die wel langetermijn-informatie gebruiken (zoals Vision Transformers of Diffusion-modellen) hebben vaak een te hoge rekencomplexiteit voor real-time toepassingen. Bestaande State Space Models (SSM's), zoals Mamba, bieden lineaire complexiteit en een groot receptief veld, maar lijden onder verlies van ruimtelijke continuïteit door hun scanningsmechanismen (zoals Hilbert-scanning) en zijn nog niet optimaal toegepast op online VSR.

Methodologie: TS-Mamba

De auteurs stellen TS-Mamba voor, een nieuw model dat gebruikmaakt van "Trajectory-aware Shifted State Space Models". De architectuur bestaat uit de volgende kerncomponenten:

1. Trajecto-gestuurde Token Selectie:

   • In plaats van alle informatie uit voorgaande frames te verwerken, construeert het model eerst trajecten binnen de video.
   • Op basis van deze trajecten worden de meest vergelijkbare tokens (features) uit de voorgaande frames geselecteerd. Dit zorgt voor een efficiënte aggregatie van langetermijn-tijdelijke informatie zonder overbodige berekeningen.
   • Een trajecto-bewuste verliesfunctie ($L_{trj}$) wordt gebruikt tijdens het trainen om de nauwkeurigheid van deze trajectgeneratie en token-selectie te waarborgen.

2. Trajectory-Aware Shifted Mamba Aggregation (TSMA):

   • Dit is het centrale aggregatiemodule dat de geselecteerde tokens samenvoegt met het huidige frame.
   • Het module gebruikt een "Scan-Shift-Scan" strategie om de inherente ruimtelijke discontinuïteiten van Hilbert-scanning in SSM's te compenseren.
   • Shifted SSM Blocks: Het model introduceert vier specifieke verschuivingsoperaties (bijv. Up, Up-Left) gecombineerd met verschillende Hilbert-scannings. Dit creëert "verschuifbare vensters" die de discontinuïteiten binnen en tussen vensters elimineren.
   • Er zijn twee parallelle takken in het TSMA-module:
     • Intra-window compensation branch: Compenseert discontinuïteit binnen lokale vensters.
     • Inter-window compensation branch: Compenseert discontinuïteit tussen verschillende vensters.
   • Door deze verschuivingen wordt de ruimtelijke continuïteit van het Mamba-model versterkt zonder de rekencomplexiteit significant te verhogen.

3. Selectieve Scanning (SS3D):

   • De tokens worden via een ruimtelijk Hilbert-gebaseerde selectieve scanning langs de tijdsdimensie verwerkt. Dit interweaves (verweeft) de geselecteerde tokens met het huidige frame, waardoor informatie over zowel ruimtelijke als tijdelijke dimensies kan interageren.

Belangrijkste Bijdragen

1. Eerste SSM-gebaseerde Online VSR: TS-Mamba is het eerste model dat State Space Models toepast op online VSR, waarbij het langetermijn-tijdelijke informatie op token-niveau aggregeert, in tegenstelling tot CNN-gebaseerde methoden die vaak beperkt zijn tot één frame.
2. Integratie van Video-Trajecten in Mamba: Voor het eerst worden video-trajecten gebruikt om de meest relevante tokens uit het verleden te selecteren, wat een efficiënte token-level aggregatie mogelijk maakt.
3. Nieuwe Shifted SSM Blocks: De auteurs ontwerpen nieuwe blokken die gebaseerd zijn op Hilbert-scanning en verschuivingsoperaties. Dit lost het probleem van ruimtelijke discontinuïteit op dat vaak voorkomt bij Mamba-modellen, waardoor lokale ruimtelijke samenhang wordt behouden.
4. Efficiëntie en Performance: Het model introduceert een trajecto-bewuste loss-functie om de selectie van tokens te optimaliseren, wat leidt tot een betere balans tussen kwaliteit en complexiteit.

Resultaten

TS-Mamba is getest op drie veelgebruikte datasets: REDS4, Vid4 en Vimeo-90K-T.

• Performance: Het model behaalt state-of-the-art (SOTA) resultaten in termen van PSNR en SSIM, vaak beter dan zes bestaande online VSR-benchmarkmodellen (zoals BasicVSR++, FDAN, KSNet, TMP).
• Complexiteit: Ondanks de verbeterde prestaties, reduceert TS-Mamba de rekencomplexiteit (gemeten in MACs) met meer dan 22,7% ten opzichte van de beste concurrenten.
• Snelheid: Het model bereikt een inferentiesnelheid die geschikt is voor real-time toepassingen (2e snelste onder de geteste online methoden, na TMP, maar met aanzienlijk lagere complexiteit).
• Ablatie Studies: Experimenten bevestigen dat zowel de trajecto-generatie, de loss-functie, als de specifieke verschuivingsoperaties (Intra/Inter-window) essentieel zijn voor de uiteindelijke prestaties.

Betekenis

Deze paper is significant omdat het een brug slaat tussen de efficiëntie van State Space Models (Mamba) en de eisen van real-time videoverwerking. Door het probleem van ruimtelijke discontinuïteit in Mamba op te lossen via slimme verschuivingen en door langetermijn-informatie selectief te benutten via trajecten, biedt TS-Mamba een nieuwe richting voor videoversterking. Het bewijst dat het mogelijk is om langetermijn-tijdelijke afhankelijkheden te modelleren zonder de hoge rekenkosten van Transformers of Diffusion-modellen, wat het zeer waardevol maakt voor applicaties met beperkte resources zoals mobiele apparaten en live-streaming.