Riepilogo Tecnico: SPARK – Spatial Policy-driven Adaptive Reinforcement Learning per la Knowledge Distillation

1. Definizione del Problema

La quantizzazione a basso bit (utilizzando 8 bit o meno) è essenziale per il dispiegamento di reti di restauro d'immagine (IR) su dispositivi con risorse limitate, come smartphone e droni. Tuttavia, la quantizzazione introduce un rumore di arrotondamento che degrada sproporzionatamente le regioni ad alta frequenza, inclusi bordi, texture fini e strutture ripetitive. Sebbene la Knowledge Distillation (KD) sia una tecnica standard per addestrare modelli leggeri supervisionandoli con un modello complesso (teacher), gli esistenti metodi di KD per l'IR applicano i segnali di distillazione in modo uniforme in tutte le posizioni spaziali.

Questo approccio uniforme è subottimale perché:

• Allocazione inefficiente del gradiente: Le regioni lisce a bassa frequenza (ad esempio, cieli uniformi) sono facili da ricostruire sia per il teacher che per lo student; forzare lo student a imitare il teacher in queste aree produce poco gradiente utile e sottrae risorse all'apprendimento nelle regioni più difficili.
• Modalità di fallimento critiche: Il rumore di quantizzazione è più dannoso nelle regioni ad alta frequenza, dove gli osservatori umani valutano la qualità. Le funzioni di perdita uniformi penalizzano efficacemente queste aree critiche.
• Euristiche statiche: I precedenti tentativi di affrontare questo problema utilizzando maschere di bordo fatte a mano o euristiche di difficoltà fisse si impegnano in un'unica nozione di "difficoltà" che rimane statica durante l'addestramento, fallendo nel adattarsi alla capacità evolutiva dello student o alle specifiche caratteristiche della degradazione.

2. Metodologia: Il Framework SPARK

Gli autori propongono SPARK (Spatial Policy-driven Adaptive Reinforcement learning for Knowledge distillation), un framework che tratta l'allocazione spaziale dello sforzo di distillazione come un problema decisionale sequenziale risolto tramite una politica di Reinforcement Learning (RL) leggera.

Componenti Core

1. Reti Teacher e Student:

   • Un Teacher Network di grande capacità, pre-addestrato e congelato, fornisce output di alta qualità ($I_t$) che fungono da segnale di supervisione morbido (soft).
   • Uno Student Network leggero (quantizzato tramite Quantization-Aware Training, QAT) impara a replicare il teacher. Durante l'addestramento, viene applicata una falsa quantizzazione per simulare il comportamento di deployment, mantenendo però i gradienti a piena precisione per l'ottimizzazione.

2. Estrattore di Caratteristiche di Difficoltà:
   Ad ogni step di addestramento, un modulo non differenziabile estrae quattro mappe di caratteristiche complementari per valutare la difficoltà di ricostruzione spaziale:

   • Varianza del Laplaciano: Cattura la nitidezza dei bordi (regioni ad alta frequenza).
   • Varianza dei Pixel: Misura la complessità della texture locale.
   • Errore di Ricostruzione dello Student: La differenza assoluta tra l'output dello student e la Ground Truth ($|Student - GT|$).
   • Gap di Conoscenza Teacher-Student: La differenza tra gli output del teacher e dello student.

3. Rete di Politica Spaziale (Agente RL):

   • Una compatta CNN (tre strati convoluzionali) elabora le caratteristiche di difficoltà su una griglia $8 \times 8$.
   • Produce una mappa di pesi spaziali stocastica ($W$) campionando da una distribuzione Gaussiana appresa ($\mathcal{N}(\mu, \sigma)$).
   • L'output è limitato nell'intervallo $[0, 1]$ per servire come mappa di pesi normalizzata.
   • Stocasticità: Il campionamento casuale permette alla politica di esplorare strategie di allocazione spaziale non banali che obiettivi puramente deterministici e supervisionati potrebbero perdere.

4. Obiettivo di Addestramento e Reward:

   • Funzione di Perdita: La perdita totale combina una standard L1 reconstruction loss (che ancora lo student alla Ground Truth) e una Spatial-Weighted KD Loss. La perdita KD è una Mean Squared Error (MSE) pesata tra lo student e il teacher, modulata dalla mappa di pesi della politica $W$.
   • Segnale di Reward: La politica è addestrata utilizzando una funzione di reward basata sul miglioramento della qualità di ricostruzione (PSNR e SSIM) rispetto a uno student di base addestrato senza pesatura spaziale.
   • Distacco del Gradiente (Gradient Detachment): Fondamentalmente, l'azione campionata (la mappa di pesi) viene staccata prima di essere applicata alla perdita KD. Ciò decupla l'apprendimento della politica dai gradienti di ricostruzione pixel-wise, permettendo alla politica di essere ottimizzata esclusivamente in base al segnale di reward globale.

