技術要約：SPARK – 知識蒸留のための空間的ポリシー駆動型適応的強化学習

1. 問題提起

低ビット量子化（8ビット以下）は、スマートフォンやドローンなどのリソース制約のあるデバイスに画像復元（IR）ネットワークをデプロイするために不可欠である。しかし、量子化は丸め誤差を導入し、エッジ、微細なテクスチャ、反復的な構造を含む高周波領域を不均衡に劣化させる。知識蒸留（KD）は、複雑な教師モデルを用いて軽量なモデルを訓練するための標準的な手法であるが、既存のIR向けKD手法は、すべての空間位置に対して一様に蒸留信号を適用する。

この一様なアプローチは、以下の理由から最適ではない：

• 勾配割り当ての非効率性： 滑らかな低周波領域（例：平坦な空）は、教師と生徒の両方にとって再構成が容易である。これらの領域で生徒に教師を模倣させることは、有用な勾配をほとんど生み出さず、より困難な領域での学習を阻害する。
• 決定的な失敗モード： 量子化ノイズは、人間の観察者が品質を判断する際に重要な高周波領域で最も大きなダメージを与える。一様な損失関数は、これらの重要な領域を実質的に過小評価してしまう。
• 静的なヒューリスティック： 手作りのエッジマスクや固定された困難度ヒューリスティックを用いてこれに対処しようとする過去の試みは、単一の「困難度」の概念に固執しており、生徒の進化する能力や特定の劣化特性に適応できない。

2. メソドロジー：SPARKフレームワーク

著者らは、空間的な蒸留努力の割り当てを、軽量な強化学習（RL）ポリシーによって解決される逐次的な意思決定問題として扱うフレームワーク、SPARK（Spatial Policy-driven Adaptive Reinforcement learning for Knowledge distillation）を提案する。

コアコンポーネント

1. 教師および生徒ネットワーク：

   • 大容量で事前学習済みの、凍結された教師ネットワークが、高品質な出力（$I_t$）を提供し、ソフトな監督信号として機能する。
   • 軽量な生徒ネットワーク（量子化認識訓練（QAT）を介して量子化される）は、教師を複製するように学習する。訓練中、展開時の挙動を模倣するために偽量子化が適用されるが、最適化のためにフルプレシジョンの勾配は保持される。

2. 困難度特徴抽出器：
   各訓練ステップにおいて、非微分モジュールが、空間的な再構成の困難度を評価するために4つの補完的な特徴マップを抽出する：

   • ラプラシアン分散（Laplacian Variance）： エッジの鋭さ（高周波領域）を捉える。
   • ピクセル分散（Pixel Variance）： 局所的なテクスチャの複雑さを測定する。
   • 生徒の再構成誤差（Student Reconstruction Error）： 生徒の出力と正解（GT）との絶対差（$|Student - GT|$）。
   • 教師と生徒の知識ギャップ（Teacher-Student Knowledge Gap）： 教師の出力と生徒の出力の差。

3. 空間的ポリシーネットワーク（RLエージェント）：

   • コンパクトなCNN（3層の畳み込み層）が、$8 \times 8$のグリッド上で困難度特徴を処理する。
   • 学習されたガウス分布（$\mathcal{N}(\mu, \sigma)$）からサンプリングすることで、確率的な空間重みマップ（$W$）を出力する。
   • 出力は$[0, 1]$にクランプされ、正規化された重みマップとして機能する。
   • 確率性： ランダムなサンプリングにより、決定論的で純粋に教師ありの目的関数では到達できない、非自明な空間割り当て戦略の探索が可能になる。

4. 訓練目的関数と報酬：

   • 損失関数： 全体の損失は、標準的なL1再構成損失（生徒を正解に繋ぎ止めるもの）と、ポリシーの重みマップ $W$ によって変調される空間重み付きKD損失を組み合わせたものである。KD損失は、教師と生徒の間の平均二乗誤差（MSE）であり、重み付けされたものである。
   • 報酬信号： ポリシーは、空間重み付けなしで訓練されたベースラインの生徒と比較した、再構成品質（PSNRおよびSSIM）の改善に基づく報酬関数を用いて訓練される。
   • 勾配の切り離し（Gradient Detachment）： 重要な点として、サンプリングされたアクション（重みマップ）は、KD損失に適用される前に切り離される。これにより、ポリシーの学習とピクセルごとの再構成勾配がデカップリング（分離）され、ポリシーをグローバルな報酬信号のみに基づいて最適化することが可能になる。

