Word Embeddings Are Steers for Language Models (2024)

Abstract: Language models (LMs) automatically learn word embeddings during pre-training on language corpora. Although word embeddings are usually interpreted as feature vectors for individual words, their roles in language model generation remain underexplored. In this work, we theoretically and empirically revisit output word embeddings and find that their linear transformations are equivalent to steering language model generation styles. We name such steers LM-Steers and find them existing in LMs of all sizes. It requires learning parameters equal to 0.2% of the original LMs’ size for steering each style. On tasks such as language model detoxification and sentiment control, LM-Steers can achieve comparable or superior performance compared with state-of-the-art controlled generation methods while maintaining a better balance with generation quality. The learned LM-Steer serves as a lens in text styles: it reveals that word embeddings are interpretable when associated with language model generations and can highlight text spans that most indicate the style differences. An LM-Steer is transferrable between different language models by an explicit form calculation. One can also continuously steer LMs simply by scaling the LM-Steer or compose multiple LM-Steers by adding their transformations. Our codes are publicly available at https://github.com/Glaciohound/LM-Steer.

## introduction in recent years, language models (lms) have significantly advanced various natural language processing (nlp) tasks such as machine translation, sentiment analysis, schema induction, summarization, and sociocultural understanding (brown et al., 2020; kojima et al.; li et al., 2023b; radford et al., 1 please be advised that this paper contains potentially controversial results and examples to some readers, included solely for research purposes to explore model capabilities. figure 1 : we find hidden steers in output word embeddings. by linearly transforming word embeddings, language model generations are "steered" toward different style polarity and levels. their output word embeddings are learned automatically to calculate word output likelihoods during pre-training on language corpora. typically, the dot product c ⊤ e v between a computed context vector and a learnable output word embedding e v for token v is usually used as the word logit. the word output probability is defined as the softmax over all word logits: empirically, we find such lm-steers exist prevalently in lms thanks to its general formulation and ignorance of model architectures. it can efficiently steer lms of all sizes, ranging from smaller gpt-2 family @xcite to modern large language models, including gpt-j @xcite and llama-2-7b @xcite . on tasks such as language model detoxification and sentiment control, it achieves comparable or superior performance compared to the state-ofthe-art controlled generation baselines. we also find multiple merits of lm-steer both as a steering tool and lens for inspecting the relation between word embeddings and language model generation. an lm-steer can highlight text spans that best indicate a style in the full text. a learned lm-steer also makes word embedding dimensions interpretable by pointing out the dimensions closest related to a style, revealing what kinds of words contribute to or contradict a style. lm-steer is both parameter efficient and data-efficient: on gpt2-large, it learns only parameters only 0.2% the size of the original model (9% of the size of lora @xcite , a parameter-efficient fine-tuning method) and able to text sentiments on dozens of sentences. a learned lm-steer is transferable to other lms with different sizes by explicit-form calculation without additional training. moreover, lm-steer theoretically enables both continuous and compositional control. this allows for dealing with diverse and nuanced situations, such as fine-grained personalized or customized generation, without re-training for each scenario. ## lm-steer: revealing hidden steers in word embeddings as a theoretical motivation, we first show an informal theorem relating output word embeddings with generation styles. we leave the formal statement as well as the proof to appendix c and only present an intuitive interpretation. theorem 1. (informal) with certain assumptions, shifting styles in language models is equivalent to a linear transformation in word embedding space. inspired by this discovery, we propose lm-steer to apply a linear transform in the output word embedding space. lm-steer is conceptually simple and straightforward to implement. an illustration of lm-steer is presented in figure 2(a) . specifically, we assume a language model with fixed parameters. we replace its each output word embeddings e v with and call the resulting language model p ϵw a "lm-steered model". here, the "steer matrix" w is the only learnable parameter in lm-steer, and ϵ is an adjustable scalar indicating the polarity and intensity of the "steering value". without loss of generality, we arbitrarily pick a small value ϵ 0 = 1e -3 as the default steering value. 2 we use p ϵw to denote the steered language model. figure 2(b , c ) shows the training and generation process of lm-steer. during training, we use the positively steered model p ϵw to fit the positively labeled texts, with maximal likelihood as the training objective. when negative texts are available, we also fit them with p -ϵw . when generating with lm-steer, the user only needs to specify a steering value ϵ and then decode the language model. more details are in appendix d. intuitively explaining, lm-steer matrix w learns to identify word embedding dimensions that are best associated with a target style and manipulate among those dimensions to achieve language rewording. in section 5.1, we show this enables an interpretation of word embeddings by analyzing the learned w . we also theoretically and empirically compare lm-steer against a simplified version, a (soft) word blacklist (swb): learning a global logit offset is applied to each token candidate after the original logits are computed. as we demonstrate in appendix i, adding a (learnable) vector to context vectors c achieves a similar effect with swb. we also further prove that lm-steer is theoretically expressive of any distribution shift, while a swb is unable to do so. in 4, we show that swb indeed yields inferior performance than lm-steer. ## lm-steers connect word embeddings with the text distribution in previous sections, lm-steer revealed the hidden biases encoded in the automatically learned word embeddings of lms. this section provides an alternative perspective, where these hidden biases serve as a lens for interpreting the connection between word embeddings and the generation distribution of lms. section 5.1 demonstrates how lm-steer ## conclusions in this work, we discover the prevalent phenomenon of word embeddings containing steers for language model generation. we demonstrate the promise and efficacy of lm-steer, a theoretically grounded, simple, and lightweight approach for the steering of generative language models. lm-steer can model various styles and achieve comparable or superior performance to baselines in language model detoxification and generation control. lm-steer also allows for continuous and compositional control and can be transferred to other language models. more importantly, it provides an interpretation of how word embeddings interplay with language model generation. so far, we have only studied output word embeddings, so it is intriguing to ask whether similar phenomena apply to other components, such as input word embeddings and hidden layers.