5. Inference:
   La rete della politica RL viene scartata dopo l'addestramento. Di conseguenza, SPARK aggiunge zero costi di inferenza e non richiede modifiche architettoniche alla rete student, rendendolo agnostico rispetto al compito di IR (applicabile a Low-Light Enhancement, Denoising e Super-Resolution).

3. Principali Contributi

• Identificazione di un Collo di Bottiglia: Il paper identifica l'uniformità spaziale degli esistenti loss di KD come un collo di bottiglia primario per la quantizzazione a basso bit, motivando la necessità di una distillazione spazialmente adattiva per una restauro robusto alla quantizzazione.
• Politica Spaziale Guidata da RL: L'introduzione di SPARK, la prima strategia di pesatura spaziale basata su RL per la KD nel basso-livello IR. Utilizza caratteristiche di difficoltà interpretabili per apprendere una mappa di pesi per regione end-to-end durante la QAT.
• Esplorazione Stocastica: Utilizzando un'azione campionata da una Gaussia stocastica con gradienti staccati, il framework consente l'esplorazione di allocazioni spaziali inaccessibili a obiettivi puramente differenziabili.
• Ampia Applicabilità: Il metodo è dimostrato efficace attraverso diversi compiti di IR (LLE, Denoising, SR) e varie architetture student senza richiedere modifiche specifiche per il compito.

4. Risultati Sperimentali

Gli autori hanno valutato SPARK su benchmark standard: LOLv1 (Low-Light Enhancement), SIDD (Image Denoising) e Urban100 (Super-Resolution).

• Performance Quantitative:

  • SPARK ha costantemente superato PTQ (Post-Training Quantization), la standard QAT e i metodi di KD allo stato dell'arte (FAKD, DCKD, SLKD) attraverso molteplici modelli student (MIRNet, DnCNN, RFDN, SRCNN).
  • LOLv1: SPARK ha raggiunto 22.01 dB PSNR (rispetto ai 21.67 dB del secondo miglior metodo, SLKD) con MIRNet come student.
  • SIDD: SPARK ha raggiunto 33.89 dB PSNR (rispetto a 32.14 dB per SLKD).
  • Urban100 (SR): Alla scala $\times 2$, SPARK ha ridotto il gap rispetto al teacher a piena precisione a soli -0.16 dB. Alla scala $\times 4$, ha superato PTK di +2.53 dB per gli studenti RFDN.
  • SPARK ha costantemente ottenuto i valori più alti di PSNR/SSIM e i valori più bassi di LPIPS.

• Performance Qualitative:

  • I confronti visivi hanno mostrato che SPARK recupera la luminosità e i dettagli strutturali in modo più accurato in scenari di scarsa illuminazione rispetto ai competitor.
  • Nel denoising, SPARK produce risultati più puliti con maggiore uniformità cromatica e bordi più nitidi, evitando il rumore residuo e la distorsione del colore visti in altri metodi.
  • Nella Super-Resolution, SPARK ricostruisce caratteristiche strutturali più fini e geometrie accurate, evitando l'over-smoothing e gli artefatti di aliasing comuni nelle approcci SOTA.

• Studi di Ablazione:

  • I confronti hanno confermato che il meccanismo di reward basato su RL è il principale motore del miglioramento, superando significativamente l'addestramento senza RL (usando solo L1 o perdita MSE).

5. Significato e Rivendicazioni

Il paper sostiene che SPARK stabilisce una base per la distillazione spazialmente adattiva nella visione a basso livello. La sua importanza risiede in:

• Efficienza: Risolve il problema del degrado indotto dalla quantizzazione nelle regioni critiche senza aggiungere overhead di inferenza.
• Adattabilità: Sostituisce le euristiche statiche e artigianali con un processo decisionale apprendibile e dipendente dallo stato che evolve man mano che la capacità dello student cambia.
• Generalizzazione: Dimostra che un framework unificato e agnostico rispetto al compito può gestire efficacemente diversi compiti di restauro (LLE, Denoising, SR) sotto significativi vincoli computazionali.

Gli autori riconoscono modestamente un limite: il segnale di reward si basa su metriche globali (PSNR/SSIM), che potrebbero non catturare completamente la qualità percettiva in tutti gli scenari di degradazione, suggerendo questo come direzione per il lavoro futuro. Tuttavia, i risultati attuali validano l'ipotesi centrale che la distillazione spazialmente adattiva sia superiore alla supervisione uniforme per l'IR a basso bit.