Technical Summary: SPARK – Spatial Policy-driven Adaptive Reinforcement Learning for Knowledge Distillation

1. Problem Statement

Low-bit quantization (using 8 bits or fewer) is essential for deploying image restoration (IR) networks on resource-constrained devices such as smartphones and drones. However, quantization introduces rounding noise that disproportionately degrades high-frequency regions, including edges, fine textures, and repetitive structures. While Knowledge Distillation (KD) is a standard technique for training lightweight models by supervising them with a complex teacher, existing KD methods for IR apply distillation signals uniformly across all spatial locations.

This uniform approach is suboptimal because:

• Inefficient Gradient Allocation: Smooth, low-frequency regions (e.g., flat skies) are easy for both teacher and student to reconstruct; forcing the student to mimic the teacher in these areas yields little useful gradient and crowds out learning in harder regions.
• Critical Failure Modes: Quantization noise is most damaging in high-frequency regions where human observers judge quality. Uniform loss functions effectively underweight these critical areas.
• Static Heuristics: Previous attempts to address this using hand-crafted edge masks or fixed difficulty heuristics commit to a single notion of "difficulty" that remains static throughout training, failing to adapt to the evolving capacity of the student or the specific characteristics of the degradation.

2. Methodology: The SPARK Framework

The authors propose SPARK (Spatial Policy-driven Adaptive Reinforcement learning for Knowledge distillation), a framework that treats the spatial allocation of distillation effort as a sequential decision problem solved via a lightweight Reinforcement Learning (RL) policy.

Core Components

1. Teacher and Student Networks:

   • A large-capacity, pre-trained, frozen Teacher Network provides high-quality outputs ($I_t$) serving as a soft supervisory signal.
   • A lightweight Student Network (quantized via Quantization-Aware Training, QAT) learns to replicate the teacher. During training, false quantization is applied to mimic deployment behavior while full-precision gradients are retained for optimization.

2. Difficulty Feature Extractor:
   At each training step, a non-differentiable module extracts four complementary feature maps to assess spatial reconstruction difficulty:

   • Laplacian Variance: Captures edge sharpness (high-frequency regions).
   • Pixel Variance: Measures local texture complexity.
   • Student Reconstruction Error: The absolute difference between the student output and the Ground Truth ($|Student - GT|$).
   • Teacher-Student Knowledge Gap: The difference between the teacher and student outputs.

3. Spatial Policy Network (RL Agent):

   • A compact CNN (three convolutional layers) processes the difficulty features on an $8 \times 8$ grid.
   • It outputs a stochastic spatial weight map ($W$) by sampling from a learned Gaussian distribution ($\mathcal{N}(\mu, \sigma)$).
   • The output is clamped to $[0, 1]$ to serve as a normalized weight map.
   • Stochasticity: The random sampling allows the policy to explore non-trivial spatial allocation strategies that deterministic, purely supervised objectives might miss.

4. Training Objective and Reward:

   • Loss Function: The total loss combines a standard L1 reconstruction loss (anchoring the student to the Ground Truth) and a Spatial-Weighted KD Loss. The KD loss is a weighted Mean Squared Error (MSE) between the student and teacher, modulated by the policy's weight map $W$.
   • Reward Signal: The policy is trained using a reward function based on the improvement in reconstruction quality (PSNR and SSIM) relative to a baseline student trained without spatial weighting.
   • Gradient Detachment: Crucially, the sampled action (the weight map) is detached before being applied to the KD loss. This decouples the policy learning from the pixel-wise reconstruction gradients, allowing the policy to be optimized solely based on the global reward signal.

5. Inference:
   The RL policy network is discarded after training. Consequently, SPARK adds zero inference cost and requires no architectural changes to the student network, making it IR task-agnostic (applicable to Low-Light Enhancement, Denoising, and Super-Resolution).

3. Key Contributions

• Identification of a Bottleneck: The paper identifies the spatial uniformity of existing KD losses as a primary bottleneck for low-bit IR, motivating spatially adaptive distillation as a necessary component for quantization-robust restoration.
• RL-Driven Spatial Policy: The introduction of SPARK, the first RL-based spatial weighting strategy for KD in low-bit IR. It utilizes interpretable difficulty features to learn a per-region weight map end-to-end during QAT.
• Stochastic Exploration: By employing a stochastic Gaussian-sampled action with detached gradients, the framework enables the exploration of spatial allocations inaccessible to purely differentiable objectives.
• Broad Applicability: The method is demonstrated to be effective across diverse IR tasks (LLE, Denoising, SR) and various student architectures without requiring task-specific modifications.

4. Experimental Results

The authors evaluated SPARK on standard benchmarks: LOLv1 (Low-Light Enhancement), SIDD (Image Denoising), and Urban100 (Super-Resolution).

• Quantitative Performance:

  • SPARK consistently outperformed PTQ (Post-Training Quantization), standard QAT, and state-of-the-art KD methods (FAKD, DCKD, SLKD) across multiple student models (MIRNet, DnCNN, RFDN, SRCNN).
  • LOLv1: SPARK achieved 22.01 dB PSNR (vs. 21.67 dB for the next best, SLKD) with MIRNet as the student.
  • SIDD: SPARK achieved 33.89 dB PSNR (vs. 32.14 dB for SLKD).
  • Urban100 (SR): At $\times 2$ scale, SPARK reduced the gap to the full-precision teacher to only -0.16 dB. At $\times 4$ scale, it outperformed PTQ by +2.53 dB for RFDN students.
  • SPARK consistently achieved the highest PSNR/SSIM and lowest LPIPS scores.

• Qualitative Performance:

  • Visual comparisons showed SPARK recovered brightness and structural detail more accurately in low-light scenarios compared to competitors.
  • In denoising, SPARK produced cleaner results with higher color uniformity and sharper edges, avoiding the residual noise and color distortion seen in other methods.
  • In Super-Resolution, SPARK reconstructed finer structural features and accurate geometry, avoiding the over-smoothing and aliasing artifacts common in SOTA approaches.

• Ablation Studies:

  • Comparisons confirmed that the RL-based reward mechanism is the primary driver of improvement, significantly outperforming training without RL (using only L1 or MSE loss).

5. Significance and Claims

The paper claims that SPARK establishes a foundation for spatially adaptive distillation in low-level vision. Its significance lies in:

• Efficiency: It solves the problem of quantization-induced degradation in critical regions without adding inference overhead.
• Adaptability: It replaces static, hand-crafted heuristics with a learnable, state-dependent decision process that evolves as the student's capacity changes.
• Generalization: It demonstrates that a unified, task-agnostic framework can effectively handle diverse restoration tasks (LLE, Denoising, SR) under significant computational constraints.

The authors modestly acknowledge a limitation: the reward signal relies on global metrics (PSNR/SSIM), which may not fully capture perceptual quality in all degradation scenarios, suggesting this as a direction for future work. However, the current results validate the core hypothesis that spatially adaptive distillation is superior to uniform supervision for low-bit IR.