Linearly Controlled Language Generation with Performative Guarantees

Authors: Emily Cheng, Marco Baroni, Carmen Amo Alonso
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: control theory, representation engineering, large language models

Abstract:
The increasing prevalence of Large Language Models (LMs) in critical applications highlights the need for controlled language generation strategies that are not only computationally efficient but that also enjoy performance guarantees. To achieve this, we use a common model of concept semantics as linearly represented in an LM’s latent space. In particular, we take the view that natural language generation traces a trajectory in this continuous semantic space, realized by the language model’s hidden activations. This view permits a control-theoretic treatment of text generation in latent space, in which we propose a lightweight, gradient-free intervention that dynamically steers trajectories away from regions corresponding to undesired meanings. Crucially, we show that this intervention, which we compute in closed form, is guaranteed (in probability) to steer the output into the allowed region. Finally, we demonstrate on a toxicity avoidance objective that the intervention steers language away from undesired content while maintaining text quality.

Decision: reject
Reviewer scores: 6, 5, 6, 3

Reviews and discussion:
<thread>
creation date:
1729700177020

thread_type:
official_review

<thread_object>
creator: reviewer_2dsm

<content>
summary:
This work proposes LiSeCo, an intervention method that applies control theory principles in the latent space of language models to dynamically steer the sequence generation to desired semantic regions, such as avoiding toxic or undesired content, while preserving text quality.

strengths:
1. The method is applicable in post-hoc adaptation, making it efficient and universally applicable to frozen language models.
2. The work offers a closed-form solution for optimal interventions in the latent space of language models and provides theoretical guarantees on the controllability of the method.
3. The method is designed to be lightweight and dynamic, eliminating the need for gradient-based methods, which reduces computational overhead during the generation process.

weaknesses:
**Theory:**
1. *Semantic Probe:* The probing classifier function $f_t$ maps the latent space to the decision space $\[ 0, 1 \]$, meaning it assigns a single scalar value to each hidden representation (1-dimension). However, line 157 and Figure 1 suggest that the latent space is mapped to a 2-dimensional vector, as indicated by  $W_t \in \mathbb{R}^2$. This creates a discrepancy in the mathematical formulation: either the definition of $f_t$ or the dimension of $W_t$ is incorrect. In general, I assume that there is no benefit in having a second dimension. The closed-form solution could be computed as a binary classification problem (using the sigmoid instead of the softmax function) as well (see 1. Question).
2. *Optimal Controller:* In essence, you assume that $f_t$ can differentiate between toxic and non-toxic inputs. So, your control setup changes toxic inputs in a way so that $f_t$ classifies them as non-toxic. First, as part of the optimization problem, 1c and 1d are not necessary as these are no constraints and rather trivial in the context of language models (see 2. Question). Second, your control setup puts a lot of weight on $f_t$, as it ultimately assumes that $f_t$ knows how to change the hidden representations so that the language model does not generate toxic sequences.

**Evaluation:**

3. *Semantic Probe:* The probing classifier was trained on the train split of the constraint dataset, and its performance was evaluated on the corresponding validation split. The evaluation is based on sequences from the constraint dataset that may not have been generated by the language model from which the hidden representations are derived. This raises concerns about the validity of the evaluation: for instance, the language model might assign lower probabilities to toxic than to non-toxic sequences, which could have a direct influence on the hidden representations. This underlying influence, which is not accounted for in the current evaluation, could be mitigated by evaluating the performance on sequences that the language model has actually generated, using your proposed automatic scorer for determining the ground truth toxicity. This would provide a more realistic assessment of the probing classifier’s ability to separate toxic from non-toxic sequences. Since your theoretical guarantees rely on this assumption that toxic and non-toxic sequences are separable by $f_t$, this should be investigated in more detail (see 3. Question).

4. *Experiments:* It is not immediately clear how the value of the free parameter $p$ should be chosen. For instance, $p = 0.01$ seems necessary to avoid generating toxic sequences with the “llama” and “pythia” models. However, for the “mistral” model, $p = 0.3$ is the best setting, while $p = 0.01$ significantly sacrifices naturalness (see 4. Question). Moreover, for the “pythia” model, sequences are still largely classified as negative, regardless of the value of $p$. Thus, while your method theoretically provides a guarantee, this guarantee holds only if the allowable region suggested by the probing classifier aligns with the region where the language model actually generates non-toxic/non-negative sequences. The empirical results suggest that this alignment may not be as strong as claimed. If the probing classifier truly had 90% accuracy, then only slightly over 1% of sequences should have scores above 0.5 when $p = 0.01$.

