Style Locality for Controllable Generation with k NN Language Models (2023)

Abstract: Recent language models have been improved by the addition of external memory. Nearest neighbor language models retrieve similar contexts to assist in word prediction. The addition of locality levels allows a model to learn how to weight neighbors based on their relative location to the current text in source documents, and have been shown to further improve model performance. Nearest neighbor models have been explored for controllable generation but have not examined the use of locality levels. We present a novel approach for this purpose and evaluate it using automatic and human evaluation on politeness, formality, supportiveness, and toxicity textual data. We find that our model is successfully able to control style and provides a better fluency-style trade-off than previous work

## introduction controllable language generation is the study of developing models to generate language with predefined characteristics, including stylistic (e.g. politeness, formality), demographic (e.g. gender, age), or content-related (e.g. keywords, entities) attributes @xcite . in recent years, chatbots have quickly gained widespread use in industries including marketing, support systems, education, healthcare, cultural heritage, and entertainment @xcite . chatbots generate language to prepare a humanlike or natural response for the user based on the intent and context information provided by the user, environment, and other factors @xcite . for a chatbot in a support system, controllable generation could be leveraged to respond more politely. in an educational context, the chatbot could give more entertaining factual answers to enhance the learning experience. recently, memory-augmented language models have been shown to substantially improve performance by memorizing rare language patterns @xcite . the use of a datastore allows the model to scale to larger text collections without additional training by adding new data to the datastore. we present a model that combines the advantages of transformer language models, locality types, and memorization to enable the controlled generation of language in particular styles. the model makes use of a k-nearest neighbor architecture to retrieve similar contexts from memory and uses specifically trained locality weights to reweight neighbors according to the desired style. in contrast to previous locality work, our locality levels are based on the style, data source, and similarity between styles, rather than the location of text within a document @xcite . we compare our model to previous work, showing that it outperforms previous memory-augmented language models on style control, while providing a favorable trade-off between style control and fluency. neural turing machines are neural networks connected to resources in external memory that they can access through attention. the combined system is comparable to a turing machine but is differentiable and can be trained by gradient descent. work from @xcite showed the model's ability to learn algorithms like copying, sorting, and associative recall. memory-based recurrent neural networks similarly have been shown to implement continuously differentiable versions of well-known data structures including stacks and queues, outperforming deep rnns on sequenceto-sequence transduction tasks @xcite , and can solve difficult pattern recognition tasks with an expandable memory @xcite . end-to-end memory networks are neural networks that have a recurrent attention model applied to a potentially huge external memory. since the model is trained end-to-end, it needs much less supervision during training and is therefore more widely applicable in realistic settings. on the text8 and penn treebank @xcite datasets, the model showed similar performance to rnns and lstms @xcite . extending language models with a continuous cache to adjust their prediction to the recent history showed improvement over previous memoryaugmented networks @xcite . the model is a more condensed version, which accesses past hidden activations through a dot product with the current hidden activation and stores them as memory. this system is effective in handling enormous memory sizes. @xcite introduced k-nearest neighbor language models (knn-lms), outperforming previous memory approaches. the model encodes context vectors and stores them as keys in a dictionary with a corresponding next word. during decoding, the model finds nearest neighbors to the current context, retrieves their values from the dictionary, and interpolates the distribution given by the nearest neighbors with the distribution given by the base language model. recently, @xcite introduced structural locality levels to knn-lms. the idea is to reweight nearest neighbors by their similarity to other text as determined by the location of the text in the structure of the document. they showed that their model outperformed @xcite on wikipedia text and java code. the locality levels for code included the project and directory that the code was in, while for wikipedia it included four levels; (1) different category and section, (2) same section -different category, (3) different section -same category, and (4) same category and section. each of these levels is assigned a weight, learned by a linear model on a sample of the data. they found that learned weights are lower for more similar levels and that the final predictions are more likely to be nearest neighbors from the same locality. controllable text generation controllable characteristics of text include stylistic ones like formality, sentiment, and politeness. they may also be the demographic characteristics of the author, such as age and gender, or keywords and entities as content @xcite . ziegler et al. (2019) showed that language models can control use of sentiment and physically descriptive language when learning from human feedback via reinforcement learning. @xcite introduced a model where attribute classifiers are used to guide generation. the gradients of these classifiers are used to push the latent space representation to control for the topical style, sentiment, and toxicity. @xcite created a generator and detector of fake news that controls generated text by conditioning on the authors, date, headline, and/or domain. other work examines control at decoding time. @xcite used a combination of lms finetuned on desirable or undesirable properties. the distribution of next word probabilities from the desired properties are added to the base lm, while the undesirable distribution is subtracted. similarly, @xcite used classifiers to determine if possible generation steps will result in text with a given stylistic attribute. @xcite used a knn language model which was modified to encode style attributes, altering the encoding space such that retrieved neighbors are more often posts of a particular style. they experimented with politeness, toxicity, and formality, finding that their architecture provides better style control and fluency for some styles over others, and that specialized datastores for each style outperformed datastores with multiple styles. ## datasets the data collected for the experiments were obtained from four different style categories; supportive, formal/informal, polite/impolite and toxic data. in subsequent experiments, we use all combined style datasets excluding neutral data except for experiments on single style models. all data is in english. supportive @xcite published the rud-dit corpus, containing 6k reddit comments and a continuous value representing the offensiveness or supportiveness of each in the range of -1 to 1. we took the top 25% of the dataset as the supportive data, representing 1.1k comments. formal / informal @xcite introduced the gyafc corpus containing ≈ 135k informal and ≈ 144k formal sentences for style transfer. the documents cover the topics entertainment and music and family and relationships. polite / impolite the politeness corpus comes from @xcite and contains 4k posts from wikipedia talk pages. each comment has a label for polite, impolite, or neutral, with 1k of each polite and impolite. toxic we use two datasets for toxicity. the wikipedia detox corpus @xcite ) contains 116k comments that were labeled for toxicity from wikipedia talk pages, with 14k labeled as toxic. the second dataset comes from the civil comments platform @xcite and contains 160k toxic comments. the dataset contains toxicity subcategories that were not used in this work. ## models our work builds off of the knn language models used in previous work. in a standard language model, we want to predict the next word w t given a context w 1 , ..., w t-1 . the language model learns the probability p(w t |w 1 , ..., w t-1 ). a knn language model requires a datastore, d, which is constructed with the encodings of context sequences and their corresponding target word. the encodings are given by a function, f , which in our work is computed with a transformer language model. the keys, k, and values, v of the datastore are then defined as (k, v ) = {(f (w 1 , ..., w t-1 ), w t )} for all context and target word pairs in d. during inference, the model must retrieve similar contexts to the current context. the current context sequence is encoded with f and used to query the datastore to retrieve the nearest neighbors using a distance function, which is the l 2 distance in our experiments. these distances are converted into a distribution by passing the distances to a softmax and aggregating across duplicate target words. finally, the model interpolates the nearest neighbor distribution with the base language model distribution with hyperparameter λ, result-ing in the final distribution p(w t |w 1 , ..., w t-1 ) = λp kn n + (1λ)p lm . xu et al. (2022) introduced locality levels on top of this model, intuiting that additional information about the documents that the text came from could be used to weight nearest neighbors differently. the information they use for this can be thought of as a kind of depth in structured documents. for example, text from the same article and section on wikipedia should probably be more similar than text from different articles or even different sections. we also use locality levels, but do not associate these levels with depth but rather with specific and general styles, and with the source of the text (e.g. wikipedia versus yahoo answers). the knn distribution can be defined formally as p kn n (w t |c t ) ∝ for keyvalue pairs (k, v) in the set of nearest neighbors, n , distance function, d, and context c @xmath0, ..., w t-1 . the addition of locality levels is a parameterization of the distance function, which reweights the distances based on the similarity of localities between the neighbor contexts and current context. this changes the exponentiated function and distribution to for key and context sharing locality level n. we use a linear function for g, as in previous work, with a bias term of zero. the function can be learned with an annotated sample of the training data and is optimized to minimize the log-likelihood of the target token. @xcite experimented with using knn language models for controllable generation and included two baselines. the first is single style models, which are trained on only one style each. this is less efficient as it requires training a separate model for each style. the other baseline is a single style datastore approach. in this method, there are separate datastores for each style rather than having all styles mixed into one datastore. intuitively, the model may benefit from having more similar neighbors in the datastore from which to interpolate, but these similar contexts may not share the desired style to a satisfactory degree. using a separate datastore for each style and restricting retrieval to the desired datastore is effectively forcing the weights of g to zero for contexts that do not share the same style locality level. ## experimental setup we first built a language model from the style dataset and compared it to a fine-tuned language model first trained on wikitext103. all models from previous work were reimplemented and finetuned on our style dataset. we compare to single style models, which are models fine-tuned on only the subset of data for that particular style, which means one model is fine-tuned for each style subset. additionally, we compare to the specific datastore setting from @xcite . this means that instead of encoding all contexts in a single datastore we instead store them separately, so that retrieval is forced to return only neighbors in that style. we describe our training of the style locality generation model and our comparison using human evaluation. experiments used an nvidia a100 gpu. the datastores for our experiments took between 8gb and 20gb disk space. aside from the datastore sizes, we use the same parameters as @xcite . we built our models off of the fairseq library @xcite . we used their wiki103 lm architecture. 1 the model has a hidden size of 1024, 16 layers, and 8 attention heads. ## conclusion we constructed a k-nearest neighbors language model that is able to leverage locality levels based on the style, data source, and similarity between styles, effectively reweighting nearest neighbors with additional available information. we discussed how our model is trained using polite, impolite, formal, informal, supportive, and toxic data. reimplementing models from previous work, we found that using style and data source gave an improved perplexity and was able to outperform the single style models and @xcite on style control. in contrast to previous work, we found that the single style datastores were only slightly better for style-control but came with a large drop in fluency, making them less favored overall. our model showed better fluency than all models except for @xcite , though our model provides a better style-fluency trade-off than previous work. our work provides initial steps toward style controllable nearest neighbors language models. future work may explore the style-fluency trade-off in more detail and the impact on downstream applications, such as with conversational agents.