mlsanityData to 7 Oct 2026

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

arXiv 2110.02523 ↗First posted 6 Oct 2021cs.CLcited by 7

Not on the map. Why?

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

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

Builds on

34 of 55 references.

References · 34

  1. 2017Attention Is All You Needcited by 59,932
  2. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  3. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,204
  4. 2019RoBERTa: A Robustly Optimized BERT Pretraining Approachcited by 13,266
  5. 2020A Simple Framework for Contrastive Learning of Visual Representationscited by 11,979
  6. 2018Representation Learning with Contrastive Predictive Codingcited by 7,662
  7. 2019Momentum Contrast for Unsupervised Visual Representation Learningcited by 7,482
  8. 2014Explaining and Harnessing Adversarial Examplescited by 7,287
  9. 2016SQuAD: 100,000+ Questions for Machine Comprehension of Textcited by 4,115
  10. 2020wav2vec 2.0: A Framework for Self-Supervised Learning of Speech Representationscited by 3,853
  11. 2017Prototypical Networks for Few-shot Learningcited by 3,001
  12. 2016Towards Evaluating the Robustness of Neural Networkscited by 2,976
22 more
  1. 2015FaceNet: A Unified Embedding for Face Recognition and Clusteringcited by 2,833
  2. 2020Supervised Contrastive Learningcited by 2,803
  3. 2019XLNet: Generalized Autoregressive Pretraining for Language Understandingcited by 2,801
  4. 2019ALBERT: A Lite BERT for Self-supervised Learning of Language Representationscited by 2,772
  5. 2017A Broad-Coverage Challenge Corpus for Sentence Understanding through Inferencecited by 2,665
  6. 2015Character-level Convolutional Networks for Text Classificationcited by 2,293
  7. 2018Unsupervised Feature Learning via Non-Parametric Instance-level Discriminationcited by 1,838
  8. 2018Learning deep representations by mutual information estimation and maximizationcited by 1,265
  9. 2018Generalized Cross Entropy Loss for Training Deep Neural Networks with Noisy Labelscited by 994
  10. 2019Learning Imbalanced Datasets with Label-Distribution-Aware Margin Losscited by 987
  11. 2019Generalization through Memorization: Nearest Neighbor Language Modelscited by 788
  12. 2019Data-Efficient Image Recognition with Contrastive Predictive Codingcited by 672
  13. 2020BERT-ATTACK: Adversarial Attack Against BERT Using BERTcited by 488
  14. 2017Time-Contrastive Networks: Self-Supervised Learning from Videocited by 453
  15. 2020Fine-Tuning Pretrained Language Models: Weight Initializations, Data Orders, and Early Stoppingcited by 377
  16. 2019On Mutual Information Maximization for Representation Learningcited by 315
  17. 2016Large-Margin Softmax Loss for Convolutional Neural Networkscited by 296
  18. 2020Supervised Contrastive Learning for Pre-trained Language Model Fine-tuningcited by 275
  19. 2020Revisiting Few-sample BERT Fine-tuningcited by 237
  20. 2018Deep k-Nearest Neighbors: Towards Confident, Interpretable and Robust Deep Learningcited by 203
  21. 2020BERT-kNN: Adding a kNN Search Component to Pretrained Language Models for Better QAcited by 36
  22. 2019Cross-Entropy Loss and Low-Rank Features Have Responsibility for Adversarial Examplescited by 17

Built on by

Cited by year

Citing papers per year: 2022: 4; 2023: 2; 2024: 0; 2025: 1.

7 papers.

  1. 2023Mask-guided BERT for Few Shot Text Classificationcited by 25
  2. 2022Relation Extraction as Open-book Examination: Retrieval-enhanced Prompt Tuningcited by 13
  3. 2022Reasoning Through Memorization: Nearest Neighbor Knowledge Graph Embeddingscited by 12
  4. 2022Feature-Level Debiased Natural Language Understandingcited by 10
  5. 2022Mere Contrastive Learning for Cross-Domain Sentiment Analysiscited by 8
  6. 2023Revisiting k-NN for Fine-tuning Pre-trained Language Modelscited by 0
  7. 2025Label Distribution Learning-Enhanced Dual-KNN for Text Classificationcited by 0