questions:
1. What is the benefit in having a 2-dimensional decision space?
2. What is the rational behind including 1c and 1d in the optimization problem?
3. How does the probing classifier perform on differentiating between toxic and non-toxic sequences that were actually generated by the language model?
4. Is there a clear guidance on choice of the free parameter $p$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729894790014

thread_type:
official_review

<thread_object>
creator: reviewer_akvp

<content>
summary:
The paper proposes to control the generation of LLMs using controls defined in their latent space. On top of these controls, the paper applies techniques borrowed from Control Theory to steer trajectories (latent representations) away from regions corresponding to undesired behaviors. 

Here is a summary of their approach: 

== Approach formulation: 

- **Given:** We are given **an already pre-trained** LLM. The goal is to modify the internal representation of the model so that output responses: (1) are **guaranteed** to never fall into a disallowed region. (2) are as close as possible to a desired region. 

- **The control mechanism:** The control mechanism alters the output vector embedding at **every layer** that guarantees (in probability) that latent trajectories remain out of the "unsafe" region. The control mechanisms here are additive, i.e., we design a controller parameter $\theta_t$ that gets added to its later embedding $x_t$, at every layer $t$. 

- **Quantifying unsafe region:** The approach also requires a "probing function" that maps latent vectors to their safety/unsafety(e.g., toxicity). In practice, building this classifier requires supervision. Specifically, any embedding that satisfies $\sigma (W^\top (x_t)) \geq p$ (I). 

- **Steering away from unsafe regions:** This additive control is designed so that it satisfies the following inequality   $\sigma (W^\top (\theta_t + x_t)) \leq p$ (II), for any $x_t$ that leads to "unsafe" generations characterized by Eq.I.  

- **Problem formulation** Now one can formally define the goal as finding a set of controls for all layers $t$: $\{ \theta_t \}_t$, by ensuring the condition in Eq.II for each layer.  

- **Simplifying the goal** Rather than solving it for **all the layers** (previous paragraph), the paper chooses a more modest goal of solving the goal for each layer. With this simplification, there is an analytical solution (Theorem 1) since the inequality of a linear separator and a nonlinearity (softmax).

strengths:
- The presented formulation seems novel to me. The math is beautiful. So kudos to the authors. 
 - The construction of using semantic probes for controlled generation is intuitive and interesting.
 - The closed-form solution renders the approach very efficient compared to other controlled generation approaches like FUDGE or PPLM which requires classifiers or gradients.
 - Overall writing is clean and clear. I was able to follow and understand them. 
(I also have extensive feedback on wherever the writing needs to improve; see them below)

weaknesses:
- I have extensive comments below where I elaborate on your writing. 

- I think your experimental results are your weakness. Part of this is writing (I have elaborated on this layer). But there are other aspects too that I am unclear about which may require more experiments or better arguments. 

- The work provides guarantee (in probability) that output activation will lie in the desired region. However, defining a proper "allowed region" in practice is challenging. In this work, defining allowed/disallowed regions is achieved by a linear probing classifier. Thus the performance guarantee lies on the classifier being accurate, which is not always the case. In particular, your "probing functions" are typically not very reliable on the edges of safe and unsafe generations. In other words, the guarantees that you give are on the model behavior, **assuming** that the probe is reliable (whether it is reliable or not, is not something that is quantified or guaranteed by your framework). For example, it may be difficult to probe whether early tokens in the generation would lead to unsafe responses or not because safety is usually defined on the sequence instead of token level. FUDGE takes care of this issue by training discriminators that predict whether a target attribute will become true in the future. However, it's unclear whether a simple linear prob can do this too. Moreover, probes are less effective in out-of-domain cases. It calls into question of the effectiveness of the proposed approach's OOD generalization capability, which is not sufficiently studied by the experiments.

- The framework assumes a binary allowed / disallowed region. This does not capture the need for more fine-grained control, such as the degree of toxicity or a spectrum of sentiments. This limitation is also reflected in the relatively toy experimental setup, where the labels are binary for the toxicity and negativity avoidance tasks.

