Regularized Training of Nearest Neighbor Language Models

TL;DR

This paper introduces a regularization technique during training of language models that enhances the effectiveness of post-hoc kNN retrieval, leading to improved language modeling performance especially for high-frequency words.

Contribution

It proposes a novel L2 regularization on activations during training to improve kNN-based language model performance, a concept not previously explored.

Findings

Regularization improves kNN classification accuracy.

Enhanced performance on WIKI-2 and WIKI-103 datasets.

Better handling of high-frequency words.

Abstract

Including memory banks in a natural language processing architecture increases model capacity by equipping it with additional data at inference time. In this paper, we build upon $k$NN-LM \citep{khandelwal20generalization}, which uses a pre-trained language model together with an exhaustive $k$NN search through the training data (memory bank) to achieve state-of-the-art results. We investigate whether we can improve the $k$NN-LM performance by instead training a LM with the knowledge that we will be using a $k$NN post-hoc. We achieved significant improvement using our method on language modeling tasks on \texttt{WIKI-2} and \texttt{WIKI-103}. The main phenomenon that we encounter is that adding a simple L2 regularization on the activations (not weights) of the model, a transformer, improves the post-hoc $k$NN classification performance. We explore some possible reasons for this improvement. In particular, we find that the added L2 regularization seems to improve the performance for high-frequency words without deteriorating the performance for low frequency ones.