Technische Samenvatting: DCVC-MB: Neurale B-Frame Videocompressie met State Space Models

Probleemstelling

Hoewel traditionele videocodecs (bijv. VVC, HEVC) aanzienlijke compressiewinsten behalen door middel van bidirectionele voorspelling (B-frames), hebben Neurale Videocodecs (NVC's) zich voornamelijk gericht op unidirectionele P-frame structuren. Bestaande pogingen om B-frames in NVC's te integreren, worstelen met twee primaire problemen:

1. Ineffectieve Fusie: Eerdere methoden slagen er niet in om bidirectionele temporele informatie efficiënt te fuseren, wat vaak leidt tot een lagere prestatie vergeleken met P-frame-only benaderingen.
2. Computationele Complexiteit: Transformer-architecturen, die in staat zijn om langetermijnafhankelijkheden te vangen, schalen kwadratisch ($O(N^2)$) met de resolutie. Dit leidt tot een verbruik van het geheugen dat de VRAM-limieten overschrijdt (bij hoge resoluties zoals 720p, 1080p) en tot prohibitieve computationele kosten, wat ze onpraktisch maakt voor real-world implementatie.

Methodologie

De auteurs stellen DCVC-MB (DCVC-Mamba) voor, een neuraal videocodec-framework dat specifiek is ontworpen voor B-frame codering met behulp van State Space Models (SSMs). De architectuur is gebaseerd op het DCVC-DC framework, maar introduceert fundamentele wijzigingen om gelijktijdige toegang tot zowel verleden als toekomstige temporele informatie mogelijk te maken.

Kerncomponenten van de Architectuur

1. Bidirectionele Feature Fusie via Mamba:

   • In tegenstelling tot eerdere modellen die een enkele Decoded Picture Buffer (DPB) uit het verleden gebruiken, houdt DCVC-MB twee DPB's bij: één van het vorige frame ($X_{t-1}$) en één van het toekomstige frame ($X_{t+1}$).
   • Deze features worden gefuseerd met behulp van een Bidirectional Mamba Block (BMB). De Mamba-architectuur biedt lineaire complexiteit ($O(N)$) vergeleken met de kwadratische complexiteit van Transformers, wat efficiënte verwerking op hoge resoluties mogelijk maakt.
   • Het fusieproces omvat het concateneren van referentieframes, het embedden van deze via een 2D CNN, het tokeniseren van de frames, en het verwerken van de sequentie door de SSM. De SSM modelleert systeemdynamica met behulp van een gediscretiseerde Ordinary Differential Equation (ODE) met leerbare parameters ($A, B, C$).
   • Om complexe bidirectionele afhankelijkheden te vangen, maakt het model gebruik van een Spatial-First Bidirectional methode, waarbij tokens in zowel voorwaartse als achterwaartse row-major scans worden verwerkt en de outputs worden opgeteld.

2. Adaptieve Latent Skipping:

   • Om de tijd voor entropiecodering te verminderen, gebruikt het framework een mechanisme voor adaptieve latent skipping. In plaats van een vaste drempelwaarde wordt een relatieve drempelstrategie gebruikt op basis van de standaarddeviatie ($\sigma$) van de latente elementen binnen een frame.
   • Latents met een $\sigma_i < \text{mean}(\sigma)$ worden overgeslagen (vervangen door een gemiddelde waarde), terwijl anderen worden gecodeerd. Dit vermindert de entropiecoderingstijd met ongeveer 9$\times$ voor P-frames en 5$\times$ voor B-frames, zonder asymmetrische computationele last te introduceren.

3. Canonical Position Embedding:

   • Om het verschil in resolutie tussen training (256$\times$256) en inferentie (tot 1080p) aan te pakken, vervangen de auteurs standaard pixel-coördinaat positie-embeddings door Canonical Position Embeddings.
   • Ruimtelijke coördinaten worden genormaliseerd naar een canoniek beeldvlak ($x, y \in [0, 1]$) in plaats van ruwe pixelwaarden te gebruiken. Dit zorgt ervoor dat de positionele harmonieken consistent blijven over verschillende resoluties, waardoor overfitting op de trainingsresolutie wordt voorkomen.

4. Optimalisaties tijdens Inferentie:

   • Open GoP: Hiermee kan het laatste B-frame van een Group of Pictures (GoP) refereren aan het I-frame van de daaropvolgende GoP, wat zorgt voor dichterbij gelegen temporele referenties.
   • Bidirectionele Codering: Codeert oneven geïndexeerde GoP's in zowel voorwaartse als achterwaartse richting, waarbij de resultaten met de laagste rate-distortion kosten worden geselecteerd. Even geïndexeerde GoP's gebruiken een gesloten configuratie om de lineaire complexiteit te behouden.

Belangrijkste Bijdragen

• Nieuwe Architectuur: Introductie van een bidirectionele feature fusie-architectuur gebaseerd op State Space Models (Mamba) die efficiënte B-frame codering mogelijk maakt op hoge resoluties waar Transformer-gebaseerde benaderingen falen door geheugenbeperkingen.
• IBP Frame Strategie: Implementatie van een basis neuraal compressiemodel dat ondersteuning biedt voor Random Access (RA) videocodering middels IBP frame-ordening.
• Efficiëntiemechanismen: Ontwikkeling van een adaptieve entropie-gebaseerde latent-skipping techniek om de entropiecodering te versnellen, en de integratie van inferentie-strategieën (Open GoP en Bidirectionele Codering) om de compressieprestaties te verbeteren.
• Resolutie-generalisatie: Het gebruik van canonical position embeddings om te zorgen dat het model effectief generaliseert van lage-resolutie trainingsdata naar hoge-resolutie inferentie.

Experimentele Resultaten

De methode is geëvalueerd op standaard datasets (HEVC, UVG, MCL-JCV) met RGB BT.709 conversie, waarbij werd vergeleken met VTM-19.0 (LDP en RA configuraties) en state-of-the-art NVC's (DCVC-DC, DCVC-FM).

• Prestaties vs. Traditionele Codecs: DCVC-MB bereikt BD-rate reducties van tot wel 30,45% ten opzichte van VTM-19.0-LDP en 1,81% ten opzichte van VTM-19.0-RA (Inter-GoP=16).
• Prestaties vs. Neurale Codecs: Het model vertoont gemiddelde BD-rate reducties van 8,98% vergeleken met eerdere neurale videocodecs (specifiek DCVC-FM) en 20,19% ten opzichte van de sterkere VTM-19.0-RA benchmark.
• Efficiëntie: De lineaire schaling van het Mamba fusion model zorgt ervoor dat de volledige B-frame pipeline binnen 13,8GB VRAM past bij 1080p, terwijl een vergelijkbare Transformer-architectuur meer dan 40GB overschrijdt.
• Ablatie-studies: Experimenten bevestigen dat het Mamba fusion model verantwoordelijk is voor het merendeel van de BD-rate winsten ($-4,55\%$), waarbij de inferentie-strategieën additieve verbeteringen leveren. Er werd aangetoond dat canonical position embeddings cruciaal zijn voor generalisatie naar onbekende resoluties.

Betekenis en Claims

Het artikel beweert dat DCVC-MB een significante vooruitgang vertegenwoordigt in neurale videocompressie door succesvol de kloof te overbruggen tussen de compressie-efficiëntie van bidirectionele voorspelling en de computationele haalbaarheid van deep learning modellen.

De auteurs stellen dat:

1. Schaalbaarheid: Door gebruik te maken van Mamba bereikt het framework praktische B-frame codering bij hoge resoluties (720p, 1080p), waarmee de geheugenbottlenecks die eerdere bidirectionele NVC's beperkten, worden overwonnen.
2. Concurrentiekracht: DCVC-MB is de eerste NVC die de sterkere VTM-19.0-RA configuratie overtreft, wat aantoont dat neurale benaderingen nu kunnen concurreren met of zelfs beter kunnen presteren dan leidende traditionele codecs in Random Access scenario's.
3. Praktische Bruikbaarheid: Door middel van adaptieve latent skipping en optimalisaties tijdens de inferentie, balanceert de methode hoge compressieprestaties met een verminderde encoding- en decodingtijd, waardoor de weg naar praktische implementatie van neurale videocompressie wordt geëffend.

De auteurs erkennen beperkingen, waarbij zij opmerken dat de prestaties kunnen verslechteren op specifieke datasets (bijv. MCL-JCV) door domeinverschillen tussen trainingsdata (natuurlijke video) en testdata (bijv. geanimeerde content met drastische beweging), en dat YUV420 optimalisatie nog steeds een uitdaging vormt vergeleken met RGB training.