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-theart 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. 1

heading: 1 Introduction
text: 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.,
1Please be advised that this paper contains potentially controversial results and examples to some readers, included solely for research purposes to explore model capabilities.
2018; OpenAI, 2023; Fung et al., 2023, 2024). Their output word embeddings are learned automatically to calculate word output likelihoods during pre-training on language corpora. Typically, the dot product c⊤ev between a computed context vector and a learnable output word embedding ev for token v is usually used as the word logit. The word output probability is defined as the softmax over all word logits:
P (v|c) = exp(c ⊤ev)∑
u∈V exp(c ⊤eu)
, (1)
where V is the whole vocabulary. While being a fundamental topic in natural language processing, previous work on interpreting them is usually focused at the word level, such as their semantic information (Şenel et al., 2018), word senses (Hewitt et al., 2023), and analogical relations (Mikolov et al., 2013; Park et al., 2017). However, as the word embeddings are optimized for generation loss during pre-training, the learned embedding space should be closely associated with LMs’ generation distributions. In this work, we propose to study the roles that word embeddings play in LM generation, which remains an underexplored topic, and analyze a simple while effective LM steering method LM-Steer.
16410
We start by examining the theoretical relation between word embeddings and LM generation distributions. We find that linear transformations in output word embeddings are equivalent to LM generation style changes. This motivates our empirical investigation on a lightweight and simple steering method, LM-Steer, to steer LM generation distribution flexibly and transparently. LM-Steer deploys a d×d learnable linear transformation W on the output word embeddings, where d is the embedding dimension. Specifically, the embeddings ev are replaced with ev + ϵWev. Here, ϵ acts as a “steering value” to provide a simple control on steering polarity and intensity. Inherently, LM-Steer discovers hidden dimensions in word embeddings that are associated with text styles. By transforming these dimensions, LM-Steer affects LMs’ interpretation of words and, consequently, generation distributions.
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 GPT2 family (Radford et al., 2019) to modern large language models, including GPT-J (Wang and Komatsuzaki, 2021) and Llama-2-7B (Touvron et al., 2023). 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 (Hu et al., 2021a), 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.

heading: 2 Related Work
text: Understanding Word Embeddings Language models learn word embeddings for individual words automatically after pre-training. To understand them, (Mikolov et al., 2013; Allen and Hospedales, 2019) discover linear translational relations among embeddings, while (Park et al., 2017; Rothe and Schütze, 2016; Ethayarajh, 2019) examine their rotational relations. Some other work examines the semantic information (Şenel et al., 2018; Lund and Burgess, 1996; Jang and Myaeng, 2017; Şenel et al., 2022; Murphy et al., 2012; Faruqui et al., 2014) or word senses (Panigrahi
et al., 2019; Hewitt et al., 2023) of individual words in their embeddings. These efforts, however, have mostly focused on and evaluated word-level interpretations of word embeddings (Chang et al., 2009) while we first investigate their relations with language model generations. Control of Language Models has been of growing interest in recent years, motivated by the increasing capabilities of LMs (Li et al., 2023a). This area originates from the need to leverage the generation capabilities of large language models while avoiding the need for time-consuming and costly retraining or fine-tuning. Attempts include applying attribute classifiers or heuristic constraints at decoding time (Kumar et al., 2022; Dathathri et al.; Liu et al., 2021; Yang and Klein, 2021), treating the generation process as an optimization problem over the embedding or token sequences (Kumar et al., 2021), or post-editing the output (Li et al., 2018). These techniques are often computationally expensive and rely on suitable external classifiers. More recently, prompting-based control for large language models has received much attention, which, however, relies on the quality and availability of large language models (Brown et al., 2020; OpenAI, 2023), and may also necessitate the deliberate training (Raffel et al., 2020; Zhou et al., 2023). It can also be challenging to design effective prompts for complex or nuanced control goals. Parameter-efficient fine-tuning such as LoRA (Hu et al., 2021a) focuses on learning lowrank approximations of model parameters. However, this cannot achieve flexible and transferrable language model steering like ours. Probably most closely related to our work are attempts to discover “steering” vectors or tokens (Subramani et al., 2022; Li and Liang, 2021), and also similar work in image generation (Jahanian et al.; Hu et al., 2021b). Different from our model, these efforts focus on other applications such as multi-task learning and sentence recovery, and the learned vectors are not shown to be transferrable or interpretable nor enable flexible control. Language Model Detoxification Motivated by the goal to address the systematic biases embedded in language models, there are efforts in conducting language model de-biasing or de-toxification (Meade et al., 2022; Kaneko et al., 2022). Approaches span all aspects of the language model pipeline. A line of work focuses on automatically obtaining cleaner data (Barikeri et al., 2021; Webster et al., 2021; Dinan et al., 2020). Another line
of work modifies the model workflow design to explicitly accommodate the bias factors (Webster et al., 2021; Schick et al., 2021; Yu et al., 2023; Omrani et al., 2023; Yang et al., 2023). The most related line of work to the herein proposed method involves manipulating embedding space such as Principle Component Analysis and Nullspace Projection (Liang et al., 2020; Bolukbasi et al., 2016; Ravfogel et al., 2020). The evaluation in these settings (Kaneko and Bollegala, 2021; Nadeem et al., 2021; Nangia et al., 2020) mostly consists of quiz-question checking for stereotypical misbeliefs. More related to our method are those mentioned in language model control (Kumar et al., 2022; Dathathri et al.; Liu et al., 2021; Yang and Klein, 2021; Kumar et al., 2021), which constrains or guides text generation according to a classifier. A unique contribution in our work is that the learned LM-Steer can be transferred to detoxify other offthe-shelf language models without a costly training process.

