LiteKD-Net: Lightweight Knowledge-Distilled Network for Mobile Image Denoising

기술 요약: 모바일 이미지 디노이징을 위한 LiteKD-Net

문제 정의
모바일 이미지 디노이징은 고품질의 복원 성능을 달면서도 온디바이스 배포를 위한 낮은 계산 비용을 유지해야 하는 이중 과제에 직면해 있다. 모바일 기기는 전문 카메라보다 작은 센서와 픽셀을 사용하므로, 본질적으로 신호 대 잡음비(SNR)가 낮고 단순한 가우시안 가법 모델로는 포착할 수 없는 복잡하고 신호 의존적인 노이즈 패턴(예: 광자 샷 노이즈, 리드 노이즈, 로우 노이즈)을 생성한다. 또한, 이러한 작업을 위한 딥러닝 모델 학습은 엄격하게 정렬된 저품질(LQ) 및 그라운드 트루스(GT) 이미지 쌍의 부족으로 인해 어려움을 겪는다. 움직임이나 카메라 내부 처리 방식의 변화로 인해 서로 다른 노이즈 특성을 가진 동일한 장면을 캡처하는 것은 매우 어렵기 때문이다. 수천만 개의 파라미터를 가진 기존의 고용량 네트워크들은 모바일 하드웨어에서 높은 지연 시간, 과도한 에너지 소비, 큰 메모리 점유율 문제를 야기하며, 클라우드 기반 처리는 지연 시간과 개인정보 보호 문제를 유발한다.

방법론
저자는 데이터 부족, 모델 효율성, 성능 회복 문제를 해결하기 위해 세 가지 핵심 구성 요소를 통합한 통합 프레임워크인 LiteKD-Net을 제안한다.

1. 노이즈 시뮬레이션 파이프라인 (NSP): 정렬된 실제 데이터 쌍의 부재를 극복하기 위해, 저자는 깨끗한 sRGB 이미지로부터 합성 훈련 데이터를 생성하도록 물리 기반 노이즈 모델(본래 DSLR 카메라용)을 변형하였다. 이 파이프라인은 sRGB 이미지를 선형화하고, 광자 샷 노이즈, Tukey–Lambda 리드 노이즈, 로우 노이즈를 시뮬레이션하며, 결정적으로 **픽셀 크로스토크(pixel crosstalk)**를 도입한다. 이 크로스토크는 특정 컨볼루션 커널을 통해 모델링되어 인접 픽셀 간의 국소적 공간 상관관계를 생성함으로써, 소형 모바일 센서에서 발견되는 전기적/광학적 누설을 모사한다. 마지막으로 양자화 노이즈를 추가하고, 이미지를 클리핑한 후 다시 sRGB로 변환한다.
2. Lite-RRDB 아키텍처: 저자는 업샘플링 모듈을 제거하여 아이덴티티 해상도(identity-resolution) 디노이징을 수행하도록 Real-ESRGAN 아키텍처를 변형하였다. 경량 학생(Student) 네트워크를 만들기 위해, 표준 잔차 밀집 블록(RRDB)을 Lite-RRDB로 교체하였다. 이 설계에서는 표준 컨볼루션을 깊이별 분리 컨볼루션(depthwise separable convolutions)(Depthwise 후 Pointwise)으로 대체하고, 밀집된 특징 결합(dense feature concatenation)을 제거하는 대신 고정된 채널 폭(64)을 가진 순차적 처리 방식을 채택하였다. 이는 잔차-내-잔차(residual-in-residual) 훈련 구조를 유지하면서 파라미터 수와 계산 복잡성을 크게 줄인다.
3. 노이즈-KD (Knowledge Distillation): 경량화된 학생 네트워크의 감소된 용량을 보완하기 위해, 특징 수준의 지식 증류 전략이 사용된다. 고용량 교사(Teacher) 네트워크(변형된 Real-ESRGAN)를 먼저 학습시킨 후 고정한다. 학생 네트워크 학습 과정에서, 학생 네트워크는 평균 제곱 오차 손실($L_{KD}$)을 사용하여 교사 네트워크의 중간 특징 맵을 일치시키도록 유도된다. 학생 네트워크는 픽셀 수준의 재구성 손실($L_1$), 지각 손실($L_{per}$), 그리고 증류 손실을 통해 공동 최적화된다. 결정적으로, 교사 네트워크는 학습 후 폐기되므로 추가적인 추론 비용을 발생시키지 않는다.

