KRA: K-Nearest Neighbor Retrieval Augmented Model for Text Classification

Abstract

Text classification is a fundamental task in natural language processing (NLP). Deep-learning-based text classification methods usually have two stages: training and inference. However, the training dataset is only used in the training stage. To make full use of the training dataset in the inference stage in order to improve model performance, we propose a k-nearest neighbors retrieval augmented method (KRA) for deep-learning-based text classification models. KRA works by first constructing a storage system that stores the embeddings of the training samples during the training stage. During the inference stage, the model retrieves the top k-nearest neighbors of the testing text from the storage. Then, we use text augmentation methods to expand the retrieved neighbors, including traditional augmentation methods and a large language model (LLM)-based method. Next, the method weights the augmented neighbors based on their distances from the target text and incorporates their labels into the inference of the final results accordingly. We evaluate our KRA method on six benchmark datasets using four commonly used deep learning models: CNN, LSTM, BERT, and RoBERTa. The results demonstrate that KRA significantly improves the classification performance of these models, with an average accuracy improvement of 0.3% for BERT and up to 0.4% for RoBERTa. These improvements highlight the effectiveness and generalizability of KRA across different models and datasets, making it a valuable enhancement for a wide range of text classification tasks.