heading: 3 LM-Steer: Revealing Hidden Steers in Word Embeddings
text: 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 ev with
e′v = ev + ϵWev = (I + ϵW )ev, (2)
and call the resulting language model PϵW a “LMSteered 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
2Using an arbitrary ϵ0 possess the same representation
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.

heading: 4 Steering Language Model Generation

heading: 4.1 Language Detoxification
text: It is known that large pre-trained LMs might generate toxic content that appears in the pre-training distribution (Sheng et al., 2019; Gehman et al., 2020), such as inaccurate information, harmful stereotypes, and unethical content. Language model detoxification is the task of mitigating or avoiding these generations in order to enable the safe usage of language models. We experiment on the GPT2 family, Pythia family (Biderman et al., 2023), GPTJ-6B, and Llama-2-7B as the backbone language models.
Setting: Following (Liu et al., 2021), we use Jigsaw Unintended Bias in Toxicity Classification Kaggle challenge3 as the training dataset. For evaluation, we use 10K nontoxic prompts from the RealToxicityPrompts dataset (Gehman et al., 2020). We randomly generate 25 sentences of up to 20 tokens using nucleus sampling (Holtzman et al.)
power as any other ϵ, as there always exists W ′ = ϵ ϵ0 W so that ϵ0W ′ = ϵW holds. 3https://bit.ly/3cvG5py
Baselines LM-Steered⊕
PPLM
0.55
with p = 0.9. Then the toxicity scores (in range [0, 1]) of generations are evaluated using Perspective API 4. Two metrics are reported: the maximal toxicity of generations on each prompt averaged across prompts (“Avg. max. toxicity”) and the averaged probability of generating > 0.5 toxicity (“Toxicity prob.”). We also evaluate generation quality in terms of fluency (perplexity score measured by a GPT2-large) and diversity (Dist-{1, 2, 3}: the portion of distinct {1, 2, 3}-grams). When decoding, we use a steering value of 5ϵ0 for generation, selected based on the balance between generation fluency and task performance on the dev set in Appendix E.
Baselines: DExperts (Liu et al., 2021) trains positive and negative label classifiers and uses the difference in two classifiers’ scores to offset the LM’s original logits. DAPT (Gururangan et al., 2020) simply further pretrains the language model on the non-toxic subset (filtered by Perspective API) of OpenWebText Corpus (OWT) (Gokaslan et al., 2019). PPLM (Dathathri et al.) learns to use the gradients of the label classifier to update the LM’s hidden representations. GeDi (Krause et al.,
4https://perspectiveapi.com
2021) is a model that uses the Bayesian rule for class-conditioned LM generation. MuCoLa (Kumar et al., 2022) models the text generation as an optimization problem regarding the classifier scores. PromptT5 (Raffel et al., 2020) T5 is a pre-trained LM optimized for prompt-based task solving, and we use “Complete this sentence so that it embodies a {positive/negative} sentiment:” to prompt T5. LoRA (Hu et al., 2021a) trains low-rank approximations of parameter matrices to achieve parameter-efficient fine-tuning. Finally, we compare with the soft blacklist baseline discussed in Section 3.
Results and Analysis: We present the results in Table 1. Despite the simple design, LM-Steer achieves the best detoxification scores on both metrics, reducing Avg. max. toxicity by > 6% absolute percentages. It is also noteworthy that LMSteer also demonstrates reasonable balance on fluency (2nd lowest perplexity score) and diversity (same-level Dist-k scores with baselines). Figure 3 further shows the detoxification versus baseline size, where LM-Steer+{GPT2 family, Pythia fam-
ily, GPT-J and Llama-2} uniformly outperforms baselines where of all sizes, where more numerical results can be found in Appendix J and M. Incorporation of LM-Steer with LoRA, instruction following, and full embedding tuning are explored in Appendix L, N and O, respectively.
Human Evaluation We compare LM-steer with LoRA, DExperts, and GPT-2 in a pairwise manner with human annotators. Specifically, we follow the practice in DExperts and ask four student human annotators to compare 50 generations from LMsteer and the baseline from 3 perspectives: detoxification, fluency, and being topical to the prompt. The results are as follows. We can see that LMsteer is ranked significantly less toxic and more topical than the baseline. It performs similarly to DExperts and GPT-2 but better than LoRA in terms of fluency.

