기술 요약: SPARK – 지식 증류를 위한 공간적 정책 기반 적응형 강화 학습

1. 문제 정의

저비트 양자화(8비트 이하 사용)는 스마트폰이나 드론과 같은 자원 제한적인 장치에 이미지 복원(IR) 네트워크를 배포하는 데 필수적입니다. 그러나 양자화는 에지, 미세한 질감, 반복적인 구조와 같은 고주파 영역을 불균형하게 저하시키는 반올림 노이즈를 유발합니다. 지식 증류(Knowledge Distillation, KD)는 복잡한 교사(Teacher) 모델이 경량화된 학생(Student) 모델을 감독하여 학습시키는 표준적인 기술이지만, IR을 위한 기존의 KD 방식은 모든 공간적 위치에 대해 증류 신호를 균일하게 적용합니다.

이러한 균일한 접근 방식은 다음과 같은 이유로 최적이 아닙니다:

• 비효율적인 그래디언트 할당: 매끄러운 저주파 영역(예: 평탄한 하늘)은 교사와 학생 모두 재구성하기 쉽습니다. 이러한 영역에서 학생이 교사를 모방하도록 강제하는 것은 유용한 그래디언트를 거의 생성하지 못하며, 오히려 어려운 영역에서의 학습을 방해합니다.
• 결정적인 실패 모드: 양자화 노이즈는 인간 관찰자가 품질을 판단하는 기준인 고주파 영역에서 가장 치명적입니다. 균일한 손실 함수는 이러한 중요한 영역의 비중을 효과적으로 과소평가합니다.
• 정적인 휴리스틱: 수작업으로 만든 에지 마스크나 고정된 난이도 휴리스틱을 사용하여 이를 해결하려는 이전의 시도들은 "난이도"에 대한 단일한 개념에 고착되어 있으며, 이는 학생의 진화하는 능력이나 특정 열화 특성에 적응하지 못합니다.

2. 방법론: SPARK 프레임워크

저자들은 SPARK(Spatial Policy-driven Adaptive Reinforcement learning for Knowledge distillation)를 제안합니다. 이 프레임워크는 공간적 증류 노력의 할당을 경량 강화 학습(RL) 정책을 통해 해결되는 순차적 의사 결정 문제로 취급합니다.

핵심 구성 요소

1. 교사 및 학생 네트워크:

   • 대용량의 사전 학습된 고정된 **교사 네트워크(Teacher Network)**는 부드러운 감독 신호 역할을 하는 고품질 출력($I_t$)을 제공합니다.
   • 경량화된 **학생 네트워크(Student Network)**는 양자화 인식 학습(Quantization-Aware Training, QAT)을 통해 교사를 복제하도록 학습됩니다. 학습 중에는 배포 동작을 모방하기 위해 가짜 양자화(False Quantization)가 적용되지만, 최적화를 위해 전정밀도(Full-precision) 그래디언트는 유지됩니다.

2. 난이도 특징 추출기 (Difficulty Feature Extractor):
   각 학습 단계에서 비미분 모듈이 공간적 재구성 난이도를 평가하기 위해 네 가지 상호 보완적인 특징 맵을 추출합니다:

   • 라플라시안 분산(Laplacian Variance): 에지의 선명도(고주파 영역)를 포착합니다.
   • 픽셀 분산(Pixel Variance): 국부적인 질감 복잡도를 측정합니다.
   • 학생 재구성 오차(Student Reconstruction Error): 학생 출력과 정답(Ground Truth, GT) 사이의 절대 차이($|Student - GT|$)입니다.
   • 교사-학생 지식 격차(Teacher-Student Knowledge Gap): 교사와 학생 출력 간의 차이입니다.

3. 공간적 정책 네트워크 (Spatial Policy Network, RL Agent):

   • 소형 CNN(3개의 컨볼루션 레이어)이 $8 \times 8$ 그리드 상의 난이도 특징을 처리합니다.
   • 학습된 가우시안 분포($\mathcal{N}(\mu, \sigma)$)에서 샘플링하여 확률적 공간 가중치 맵($W$)을 출력합니다.
   • 출력값은 $[0, 1]$로 클램핑되어 정규화된 가중치 맵 역할을 합니다.
   • 확률성(Stochasticity): 무작위 샘플링을 통해 정책은 결정론적이고 순수하게 지도 학습적인 목적 함수가 놓칠 수 있는 비자명한 공간 할당 전략을 탐색할 수 있습니다.