-  In practice, you need to sacrifice your guarantees to maintain naturalness; so there is a clear limit on how much you can control the output safety. 

 -  While you're solving layer-wise guarantees, only the guarantee on the final layer is going to matter to the user. So one is left wondering whether this layer-wise control is necessary or one can just do it for the final layer. Or perhaps apply a stronger control (higher p) for middle layers, but lower it for the final layer to maintain naturalness? These clearly require more experimental investigation.

- Unclear whether the proposed LiSeCo approach can achieve additional gains on top of instruction tuning or conventional safety alignment. Because of the lightweight nature of the proposed approach, only moderate improvements on safety can still be an helpful addition on safety-aligned models.
 - There is limited analysis of experimental results.  (a) For example, LiSeCo (0.01) outperforms FUDGE only on the Mistral model. However, Mistral+LiSeCo (0.01) shows very low naturalness (Figure 2). Maybe the output has already degraded to low quality texts. Further explanations are needed. (b) It would be interesting to conduct analysis on what is the probe’s effect on the distribution of next tokens. Is there a clear reduction of the probability of toxicity-related tokens? (c) It would be helpful to discuss the trade-off between controllability and latency a bit more (or perhaps more clearly?). It seems that LiSeCo attains worse controllability than FUDGE in many cases but is much more efficient. I have more comments on this below.

questions:
Questions are listed in no particular order: 

 - In the "Semantic Probe" paragraph, did you mean to define $f_t=\sigma(W_t^\top x_t)$? This seems to be implied but not explicitly given. 

 - Are your control mechanisms defined for each input or for a collection of inputs? The paragraph in S3 is not explicit, but since it defines R_t for any x ($Rt = \{ x | \sigma(W^T x) > p\}$), it makes it seem like we define the control for **any** input. However, in S4, it seems like we're fixing the input. For instance, "prompt sequence" (Eq1d) and also specific choices of x_t in Theorem 1. Given the statement of your theorem, I think the latter is correct (a control per prompt), but please make it explicit upfront in S3.  