heading: 4.2 Sentiment Control
text: We also evaluate LM-Steer’s performance on an extensively studied generation task controlled by sentiment. This ability can be found useful when tailor-
ing persuasive and emotionally appealing messages to specific target audiences in marketing or advertising or to create personalized and engaging user experiences in chatbot systems.
Setting: We follow the setting in (Liu et al., 2021) and use Stanford Sentiment Treebank (SST5) (Socher et al., 2013) as training data, where we use texts with labels 1∼2 as negative samples, and those with 4∼5 labels as positive samples. For evaluation, we use the HuggingFace’s sentiment classifier (Wolf et al., 2020). The generation prompts are a subset of the OpenWebText Corpus filtered by the sentiment analysis classifier. Models are applied on these prompts 25 times to generate up to 20 tokens. We then measure the average percentage of posi-
tive generations for each prompt as the “Positivity” score. Similar to the detoxification task, we use 5ϵ0 for positive sentiment and −5ϵ0 for negative sentiment control.
Baselines: In addition to the baselines used in detoxification, two variants of DExperts, DExperts (pos) and DExperts (neg), which only use one of the two classifiers for guiding generation, are also listed.
Results: Table 3 presents the full results. LMSteer, despite a much simpler and smaller model, takes 1st place on the positive side and 2nd or 3rd place on the negative side and achieves a reasonable balance on fluency and diversity.

heading: 4.3 Continuous and Compositional Control
text: The conceptually simple design of LM-Steer makes it an architecture-agnostic plug-in to diverse language models. We demonstrate that LM-Steer maintains a linearity guarantee, which enables continuous and compositional control. More specifically, our model allows for interpolation and extrapolation on the steering spectrum by simply interpolating and extrapolating the steering value. Moreover, if two LM-Steers ϵ1W1, ϵ2W2 are learned on potentially different tasks, their effect can be combined by decoding with Pϵ1W1+ϵ2W2 .
In Figure 4a, we plot the distribution shift when adjusting sentiment steer ϵ. We also curve the maximal likelihood estimated Beta distribution. In Figure 4b, we observe that LM-Steer can compositionally control sentiment and toxicity, even though there exists a mutual influence between these two factors (e.g., a negative sentiment might also lead to more toxic comments). Table 5 also provides an example of how the generated sentence is continuously steered from toxic to non-toxic, demonstrating a simple fine-grained control on the toxicity level. When the steering value increases from negative to positive, both the number and the intensity
of toxic words (bolded in the table) decrease.

heading: 4.4 Efficiency
text: Thanks to its simple design, LM-Steer enjoys efficiency in multiple perspectives. We vary the detoxification dataset size from 30 to 10k and measure LM-Steer’s performance in Figure 5(b). We see that as few as 30 data points still enable LM-Steer to achieve high detoxification scores (0.322). When dataset size exceeds 3k LM-Steer acquires a good balance between detoxification and generation quality. We also show decoding time and parameter efficiency in Table 4, where our model only uses 1% of the baseline’s learnable parameter size and uses a low computation overhead during decoding.

heading: 5 LM-Steers Connect Word Embeddings with the Text Distribution
text: 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
0.2
0.275
0.35
0.425
0.5
gpt2 gpt2-medium gpt2-xl gpt-j-6b
Transferred
Avg. Max. Toxicity ( )↓ Original Trained
(124M) (355M) (1.5B) (2.7B)
(a) Transferring an LM-Steer to other LMs with explicit-form calculation. The transferred LM-Steer maintains the detoxification ability partially.
30 100 300 1k 3k 10k full 0.1
1
10
100
1000
0.2
0.28
0.36
Perplexity Toxicity
0.44
0.52
(b) LM-Steer is capable of learning from only dozens of data samples to achieve decent detoxification. More data are beneficial for fluency.
decomposes word embeddings into interpretable dimensions closely related to word selection in a particular distribution. Section 5.2 shows that LM-Steer can highlight indicative words. Finally, in light of the hypothesis that different LMs have correlated word embedding spaces, Section 5.3 illustrates how LM-Steer can be transferred between LMs with an explicit expression.

heading: 5.1 Interpreting Word Embeddings
text: LM-Steer provides a lens on how word embeddings correlate with LM word embeddings: what word dimensions contribute to or contrast to a specific style. In the detoxification experiment, we conduct an SVD decomposition of the learned W . Among S, V,D, the D component can be interpreted as a ranked list of the most “magnified” row dimensions in the transformation W . We then take its first 9 rows and list the most influenced words in Table 6. Dimensions 2, 4, and 6 are filtered out as they only match non-English tokens. Although offensive to read, this table helps us understand what kind of words are most related to toxicity and thus suppressed by LM-Steer in a generation. More details are explained in Appendix G.

