[논문 리뷰] Rethinking Transformer-Based Blind-Spot Network for Self-Supervised Image Denoising

Junyi Li, Zhilu Zhang외 1명|arXiv (Cornell University)|2024. 04. 11.

Image and Signal Denoising Methods인용 수 6

한 줄 요약

트랜스포머 기반의 블라인드-스팟 네트워크( TBSN )를 도입하여 자가지도(Self-supervised) 이미지 디노이징을 수행하며, 마스킹된 윈도우 주의(attention)와 그룹화된 채널별 주의(attention)로 블라인드-스팟 제약을 만족시키고 receptive field를 확장한다; 지식 증류를 통한 U-Net 학생 모델(TBSN2UNet)으로 추론 효율성을 높인다.

ABSTRACT

Blind-spot networks (BSN) have been prevalent neural architectures in self-supervised image denoising (SSID). However, most existing BSNs are conducted with convolution layers. Although transformers have shown the potential to overcome the limitations of convolutions in many image restoration tasks, the attention mechanisms may violate the blind-spot requirement, thereby restricting their applicability in BSN. To this end, we propose to analyze and redesign the channel and spatial attentions to meet the blind-spot requirement. Specifically, channel self-attention may leak the blind-spot information in multi-scale architectures, since the downsampling shuffles the spatial feature into channel dimensions. To alleviate this problem, we divide the channel into several groups and perform channel attention separately. For spatial selfattention, we apply an elaborate mask to the attention matrix to restrict and mimic the receptive field of dilated convolution. Based on the redesigned channel and window attentions, we build a Transformer-based Blind-Spot Network (TBSN), which shows strong local fitting and global perspective abilities. Furthermore, we introduce a knowledge distillation strategy that distills TBSN into smaller denoisers to improve computational efficiency while maintaining performance. Extensive experiments on real-world image denoising datasets show that TBSN largely extends the receptive field and exhibits favorable performance against state-of-theart SSID methods.

연구 동기 및 목표

• Transformer의 능력을 활용하되 블라인드-스팟 제약을 유지하여 SSID를 motivate하고 개선한다.
• 실제 노이즈 패턴에 대해 수용 가능한 receptive field를 확장하는 트랜스포머 기반 블라인드-스팟 네트워크(TBSN)를 설계한다.
• 채널 주의에서 정보 누수를 방지하기 위해 채널을 그룹화하고 그룹 내에서 주의를 적용한다.
• 지식 증류 전략을 통해 효율적인 추론을 위한 U-Net 학생 모델(TBSN2UNet)을 생성하는 실용성을 높인다.

제안 방법

• 학습된 주의 마스크를 가진 마스크드 윈도우 기반 자기-주의(M-WSA)를 도입하여 주의가 짝수 좌표 위치로만 제한되도록 하여 확산된 확장(convolution) 효과를 구현한다.
• 채널 수가 공간 해상도보다 클 때 블라인드-스팟 정보 누출을 막기 위해 채널을 작은 그룹으로 처리하는 그룹화된 채널-별 자기-주의(G-CSA)를 도입한다.
• M-WSA, G-CSA 및 FFN을 결합한 확장된 트랜스포머 주의 블록(DTAB)을 구성하고 이를 엔코더-디코더 U-Net 내의 확장된 트랜스포머 아키텍처로 적용하여 SSID를 수행한다.
• 노이즈 상관을 깨뜨리면서도 블라인드-스팟 무결성을 유지하기 위해 학습 및 추론 단계에서 비대칭 인자(PD)로 픽셀 셔플 다운샘플링을 적용한다.
• 사전 학습된 TBSN이 교사 역할을 하고, 효율적인 추론을 위한 컴팩트한 U-Net 학생 모델(TBSN2UNet)을 학습시키는 지식 증류 체계를 제안한다.
• 실제 노이즈 제거 벤치마크 SIDD 및 DND에서 평가하여 최첨단 SSID 방법들과 비교한다.

실험 결과

연구 질문

• RQ1 Transformer 기반 연산자를 SSID의 블라인드-스팟 요구를 만족하도록 재설계할 수 있는가?
• RQ2 공간적/채널 자기-주의가 블라인드-스팟 무결성과 디노이징 성능에 어떤 영향을 미치는가?
• RQ3 TBSN을 더 작은 U-Net으로 증류하였을 때 컴퓨팅 비용이 감소하면서 성능이 유지되는가?

주요 결과

• TBSN은 자기지도 방법들 중에서 SIDD 및 DND 벤치마크에서 최첨단 SSID 성능을 달성한다.
• 마스크드 윈도우 기반 자기-주의(M-WSA)는 블라인드-스팟 제약을 준수하면서 국부 수용필드(receptive field)를 확장하여 디노이징 정확도를 향상시킨다.
• 그룹화된 채널-별 자기-주의(G-CSA)는 다중 스케일 아키텍처에서 블라인드-스팟 정보 누출을 방지하고 성능을 유지한다.
• DTAB은 로컬 및 전역 특징의 보완적 융합을 가능하게 하여 수용영역의 확장을 크게 하고 PSNR 이득을 준다.
• 지식 증류로 TBSN2UNet을 얻으면 교사 모델과 유사한 성능을 보이면서 추론 효율이 크게 향상된다.
• TBSN은 여러 선행 SSID 방법들을 능가하고 실제 데이터셋에서 지도 학습 대비 접근에 근접한 성능을 보인다.