[Paper Review] Rethinking Nearest Neighbors for Visual Classification

Menglin Jia, Bor-Chun Chen+4 more|arXiv (Cornell University)|Dec 15, 2021

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper proposes a simple yet effective method to enhance modern neural network classifiers by integrating k-Nearest Neighbors (k-NN) during both training and inference. By using k-NN predictions to identify hard examples during training and linearly interpolating k-NN and model probabilities at test time, the approach achieves state-of-the-art performance on diverse visual classification tasks, especially in low-data regimes and fine-grained recognition, outperforming standard linear classifiers and demonstrating the enduring value of classical methods in the deep learning era.

ABSTRACT

Neural network classifiers have become the de-facto choice for current "pre-train then fine-tune" paradigms of visual classification. In this paper, we investigate k-Nearest-Neighbor (k-NN) classifiers, a classical model-free learning method from the pre-deep learning era, as an augmentation to modern neural network based approaches. As a lazy learning method, k-NN simply aggregates the distance between the test image and top-k neighbors in a training set. We adopt k-NN with pre-trained visual representations produced by either supervised or self-supervised methods in two steps: (1) Leverage k-NN predicted probabilities as indications for easy vs. hard examples during training. (2) Linearly interpolate the k-NN predicted distribution with that of the augmented classifier. Via extensive experiments on a wide range of classification tasks, our study reveals the generality and flexibility of k-NN integration with additional insights: (1) k-NN achieves competitive results, sometimes even outperforming a standard linear classifier. (2) Incorporating k-NN is especially beneficial for tasks where parametric classifiers perform poorly and / or in low-data regimes. We hope these discoveries will encourage people to rethink the role of pre-deep learning, classical methods in computer vision. Our code is available at: https://github.com/KMnP/nn-revisit.

Motivation & Objective

• To re-evaluate the role of classical k-NN classifiers in modern deep learning pipelines for visual classification.
• To investigate whether k-NN can complement modern neural network models, especially in low-data or fine-grained recognition scenarios.
• To develop a unified framework that integrates k-NN during both training and inference to improve model robustness and accuracy.
• To demonstrate that k-NN, despite being non-parametric and 'lazy', remains highly effective when combined with pre-trained deep features.
• To inspire renewed interest in pre-deep learning methods by showing their complementary strengths with modern neural networks.

Proposed method

• During training, the method uses k-NN predictions to identify hard examples and scales the cross-entropy loss based on k-NN confidence, forcing the model to focus on difficult samples.
• At inference time, the final prediction is obtained by linearly interpolating the model’s output with the k-NN predicted probability distribution.
• The k-NN classifier operates directly on pre-trained visual features from either supervised or self-supervised models, without requiring retraining.
• The method is applied across various backbones (ResNet, ViT, Swin) and pre-training objectives (ImageNet, MoCo, DINO, Barlow Twins), showing broad compatibility.
• The interpolation coefficient is tuned on validation data to balance model and k-NN predictions, improving generalization.
• The approach is evaluated under standard linear evaluation and end-to-end fine-tuning protocols across multiple benchmarks.

Experimental results

Research questions

• RQ1 Can k-NN improve performance when integrated into modern deep learning pipelines for visual classification?
• RQ2 Does k-NN provide particular benefits in low-data regimes or for fine-grained recognition tasks?
• RQ3 Is training-time integration of k-NN predictions more effective than inference-only integration?
• RQ4 How does k-NN compare to standard linear classifiers when used with the same features?
• RQ5 Can k-NN serve as a robust, model-agnostic component across diverse pre-trained representations and architectures?

Key findings

• k-NN integration improves performance across all evaluated datasets, including ImageNet, fine-grained benchmarks, and the NeWT benchmark, with consistent gains.
• On the ImageNet linear evaluation protocol, the method achieves top-1 accuracy of 80.2% using ViT-S/16 features, outperforming the standard linear classifier.
• For low-data regimes, such as in the iNaturalist2021-mini and NeWT benchmarks, k-NN integration leads to significant improvements, with gains of up to 5.2% in mean accuracy.
• The method achieves state-of-the-art results on the Caltech-UCSD Birds-200-2011 and Stanford Dogs datasets, surpassing previous SOTA by over 1% in top-1 accuracy.
• k-NN consistently outperforms or matches linear classifiers across all backbone and pre-training combinations evaluated, including ResNet-50 and ViT-B/16.
• The ablation study confirms that joint training and inference integration is more effective than inference-only integration, with a 1.5% average gain on ImageNet.