Paper: Taming AI Bots: Controllability of Neural States in Large Language Models (ICLR-2024)

Reviewer comment (Theory):
The biggest weakness of this work is the strong assumption in Postulate 1 about the bijective/invertible nature of the map $F_w$. This assumption almost gives away the controllability conclusion: if there is an invertible map, then it at least intuitively follows very clearly how such a system would be controllable.

Author response:
Indeed, we acknowledge that Postulate 1 is a limitation of this paper. However, the assumption are actually not that strong: that current models are well-trained according to our definition can be tested by measuring  perplexity, and qualitatively one would be hard-pressed to argue that at least some current LLMs generate sentences that are statistically distinguishable from human-generated training data, which presumably was written to convey some meaning. We could have made the claim seemingly rigorous by enumeration arguments in the style of old theory papers, so the conclusion would be formally proven. However, we are unsatisfied with this kind of argument which is why we state the proposition as a postulate. By doing so, we  we hope to open a discussion, rather than close it, on this topic that we concur is an important one.

Reviewer comment (Presentation):
While the primary contribution of this work is theoretical, the authors do conduct experiments and present results in Appendix A. Even if keeping results in appendices, the authors should at least provide a sentence or two in the main paper highlighting high-level trends of the experiment results.

Author response:
Thank you, will do.

Reviewer comment (Theory):
In Section 3, the authors describe how, "For a sufficiently high temperature parameter $T$..." random sequences can be generated from any start token. The reviewer questions whether this is true for any $T > 0$ (barring numerical precision errors from very small numbers) and does not see why one would need a larger value of $T$ for this theoretical claim, while recognizing that in practice higher $T$ would "spread out" faster to cover more random sequences.

Author response:
Yes, we could just say $T>0$. Thanks for pointing that out.

Reviewer comment (Theory):
The reviewer is concerned about Postulate 1 and asks what evidence or arguments the authors have supporting this assumption.

Author response:
The fact that there exist LLMs that are well-trained according to our definition should not be very controversial, as one could be convinced by simply interacting with some of the best models, or measuring their perplexity on any benchmark validation set. There are also quantitative empirical studies establishing that the best models generate sentences that are indistinguishable from human-generated, with high-frequency/probability. As for the rest of the claim, that follows from the complexity of the trained map coupled with the finiteness of the context. Also empirical evidence is the practice of jailbreaking and adversarial prompting, as well as the massive effort that goes into preventing such in the area of Responsible AI.  However, we recognize that even taking all this into account still does not make a proof.  In our view, though, the framing of the problem in the language of controls still would be beneficial to the community.

Reviewer comment (Theory):
Definition of equivalence seems to be insufficient and lacks important components.

Author response:
Any model trained on natural sentences determines a partition of the space of input sentences, including gibberish, indeed. This is why we restrict the model to operate in the quotient space (10), so input gibberish is excluded from the analysis. In theory, one can of course feed gibberish to a model, but that does not have any meaning in our definition since it is outside $\sigma({\cal I})$. In practice, one can easily exclude gibberish sequences by filtering the input, as done to avoid jailbreaking.

Reviewer comment (Theory):
Definition of meaning seems to defeat the whole purpose of this paper, as it allows for any gibberish/random sequence of tokens to still induce a meaning and possibly a set of many other gibberish sequences to be in their equivalent class.

Author response:
Any model trained on natural sentences determines a partition of the space of input sentences, including gibberish, indeed. This is why we restrict the model to operate in the quotient space (10), so input gibberish is excluded from the analysis. In theory, one can of course feed gibberish to a model, but that does not have any meaning in our definition since it is outside $\sigma({\cal I})$. In practice, one can easily exclude gibberish sequences by filtering the input, as done to avoid jailbreaking.

Reviewer comment (Novelty):
How to find a discriminant for meaning is left out as the authors explicitly assume that “the mechanism is provided by human annotators and other providers of training data.” While the authors emphasize in the introduction that such information can be used in the LLM training without external reward model: “This observation shows that sentence-level annotations can be incorporated directly into the trained model without the need for any external reward model nor external policy model, simply by sentence-level feedback,” I do not see the advantage of this approach over using the very same data to train a reward model and use that either during the training (as in RLHF) or as an augmentation (as in Rectification method), the latter indeed provides quite strong theoretical guarantees. If this secondary goal is valid, the paper requires sufficient reasoning for why the presented approach is superior. If not, the presentation requires to change and reflect only the controllability analysis.

