[논문 리뷰] Neural Nearest Neighbors Networks

Tobias Plötz, Stefan Roth|TUbilio (Technical University of Darmstadt)|2018. 10. 30.

Image and Signal Denoising Methods참고 문헌 36인용 수 30

한 줄 요약

이 논문은 비차별 가능한 k-최근접 이웃(KNN) 선택 규칙의 미분 가능이고 연속적인 근사를 제안하여 비국소 신경망의 엔드 투 엔드 학습을 가능하게 한다. 이 방법은 학습 가능한 특징 공간을 통해 데이터의 자기 유사성을 활용하는 신경적 최근접이웃 블록(N³ 블록)을 도입하며, 이미지 복원 및 초해상도 분석에서 CNN 및 비국소 기반 모델을 능가하고, Urban100 및 BSD100 데이터셋에서 최신 기준 성능을 달성한다.

ABSTRACT

Non-local methods exploiting the self-similarity of natural signals have been well studied, for example in image analysis and restoration. Existing approaches, however, rely on k-nearest neighbors (KNN) matching in a fixed feature space. The main hurdle in optimizing this feature space w.r.t. application performance is the non-differentiability of the KNN selection rule. To overcome this, we propose a continuous deterministic relaxation of KNN selection that maintains differentiability w.r.t. pairwise distances, but retains the original KNN as the limit of a temperature parameter approaching zero. To exploit our relaxation, we propose the neural nearest neighbors block (N3 block), a novel non-local processing layer that leverages the principle of self-similarity and can be used as building block in modern neural network architectures. We show its effectiveness for the set reasoning task of correspondence classification as well as for image restoration, including image denoising and single image super-resolution, where we outperform strong convolutional neural network (CNN) baselines and recent non-local models that rely on KNN selection in hand-chosen features spaces.

연구 동기 및 목표

• 신경망 내 KNN 선택의 비차별성 문제를 해결하여 엔드 투 엔드 학습이 가능하도록 하는 것.
• 고정된 수작업 특징에 의존하는 것이 아니라, KNN 매칭에 사용되는 특징 공간을 최적화할 수 있도록 하는 것.
• 온도 매개변수 값이 0에 수렴할 때 원래 KNN 동작을 유지하는, 미분 가능한 연속적 근사를 개발하는 것.
• 제안된 N³ 블록이 이미지 복원 및 집합 추론 작업에서 성능 향상에 얼마나 효과적인지 입증하는 것.

제안 방법

• 쌍별 거리에 대해 온도 제어 소프트맥스를 사용하여 KNN 선택의 연속적 결정론적 근사를 제안하며, 선택 과정을 통해 역전파가 가능하도록 한다.
• 온도 매개변수가 선택의 날카기 정도를 제어하는, 미분 가능한 이웃 가중치 부여 방식을 정의한다.
• 특징 맵에 대해 이완된 KNN 연산을 적용하는 학습 가능한 신경망 레이어인 N³ 블록을 도입하며, 유사한 패치들로부터 정보를 미분 가능한 방식으로 통합한다.
• 온도 매개변수를 통해 고온일 경우 균일 평균(고온)에서 저온일 경우 딱딱한 KNN 선택(저온)으로의 보간이 가능하며, 최적화를 위한 기울기가 제공된다.
• 이완된 KNN 연산을 신뢰할 수 있는 어텐션 유사 메커니즘으로 활용하며, 표준 소프트 어텐션을 일반화할 수 있고, 깊은 네트워크에 중첩하여 사용할 수 있다.
• 전체 네트워크를 엔드 투 엔드로 학습하며, 역전파를 통해 특징 추출기, 거리 측정 기준, 이웃 집계를 동시에 최적화한다.

실험 결과

연구 질문

• RQ1 KNN 선택을 미분 가능하게 만들 수 있을까? 이를 통해 딥 신경망 내 엔드 투 엔드 학습이 가능해지는가?
• RQ2 KNN 매칭을 위한 특징 공간을 최적화하면 이미지 복원 작업 성능 향상에 기여하는가?
• RQ3 고정된 특징을 사용하는 비국소 방법과 강력한 CNN 기반 모델을 뛰어넘는, KNN의 미분 가능한 근사가 이미지 복원 및 초해상도 분석에서 성능을 높일 수 있는가?
• RQ4 깊은 네트워크의 특징 계층에서 학습된 온도 및 이웃 가중치는 어떻게 변화하는가?
• RQ5 N³ 블록은 대응 분류와 같은 집합 기반 데이터에 효과적으로 적용될 수 있는가?

주요 결과

• Urban100에서 ×2 초해상도 복원 시 N³ 블록은 PSNR 30.80 dB를 달성하여 강력한 VDSR 기반 모델(30.76 dB)과 다른 비국소 방법을 능가한다.
• Urban100에서 ×4 초해상도 복원 시 N³ 넷은 25.23 dB를 기록하며, VDSR(25.18 dB)와 WSD-SR(25.16 dB)를 모두 앞서며, 더 선명한 시각적 결과를 제공한다.
• σ=25인 가우시안 노이즈 제거 작업에서, N³ 블록은 기준 DnCNN의 PSNR 29.74 dB에서 29.99 dB로 향상되었으며, 더 큰 패치로 학습한 경우에도 성능 향상이 유지된다.
• 네트워크는 초기 레이어에서 더 날카운 선택(낮은 온도)을 사용하며 평균 상위 가중치 ≈ 0.21을, 더 깊은 레이어에서는 더 균일한 가중치 부여(≈ 0.04)를 사용함으로써 계층적 특징 활용을 학습하고 있음을 보여준다.
• 전체 N³ 모델의 런타임 오버헤드는 DnCNN의 3.5배이지만, 유사한 오버헤드를 가진 DnCNN 및 KNN 기반 모델보다 유의미하게 높은 성능을 달성한다.
• 대응 분류 작업에서 최근 신경망 기반 모델에 단순히 N³ 블록을 추가하기만 해도 성능 향상이 크게 발생함을 보여주며, 집합 기반 작업에 대한 일반화 능력을 입증한다.