[논문 리뷰] Masked Image Training for Generalizable Deep Image Denoising (2023 CVPR)

Introduction

최근 이미지 노이즈 제거 (denoising) 분야에서는 대부분 Deep Learning을 사용한다. 하지만 Deep Learning을 사용하여 denoising을 학습할 때 문제점이 학습 데이터에 존재하는 노이즈 분포는 잘 제거하지만 학습 데이터에 없는 노이즈에 대해서는 잘 제거를 못한다는 것이다. 기존에 denoising을 학습할 때 clean image에 가우시안 노이즈를 더해서 노이즈 이미지를 만들어 학습에 사용을 많이 하는데 이렇게 학습한 네트워크는 ISP process에 의한 노이즈나, salt pepper 노이즈, poisson 노이즈 등등 가우시안 분포와 다른 노이즈는 잘 제거하지 못한다. 이 논문에서는 이런 문제점을 mask training이라는 방법으로 해결하고자 했다. 

위의 그림의 왼쪽같이 학습 데이터에 있는 In-distribution 노이즈는 기존 네트워크가 잘 제거하지만 오른쪽처럼 학습 데이터에 없는 out-of-distribution 노이즈가 들어오면 네트워크가 잘 제거하지 못한다. 이 논문에서 제시한 방법을 사용하면 오른쪽 ours처럼 out-of-distribution에서 조금 더 좋은 결과를 얻을 수 있다. 

Method

구체적인 방법 설명이전에 motivation을 좀 더 자세히 설명한다. 

Motivation

기본적으로 우리가 denoising을 학습한 네트워크에게 기대하는 것은 clean하고 natural한 이미지를 생성하는 것이다. 이렇게 학습이 됐다고 가정하면 어떠한 노이즈나 다른 degradation이 들어와도 clean한 이미지를 생성할 수 있을 것이다. 하지만 denoising을 학습을 시키면 네트워크는 clean한 이미지를 생성하는것을 목표로 하는 것이 아니라 학습 데이터에 있는 노이즈 분포를 학습하고 노이즈를 찾아내는 쪽으로 학습이 된다. 이 때문에 학습 데이터에 없는 노이즈가 들어오면 clean한 이미지를 생성해내지 못하는 경향이 생긴다. 이 논문에서는 이러한 경향을 파악해 아래 그림으로 보여주고 이 논문에서 제안한 방법은 네트워크가 기본 이미지에 더 집중을 잘 할 수 있음을 보여준다. 

위의 figure에 대해 설명을 하자면 기본적으로 왼쪽같은 자연스럽지 않은 이미지에 가우시안 노이즈를 더해서 denoising 네트워크를 학습시키고 오른쪽같은 natural한 이미지에서 테스트를 진행한것이다. 만약 네트워크가 노이즈분포가 아닌 clean 이미지에 더 집중을 해서 denoising을 한다면 다른 natural한 이미지에서 denoising을 할 때 성능이 떨어져야할 것이다. 하지만 위의 그림처럼 unnatrual 이미지에서 학습을 한 네트워크가 natural한 이미지에서도 denoising을 잘 한다. 이는 네트워크가 노이즈 분포에 집중하고 이미지에는 집중을 잘 하지 않았음을 보여준다. 하지만 이 논문에서 제안한 방법을 사용하면 training image의 성질이 테스트에서도 나타나고 이는 즉 이미지 자체에 네트워크가 조금 더 집중을 잘 했음을 나타낸다. 

Masked Training 

기본적으로 method는 굉장히 simple하다. input을 1×1 convolution으로 채널을 늘려주고 거기에 mask를 random으로 mask를 씌우는 input mask가 있고 중간 attention block에서 feature에 mask를 씌우는 attention mask가 있다. 아래 그림이 잘 나타내주고 있다. 