Author response:
We do not claim that forgoing an external reward model is superior, just that it is sufficient. Indeed the reviewer is correct that, if one can afford it, direct augmentation or RLHF with an externally-trained reward model (which then becomes the origin of meanings) is also possible, and possibly better. All we are saying is that an external reward model is not necessary, because one can use the model itself -- as has been established in various alternatives to RLHF, such as Direct Preference Optimization (DPO). We were not aware of the "rectification method", as discussed below, since that appeared after our paper was completed, but we will include it as an alternative way of endowing the trained model with meanings.

Reviewer comment (Novelty):
Last paragraph of introduction: “we hope to encourage the design of actual new control mechanism” --> It looks like the authors are not aware of the [*Rectification* method](https://arxiv.org/abs/2302.14003), which is indeed a formal control paradigm for avoidance any definable “meaning,” as in the terminology of this paper. See the point below.

Author response:
Thanks for the useful reference and the thorough description of that work, which we were not aware of since it appeared on ArXiv after our work was completed. We will add discussion as it certainly sounds pertinent, thank you for the pointer and the explanation.

Reviewer comment (Novelty):
Let me first emphasize that I am aware of the differences between these two works (no need to tell me they say X while we say Y). My main point is to help you broaden your research and make the current paper be more comprehensive and useful for the readers, also this is aligned with your last paragraph of the introduction: The concept of probabilistic avoidance, which is akin to *stochastic* controllability but in preventing a meaning from being reached is [studied before](https://arxiv.org/abs/2302.14003). Even though the authors of that paper did not call it controllability (they call it security condition which also comes [from an older paper](http://proceedings.mlr.press/v97/fatemi19a/fatemi19a.pdf)), but it is indeed stochastic controllability in the context of avoidance. To use the terminology of this submission, what that work shows mathematically is that the LLM’s output probabilities can be minimally corrected and that will guarantee the avoidance of any prescribed meaning of interest (the one under their study was toxicity, but essentially the same analysis is applicable to any other meaning). Importantly, such results show that a certain optimal value function corresponding to reaching the prescribed meaning (with reward of -1 and no discount) is key at least in principle for the avoidance, which also highlights the sufficiency of RL as the core learning paradigm for closed-loop controlling/steering an LLM.

Author response:
Thanks for the useful reference and the thorough description of that work, which we were not aware of since it appeared on ArXiv after our work was completed. We will add discussion as it certainly sounds pertinent, thank you for the pointer and the explanation.

Reviewer comment (Theory):
Section Preliminaries, first part --> An LLM in its standard form provides a simplex over token space, not an embedding + a metric. How can an LLM be a discriminant?

Author response:
An LLM is trained so that the vector of logits, or the corresponding soft-max, approximates the (log-)posterior of the next token, which is the optimal (Bayesian) discriminant. So, if you view the LLM as a map from an input sequence of discrete tokens to the softmax vector, it is indeed a discriminant, and since it is trained with the (exponential of the) inner product of the embedding with the one-hot encoding of the next token, the inner product (metric) is built into the embedding space by design of the training process.

Reviewer comment (Writing):
Section Preliminaries, second part --> The definition of equivalence is quite confusing. It looks like that equivalence requires both discriminant and a set of classes, no matter how they are defined (does it?). However, the presented definition only uses discrimanent, which seems incomplete. In that case, $x^1 \stackrel{\phi}{\sim} x^2$ is ambiguous.

Author response:
The definition requires either a discriminant or a set of classes, since one determines the other. If given a discriminant, one can define partitions using simple order relations or rules (e.g. thresholds), and given a set of classes one can define a discriminant as a (piecewise constant) function.

Reviewer comment (Theory):
Definition 1 --> [related to the point above] without a set of classes (meanings), the definition is incomplete. Indeed, $\phi$ only maps a given sentence to a metric space, that is a vector $\in \mathbb{R}^K$ (plus a given metric), nothing more. One needs to have (i) a set of classes in addition to (ii) a criteria (like argmax in your example) to define equivalence.

Author response:
That is technically correct, and we will clarify in the narrative: We leave (ii) open as a design choice, since typically fine-tuning is done with a variety of ranking (with argmax), or alignment (penalizing the inner product directly) or scoring.  In all cases, the decision rule is based on the embedding and the metric: During pre-training, the classes are represented by the one-hot encoding of the next token, and the inner product with the discriminant (logit vector) is explicit in the empirical cross-entropy loss. During fine-tuning, either through ranking or alignment, the same inner product is used either to rank options (usually expressed in context), to align with the soft-max of the corresponding tokens (1, 2, 3, etc.), or the embedding of natural sentences and generated sentences are compared, again with the same inner product. In all cases, a vector plus an inner product (metric) is sufficient to determine a classifier (discriminant vector).

Reviewer comment (Theory):
By extension, $[x]$ does not provide a comprehensible set by just using $\phi$. Intuitively, one can ask, given another vector $y$, then $<\phi(x), \phi(y)>$ must be what in order for $y$ to be part of $[x]$?

Author response:
As a simple example, consider the meaning "toxic". Although the same sentences can be deemed toxic by one annotator and tame by another, in the end for any chosen threshold the classifier will partition the space of input sentences into two disjoint regions. Then, for a given $x$, the equivalence class $[x]$ is any sentence $y$ whose discriminant falls in the same partition as $x$, which can be determined by the inner product $\langle \phi(x), \phi(y) \rangle$, or the corresponding softmax; since in this binary classification, $P_w(y|x) > 0.5$ if $y \in [x]$, and $P_w(y|x) < 0.5$ otherwise. (If it happens to be exactly $0.5$ up to the numerical resolution, then it is ambiguous). One can do the same with respect to an arbitrary number of classes, each represented by any element (prototype) in the class. Then, belonging to one's class is equivalent to a nearest neighbor rule in the learned feature space $\phi_w(X)$ for $X$ the set of all input sequences.

Reviewer comment (Theory):
Even if the above issue is resolved, the provided definition of meaning lacks any connection to linguistic meanings, as any random sequence of tokens will induce its own $[x]$. This conflicts with the authors initial claim that they would like to distinguish between a meaningful sentence and a random sequence (first line of page 2). Importantly, this is totally separate from the question of what a good $\phi$ is, and it hold from *any* given $\phi$.

Author response:
The apparent conflict is vacated by the fact that random sequences are not in $\sigma({\cal I})$. There are connections with distributional theories in linguistics, since one can convert sets and partitions to conditional probabilities and vice-versa, as others have shown. Our definition in terms of equivalence classes is also implicit in semantic paraphrasing in linguistics.  Since there is no universal notion of meaning that we know of, data collection yields a distribution of expressions, which is captured by a trained model, which can then use it to partition the input space. Different data, and different models (and different individual humans) all maintain different meaning representations. We therefore limit our claims to the system of meanings of a specific trained model.

Reviewer comment (Theory):
As for selecting $\phi$, the authors say “Here, we assume that the mechanism is provided by human annotators and other providers of training data.” How is this different from an external reward model, as the authors named as a disadvantage of RL-based methods? (Note: the reward is defined for a complete text not a given point at the middle, so these seem to be exactly the same.)

Author response:
If an external reward model exists, one can use it to induce a meaning structure, in perfect harmony with our framework. The disadvantage is having to rely on an external model that defines meaning, which then prompts the question of how that model acquires meaning in the first place. Since that model is also trained with human utterances and/or human supervision, then even if using an external reward model, the mechanism by which meanings are ultimately transferred to the trained model is still the same that we describe, originating with human annotators and content providers, then transferred to the trained model during the learning process (directly in SFT, indirectly through the external reward model in RLHF). We will clarify this further and add reference to the possibility of using an external reward model if one so desires.

Reviewer comment (Theory):
Eq 1 --> This is confusing. The range of $\phi$ is $\mathbb{R}^K$, where $K$ was the dimension of matric space, yet the output of $\phi$ is a probability distribution (which is a simplex).

Author response:
$K$ is the number of tokens in the dictionary (so, the discriminant has $K$ components), and $M$ is the dimension of the vectorized embedding of the tokens. The metric space is where the logits live, which is a continuous space of dimension equal to the number of hypotheses, or tokens, which is $K$. This should be clear in equations (4), (5), and (6) which are devoted to stress the point that the same entities (elements of the dictionary) are represented differently at different stages of the LLM.

Reviewer comment (Theory):
The characterization of meanings seems quite subjective. The authors mentioned that their characterization of meanings is compatible with the deflationary theories in epistemology. It seems that this explanation itself does not justify why such a characterization is meaningful for studying the controllability of LLMs in general.

Author response:
Indeed, our definition of meaning is specific to a discriminant function $\phi(\cdot)$, which in the case of an LLM is implemented by the trained model $\phi_w(\cdot)$. Different models partition the data differently, although models trained on the same or similar data will likely yield similar partitions. We do not believe in a universal notion of meaning, and in any case our scope is limited to how meanings are represented by LLMs.

It is not essential that our definition aligns with any particular epistemological theory, but we note that, once mapped to LLMs,  some theories boil down to representing meanings as equivalence classes of expressions, and then differ on their genesis. In our case, such classes are determined by the state of the trained model, which is what we want to control.

Reviewer comment (Theory):
The authors claim that their conditions are largely met by today’s LLMs. More justifications are needed. This also seems a hand-waving statement. "Largely met" means "not always met"?

Author response:
Met by some, but not all, trained LLMs. LLMs by design represent a joint distribution on a sequence of tokens, but how well that approximates the empirical distribution of natural sentences is a matter of approximation rather than constraint satisfaction. In this context, by "largely" we mean that perplexity of trained models is low enough that generated sentences are indistinguishable from human generated ones. This may not be true of all models, so we use "largely" to indicate that the conditions are valid only for some models.

Reviewer comment (Theory):
Does Definition 1 assume some sort of underlying classifier such that the equivalent classes can be defined? What is that specific classifier? The authors mentioned "the mechanism is provided by human annotators and other providers of training data." This is very confusing. I don't understand what exactly this phi is.

Author response:
In the definition, $\phi$ can be any discriminant vector. In a trained model, $\phi = \phi_w$ is the soft-max vector after supervised fine-tuning (and optionally RLHF). Since that vector is a function of the training data, both in pre-training and fine-tuning, and the training data is provided by human annotators and authors, ultimately the source of meanings is latent in the expressions, ranking, or scoring humans have provided for training.

Reviewer comment (Theory):
Equation (6) is very confusing. I mean, in Equation (4), y is sampled from the softmax operation. Then all of a sudden, it becomes an additive noise? It seems that the noise n_t depends on x_{1:t}?

Author response:
Indeed the fact that the same token can be represented in different ways can be confusing at first:  Tokens can be represented as elements of a discrete dictionary after sampling, as in (4), or as discriminant vectors before sampling, as in (5). In the latter case, the effect of different selections in the (previous round of) sampling results in a perturbation of the embedding vector of the next token, which is represented in (6) as an additive perturbation since embedding vectors (logits) are  elements of a vector space. In general, it is true that the ``noise'' $n_t$ is dependent on the previous history, $x_{1:t-1}$, but in complex ways that can be modeled statistically.

Reviewer comment (Theory):
Regarding Definition 3, an LLM is well-trained if theta is any small positive number?

Author response:
Yes, although different users, or use cases, may have different standards on what constitutes a well-trained model, reflected in the perplexity thresholds that are considered acceptable.

Reviewer comment (Theory):
What does Postulate 1 really mean? How to justify this?

Author response:
It means that one can find a prompt that causes the model to utter any meaningful sentence with high probability. It could be stated as a theorem since the context is finite and one could trivially find the prompt by enumeration, but since today's models have tens of thousands of tokens as context, enumeration is not viable in practice. So, we state it as a postulate in hope that others will find better methods -- analytical or empirical -- to prove (or disprove) the claim. It is admittedly a weakness of our analysis; yet framing the problem as controllability of trajectories in a quotient space determined by the dynamical model itself is novel to the best of our knowledge, and may spark interest in communities beyond NLP.

Reviewer comment (Theory):
How can we justify the assumptions in Theorem 2?

Author response:
Using Postulate 1: if one accepts the enumeration argument, they can consider the conditions of Theorem 2 met.