heading: 5.2 Highlighting Keywords in Styled Texts
text: LM-Steer also enables automatically pointing out what specific words are most indicative of the style in a given sentence. We conjecture that these distinctive words contribute most to the likelihood difference between a task-specific LM and a domain-general LM. In Table 7, we list a few toxic sentences and the most indicative text spans highlighted by LM-Steer. To acquire such text spans, we calculate logPϵW (vi| · · · vi−1) − logP0(vi| · · · vi−1) which is the difference of loglikelihoods of each token under steered (PϵW ) and the original language model (P0). Then, we deploy a dynamic programming algorithm to obtain the continuous sub-sequence under length 5 with the highest cumulative difference. We list highlighted words in toxic prompts in the RealToxicityPrompts dataset in Table 7, which shows insulting, cursing, controversial, and sexually explicit words in each sentence.
There’s another controversial Hollywood racial decision that Stacey Dash is sinking her teeth into.
The UFC champ then suggested Justino is a longtime PED user with her most d**ning comments.
But I really have a question for you: Why would I go on a game show and play into the bulls**t allowing myself to be ranked by some fake competition?
I think sexism prevents this from being a real win for fat people.
If they want to be fair and non hypocritical idiots they should.
Table 7: Toxic sentences with toxic keywords highlighted by LM-Steer after training detoxification on GPT2-Large.

heading: 5.3 Transfering LM-Steer Between Models
text: A much-desired property of LM-Steer, because of its theoretical soundness, is its transferability to other language models. Details and derivations of LM-Steer transfer are in Appendix F. Intuitively speaking, the original logit c⊤ev can be understood as a similarity or matching metric between context vector c and word embedding ev. In LM-Steer, the logit is offset by ϵ times c⊤Wev, which is also a bilinear similarity. To transform this LMSteer to another language model, we need to map the context vectors and word embeddings between word embedding spaces ev = He′v
c⊤Wev = (Hc′)⊤W (He′v) = c ′⊤(H⊤WH)e′v
(3)
We work by first identifying a linear mapping H from target LM word embeddings to source LM word embeddings. Then, the matrix H⊤WH can be inserted into the target LM as LM-Steer. This is motivated by prior work on the linear mapping between word embeddings from different models (Li et al., 2021). Finally, the calculated steering matrix is directly applied to the target LM. Figure 5(a) shows the performance after we transfer the LMSteer learned on GPT2-large to LMs of other sizes, ranging from gpt2 (124M) to GPT-J-6B (6B). We can see a uniform improvement in transferred LMSteers, with GPT2 and GPT2-medium getting similar scores (0.307 and 0.308) to the best baseline (DExperts).

heading: 6 Conclusions
text: 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. LMSteer 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.
Limitations
One limitation of LM-Steer is that it works on word embeddings and focuses on conditions related to wording. This restricts its capability to deal with more complex tasks, such as syntactic trees or persuasive techniques that involve logical reasoning. Additionally, our model is dependent on word embeddings, so the model cannot work with language model APIs that do not provide direct access to these embeddings.

heading: Acknowledgement
text: This research is partially supported by U.S. DARPA KAIROS Program No. FA875019-2-1004, DARPA INCAS Program No. HR001121C0165, U.S. DARPA SemaFor Program No. HR001120C0123, DARPA MIPs Program No. HR00112290105, and DARPA ITM Program No. FA8650-23-C-7316. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of DARPA or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes, notwithstanding any copyright annotation therein.

heading: A Broader Impacts
text: The intended use of this work is to contribute to advancements in fine-grained and efficient control of language generation in AI, with experiments shown on sentiment modulation, political stance adjustment, and language detoxification. We do not aim to create a tool for manipulating public opinion or promoting specific political ideologies, but instead to provide methods for enhancing the reasoning interpretability, and safety of AI language models. Our techniques offer the potential for fine-tuned sentiment adjustment and toxic content mitigation, thereby contributing to more reliable, unbiased, and respectful language generation systems. We would like to emphasize that on the problem of language model toxicity, we limit our model to modeling detoxification only. This encourages positive and socially beneficial usage of our model as well as general language models.

