Technische Samenvatting: SPARK – Spatial Policy-driven Adaptive Reinforcement Learning voor Knowledge Distillation

1. Probleemstelling

Low-bit kwantisatie (met 8 bits of minder) is essentieel voor het implementeren van image restoration (IR) netwerken op middelen-beperkte apparaten zoals smartphones en drones. Kwantisatie introduceert echter afrondingsruis die hoogfrequente regio's, inclusief randen, fijne texturen en repetitieve structuren, onevenredig veel verslechtert. Hoewel Knowledge Distillation (KD) een standaardtechniek is voor het trainen van lichtgewicht modellen door ze te superviseren met een complexe teacher, passen bestaande KD-methoden voor IR de distillatiesignalen uniform toe op alle ruimtelijke locaties.

Deze uniforme aanpak is suboptimaal omdat:

• Inefficiënte Gradiëntallocatie: Gladde, laagfrequente regio's (bijv. egale luchten) zijn gemakkelijk te reconstrueren door zowel de teacher als de student; het dwingen van de student om de teacher in deze gebieden te imiteren levert weinig nuttige gradiënten op en verdringt het leren in moeilijkere regio's.
• Kritieke Foutmodi: Kwantisatieruis is het meest schadelijk in hoogfrequente regio's waar menselijke waarnemers de kwaliteit beoordelen. Uniforme verliesfuncties wegen deze kritieke gebieden effectief onder.
• Statische Heuristieken: Eerdere pogingen om dit aan te pakken met handmatige edge-masks of vaste moeilijkheidsgraad-heuristieken committeren zich aan één enkele definitie van "moeilijkheid" die statisch blijft tijdens de training, waardoor het niet kan adapteren aan de evoluerende capaciteit van de student of de specifieke kenmerken van de degradatie.

2. Methodologie: Het SPARK-framework

De auteurs stellen SPARK voor (Spatial Policy-driven Adaptive Reinforcement learning voor Knowledge distillation), een framework dat de ruimtelijke allocatie van distillatie-inspanning behandelt als een sequentieel beslissingsprobleem dat wordt opgelost via een lichtgewicht Reinforcement Learning (RL) policy.

Kerncomponenten

1. Teacher en Student Netwerken:

   • Een grote, vooraf getrainde, bevroren Teacher Network levert hoogwaardige outputs ($I_t$) die dienen als een soft supervisory signaal.
   • Een lichtgewicht Student Network (gekwantiseerd via Quantization-Aware Training, QAT) leert de teacher te repliceren. Tijdens de training wordt valse kwantisatie toegepast om het deployment-gedrag na te bootsen, terwijl full-precision gradiënten behouden blijven voor optimalisatie.

