Neural Nearest Neighbors Networks

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.

to appear at NIPS*2018, code available at https://github.com/visinf/n3net/

References in corpus (8)

• Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising
• Categorical Reparameterization with Gumbel-Softmax
• Multi-Scale Context Aggregation by Dilated Convolutions
• Non-Local Recurrent Network for Image Restoration
• Variational inference for Monte Carlo objectives
• Nonlocality-Reinforced Convolutional Neural Networks for Image Denoising
• Deep Mean-Shift Priors for Image Restoration
• A Trilateral Weighted Sparse Coding Scheme for Real-World Image Denoising

Cited by in corpus (2)

• SwinIR: Image Restoration Using Swin Transformer
• GRDN:Grouped Residual Dense Network for Real Image Denoising and GAN-based Real-world Noise Modeling