Nearest Neighbor Language Models for Stylistic Controllable Generation (2022)

Abstract: Recent language modeling performance has been greatly improved by the use of external memory. This memory encodes the context so that similar contexts can be recalled during decoding. This similarity depends on how the model learns to encode context, which can be altered to include other attributes, such as style. We construct and evaluate an architecture for this purpose, using corpora annotated for politeness, formality, and toxicity. Through extensive experiments and human evaluation we demonstrate the potential of our method to generate text while controlling style. We find that style-specific datastores improve generation performance, though results vary greatly across styles, and the effect of pretraining data and specific styles should be explored in future work.

## introduction language models with external memory, like @xcite 's recent k-nearest neigbour language model (knn-lm), have demonstrated impressive predictive performance. great reductions in perplexity are achieved through storing the encoding of contexts from the training data. a sequence of tokens is encoded by the model and stored as a key, which is paired with a value representing the next token in the sequence. during decoding, similar contexts are recalled based on their key similarity, and values are interpolated with the base language model's predictions. in this work, we augment the encoding with stylistic attributes, such that the keys are more heavily influenced by the style of the encoded text. by explicitly encoding the style, the similarity is more strongly affected by the stylistic attributes than previous models. when decoding, we can then provide a style (e.g. polite or formal) and the most similar contexts are both relevant in content and more likely to conform to the provided style. the ## related work recent approaches to controllable generation have included fine-tuning large models, such as @xcite , who condition on 50 control codes during training, which represent different styles, topics, and languages. other approaches avoid retraining by modifying the predictions only at decoding time. the fudge model predicts for a sequence, the likelihood that possible generation steps will result in a sequence that satisfies a given constraint @xcite . the dexperts model alters the probability distribution of a language model (lm) based on the predictions of other lms that are fine-tuned on specific desired or undesired attributes @xcite . @xcite and others similarly modify gradients directly during prediction. khandelwal et al.'s k-nearest neigbour language model is based on the groundwork laid by previous work that augmented language models with a cache memory of recent observations. @xcite captured local context of up to a few thousands tokens to improve predictions. @xcite expanded upon this idea by storing all past hidden activations in a memory. khandelwal et al. then replaced the recurrent network with a transformer to better model long-term dependencies. aspects of knn-lms have been improved upon, in terms of performance, efficiency, and additional functionality. even though knn-lms achieve state-of-the-art predictive performance, the retrieval operation is very computationally expensive. @xcite and @xcite explored techniques to improve inference speed, such as compressing embeddings, or training an additional model to dynamically disable retrieval for predictions where the datastore is unlikely to improve the result. @xcite implemented a similar model and focused on improving scalability. @xcite extended the model with a gating mechanism that learns to combine short and long term memory with local context. @xcite improved performance by leveraging structural locality features such as topic clusters in text or project hierarchies in source code repositories. @xcite also extended their model for use in machine translation @xcite , which has also received efficiency improvements @xcite . ## methodology the main goal of our work is to expand knn-lm functionality and we build off of @xcite , which we will refer to as the baseline architecture. we modified this architecture to accept additional style attributes as input, and concatenate these to the input text encoding. this has the effect of modifying the embedding space such that it encodes both semantic and stylistic properties (see appendix b). we will refer to this as the style architecture. after the input style and context are encoded, they are stored in the datastore. we experimented with separating datastores based on the distribution of style values in the dataset. for this part we take subsets with specific style values (e.g. only toxic or only polite) from the datasets and construct datastores containing only data from those subsets. we refer to this as separate datastores. datastores containing examples of different style are referred to as mixed datastores. ## datasets we use 4 datasets, 3 of which contain style attributes. unless stated otherwise, we use the default splits provided by the original work. wikitext-103 wikitext-103 (wt103) is a collection of good and featured wikipedia articles @xcite . it is provided in a tokenized form, with case, punctuation and numbers, totalling 103 b tokens. for better compatibility with our tokenization of other datasets we modified wt103's tokenization by replacing all occurrences of *n 't (e.g. in "can't") with * n't. politeness the stanford politeness corpus (spc) consists of 11k utterances from stackexchange and wikipedia talk pages, annotated with politeness scores, which we use as style attributes @xcite . we created an 85-7.5-7.5 split for training, validation, and test. formality grammarly's yahoo! answers formality corpus (gyafc) is the largest available corpus for formality style transfer, containing about 110k formal/informal sentence pairs @xcite . it is divided into the domains entertainment & music and family & relationships, which makes it suitable for training and evaluation with in/out-of-domain data. we do not use the parallel nature of the corpus, but assign each sentence a style attribute (-1 for informal-and 1 for formal sentences) and re-split the training subset to obtain an 80-10-10 train/validation/test split. toxicity for toxicity we use the jigsaw unintended bias in toxicity classification dataset, which contains comments from the civil comments platform annotated with several binary toxicity labels representing types of toxicity (e.g. insult, identity attack, sexually explicit) @xcite . we use only the toxic label as a style attribute. we further use the real toxicity prompts dataset for human evaluation @xcite . this dataset contains the beginning of sentences and has been used to test if models can continue generation without toxicity. ## experiments in preliminary experiments we examined the effect of float precision on model performance. using wt103 and the setup of @xcite we found that using half precision halves inference time without hurting perplexity, and reduces inference time for the knn-lm with a negligible 0.25 increase in perplexity. in the following sections we use half-precision to speed up the experimentation process. in the main experiments we first examine the impact on perplexity when incorporating style values into the model. then we compare our method to the baseline through human evaluation. ## conclusion we examined the use of knn language models for controllable stylistic generation using politeness, formality, and toxicity as target styles. our findings show that simply encoding style in the architecture improved perplexity of the language model. a human evaluation further showed that specific datastores for target styles outperform the standard mixed datastore, and that our model generally outperformed the baseline knn model in terms of fluency and style control, though results on specific styles varied. future work is needed to fully understand the effect of pretraining and benefits of the model variants for specific styles and should also consider comparisons to other controllable generation models, such as @xcite . our code is available on github at https:// github.com/d8xa/style-knnlm .