[论文解读] Neural Nearest Neighbors Networks

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

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匹配的特征空间的优化，而非依赖于固定的手工设计特征。
• 开发一种可微分的、连续的KNN松弛方法，使其在温度参数趋近于零时保持原始KNN的行为。
• 展示所提出的N³块在图像恢复和集合推理任务中提升性能的有效性。

提出的方法

• 提出一种基于温度控制的softmax在成对距离上的连续确定性松弛，实现对选择过程的反向传播。
• 定义一种可微分的邻居加权方案，其中邻居权重通过距离上的软最小值计算，温度参数控制选择的锐度。
• 引入N³块，一种可学习的神经网络层，将松弛后的KNN操作应用于特征图，以可微方式聚合来自相似块的信息。
• 采用温度参数，实现从均匀平均（高温）到硬KNN选择（低温）的插值，同时支持梯度优化。
• 将松弛后的KNN操作用作可微分的注意力机制，可泛化标准软注意力，并可堆叠于深层网络中。
• 采用端到端训练方式，联合优化特征提取器、距离度量和邻居聚合，通过反向传播实现。

实验结果

研究问题

• RQ1 我们能否使KNN选择可微分，以支持深度神经网络中的端到端训练？
• RQ2 对KNN匹配的特征空间进行优化是否能提升图像恢复任务中的性能？
• RQ3 一种可微分的KNN松弛方法是否能在图像去噪和超分辨率任务中超越固定特征的非局部方法和强大的CNN基线模型？
• RQ4 在深层网络中，学习到的温度和邻居加权在不同特征层次上如何变化？
• RQ5 N³块能否有效应用于集合值数据，如对应关系分类任务？

主要发现

• 在Urban100数据集上，×2超分辨率任务中N³块实现了30.80 dB的PSNR，优于强基线模型VDSR（30.76 dB）及其他非局部方法。
• 在Urban100数据集上，×4超分辨率任务中，N³网络实现了25.23 dB的PSNR，超过VDSR（25.18 dB）和WSD-SR（25.16 dB），且视觉结果更清晰。
• 在σ=25的高斯去噪任务中，N³块将PSNR从基线DnCNN的29.74 dB提升至29.99 dB，即使在更大图像块上训练也表现优异。
• 网络学习到在浅层使用更锐利的（温度更低）选择方式（平均最高权重≈0.21），而在深层使用更均匀的加权方式（≈0.04），表明其具有分层特征利用特性。
• N³完整模型的运行时开销为DnCNN的3.5倍，但性能显著优于DnCNN和KNN基线模型，且开销相近。
• 在对应关系分类任务中，仅将N³块添加到近期神经网络基线模型中，即可获得显著性能提升，证明其在集合类任务中的通用性。