- Proof of Theorem, you open by stating KKT conditions. Here you need to first define what lambda is. I suggest you reorder so that the definition of Lagrangian appears first, then define lambda, and then show the KKT condition (and say why it's important to check these for those who may not be familiar with the constrained optimization via Lagrangians)

- In Eq C.8, how did we obtain $w_1$ and $w_2$ from $W$? And from there, how do we define $w$? Seems like you're doing some reparameterizations without clearly stating it.

- [Potentially related to the previous point] The term "the difference of the columns of $W =: [W_1, W_2]$" (Theorem 1) is interesting but also a bit surprising. Intuitively why does it emerge in your solution? Also, am I interpreting  $W =: [W_1, W_2]$ correctly when we split the matrix into two vertical matrices? (i.e., we don't imply that $W$ has only two columns). 

- In your optimization formulation in Eq1, you're missing additional term(s) that specify the domain of $x_t$ and their membership to the unsafe region: $\sigma(W^\top x_t) > p$ (not just any point) similar to the condition in Eq. 3a. 

- While you're solving layer-wise guarantees, only the guarantee on the final layer is going to matter to the user. I think you want to articulate it clearly. Perhaps the only rationale for intermediate layer adaptation is to simplify the achievement of safety in the final layer, which is a hypothesis that needs to be verified. 

> We consider instruction-tuning, which relies on extensive LM finetuning, to be a target, or “upperbound", baseline. 

Why are these considered as an "upperbounds"? If you're using the same set of supervision used for your reward models, the direct comparison to these supervised/tuned models is a fair game. 

- Have you considered trying other additive tuning approaches, say, LoRA or BitFit? They're other additive adaptations that I'm forgetting here. 

> FUDGE has the potential to incur a high inference latency compared to other methods
>Of the methods tested, FUDGE has the highest latency by several orders of magnitude, around 3 seconds compared to other baseline 

But these k sequences can be parallelized since they're independent of each other; right? 

> This comes at the benefit of directly optimizing for the evaluator. For this reason, we consider FUDGE to be an upper-bound baseline, given that other steering methods do not have access to the ground-truth scorer.

Your approach is also using the probe's feedback; so not sure if I buy this. If I am missing something please elaborate.

> The fact that probes attain high accuracies of ∼90% confirms that the disallowed toxic (negative) regions Rt are linearly decodable with high probability

This is misleading. Given that you apply your intervention per layer, one needs to look at the layer-wise accuracy. And we see that 85% holds typically for layer ten or higher. Instead of Table 1, I'd actually include your figures in the appendix which, IMO, are easier to parse. 

- One odd thing about Fig2 is that, upon applying intervention, "naturalness" (fluency) gets higher than "no-control" (the green box); for example, top-right subfigure for llama. Given that the base model (no intervention) is optimized for fluency, this is clearly signals noise in your evaluation. 

- Consider categorizing your baselines into two: (1) tuning-base (training, LoRA, etc.); (2) inference-time steering (FUDGE, yours, etc.). And right off the bat, you want to say that category (1) is not easily extensible. The discussion of category (2) is tricky. Currently, your arguments rely on the slow speed of FUDGE, but that’s a weak argument (FUDGE can be implemented faster, I think, as I mentioned above). Plus, there are (many?) other adaptation-based that you do not compare with. I think your strengths here are: (1) your proposal is principled and allows for guarantees (even though, in practice, you need to sacrifice this guarantee; otherwise, you miss out on naturalness, which defeats the whole purpose); (2) it requires producing distribution over the words in the vocabulary. Yours is just an embedding projection.

> modifying LM weights and can be extremely energy-intensive, **needing > 4 orders of magnitude more data**

I don’t think you show any result on this (comparison with tuning LM, as a function of its supervision size). 

- It’s worth viewing your EQ1 in terms of (and comparing against) other existing known frameworks, such as SVM https://en.wikipedia.org/wiki/Support_vector_machine  

- Lastly, I think the “control-theoretic” framing of the paper might be a bit of a stretch. Ultimately, Eq1 is the most general form of your problem and Eq2 is what you actually solve. If KKT condition and interpretation as the Bellman equation constitute a “control paper,” all the alignment papers that also use some variant of RL or sequence prediction (and are most of AI now) are also “control papers.” I acknowledge that it’s a bit murky what constitutes a “control” paper; so I won’t hold this against you. 

=== Minor

- Fig2: in the caption, I’d say what 0.01, 0.1, … in the figure are (I think they’re $p$). I’d even revise the legend to say `LiSeCo(p= … )` for absolute clarity

- I don't see where you define $\mathcal{R}_t^C$ (just under the Theorem). Previously, we only defined $\mathcal{R}_t$. 

 - Expressions such as $p(unsafe) < p$ (under Theorem 1) are confusing. I suggest you choose a different notation for probability $P(.)$ so that it does not get confused with your design parameter $p$. 

 - In various places there is a subscript $2$ attached to $\sigma$ (the softmax function). Unsure what that means. I don't think it's a standard notation. 

- Figure 2 and 3 are not colorblind-friendly and colors are Figure 2 color is hard to distinguish.

> Activation Addition (ActAdd) Turner et al. (2023)

Use `\citep{}`

= Missing citations: 

There is a ton of literature on controlled generation (https://arxiv.org/abs/2202.11705, https://arxiv.org/abs/2205.14217, https://arxiv.org/abs/2112.08726, https://openreview.net/forum?id=kTy7bbm-4I4) that try to solve a similar problem. I have included a few below, but there is a ton that I am missing.  
As far as I know, these works do not provide guarantees that you do, but I don’t think it’s helpful for our collective progress to ignore them.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730675444222

thread_type:
official_review

<thread_object>
creator: reviewer_jnbs

<content>
summary:
This paper presents a novel approach, Linear Semantic Control (LiSeCo), for controlled language generation within Large Language Models (LLMs). The method employs a control-theoretic approach, optimizing latent space interventions to guide LLM outputs away from undesired regions (which are found by probing), particularly demonstrated for toxicity avoidance. Unlike other approaches that rely on complex or computationally heavy adjustments, LiSeCo operates with low latency through a closed-form solution and provides probabilistic performance guarantees.

strengths:
1. This paper is well written, easy to understand and the proposed method is not complicated, it just intervenes on the hidden representations. 

2. The proposed method is empirically effective without inducing much computational overhead.

3. The empirical results are backed up with theoretical guarantees.

weaknesses:
1. The paper mentions that small modifications in latent space can have unpredictable outcomes (similar to the "butterfly effect"), yet there are no experiments analyzing such unintended consequences in practice. A detailed analysis on potential adverse effects of these interventions, especially when they involve semantically sensitive regions of latent space, would lend weight to the discussion and address the potential risks in practical applications. It would be crucial to see that the proposed method maintains / minimally decreases the quality of the text. 

2. The applicability - there are certain attributes that are probably more easily probed (e.g. in this work toxicity - since this is associated with some of the toxic words). But there are also attributes that are not (e.g. tone, formality…). Extending this control to more nuanced tasks would really show the effectiveness / limitation of this approach.

If the author can resolve my concerns here, I would be happy to raise my score.

questions:
035 - 040: There are also decoding time algorithms for controllable generations, which are omitted in this paragraph.

042:  For instance, while knowledge editing provides an efficient alternative to exhaustive retraining, it poses risks akin to the butterfly effect: minor adjustments could lead to unintended consequences - this is not backed up with citations.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730883918776

thread_type:
official_review

<thread_object>
creator: reviewer_ucxh

<content>
summary:
The paper proposed an interesting, gradient-free method to control the language generation. The authors introduce a method called LiSeCo, where they first train layer-specific probers on data within the "unsafe" region and then apply these probers at inference using a closed-form solution.

strengths:
1. The paper is well organized and clearly presented.
2. The author used an interesting locally-optimized close-form solution to significantly speed up the inference process.
3. The performance of semantic control and quality maintenance looks good.

weaknesses:
1.  The tasks selected by the authors appear somewhat too old and easy to me. It would be beneficial to see the method applied to more recent and challenging tasks, particularly those related to safety. For example, datasets like UltraSafety (https://huggingface.co/datasets/openbmb/UltraSafety) or HarmfulQA (https://huggingface.co/datasets/declare-lab/HarmfulQA) could provide a better evaluation of the method's effectiveness in handling harmful or unsafe content. Tasks like IMDB reviews are too straightforward for current large language models and may not adequately showcase the capabilities of the proposed approach.

2. The evaluation could be significantly improved. The authors rely on automatic evaluations using relatively small models (such as RoBERTa-based models) and self-assessments. This approach may not offer a comprehensive understanding of the method's performance. A common practice is to employ more advanced models like GPT-4 for evaluation, especially since the models discussed in the paper are around 7B parameters. Utilizing GPT-4 could provide a more robust and thorough assessment.

3. The colors in Figure 2 is hard to read but it is ok.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1732786235272

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Experimental and theory update

comment:
Our current work experiments on contrained generation for two concepts: _toxicity and negativity._

We extended our approach to account for features that can be continuously tuned. Our theory now tackles cases where, instead of eliminating a specific feature on a text, it keeps it within a range, or controls the degree to which a specific property is present/missing. This enables more controllability on the generation. We refer to this extension as "continuous tuning" in the manuscript. We provide the theoretical results and mathematical proofs that ensure guarantee satisfaction in Appendix E. 

We also provide preliminary experimental results that demonstrate the effectiveness of this extension through a formality use-case, where the degree of formality is set implicitly by LiSeCo $p$. Hence, we find this to be a good use-case to test the validity of our method at controlling the degree of a specific feature (in this case formality), as opposed to eliminating it completely. Our experimental setup and results are reported in Appendix F. We find that LiSeCo, better than other baselines, is able to vary the formality without sacrificing text quality.

Please let us know if anything remains unresolved about the theoretical and experimental extension to our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a method for controllable text generation. The approach first trains a linear probe to classify attributes from hidden states, and then derives a closed-form solution for perturbing those states to ensure that the predicted attribute probabilities are within a certain thresholds. Experiments on toxicity and sentiment control demonstrate this method.

Strengths:
1. The method has a closed-form solution and only requires training the probes, and is therefore efficient.

Weaknesses:
1. Reviewers raised a concern that this work only conducted experiments on toxicity and sentiment control and is not shown on more applications such as safety in language models.

Overall, the idea proposed in this paper is very interesting and the proposed solution is simple and efficient. However, reviewers were concerned about whether the proposed method works in broader applications such as safety in language models. Therefore, I'm recommending rejection for the current paper, but I wouldn't mind if the paper gets accepted.