2. Difficulty Feature Extractor:
   Bij elke trainingsstap extraheert een niet-differentieerbare module vier complementaire feature maps om de ruimtelijke reconstructie-moeilijkheid te beoordelen:

   • Laplacian Variance: Legt de scherpte van randen vast (hoogfrequente regio's).
   • Pixel Variance: Meet lokale textuurcomplexiteit.
   • Student Reconstruction Error: Het absolute verschil tussen de student output en de Ground Truth ($|Student - GT|$).
   • Teacher-Student Knowledge Gap: Het verschil tussen de teacher en student outputs.

3. Spatial Policy Network (RL Agent):

   • Een compacte CNN (drie convolutionele lagen) verwerkt de difficulty features op een $8 \times 8$ grid.
   • Het produceert een stochastische ruimtelijke gewichtsmap ($W$) door te samplen uit een geleerde Gaussische distributie ($\mathcal{N}(\mu, \sigma)$).
   • De output wordt afgekapt naar $[0, 1]$ om te dienen als een genormaliseerde gewichtsmap.
   • Stochasticiteit: Het willekeurige samplen maakt de policy in staat om niet-triviale ruimtelijke allocatiestrategieën te verkennen die deterministische, puur gesuperviseerde doelstellingen mogelijk zouden missen.

4. Trainingsdoelstelling en Beloning:

   • Verliesfunctie: De totale loss combineert een standaard L1-reconstructieverlies (dat de student verankert aan de Ground Truth) met een Spatial-Weighted KD Loss. De KD loss is een gewogen Mean Squared Error (MSE) tussen de student en de teacher, gemoduleerd door de gewichtsmap $W$ van de policy.
   • Beloningssignaal: De policy wordt getraind met een beloningsfunctie gebaseerd op de verbetering in reconstructiekwaliteit (PSNR en SSIM) ten opzichte van een baseline student die zonder ruimtelijke weging is getraind.
   • Gradient Detachment: Cruciaal is dat de gesamplede actie (de gewichtsmap) wordt losgekoppeld (detached) voordat deze wordt toegepast op de KD loss. Dit ontkoppelt het leren van de policy van de pixel-wise reconstructie gradiënten, waardoor de policy uitsluitend geoptimaliseerd kan worden op basis van het globale beloningssignaal.

5. Inference:
   De RL policy network wordt na de training weggegooid. Bijgevolg voegt SPARK nul inference kosten toe en vereist het geen architecturale wijzigingen aan het student netwerk, waardoor het IR-taak-agnostisch is (toe te passen op Low-Light Enhancement, Denoising en Super-Resolution).

3. Belangrijkste Bijdragen

• Identificatie van een Bottleneck: Het artikel identificeert de ruimtelijke uniformiteit van bestaande KD-losses als een primaire bottleneck voor low-bit IR, wat de noodzaak voor ruimtelijk adaptieve distillatie voor kwantisatie-robuuste restauratie motiveert.
• RL-gestuurde Spatial Policy: De introductie van SPARK, de eerste RL-gebaseerde ruimtelijke wegingstrategie voor KD in low-bit IR. Het maakt gebruik van interpreteerbare difficulty features om end-to-end een per-regio gewichtsmap te leren tijdens QAT.
• Stochastische Exploratie: Door een stochastische Gaussische sampling met gedetacheerde gradiënten te gebruiken, stelt het framework de exploratie van ruimtelijke allocaties mogelijk die ontoegankelijk zijn voor puur differentieerbare doelstellingen.
• Brede Toepasbaarheid: De methode wordt effectief aangetoond over diverse IR-taken (LLE, Denoising, SR) en verschillende student architecturen zonder dat er taakspecifieke aanpassingen nodig zijn.

4. Experimentele Resultaten

De auteurs evalueerden SPARK op standaard benchmarks: LOLv1 (Low-Light Enhancement), SIDD (Image Denoising) en Urban100 (Super-Resolution).

• Kwantitatieve Prestaties:

  • SPARK presteerde consequent beter dan PTQ (Post-Training Quantization), standaard QAT en state-of-the-art KD-methoden (FAKD, DCKD, SLKD) over meerdere student modellen (MIRNet, DnCNN, RFDN, SRCNN).
  • LOLv1: SPARK bereikte 22.01 dB PSNR (vs. 21.67 dB voor de op één na beste, SLKD) met MIRNet als student.
  • SIDD: SPARK bereikte 33.89 dB PSNR (vs. 32.14 dB voor SLKD).
  • Urban100 (SR): Bij $\times 2$ schaal reduceerde SPARK de kloof tot de full-precision teacher tot slechts -0.16 dB. Bij $\times 4$ schaal presteerde het +2.53 dB beter dan PTQ voor RFDN studenten.
  • SPARK behaalde consequent de hoogste PSNR/SSIM en laagste LPIPS scores.

• Kwalitatieve Prestaties:

  • Visuele vergelijkingen toonden aan dat SPARK helderheid en structurele details nauwkeuriger herstelt in situaties met weinig licht vergeleken met concurrenten.
  • In denoising produceerde SPARK schonere resultaten met hogere kleuruniformiteit en scherpere randen, waarbij de residuele ruis en kleurvervorming die bij andere methoden worden gezien, werd vermeden.
  • In Super-Resolution reconstrueerde SPARK fijnere structurele kenmerken en accurate geometrie, waarbij de overmatige smoothing en aliasing artefacten die gebruikelijk zijn bij SOTA-benaderingen werden vermeden.

• Ablatie Studies:

  • Vergelijkingen bevestigden dat het RL-gebaseerde beloningsmechanisme de primaire drijfveer van verbetering is, wat aanzienlijk beter presteert dan training zonder RL (gebruikmakend van alleen L1 of MSE loss).

5. Betekenis en Claims

Het artikel beweert dat SPARK een fundament legt voor spatiaal adaptieve distillatie in low-level vision. De betekenis ligt in:

• Efficiëntie: Het lost het probleem van kwantisatie-geïnduceerde degradatie in kritieke regio's op zonder inference overhead toe te voegen.
• Adaptiviteit: Het vervangt statische, handmatige heuristieken door een leerbaar, state-afhankelijk beslissingsproces dat evolueert naarmate de capaciteit van de student verandert.
• Generalisatie: Het demonstreert dat een verenigd, taak-agnostisch framework effectief verschillende restauratie-taken (LLE, Denoising, SR) kan afhandelen onder aanzienlijke computationele beperkingen.

De auteurs erkennen bescheiden een beperking: het beloningssignaal leunt op globale metrieken (PSNR/SSIM), die de perceptuele kwaliteit in alle degradatiescenario's mogelijk niet volledig kunnen vangen, wat dit suggereert als een richting voor toekomstig werk. De huidige resultaten valideren echter de kernhypothese dat ruimtelijk adaptieve distillatie superieur is aan uniforme supervisie voor low-bit IR.