heading: B Formal Statement of Theorem 1
text: Hidden Markov Models (HMMs) (Baum et al., 1970) is a widely used framework for analyzing discrete stochastic processes. Because of its generality (modeling arbitrary distributions), intuitiveness, and interpretability (containing a structured state space), it has long been used as a primary choice when modeling language distribution. is a discrete stochastic process with a set of n states S and a set of m observations or emissions O, with arbitrary indexing of S and O. The time step t = 0 distribution is determined by initial state distribution s0 ∼ π. For each later time step t ≥ 1, the state transition probabilities are represented by a matrix T, where T (s, s′) = P (st+1 = s′|st = s) denotes the probability of transitioning from state s to state s′. At each time step one observation ot is emitted, with the emission probabilities represented by a matrix B, with B(s, o) = P (ot = o|st = s). A sequence of observations can be denoted as o = {o1, o2, . . . , oT }. The probability distribution over sequences o then follows formula:
P (o1, · · · , oT ;π)
= π⊤ ( T−1∏
t=0
diag(p(ot))T ) p(oT ), (4)
where p(o) is a |S|-dim vector indicating P (o | s) for all states s ∈ S.
Language Models In generative language models, the sequence is generated word-by-word by a conditional probability P (ot | o1, · · · , ot−1). The common technique to model this probability is to first calculate the inner product between a contextual vector c(o1, · · · , ot−1) and word embeddings E = (eo, · · · ) ∈ Rd×|O|, namely, l = c(o1, · · · , ot−1)⊤E. Here, l is known as the word logits, which then usually passes through a softmax operator to get a distribution over words. For simplicity of analysis, in this work, we assume a linear formulation and let conditional probability P (ot|o1, · · · , ot−1) = c(o1, · · · , ot−1)⊤eot . By the chain rule, multiplying the conditional probabilities will give us the full probability: ∏T t=1 P (ot | o1, · · · , ot−1) = P (o1, · · · , oT ). We are then interested in the situation where a language model is good enough to represent an equivalent distribution with HMM.
In this study, we aim to model the influence of conditions in text generation. This section describes how we incorporate conditions in HMMs. Conventionally, people assume a d-dimensional state representation ϕs for every state s, and ddimensional ψo for each observation o, so that they can compute the probabilities T (s, s′) = ϕ⊤s Aϕ ′ s, B(s, o) = ϕ⊤s ψo and π(s) = ϕ ⊤ π ϕs for some ϕπ. We also use matrices Φ,Ψ to denote the stacked representations Φ = (ϕs|s ∈ S),Ψ = (ψo|o ∈ O). Here we introduce an additional condition component in state representations, so that ϕs can be par-
titioned into two sub-vectors: ϕs = ( ϕs,semantic ϕs,condition ) . Here ϕs,semantic ∈ Rds represents the ds-dim semantic information, and ϕs,condition ∈ Rdc the dcdim condition information related to state s. Then we assume that the transition probability T (s, s′) comes from both semantic relations and conditional similarities between s′ and s: T (s, s′) = ϕ⊤s,semanticA
′ϕs′,semantic + ϕ⊤s,conditionϕs′,condition. We also make the following assumptions regard-
ing the state representations:
Assumption 1. State representations ϕ also satisfy the following properties:
1. Values for each dimension are uniformly normalized to a constant: ∀i ∈ [1..d],∑s∈S ϕ2s,i = C.
2. Dimensions are linearly independent: ∀i, j ∈ [1..d] and i ̸= j,∑h∈H ϕh,iϕh,j = 0.
3. Dimensions are also conditionally independent: if i, j ∈ [1..d], k ∈ [ds + 1..d] are not all the
same, ∑
s∈S ϕs,iϕs,jϕs,k = 0.
The validity of the assumption is discussed in Appendix H. Then, we present the result below, revealing that shifting between conditions is equivalent to a linear transformation in word embedding space.
Theorem 2. Assume assumption 1 holds. Suppose there are two initial distributions π = ϕ⊤πΦ, π
′ = ϕ⊤π′Φ, so that ϕπ and ϕπ′ only differ in their condition-parts: ϕπ,semantic = ϕπ′,semantic. Also, suppose the elements in ϕπ,condition are non-zero. Then there exists a matrix W so that, by transforming word embeddings from E to WE, the LM which originally simulates the text distribution starting with π will now turn to be equivalent to a distribution initiating from π′.

heading: C Formal Statement and Proof of
text: Theorem 2
To prove Theorem 2, we start by claiming a construction of matrix W . Then we prove that when assumptions 1 hold, W can change each conditional likelihood function from p(vi | o1, ·, oi−1, π) to p(vi | o1, ·, oi−1, π′) up to a constant factor. Finally, by chaining the conditional likelihoods, we see that W can change the sentence-level probability distribution of the HMM from π-initialization to π′-initialization.
Assuming full column-rank for E and p(o), we have the following connection between LM and HMM:
Proposition 1. There exist projection matrices R1 and R2 so that R⊤1 R2 = In and
c(o1, · · · , ot−1)⊤
=
  π ⊤∏t−1 t′=1 diag(p(o ′ t))T
π⊤ (∏t−2
t′=1 diag(p(o ′ t))T ) p(ot−1)
 R⊤1 ,
