LM-Switch: Transforming Word Embedding Space for Flexible Language Model Steering

We thank the reviewer for the time and review and are glad to see that the reviewer appreciates the “versatile tool” brought by LM-Switch and the clearly demonstrated motivation. Below is a detailed response to the comments.

Table Representation

Thanks for pointing out a mismarked score in Table 2. Fluency and diversity metrics are not bolded in Tables 1 and 2 following practices in prior work, as they are not the main metric in two tasks.

Figure 2a & 2b Captions

We increased the margin between the caption boxes in Figures 2a and 2b to make them clearer.

Backbone Comparisons

The reviewer raises a valuable concern for discussion. We added the backbone model performance in the revision. For the embedding training, we would like to point out that due to the large size of vocabulary, word embeddings with size $n_{vocab} \times d$ are much larger in size than LM-Switch as well as LoRA. We list its performance below. It shows comparable performance with Soft-Blacklist but is still inferior to LM-Switch. The results are added to Appendix R.

Citation to DExperts

Thanks for pointing out the errors in citation and we corrected them in the revision.

HMM Formulation

Thanks for this valuable question. This work aims to provide theoretical motivations and understanding of LM-Switch. We select HMM as a convenient analysis tool because of its universality in modeling sequential data and wide applications in various domains. The same theorem can also migrate to multiple other frameworks, such as the Observable Operator Model (OOM) or Linear Dynamic System (LDS), with only minor modifications. We also call for new theoretical frameworks that better model the linguistic data in future works. Secondly, Assumption 1 we made in the paper is extended from properties of independent and normalized random variables. We also discuss their validity in Appendix J.

Do We Freeze Other Parameters

Yes, we freeze all other parameters in the language model, as explained in the abstract and Section 3.3.

Task Selection

In this work, we train LM-Switch to fit a steering spectrum or dimension. To this end, we should specify two points to define the steering dimension better and contrast the differences between the two directions. In detoxification, this dimension is defined by toxic v.s. non-toxic sentences, and in sentiment control, the dimension is defined by positive and negative sentences. If we only use the train on positive texts, we observe that required to generate in the opposite direction, LM-Switch simply avoids words and topics as a shortcut solution, which yields inferior performance.

Finally, we hope that this addresses your concerns. If you find our response useful in answering your questions, we kindly request you to consider updating the reviews.

Behavior on language generation: comparison with prompting

The reviewer raises a valuable question regarding LM-Switch’s efficacy and generation quality under prompting. We let GPT-J generate 25 tokens under instructions (i.e., prompting), and compare its generation with and without LM-Switch. For prompts, we let ChatGPT write 10 similar instructions to the reviewer’s suggestion. We see that LM-Switch can still steer generation even under a positive prompt while maintaining generation fluency. These results are added to Appendix Q.

Example generations (one each for the first 4 prompts)

Why HMM

Thank you for this inspiring comment on our theoretical motivation. About the selection of HMMs as a framework, this work aims to provide theoretical motivations and understanding of LM-Switch. We select HMM as a convenient analysis tool because of its universality in modeling sequential data and wide applications in various domains. The same theorem can also migrate to multiple other frameworks, such as the Observable Operator Model (OOM) or Linear Dynamic System (LDS), with only minor modifications. We also call for new theoretical frameworks that better model the linguistic data in future works. Secondly, Assumption 1 we made in the paper is extended from properties of independent and normalized random variables. We also discuss their validity in Appendix J.

Finally, we hope that this addresses your concerns. If you find our response useful in answering your questions, we kindly request you to consider updating the reviews.

We are pleased to see the reviewer appreciate the proposed method as “simple to implement” and parameter efficient. We thank them for the detailed reviews. Regarding the comments:

Non-binary tasks

The reviewer raises an interesting concern. First, the hyperparameter $\epsilon$ was selected based on a validation set in Appendix E. Second, about the selection of tasks, this work tackles the steering of language models. The steering problem requires a spectrum or dimension to define the problem properly. Additionally, we demonstrate the non-binary tasks via compositionality of LM-Switch in Section 4.2. Therefore, different sub-dimensions can be composed by LM-Switch to achieve complex control (e.g., to “generate non-toxic and positive comments”).

Baselines for toxicity

The soft-blacklist method can be seen as a simplified version of LM-Switch. It also serves as a surprisingly simple but useful baseline yet unexplored by prior papers on the detoxification task. The fact that its performance is reasonable while still inferior to LM-Switch implies the powerfulness of this thread of work by exploiting word embeddings. As for the backbone choices, in the revision, we added the backbone model for a fairer pairwise comparison with our model.

No human evaluation

