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.

• The paper introduces a continuous, differentiable relaxation of the KNN rule, enabling effective backpropagation via pairwise distances.
• It presents the neural nearest neighbors block (N block) that facilitates end-to-end training by leveraging self-similarity for non-local processing.
• Empirical results show that adding the N block significantly improves performance in image denoising, super-resolution, and correspondence classification tasks.

Overview of Neural Nearest Neighbors Networks

The paper "Neural Nearest Neighbors Networks" presents a novel approach to integrating non-local processing techniques within neural network architectures. Non-local methods have long exploited the self-similarity inherent in natural signals, particularly in areas such as image restoration and analysis. Traditional approaches heavily rely on $k$-nearest neighbors (KNN) matching within a pre-defined feature space. This strategy, however, encounters a significant challenge due to the non-differentiability of the KNN selection rule, which hinders the optimization of feature space based on application performance.

To address this issue, the authors propose a continuous and deterministic relaxation of the KNN selection rule that preserves differentiability with respect to pairwise distances. The original KNN behavior is preserved as a special case when the temperature parameter used in the relaxation approaches zero. This relaxation is operationalized through the introduction of a novel architectural component in neural networks, termed as the neural nearest neighbors block (N block). This block is designed to facilitate non-local processing by leveraging the principle of self-similarity, thus enhancing modern neural network architectures.

Technical Contributions

The paper makes three notable contributions:

1. Relaxation of the KNN Rule: The authors introduce a continuous, deterministic relaxation of the KNN rule, capable of backpropagation via pairwise distances. This relaxation is controlled by a temperature parameter, which is trainable and allows for a smooth transition between a uniform weighting scheme and the hard KNN selection.
2. Neural Nearest Neighbors Block: From the proposed relaxation, a neural network layer, named the neural nearest neighbors block (N block), is developed. This layer supports end-to-end trainable non-local processing through the use of self-similarity without requiring a fixed matching space.
3. Empirical Demonstration: The paper demonstrates substantial improvements over existing models in various tasks, such as image denoising, single-image super-resolution, and correspondence classification. The addition of the N block to establish a neural nearest neighbors network (N Net) resulted in these improvements, thus validating the effectiveness of optimizing feature spaces for non-local processing.

Experimental Validation

The proposed N Net exhibits superior performance across several image restoration tasks. Specifically, it significantly improves accuracy in image denoising and super-resolution tasks by combining the strengths of traditional non-local methods with the adaptability of neural networks. For correspondence classification, which involves operating on set-valued data, the augmentation of a neural network with the N block demonstrated increased classification accuracy, showcasing the block's effectiveness in enhancing feature representation and set reasoning tasks.

Practical and Theoretical Implications

From a practical perspective, the architecture facilitates building more robust neural networks capable of capturing non-local dependencies, which are crucial for tasks involving high degrees of self-similarity, such as those in image processing domains. Theoretically, it also opens avenues for exploring further relaxations of other non-differentiable operations in neural networks, potentially broadening the scope of end-to-end trainable models.

Future Directions

The work suggests several potential future research directions. For instance, exploring the applicability of the N block beyond image-related tasks into other domains such as text or sequential data, where non-local dependencies and self-similarity may also be present. Moreover, optimizing the temperature parameter of the relaxation based on specific characteristics of different data domains or tasks could further enhance the block's adaptability and performance.

In conclusion, this paper systematically addresses the challenge of incorporating non-local processing into neural network frameworks by introducing a differentiable relaxation of the KNN rule, which enables joint optimization of feature spaces. The proposed neural nearest neighbors networks effectively leverage the strengths of traditional image processing techniques alongside modern neural architectures, demonstrating substantial improvements in both synthetic and real-world applications.