PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Nearest Neighbors Networks

Tobias Plötz, Stefan Roth

arXiv:1810.12575v1cs.CVcs.LG

TL;DR

Non-local methods traditionally depend on hand-defined feature spaces and non-differentiable KNN selection. This paper introduces a continuous deterministic relaxation that preserves hard KNN as temperature approaches zero and builds the N3 block for trainable non-local processing, reporting improvements across correspondence classification and image-restoration tasks.

  • Problem

    Existing non-local methods use KNN matching in hand-defined feature spaces, while KNN selection is non-differentiable with respect to the distances needed for task-driven optimization.

  • Method

    The paper replaces discrete KNN selection with a continuous deterministic temperature-controlled relaxation and integrates it into the neural nearest neighbors block for end-to-end non-local processing.

  • Results

    The N3 block improves correspondence classification, image denoising, and single-image super-resolution, outperforming strong CNN and non-local baselines.

  • Takeaways & Limitations

    N3 provides a general neural-network building block that combines self-similarity-based non-local processing with differentiable neighbor selection.

Abstract

from arXiv · show

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.

1 Introduction

Local CNNs trade receptive-field size against localization, while non-local methods exploit recurring image structures but traditionally use non-differentiable KNN matching. The paper introduces a differentiable continuous relaxation and an N3 block for trainable non-local processing.

  • Motivation: CNN architectures primarily use local convolutions and stack blocks to obtain broader receptive fields for dense prediction.Dilated convolutions and hourglass architectures increase context but trade context size for localization accuracy.
  • Motivation: Non-local restoration methods exploit self-similarity by aggregating matching patches selected using hand-crafted similarity measures.Examples include non-local means and BM3D, which aggregate information across an image to restore local patches.
  • Contributions: The paper addresses KNN selection’s non-differentiability with a continuous deterministic relaxation differentiable with respect to pairwise distances and temperature.The relaxation interpolates between more uniform weighting and original KNN selection as temperature changes.
  • Contributions: The N3 block enables end-to-end trainable non-local processing based on self-similarity and improves correspondence classification, denoising, and single-image super-resolution.The authors report improvements over strong CNN and non-local baselines across these tasks.

2 Related Work

Related work spans classical self-similarity-based restoration, CNN denoising, and non-local neural networks. The paper distinguishes its neighbor-set relaxation from methods that learn features but aggregate only a single weighted average.

  • Non-local image restoration: Classical non-local restoration methods exploit repeated structures within images as a prior for recovering local patches.Non-local means weights noisy pixels by neighborhood similarity, while BM3D transforms stacks of matching patches and applies shrinkage.
  • Non-local neural networks: Non-local neural networks extend context for vision tasks, but their aggregation is restricted to a single weighted average of features.The paper presents its differentiable nearest-neighbor selection as a generalization that can recover weighted averaging when k=1.
  • Distinction: The paper’s key distinction is producing a differentiable set of neighbors rather than only aggregating neighbor labels or features.This distinction separates the approach from differentiable KNN classifiers and soft-attention-style architectures.
  • Image denoising: CNN-based denoising methods have achieved strong accuracy, alongside classical restoration approaches.Denoising also matters because image noise degrades downstream vision tasks and serves as a component of broader restoration pipelines.

3 Differentiable k-Nearest Neighbors

The method represents KNN as sequential temperature-controlled categorical sampling, then replaces discrete samples with continuous expectations. As temperature approaches zero, the resulting deterministic selection converges to hard KNN while remaining differentiable.

  • KNN formulation: KNN selects the first k database items nearest to query q under a distance metric, but its deterministic selection rule is not differentiable.The paper therefore constructs a parametric stochastic process before deriving continuous relaxations.
  • Stochastic relaxation: The first neighbor is sampled from a categorical distribution whose logits are negative query-item distances scaled by temperature.As temperature tends to zero, the distribution concentrates on the closest database item.
  • Stochastic relaxation: Further neighbors are sampled iteratively by setting the previously selected logit to negative infinity, preventing duplicate selections.Each updated set of logits defines the categorical distribution for the next index.
  • Continuous relaxation: The continuous relaxation replaces one-hot weight vectors with continuous expectations, yielding deterministic neighbor representations.Expected weight vectors are used to update logits and compute subsequent continuous neighbors.
  • Limit behavior: As t →0, every expected neighbor converges to the corresponding hard KNN item, so hard KNN remains the relaxation’s limit case.The convergence follows inductively from the first-neighbor limit and the sequential logit updates.
  • Properties: The continuous neighbors are differentiable with respect to pairwise distances and temperature, which may be learned globally or per query.The method effectively computes k weighted averages, with k=1 recovering a single weighted average.