eo = R2p(o).
We first construct a helper matrix W ′ =( Ids 0 0 Λ ) so that Λ is diagonal and W ′ϕinit = ϕ′init. Such a solution exists as we assume ϕinit,condition contains only non-zero values. Then, we can construct the matrix W as W = R+1 Φ ⊤W ′ΦR+2 ⊤, where R+1 , R + 2 are pseudoinverse matrices of R1, R2, respectively. Lemma 3. TΦ⊤W ′Φ = Φ⊤W ′ΦT .
Proof. First, it is easy to see that, by Assumption 1.1 and Assumption 1.2, the representation matrix Φ is row-orthonormal to constant C2:
ΦΦT = C2Id
.
Then we have the following proof:
TΦ⊤W ′Φ =Φ⊤ ( Ts 0 0 Idc ) ΦΦ⊤W ′Φ
=C2Φ ⊤ ( Ts 0 0 Idc )( Ids 0 0 Λ ) Φ =C2Φ ⊤ ( Ts 0 0 Λ ) Φ =C2Φ ⊤ ( Ids 0 0 Λ )( Ts 0 0 Idc ) Φ =Φ⊤W ′ΦΦ⊤ ( Ts 0 0 Idc ) Φ =W ′T
Lemma 4. ∀v ∈ V , we have that, Φdiag(p(o))Φ⊤W ′Φ =W ′Φdiag(p(o)).
Proof. To prove this, we first prove that
Φdiag(p(o))Φ⊤ has the form ( A 0 0 Λ′ ) , where Λ′ is also diagonal. This is equivalent to saying that, for any two one-hot vectors e(i), e(j), if i ∈ [ds + 1..d] or j ∈ [ds + 1..d], then
e⊤i Φ diag(p(o))Φe ⊤ j
= ∑
h∈H ϕh,iϕh,jp(v | h) = fv(i, j)1(i = j). (5)
For any i ̸= j, ∑
h∈H ϕh,iϕh,jp(v | h)
= ∑
h∈H ϕh,iϕh,j
∑
k∈[1..d] ϕh,kθv, k
= ∑
k∈[1..d] θv, k
∑ h∈H ϕh,iϕh,jϕh,k
= ∑
k ̸∈{i,j} θv,k
∑ h∈H ϕh,iϕh,jϕh,k
+ θv,i ∑
h∈H ϕ2h,iϕh,j
+ θv,j ∑
h∈H ϕh,iϕ
2 h,j
(Asm. 1.3) =0 + θv,i ∑
h∈H ϕh,iϕh,j
(Asm. 1.2) =0 + 0
=0
We then have the following:
Φdiag(p(o))Φ⊤W ′Φ
= ( A 0 0 Λ′ ) W ′Φ
= ( A 0 0 Λ′Λ ) Φ =W ′ ( A 0 0 Λ′ ) Φ =W ′ΦΦ⊤ ( A 0 0 Λ′ ) Φ =W ′Φdiag(p(o))
By combining Lemma 3 and 4, we have the following lemma: Lemma 5.
Tdiag(p(o))Φ⊤W ′Φ = Φ⊤W ′ΦTdiag(p(o))
Proof.
Tdiag(p(o))Φ⊤W ′Φ
=Φ⊤ ( Ts 0 0 Idc ) Φdiag(p(o))Φ⊤W ′Φ =Φ⊤ ( Ts 0 0 Idc ) W ′Φdiag(p(o)) =Φ⊤W ′ΦTdiag(p(o))
Finally, when we apply Lemma 3 and 5 to the language model formulation, we can see that the conditional likelihood function has been steered to:
pW (vi | o1, · · · , oi−1;π) =c(o1, · · · , oi−1;π)WE = π⊤Tdiag(o1)T · · ·Tdiag(oi−1)TR⊤1
π⊤Tdiag(o1)T · · ·Tp(oi−1) WR2PO
= π⊤Tdiag(o1)T · · ·Tdiag(oi−1)TΦ⊤W ′Φ
π⊤Tdiag(o1)T · · ·Tp(oi−1) PO
(Lemma 3)
= π⊤Tdiag(o1)T · · ·Tdiag(oi−1)Φ⊤W ′ΦT
π⊤Tdiag(o1)T · · ·Tp(oi−1) PO
(by Lemma 5)
= πΦ⊤W ′Φ⊤Tdiag(o1)T · · ·Tdiag(oi−1)T
π⊤Tdiag(o1)T · · ·Tp(oi−1) PO
= ϕinitW ′ΦTdiag(o1)T · · ·Tdiag(oi−1)T π⊤Tdiag(o1)T · · ·Tp(oi−1) PO = ϕ′initΦTdiag(o1)T · · ·Tdiag(oi−1)T
π⊤Tdiag(o1)T · · ·Tp(oi−1) PO
(by definition)
= π′Tdiag(o1)T · · ·Tdiag(oi−1)T π⊤Tdiag(o1)T · · ·Tp(oi−1) PO
∝ c(o1, · · · , oi−1;π′)E =p(oi | o1, · · · , oi−1;π′)
Therefore, the steered conditional likelihood is equivalent to an HMM initiating from π′ (up to a normalization constant over vocabulary O). By chaining the conditional likelihood functions, it is easy to see that the actual output sequence distribution is now:
pW,normalized(o1, · · · , oL;π)
= L∏
i=1
normalizeO(pW (vi | o1, · · · , oi−1;π))
= L∏
i=1
p(oi | o1, · · · , oi−1;π′)
=p(o1, · · · , oL;π′)
This concludes our proof to Theorem 2.
D Implementation Details
In this paper, we leverage the HuggingFace package5 and its model checkpoints. To implement LMSteer, we simply wrap the self.forward function of language model transformer’s lm_head, and inject the computation formula of LM-Steer. In specific, the token logits are replaced w ith c⊤(I + ϵW )eo, we change the computation order and first compute c′ = c + ϵWc), then compute c′⊤eo. We find that this increases computational efficiency in practice and avoids the problem caused by many Transformers sharing input and output word embedding parameters in storage. Another trick we applied in experiments is that, as there is a systematical distribution shift between pre-training corpus and domain-specific dataset (such as detoxification dataset and reviews), we add one “dummy” steer Wdummy to fill this overall distribution gap. Therefore, for positive label training, we use model Pϵ0(W+Wdummy), and for negative label training, we use model Pϵ0(−W+Wdummy). This is where the 3M parameters come from in Table 4.
For optimization, we use Adam optimizer (Kingma and Ba, 2014) with a 1e-2 learning rate and train for 1k steps. The steer matrix W is initialized with a Gaussian distribution of 0 mean and 1e-3 variance. Across all experiments, we run three initial seeds of 0, 1, and 2 for training. When required to generate 25 sentences on each prompt, we use random seeds 0, 1, 2, ..., 24. Our hardware is one single Tesla V-100 GPU with 16GB CUDA memory.