주요 기여

• 노이즈 시뮬레이션 파이프라인 (NSP): 모바일 특화 노이즈 시뮬레이션의 격차를 해결하기 위해 픽셀 크로스토크를 명시적으로 모델링하여 완벽하게 정렬된 GT 및 LQ 데이터를 생성하는 새로운 파이프라인을 제공한다.
• Lite-RRDB 모듈: 표준 컨볼루션을 깊이별 분리 컨볼루션으로 대체하고 밀집 결합을 제거하여 모델 크기와 FLOPs를 획기적으로 줄이는 경량 아키텍처 블록이다.
• 노이즈-KD 전략: 고정된 무거운 교사로부터 경량 학생에게 복원 능력을 전달하는 특징 수준의 증류 방법을 제공하며, 이때 추론 파라미터나 연산량은 증가하지 않는다.
• 포괄적 평가: 모델을 실제 데이터셋에서 평가하여 복원 품질과 계산 효율성 사이의 강력한 트레이드오프를 입증하였다.

실험 결과
실제 데이터셋(Mobile AI Denoising Dataset 및 MIDD)에 대한 실험 결과, LiteKD-Net은 경쟁력 있는 이미지 품질을 유지하면서도 계산 지표를 크게 감소시켰다.

• 효율성: 베이스라인인 Real-ESRGAN과 비교했을 때, LiteKD-Net은 모델 크기를 약 10배(16.70M에서 1.67M 파라미터로) 줄였으며, MACs와 FLOPs는 4배 감소시켰다.
• 속도: 추론 실행 시간은 거의 2배 단축되었으며(256×256 입력 기준 ~148ms에서 ~83ms로), 이를 통해 더 높은 초당 프레임 수(FPS)를 달행하였다(11.98 FPS vs. 6.72 FPS).
• 품질: 경량화된 학생 네트워크 단독으로는 무거운 교사에 비해 품질 저하가 나타나지만, Noise-KD를 추가함으로써 성능을 회복하였다. LiteKD-Net은 무거운 Real-ESRGAN 베이스라인에 매우 근접한 PSNR 및 SSIM 점수를 달성하였으며, 효율성 측면에서는 SwinIR을 능가한다. 지각적 지표(LPIPS, DISTS, MUSIQ) 측면에서도 경쟁력 있는 성능을 보이며, 비교된 모델 중 가장 높은 MUSIQ 점수를 기록하였다.

의의 및 주장
본 논문은 LiteKD-Net이 모바일 이미지 디노이징을 위해 복원 품질과 계산 효율성 사이의 훌륭한 트레이드오프를 제공한다고 주장한다. 물리 기반 노이즈 시뮬레이션 파이프라인과 경량 아키텍처 및 지식 증류를 통합함으로써, 저자는 대규모 네트워크의 과도한 비용이나 클라우드 처리의 지연 없이도 고품질 디노이징 모델을 모바일 기기에 배포하는 것이 가능함을 입증하였다. 결과는 제안된 방법이 모바일 하드웨어의 엄격한 자원 제약을 준수하면서도 무거운 모델의 복원 품질을 효과적으로 회복한다는 것을 보여준다. 저자는 본인의 시뮬레이션이 true RAW 데이터가 아닌 선형화된 sRGB 이미지를 대상으로 작동하므로, 파이프라인이 전체 카메라 이미징 체인(예: 데모자이킹, 톤 매핑)을 모델링하지는 못하지만, sRGB 기반 디노이징을 위한 정렬된 훈련 데이터를 생성하는 데에는 여전히 효과적이라고 겸허히 언급하였다.