4 Neural Nearest Neighbors Block

The N3 block integrates continuous, differentiable nearest-neighbor selection into feed-forward networks using self-similar features. It computes embeddings, distances, and temperatures to form neighbor volumes that are concatenated with the input features, with image-specific patch processing for broader context.

  • An N3 block derives query and database items from the same features, integrating differentiable nearest-neighbor selection into feed-forward networks.This instantiates non-local processing through the principle of self-similarity.
  • Embedding network: The embedding network computes feature embeddings and pairwise distances, while a separate branch predicts an item-specific temperature tensor.For image data, CNNs parameterize the embedding and temperature functions; Euclidean distance is used in the considered tasks.
  • Embedding network: Potential neighbors can be restricted to a local subset of items, allowing the N3 block to scale linearly rather than quadratically with the number of items.For images, this subset can be all patches within a specified local region.
  • Continuous nearest neighbors selection: The block computes k continuous nearest-neighbor feature volumes from distances and temperatures, then concatenates them with the original features instead of reducing them early.Concatenation delays feature fusion while preserving equal dimensionality between the input and neighbor volumes.
  • N3 block for image data: For image data, the N3 block operates on image patches rather than pixels to provide broader local context for matching and aggregation.This is a technical modification of the generic N3 block for image restoration.

5 Experiments

Experiments evaluate N3Net for image denoising, real-image denoising, single-image super-resolution, and correspondence classification. Across these tasks, integrating differentiable non-local processing improves strong CNN baselines and competing non-local methods.

  • Image denoising: N3Net combines DnCNN-based local processing with N3 blocks that compute learned non-local neighbor volumes.The denoising architecture uses three DnCNNs with six blocks each, while N3 blocks extract patches and match them within local image regions.
  • Ablation study: On Urban100 at σ = 25, learned patch embeddings in N3Net clearly improve over DnCNN and KNN baselines, while the lighter model exceeds DnCNN with fewer layers and parameters.The lighter model uses 12 rather than 17 layers and 427k rather than 556k parameters.
  • Comparison to the state of the art: +0.54dB at σ = 50 and +0.79dB at σ = 70 on Urban100 are reported gains over DnCNN, with N3Net outperforming local and trainable non-local models across tested denoising settings.The comparison includes DnCNN, RED30, FFDNet, NLNet, UNLNet, and NN3D; the post-hoc NN3D step performs less consistently across datasets.
  • Real image denoising: N3Net sets a new state of the art on raw-image Darmstadt denoising and surpasses the cited leading methods on developed sRGB images.The comparison reports significant margins over DnCNN and BM3D for raw images and higher PSNR than TWSC and CBDNet for sRGB images.
  • Single image super-resolution: For Set5 super-resolution, N3Net gains +0.15dB over VDSR at ×4 and +0.36dB over WSD-SR at ×2.N3Net consistently improves over VDSR across super-resolution factors, while SelfEx and WSD-SR perform worse than N3Net.
  • Correspondence classification: Correspondence classification improves by 10% to 30% in MAP when an N3 block is inserted into CNNet across the evaluated settings.The settings test outdoor generalization between St. Peter and Reichstag and indoor training and testing on Brown.

6 Conclusion

The paper introduces differentiable nearest-neighbor selection and the N3 block, demonstrating benefits across denoising, super-resolution, and correspondence classification.

  • The continuous KNN relaxation remains differentiable with respect to pairwise distances and supports integration into neural networks through the N3 block.
  • N3 blocks improve image denoising, single-image super-resolution, and correspondence classification against CNN-based and non-local approaches.
Loading 1810.12575v1…