We thank the constructive concern raised by the reviewer. We conduct a human evaluation of the detoxification task and compare LM-Switch with LoRA, DExperts, and GPT-2 in a pairwise manner. Specifically, we follow the practice in DExperts and ask 4 human annotators to compare 50 generations from LM-Switch and the baseline from 3 perspectives: detoxification, fluency, and being topical to the prompt. The results are as follows. We can see that LM-Switch is ranked significantly less toxic and more topical than the baseline. It performs similarly to DExperts and GPT-2 but better than LoRA on fluency. These results are added to Section 4.1.

Parameter efficiency

As explained in Section 1, LM-Switch only uses $d \times d$ parameters. On larger models like GPT-J 6B, as they typically have 4k word embedding dimensions, the trainable parameters are only 16.7M, which is slightly smaller than a rank-32 LoRA on the same backbone model (18.4M). Its results with LM-Switch are listed below. These results are added to Appendix P.

We thank the reviewer for the inspiring comments and are pleased to see that the reviewer appreciates the methods as being “principled” and the comparison “comprehensive”.

Hyperparameter

As the method is lightweight, the only hyperparameter is the switch value. We use $\pm 5 \epsilon_0$ universally in evaluation, which is selected based on performance on the validation set on the detoxification task. For more details, we kindly refer the reviewer to Appendix E.

Statistical significance

We thank the constructive suggestions raised by the reviewer. We headed to conduct a statistical significance test on the detoxification task using 3 random seeds. Specifically, the standard deviation of the scores is shown in the following table, which we also updated in the revision:

Incorporating in other methods

Thanks to its theoretical foundations, LM-Switch is orthogonal to other methods and is intuitively compostable with other control methods. As an example study, we select LoRA to combine with LM-Switch on detoxification with GPT2-large as the backbone. The result is as follows, which is added to Appendix O. Combining LoRA with LM-Switch produces a better detoxification score than LoRA alone, at the cost of a degraded quality score.

We thank the reviewer for the detailed and constructive comments. We are glad that the reviewer finds the paper “well-organized” and the results “promising”. Regarding your comments:

GPT-2 Sizes

The selection of using different sizes of GPT2 is for a fairer comparison with baselines of different backbones. Following DExperts’ practice, we evaluate using backbones of GPT2-base to GPT-large so that pairwise comparison with baselines is possible. In the revision, we added the backbone size of baselines in Table 1, so that such a comparison is more straightforward.

LoRA details

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 Lee et al. (https://arxiv.org/pdf/2308.07317), we set the alpha scalar equal to the rank. We report LoRA’s performance below. We see that LoRA’s performance remains stable to the rank and still uniformly performs worse than our model. We add this discussion to the revision Appendix N.

Evaluation Tasks

In this work, we focus on the steering of language models. This requires a continuous spectrum or dimension for defining the task and achieving fine-grained steering like continuous and compositional control. To fit the task of steering better, LM-Switch shows the capability of compositional control in Section 4.2 so that sub-dimensions of complex control tasks can be composed to achieve multi-dimension steering (e.g., to “generate non-toxic and positive comments”).

MuCoLa

Thanks for pointing out this concern. The original MuCoLa paper was not evaluated in the same setting as ours, and the MoCuLa model has a prolonged decoding process (taking ~ 3000 hours or 5 months on an A100 GPU to finish the sentiment task). Therefore, this model is hard to compare fairly with other models with academic resources. Their own paper only evaluated a subset of the sentiment dataset.

Compositionality

For compositionality, we add the control matrices from both the toxicity and sentiment tasks. In Figure 2b, when varying the sentiment control scalar (along the long axis), we see a color change indicating prone-negative (blue) switching to prone-positive (red). Independently, when toxicity control increases (the short axis), the height indicates the toxicity level from high to low.

HMM

The reviewer raised a valuable question for discussion. About the selection of HMMs as a framework, this work aims to provide theoretical motivations and understanding of LM-Switch. We select HMM as a convenient analysis tool because of its universality in modeling sequential data and wide applications in various domains. The same theorem can also migrate to multiple other frameworks, such as the Observable Operator Model (OOM) or Linear Dynamic System (LDS), with only minor modifications. We also call for new theoretical frameworks that better model the linguistic data in future works. Secondly, Assumption 1 we made in the paper is extended from properties of independent and normalized random variables. We also discuss their validity in Appendix J. We agree that the length of such discussions can be reduced to better balance theoretical and experimental components. We made this change in the revision.

Meaning of $h$

Thanks for pointing out a typo where $h$ should written as $s$. We corrected this in the revision.

We hope this addresses your concerns and that you could kindly consider updating the review scores.