4. 학습 목적 함수 및 보상:

   • 손실 함수: 전체 손실은 표준 L1 재구성 손실(학생을 정답에 고정)과 **공간 가중 증류 손실(Spatial-Weighted KD Loss)**의 결합으로 구성됩니다. KD 손실은 정책의 가중치 맵 $W$에 의해 조절되는 교사와 학생 간의 가중 평균 제곱 오차(MSE)입니다.
   • 보상 신호: 정책은 공간 가중치를 적용하지 않고 학습된 베이스라인 학생 대비 재구성 품질(PSNR 및 SSIM)의 개선도를 기반으로 한 보상 함수를 사용하여 학습됩니다.
   • 그래디언트 분리(Gradient Detachment): 결정적으로, 샘플링된 액션(가중치 맵)은 KD 손실에 적용되기 전에 분리(Detached)됩니다. 이는 정책 학습을 픽셀 단위의 재구성 그래디언트로부터 분리하여, 정책이 오직 글로벌 보상 신호에만 기반하여 최적화될 수 있도록 합니다.

5. 추론(Inference):
   RL 정책 네트워크는 학습 후 폐기됩니다. 결과적으로 SPARK는 제로(Zero) 추론 비용을 추가하며 학생 네트워크에 대한 구조적 변경을 요구하지 않으므로, IR 작업 불가지론적(Task-agnostic)입니다(저조도 개선, 디노이징, 초해상도에 적용 가능).

3. 주요 기여

• 병목 현상 식별: 본 논문은 기존 KD 손실의 공간적 균일성이 저비트 IR의 주요 병목 현상임을 식별하고, 공간 적응형 증류를 양자화에 강건한 복원을 위한 필수 구성 요소로 제안했습니다.
• RL 기반 공간 정책: 저비트 IR을 위한 KD에서 최초로 RL 기반 공간 가중 전략인 SPARK를 도입했습니다. 이는 해석 가능한 난이도 특징을 활용하여 QAT 과정 중 end-to-end로 학습 가능한 per-region 가중 맵을 학습합니다.
• 확률적 탐색: 결정론적인 미분 가능 목적 함수로는 도달할 수 없는 공간 할당을 가능하게 하기 위해, 그래디언트가 분리된 확률적 가우시안 샘플링 액션을 채택했습니다.
• 광범위한 적용성: 다양한 IR 작업(LLE, Denoising, SR) 및 다양한 학생 아키텍처에 대해 별도의 작업 특정 수정 없이 효과적임을 입증했습니다.

4. 실험 결과

저자들은 표준 벤치마크인 LOLv1(저조도 개선), SIDD(이미지 디노이징), Urban100(초해상도)에서 SPARK를 평가했습니다.

• 정량적 성능:

  • SPARK는 여러 학생 모델(MIRNet, DnCNN, RFDN, SRCNN)에 대해 PTQ(Post-Training Quantization), 표준 QAT, 그리고 최첨단(SOTA) KD 방법들(FAKD, DCKD, SLKD)보다 일관되게 우수한 성능을 보였습니다.
  • LOLv1: MIRNet를 학생으로 사용했을 때, SPARK는 22.01 dB PSNR을 달ach했습니다(차순위인 SLKD의 21.67 dB 대비).
  • SIDD: SPARK는 33.89 dB PSSR을 달성했습니다(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는 다른 방법들에서 보이는 잔류 노이즈나 색상 왜곡을 피하면서, 더 깨끗한 결과와 높은 색상 균일성 및 날카로운 에지를 생성했습니다.
  • 초해상도(SR)에서 SPARK는 미세한 구조적 특징과 정확한 기하학적 형태를 재구성하여, SOTA 방식에서 흔히 나타나는 과도한 평활화(Over-smoothing) 및 에일리어싱 아티팩트를 방지했습니다.

• 절제 연구(Ablation Studies):

  • 비교 연구를 통해 RL 기반 보상 메커니즘이 성능 향상의 주요 동력임을 확인했으며, 이는 RL 없이 학습했을 때(L1 또는 MSE 손실만 사용)보다 유의미하게 뛰어난 성능을 보였습니다.

5. 의의 및 주장

본 논문은 SPARK가 공간 적응형 증류를 위한 토대를 마련한다고 주장합니다. 그 의의는 다음과 같습니다:

• 효율성: 계산 제약이 큰 상황에서 핵심 영역의 양자화 유도 저하 문제를 추론 오버헤드 없이 해결합니다.
• 적응성: 정적인 수작업 휴리스틱을 학생의 능력이 변화함에 따라 진화하는 학습 가능한 상태 의존적 결정 프로세스로 대체합니다.
• 일반화: 통합된 작업 불가지론적 프레임워크가 상당한 계산 제약 하에서도 다양한 복원 작업을 효과적으로 처리할 수 있음을 입증했습니다.

저자들은 보상 신호가 글로벌 지표(PSNR/SSIM)에 의존하여 모든 저하 시나리오에서 지각적 품질을 완전히 포착하지 못할 수 있다는 한계를 겸허히 인정하며, 이를 향후 연구 과제로 제시했습니다. 그러나 현재의 결과는 공간 적응형 증류가 저비트 IR에서 균일한 감독보다 우월하다는 핵심 가설을 입증합니다.