heading: E Hyperparameter Selection
text: We select the decoding hyper-parameter based on a balance of task performance and generation quality on the detoxification task’s dev set. The scores are listed in the table below.
When gradually increasing the steering value, the detoxification success rate increases while generation fluency decreases. To better balance the two ends, we select 5ϵ0 for downstream evaluation, as it does not compromise perplexity too much while achieving decent task performance.

heading: F Details of Transferring LM-Steer to Other Language Models
text: To transfer an LM-Steer from one LM M1 to another LM M2, we notice that LM-Steer essentially
5https://huggingface.co
adds one term c⊤Weo to the logits, where both c and eo can be viewed as residing in word embedding space. Therefore, W can be considered as a similarity matrix in M1’s word embedding space. To use W in M2, we propose to map M2’s word embedding space to that of M1 before using W as usual. The process works in 2 steps.
First, we identify a linear mapping from W2 to W1’s word embedding space. We start with building a list of anchor words. Specifically, we select the top 4k words shared by both vocabularies. We denote the token embedding matrices as E′1, E ′ 2 respectively. Then, we initialize a mapping H with a Gaussian distribution of 1e-3 initial variance, and we apply Adam optimizer 0.01 learning rate for 5k steps. Secondly, After acquiring the mapping matrix H , we map both the context and embedding vectors to Hc and Heo, respectively. So the additive term for LM M2 is now c⊤H⊤WHeo, which is equivalent to using a steer matrix H⊤WH for the LM M2.
This mapping process is not precise, as word embeddings between LMS are not linearly associated. So we observe an increased instability in generation if we use large ϵ. Therefore, we reduce the steering value to 0.1 of its original scale, 0.5ϵ0 for generation. This is the setting for getting results in Figure 5a.

heading: G Details of Investigating Interpretability
text: In Section 5.1, we interpret the weights learned in LM-Steer and list discovered keywords in Table 6. A detailed description of getting these results is as follows. First, we conduct SVD decomposition of steer matrix W . The resulting D matrix can then
be interpreted as a ranked list of significant row vectors. We take the first 9 rows and compute their dot products with word embeddings. As the row vector does not tell us which of the 2 directions indicates an increased probability, we select 20 tokens with top dot product and 20 tokens with bottom dot product as two candidate groups. Each group is concatenated to a text sequence and passed to Perspective API, and the group with a larger toxicity value is considered true “keywords”. If, however, Perspective API recognizes the language as not English, which happened to rows No. 2, 4, and 6, then we discard this row as they contain mostly symbols and non-English words. Finally, we filter out suffix tokens, and the remaining keywords are listed in Table 6.
H Validity of Assumptions
To verify the validity of the assumptions, we did an experiment for searching for valid HMMs while satisfying the assumption 1. It is trivial to construct valid Ψ as long as a valid Φ can be found. So specifically, we set ds = 20 and dc = 1 to represent a one-conditional HMM. We let n = 200 and randomly initialized Φ with Gaussian distribution with variance 1e-3. Then we construct the following objective function
L = Lnorm + Ldist + Lindependence + Lconditional,
where
Lnorm = ∑
i
( ∑
s∈S ϕ2s,i −
1
dn
∑
s∈S,i′ ϕ2s,i′) 2
Ldist = ∑
s,s′ max(−T (s, s′), 0) + ∑
s
( ∑
s′ T (s, s′)− 1)2
Lindependence = ∑ i ̸=j ( ∑ s ϕs,iϕs,j) 2
Lconditional = ∑
i,j,knot one value,k∈[dc+1,d] + ( ∑
s
ϕs,iϕs,jϕs,k) 2
Generally, this objection characterizes the derivation of Φ from the assumptions. We use Adam optimizer with learning rate 1e-3, and ReduceLROnPlateau 6 with patience 100 and reduce factor 0.5. The optimization process lasts 500k steps, starting from random seeds 0, 1, 2, and 3. On all random seeds, the objective function reduces from greater than 1 to less than 1e-5. This indicates that valid HMM solutions satisfy the assumption.