5. 推論：
   RLポリシーネットワークは訓練後に破棄される。したがって、SPARKはゼロの推論コストであり、生徒ネットワークのアーキテクチャ変更も必要なく、IRタスクに依存しない（低照度強調、デノイジング、超解像に適用可能）。

3. 主な貢献

• ボトルネックの特定： 本論文は、既存のKD損失における空間的一様性が、低ビットIRにおける主要なボトルネックであることを特定し、空間適応的な蒸留が量子化に強い復元に必要であることを動機づけた。
• RL駆動型空間ポリシー： 低ビットIRにおけるKDのための、最初のRLベースの空間重み付け戦略であるSPARKを導入した。これは、解釈可能な困難度特徴を利用して、QAT中にエンドツーエンドで学習可能な、領域ごとの重みマップを利用する。
• 確率的探索： 切り離された勾配を持つガウス分布サンプリングによるアクションを採用することで、純粋に微分可能な目的関数ではアクセスできない空間割り当ての探索を可能にした。
• 幅広い適用性： 本手法は、多様なIRタスク（LLE、デノイジング、SR）および様々な生徒アーキテクチャに対して効果的であることが示されている。

4. 実験結果

著者らは、標準的なベンチマークである LOLv1（低照度強調）、SIDD（画像デノイジング）、Urban100（超解像）を用いてSPARKを評価した。

• 定量的性能：

  • SPARKは、複数の生徒モデル（MIRNet, DnCNN, RFDN, SRCNN）において、PTQ（事後訓練量子化）、標準的なQAT、および最先端のKD手法（FAKD, DCKD, SLKD）を一貫して上回った。
  • LOLv1： MIRNetを生徒とした場合、SPARKは22.01 dBのPSNRを達成した（次に優れたSLKDの21.67 dBに対し）。
  • SIDD： SPARKは33.89 dBのPSNRを達成した（SLKDの32.14 dBに対し）。
  • Urban100 (SR)： $\times 2$スケールにおいて、SPARKはフルプレシジョンの教師との差をわずか-0.16 dBに縮小した。$\times 4$スケールでは、RFDN生徒においてPTQを+2.53 dB上回った。
  • SPARKは一貫して最高のPSNR/SSIMおよび最低のLPIPSスコアを達成した。

• 定性的性能：

  • 視覚的な比較により、SPARKは競合手法と比較して、低照度シナリオにおいて明るさと構造の詳細をより正確に復元していることが示された。
  • デノイジングにおいて、SPARKは他の手法で見られる残留ノイズや色の歪みを避け、より高い色の均一性と鋭いエッジを持つ、よりクリーンな結果を生成した。
  • 超解像において、SPARKは、SOTAのアプローチで一般的な平滑化やエイリアシングのアーティファクトを避け、より微細な構造的特徴と正確な幾何学形状を再構成した。

• アブレーション研究：

  • 比較により、RLベースの報酬メカニズムが改善の主要な原動力であり、RLなしでの訓練（L1またはMSE損失のみを使用）を大幅に上回ることが確認された。

5. 意義と主張

本論文は、SPARKが空間適応的な蒸留の基礎を確立すると主張している。その意義は以下の通りである：

• 効率性： 決定的な領域における量子化誘発の劣化を、推論オーバーヘッドを追加することなく解決する。
• 適応性： 静的な手作りのヒューリスティックを、生徒の能力の変化に応じて進化する、学習可能で状態依存的な意思決定プロセスに置き換える。
• 汎用性： 統一されたタスクに依存しないフレームワークが、大幅な計算制約の下で、多様な復元タスク（LLE、デノイジング、SR）を効果的に処理できることを示した。

著者らは、報酬信号がグローバルな指標（PSNR/SSIM）に依存しており、すべての劣化シナリオにおいて知覚的な品質を完全に捉えきれない可能性があることを控えめに認め、これを今後の研究方向として示唆している。しかし、現在の結果は、空間適応的な蒸留が低ビットIRにおいて一様な監督よりも優れているという核心的な仮説を検証している。