attention mask를 도입한 이유에 대해서는 input에 대해서만 mask를 씌우는것으로는 문제가 생기기 때문이라고 설명하는데 아래 그림을 보면 attention mask 없이 input mask만 사용한 경우에는 histogram과 이미지에서도 모두 볼 수 있듯이 original보다 밝기가 상당히 밝아진다. 이는 training과 test 이미지가 완전히 같지 않기 때문에 input에만 mask를 씌워서 training을 하게 되면 네트워크가 이미지의 밝기를 키우는 쪽으로 학습을 해버릴 수 있기 때문이라고 한다. 따라서 중간 attention mask에도 mask를 같이 씌워서 학습을 하면 네트워크가 조금 더 이미지의 context 생성에 집중을 할 수 있게 된다고 한다. 

Experiments

보여주고자 하는것이 denoising의 generalization 성능이기 때문에 학습은 가우시안 노이즈로만 학습을 하고 test에는 speckle noise, poisson noise, spatially-correlated noise, salt&pepper noise, ISP noise, mixture noise로 테스트 결과를 보여준다. training에 사용한 이미지들은 DIV2K, Flickr2K, BSD500, WED 데이터셋을 섞어서 사용하고 테스트에는 CBSD68, Kodak24, McMaster, Urban100 이미지들을 사용한다. 추가로 real noise에서의 성능도 보여주기 위해 the Smartphone Image Denoising Dataset (SIDD) 와 Monte Carlo (MC) rendered image noise 데이터셋도 사용해서 테스트 결과를 보여준다. 

아래 결과는 모든 네트워크를 가우시안 노이즈로만 학습시킨 후 다양한 노이즈에 대한 결과를 측정한 것이고 mask training (ours) 방법이 모든 노이즈에서 가장 generalization 성능이 좋음을 보여준다. 

아래는 이에 대해 visualization을 한 것이다. 다른 방법들은 노이즈를 지우는데 거의 다 실패했는데 mask training을 한 네트워크는 성공적으로 노이즈를 제거한 것을 볼 수 있다. 가우시안 노이즈로만 학습시켰는데 이 정도의 결과는 신기할 정도 인 것 같다. 

ISP noise에 대해서도 측정을 했는데 PSNR에서는 DnCNN과 비슷한 수치가 나왔지만 LPIPS와 SSIM에서 많이 높은 수치가 나왔다. visualization 결과에서도 다른 방법에 비해 노이즈를 잘 제거한 것을 볼 수 있다. 

또한 이렇게 학습시킨 네트워크가 어떻게 작동하는지 조금 더 알아보기 위해 centered kernel alignment (CKA) 방법으로 결과를 분석했다. 아래 그림의 오른쪽 그림은 가우시안 노이즈와 poisson 노이즈에 대해서 네트워크 layer의 correlation을 측정한 것인데 mask triaining을 적용한 네트워크는 deep한 layer에서도 correlation이 높게 나타났으나 baseline은 deep layer에서 correlaiton이 상당히 떨어진다. 즉 possion 노이즈와 gaussian 노이즈에 대해 매우 다르게 작동한다는 의미이다. 즉, mask training이 generalization 성능에 많은 영향을 끼쳤음을 확인할 수 있다. 

마지막으로 baseline과 mask training 방법의 test 결과 curve를 나타낸 것이다. baseline 방법은 training을 할 수록 다른 노이즈 타입들에서 성능이 떨어지는 반면 mask training 방법은 training을 할 수록 다른 노이즈 타입에서도 성능이 같이 오른다. 

Conclusion & limitation

간단한 mask training 방법으로 denoising에서 generalization 성능을 올렸다. 하지만 mask training이라는 방법 자체가 input의 많은 정보를 가리기 때문에 정보 손실이 필수적으로 생기는 한계점이 있다. 

아래 그림처럼 mask ratio를 올리면 올릴 수록 다른 type의 노이즈에서는 성능이 오르지만 training에 쓰인 가우시안 노이즈에 대해서는 성능이 조금씩 떨어진다. mask가 많아질수록 information 손실이 커지기 때문에 사실 당연한 결과이고 이런 trade-off를 완화할 방법을 더 연구할 필요가 있어 보인다.