heading: I Comparison with A (Soft) Word Blacklist
text: First, we explain that a control vector is equivalent to a SWB. This is because by adding a vector to context c′ = c + ϵw, we are equivalently adding a logit bias to each word: c′⊤E = c⊤E + ϵw⊤E, where w⊤E is the static logit bias vector for each word. Then we point out that theoretically, LMSteer is more expressive in representing sequence distributions than them, since LM-Steer’s formulation c⊤(I + ϵW )E can let different words be preferred in different contexts c. Theoretically, there exists a LM-Steer for steering between any two finite-length distributions (with proof below). SWB cannot achieve this (a counterexample below). Intuitively speaking, a blacklist or whitelist uniformly applied at all positions cannot possibly achieve flexible control over the complex language distribution without hurting the generation quality.
I.1 Formal statement of the universality of LM-Steer
Let D1 and D2 be two finite-length finitevocabulary sequence distributions, there exists a
6https://pytorch.org/docs/stable/generated/ torch.optim.lr_scheduler.ReduceLROnPlateau.html
context vector function c(o1, o2, · · · , oi), a word embedding E and a matrix W so that an LM-Steer with −W and +W represents distribution D1 and D2 respectively.
Proof. We prove the existence by construction. Let I(o1, o2, · · · , oi) ∈ N be an arbitrary indexing function for all subsequences. With bounded subsequence length, I values are also bounded by a number d. We let c(o1, · · · , oi) =∑
o onehot(I(o1, o2, · · · , oi, o)) ∈ Rd. For each token o, the word embeddings is eo such that, for any subsequence (o1, o2, · · · , oi), eo[I(o1, o2, · · · , oi, o))] = (D1+D2)(o|o1,o2,··· ,oi) 2 . Then W is as follows: for any (o1, o2, · · · , oi) and token o, W [I(o1, o2, · · · , oi, o), I(o1, o2, · · · , oi, o)] = (D2−D1)(o|o1,o2,··· ,oi) (D1+D2)(o|o1,o2,··· ,oi) . It is diagonal. We omit verification due to space limits.
I.2 3.2 Construction of a counterexample for SWB
Let D1 be a one-point distribution on sequence “AB”, and D2 be a uniform distribution on sequences “BA”, “AB”. For any language model representing D1, there does not exist an SWB that can convert the language model into D2. Verification of this counterexample is trivial, and we omit it due to space limits.

heading: J Results of LM-Steer on Pythia Family
text: Pythia 7 is a family of causal language models developed by EleutherAI. Raining in size from 14M to 2.8B, these models provide an excellent testbed for evaluating the effect of language model sizes. The table below shows the performance of LMSteer applied to the Pythia language models. We can see a trend of better fluency but decreasing detoxification when the model size increases, indicating a greater controlling difficulty and better base quality of larger language models.

heading: K LoRA Configuration
text: Thanks for suggesting a more comprehensive evaluation. We use the Huggingface Transformers’ default implementation of LoRA on GPT-2-large to align with standard practices in the field and compare with our method. To test the effect of configurations, we iterate the LoRA rank over 8, 16, 32, and 64. Following the practice of (Lee et al., 2023),
7https://github.com/EleutherAI/pythia
we set the alpha scalar equal to the rank. We report LoRA’s performance below.
L Incorporating in LoRA
Thanks to its theoretical foundations, LM-Steer is orthogonal to other methods and is intuitively compostable with other control methods. As an example study, we select LoRA to combine with LM-Steer on detoxification with GPT2-large as the backbone. The result is as follows. Combining LoRA with LM-Steer produces a better detoxification score than LoRA alone (although not as good as LM-Steer alone), at the cost of a degraded quality score.

heading: M Results on GPT-J-6B
text: We go on and evaluate LM-Steer on GPT-J-6B model. We use the same evaluation setup as in the main body. The results are shown in Table 12. We observe that LM-Steer is able to reduce the toxicity of the generations while maintaining the perplexity and the fluency of the generations. We also observe that the diversity of the generations is not affected by LM-Steer.

heading: N Effect of LM-Steer on Instruction Following
text: We study the effect of LM-Steer on the LMs’ performance under prompts. For prompts, we let ChatGPT write 10 strongly positively biased prompts in Table 13. Then, we let GPT-J generate 25 tokens under instructions (i.e., prompting) and compare its generation with and without LM-Steer.
We see that LM-Steer can still steer generation even under a positive prompt while maintaining generation fluency. Without LM-Steer, the sentiment is 94.00. With a positive LM-Steer, the
average sentiment increases to 97.20. With negative LM-Steer steering, the average sentiment decreases significantly to 82.80. Example generation are listed in Table 14 and 15.

heading: O Embedding Tuning
text: We also consider the possibility of tuning the word embeddings of the backbone model. We use the same training procedure as LM-Steer, but instead of tuning the last layer, we tune the word embeddings. We use the same hyperparameters as LMSteer. The results are shown in Table 16. We observe that the performance is comparable to SoftBlacklist, but is still inferior to LM-Steer. We believe that this is because the word embeddings are shared across all the layers, and tuning the word embeddings may cause the model to forget the knowledge learned from the pre-training. We will leave the investigation of this direction for future work. Embedding-tuning is also more expensive than LM-Steer, as it requires tuning the word embeddings for the entire vocabulary, while LM-Steer only requires tuning the last layer.

heading: Future Work
text: We will leave the investigation of this direction for future work. Embedding-tuning is also more expensive than LM-Steer, as it requires tuning the word embeddings for the entire vocabulary, while LM-Steer only requires tuning the last layer