# KNN-BERT: Fine-Tuning Pre-Trained Models with KNN Classifier

arXiv: https://arxiv.org/abs/2110.02523 · HTML: https://mlsanity.com/paper/2110.02523/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 55 references that are in the map.
- First posted on arXiv 6 Oct 2021.

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representations](https://mlsanity.com/paper/2006.11477/) (2020)
- [Prototypical Networks for Few-shot Learning](https://mlsanity.com/paper/1703.05175/) (2017)
- [Towards Evaluating the Robustness of Neural Networks](https://mlsanity.com/paper/1608.04644/) (2016)
- [FaceNet: A Unified Embedding for Face Recognition and Clustering](https://mlsanity.com/paper/1503.03832/) (2015)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Unsupervised Feature Learning via Non-Parametric Instance-level Discrimination](https://mlsanity.com/paper/1805.01978/) (2018)
- [Learning deep representations by mutual information estimation and maximization](https://mlsanity.com/paper/1808.06670/) (2018)
- [Generalized Cross Entropy Loss for Training Deep Neural Networks with Noisy Labels](https://mlsanity.com/paper/1805.07836/) (2018)
- [Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss](https://mlsanity.com/paper/1906.07413/) (2019)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Data-Efficient Image Recognition with Contrastive Predictive Coding](https://mlsanity.com/paper/1905.09272/) (2019)
- [BERT-ATTACK: Adversarial Attack Against BERT Using BERT](https://mlsanity.com/paper/2004.09984/) (2020)
- [Time-Contrastive Networks: Self-Supervised Learning from Video](https://mlsanity.com/paper/1704.06888/) (2017)
- [Fine-Tuning Pretrained Language Models: Weight Initializations, Data Orders, and Early Stopping](https://mlsanity.com/paper/2002.06305/) (2020)
- [On Mutual Information Maximization for Representation Learning](https://mlsanity.com/paper/1907.13625/) (2019)
- [Large-Margin Softmax Loss for Convolutional Neural Networks](https://mlsanity.com/paper/1612.02295/) (2016)
- [Supervised Contrastive Learning for Pre-trained Language Model Fine-tuning](https://mlsanity.com/paper/2011.01403/) (2020)

## Built on by (most cited first, 7 of 7)

- [Mask-guided BERT for Few Shot Text Classification](https://mlsanity.com/paper/2302.10447/) (2023)
- [Relation Extraction as Open-book Examination: Retrieval-enhanced Prompt Tuning](https://mlsanity.com/paper/2205.02355/) (2022)
- [Reasoning Through Memorization: Nearest Neighbor Knowledge Graph Embeddings](https://mlsanity.com/paper/2201.05575/) (2022)
- [Feature-Level Debiased Natural Language Understanding](https://mlsanity.com/paper/2212.05421/) (2022)
- [Mere Contrastive Learning for Cross-Domain Sentiment Analysis](https://mlsanity.com/paper/2208.08678/) (2022)
- [Revisiting k-NN for Fine-tuning Pre-trained Language Models](https://mlsanity.com/paper/2304.09058/) (2023)
- [Label Distribution Learning-Enhanced Dual-KNN for Text Classification](https://mlsanity.com/paper/2503.04869/) (2025)

## Abstract (arXiv)

Pre-trained models are widely used in fine-tuning downstream tasks with linear classifiers optimized by the cross-entropy loss, which might face robustness and stability problems. These problems can be improved by learning representations that focus on similarities in the same class and contradictions in different classes when making predictions. In this paper, we utilize the K-Nearest Neighbors Classifier in pre-trained model fine-tuning. For this KNN classifier, we introduce a supervised momentum contrastive learning framework to learn the clustered representations of the supervised downstream tasks. Extensive experiments on text classification tasks and robustness tests show that by incorporating KNNs with the traditional fine-tuning process, we can obtain significant improvements on the clean accuracy in both rich-source and few-shot settings and can improve the robustness against adversarial attacks. footnoteall codes is available at https://github.com/LinyangLee/KNN-BERT