SPARK: Spatial Policy-driven Adaptive Reinforcement learning for Knowledge distillation

Abstract: Low-bit quantization enables deployment of image restoration (IR) networks on resource-constrained devices, but introduces rounding noise that disproportionately degrades high-frequency regions such as edges and fine textures. Existing knowledge distillation (KD) methods apply distillation signals uniformly across all spatial locations, overlooking the varying reconstruction difficulty across image regions. To address this, we propose SPARK (Spatial Policy-driven Adaptive Reinforcement Learning for Knowledge Distillation), a framework that adaptively allocates distillation effort using a lightweight reinforcement learning (RL) policy network. At each training step, a difficulty feature extractor computes four signals, namely Laplacian variance, pixel variance, student reconstruction error, and teacher-student knowledge gap, which are fed into a compact policy CNN that produces a stochastic spatial weight map to modulate the KD loss during quantization-aware training (QAT). SPARK is IR task-agnostic, adds no inference cost, and integrates into any existing QAT pipeline without architectural changes. Experiments on benchmark datasets demonstrate that SPARK consistently outperforms PTQ, QAT, and state-of-the-art (SOTA) KD approaches across multiple student architectures, achieving reconstruction quality closest to the full-precision teacher under significant computational constraints.

• The paper presents SPARK, which leverages a reinforcement learning-driven policy to adaptively weight knowledge distillation for low-bit image restoration.
• It uses interpretable spatial maps and RL to focus on challenging regions, leading to significant improvements in PSNR, SSIM, and LPIPS metrics.
• The approach integrates with quantization-aware training, outperforming static KD methods while incurring no extra inference overhead.

Spatial Policy-driven Adaptive RL for Knowledge Distillation in Low-Bit Image Restoration

Motivation and Problem Formulation

Low-bit quantization has become imperative for deploying image restoration (IR) networks in resource-constrained environments, enabling significant reduction in memory footprint, inference latency, and energy consumption. However, quantization-induced rounding noise is highly detrimental to high-frequency regions (edges, textures), degrading visual fidelity—particularly for IR tasks where pixel-wise errors are directly observable, unlike high-level vision tasks. Conventional knowledge distillation (KD) for IR applies supervision uniformly across spatial regions, disregarding heterogeneous region-wise reconstruction difficulty. This results in suboptimal allocation of student network capacity, especially under stringent quantization constraints. Static heuristics or hand-crafted masks (edge, gradient, fixed difficulty) do not adapt to evolving student-teacher dynamics during training and fail to address this bottleneck.

Framework Description

SPARK introduces a spatial policy-driven adaptive reinforcement learning (RL) scheme to address the nonuniform spatial complexity of IR tasks in the KD paradigm. The approach is IR task-agnostic and compatible with quantization-aware training (QAT) pipelines, adding no inference overhead.

Feature Extraction and Policy Network

SPARK employs a difficulty feature extractor that generates four interpretable spatial maps:

• Laplacian variance (edge complexity)
• Pixel variance (texture complexity)
• Student reconstruction error (per-pixel loss vs. ground truth)
• Teacher-student knowledge gap (difference in outputs)

These are concatenated and fed to a lightweight policy network (three-layer CNN), which outputs a stochastic spatial weight map via Gaussian sampling and clamping to $[0,1]$. The weight map modulates the distillation loss, prioritizing difficult regions dynamically. The policy network is optimized via RL, using a reward function that linearly combines improvements in PSNR and SSIM relative to uniform weighting, ensuring exploration of spatial allocation strategies beyond the reach of gradient-based objectives.

Training Objectives

SPARK’s total loss comprises two terms:

• Pixel-wise L1 reconstruction loss (student vs. ground truth)
• Spatially-weighted MSE KD loss (student vs. teacher, modulated by learned spatial map)

The stochastic policy output is detached before KD loss calculation, enforcing policy updates solely from reward signals, not pixel-wise gradients. This separation enables targeted allocation of distillation "pressure" and avoids overfitting to trivial regions.

Experimental Evaluation

SPARK is evaluated across LLE (LOLv1), denoising (SIDD), and super-resolution (Urban100), employing varied teacher-student pairs (MIRNet, DnCNN, RFDN, SRCNN) quantized to INT8. Strong quantitative results are observed:

• LLE (LOLv1): SPARK achieves 22.01 dB PSNR, 0.785 SSIM, 0.176 LPIPS—closest to the full-precision teacher and superior to PTQ/QAT and SOTA KD such as FAKD, DCKD, SLKD.
• Denoising (SIDD): SPARK delivers 33.89 dB PSNR, 0.807 SSIM, 0.320 LPIPS, substantially outperforming uniform KD and PTQ (-5.47 dB vs. teacher).
• SR (Urban100): On ×2/×4 scales, SPARK reduces the gap with teacher to -0.16 dB PSNR for RFDN; always secures highest PSNR, SSIM, and lowest LPIPS across student architectures.

Qualitative results reveal visually perceptible improvements in edge detail, color uniformity, and noise reduction—especially in regions typically impaired by quantization.

Ablation studies confirm that RL-driven policy weighting offers consistent improvements over L1 and MSE loss baselines. The reward mechanism is demonstrated to be the main source of gains, guiding the student network towards optimal use of representational capacity in challenging regions.

Practical and Theoretical Implications

SPARK validates spatially adaptive distillation as a critical missing component for quantization-robust IR, providing a principled solution to region-wise learning allocation. The policy leverages interpretable features and stochasticity, enabling principled exploration and rational region prioritization. As the policy network is discarded post-training, there is no inference cost, preserving deployment efficiency.

Theoretically, the separation between reward-driven policy updates and differentiable losses expands the optimization landscape, enabling strategies inaccessible to purely supervised learning. This advances the use of RL in low-level vision, shifting its role from architecture/design selection to direct supervision modulation.

Limitations and Future Perspectives

The reward function is tied to aggregate metrics (PSNR, SSIM), which may not always reflect perceptual quality in cases of severe or atypical degradations. Integration of learned perceptual metrics or region-specific objectives could further enhance adaptability. Extending SPARK to other modalities, temporal data, or higher-order spatial relationships remains an open direction. To generalize further, exploration of multi-agent RL policies or hierarchical spatial weighting could accelerate convergence and improve interpretability.

Conclusion

SPARK establishes spatial policy-driven adaptive RL as a robust mechanism for knowledge distillation in low-bit IR networks, outperforming uniform KD and SOTA baselines in both quantitative and qualitative terms. The approach is compatible with existing QAT and KD pipelines, and opens a new direction for RL-guided compression and supervision in low-level vision tasks (2606.15243).