Summary
The paper introduces an alternative procedure for LLM alignment that does not fine-tune LLM weights, but instead learns a separate value function that is used to update hidden states. The value function is learned using a variation of temporal difference, then applied at inference time to modify hidden states by gradient ascent, maximizing the predicted state value. Authors evaluate their approach with multiple 7B LLMs on HH-RLHF data, comparing against both RLHF and training-free baselines. The paper also analyzes OOD generalization to HarmfulQA.
Strengths
- Authors propose an interesting approach to that can be used to alter LLM behavior in general
- When experimenting with HH-RLHF dataset, authors evaluate against multiple types of baselines and provide additional analysis that was interesting to read
- The paper is generally well-written and easy to follow
- Authors made the code available, in a (mostly) serviceable state
Weaknesses
**1a. Motivation for the choice of baselines.**
In your work, you cite, among others, ARGS[26], DeAL [22], Value Augmented Sampling [21] that also learn value functions and use them to steer model outputs (in other ways), but, to the best of my knowledge, you do not compare against them as baselines, instead choosing a relatively older work on controlled decoding. While [21] may be dismissed as concurrent work, the other works appear to be a relevant alternative and it is not clear why they were not chosen as baselines.
If there is a reason why these works will, beyond reasonable doubt, fail at the task that you evaluate on, I would recommend that you explain this in the paper. If there is no such reason, the paper would benefit from comparing against them.
**1b. Motivation for the choice of models**
Your paper focuses on Llama, Vicuna and Falcon models, of the 7B variety. While these are indeed LLMs, the original Llama was released circa 1.5 years ago and since then, LLMs improved **significantly** across tasks.
Picking older LLMs appears counterintuitive, as their generally worse quality makes it harder to measure possible drawdowns introduced by LLM alignment.
If you have a reason for choosing these models, please explain why you focus on older LLMs as compared to, for example, Llama 3 8B (or 70B), Qwen2, Gemma or other models near the top of https://huggingface.co/spaces/open-llm-leaderboard/open_llm_leaderboard . If there is no such reason, the paper would benefit from switching to more accurate models.
**2. Inference time exploration**
LLM use cases are often sensitive to inference throughput (tokens per second) and latency (time to first / next token).
To the best of my understanding, RE-Control introduces an iterative optimization step to each forward pass during autoregressive inference. Depending on the configuration, this may result in a significant slowdown, which may limit the practical significance of your approach.
I would argue that the work would benefit from analyzing this difference in speed in different settings (e.g. single-sequence vs batch inference, etc).
**3. Main experiments are limited to one dataset and relatively small past generation LLMs, ranked by GPT-4**
This is definitely not a fault on authors' side, but the paper makes its main conclusions based on 7B models, using reward functions trained on a single dataset. This could result in accidental false conclusions if it turns out that, for instance, RE-Control harms the quality of stronger models or if it is somehow implicitly overfitting on on GPT4 opinions.
The standard way to minimize this risk is to diversify the experiments: try alternative alignment datasets (e.g. webgpt_comparisons, oasst1, etc), try larger models (llama-3 70B), introduce human rankings in some setups, etc. I understand that not all of these evaluations may be available to the authors, but for a NeurIPS publication, I would expect more variation in the experiments and, if there is a confounder that could not be eliminated (e.g. using GPT4 and no human eval), it should be stated among the limitations section.
Questions
**Questions on the definition of state**
To the best of my (possibly wrong) understanding, when you apply Bellman equation, you assume that the dynamic system's state satisfies Markov assumption. [If not, please explain why not]
Since LLMs use attention to previous hidden states, hidden vector for a specific state do not satisfy Markov assumption, since LLM's next token probability depends not only on them, but on a more distant past as well. In contrast, a fully markovian state would need to contain all previous hidden vectors, or the current hidden vectors and all past KV projections, or a sequence of all previous tokens (no hidden vectors necessary).
In other words, **when you define V(s), does s refer to just the current token's hiddens or a full state with Markov assumption?**
If you mean the latter state, then the test-time intervention (S4.4) needs to modify all previous hidden states of an LLM. This is important because modifying past hidden states may result in a very inefficient LLM inference algorithm.
If only the current state, you seem to apply policy iteration (S4.2-4.3) to a non-markov state. Please explain how you make sure that this algorithm still has the guarantees of optimal policy. If it doesn't, please clearly explain that the algorithm is a heuristic inspired by PI rather than actual PI.
### On reproducibility
To reiterate, the fact that you publish the code is great. None of my complaints below affected the final score.
The codebase lacks library versions (requirements.txt / dockerfile / list them in the readme), which makes it difficult to reproduce, especially in the future. While I ultimately managed to run the code by choosing the libraries with an educated guess (and minor modifications to the code), I am still not sure if I got the method to work "as intended" and not introduce silent errors.
For legal reasons, it would be best to direct the users to a version of Llama 7B that contains its original license, at least in the final version of the paper.
Using GPT-4 opinion means that the experiments would be difficult to reproduce after it is cycled ou
### Typos / minor:
> L16 LLama
The capitalization for the first version was LLaMA, second and third are Llama.
> supplementary code: intervented_model
you may have meant “intervened”
Limitations
The "Limitations and future work" appendix can be significantly improved. Currently, it focuses on future work and omits some limitations of the experiments, such as:
- using GPT-4 as the primary metric will make the results irreproducible once OpenAI cycles out GPT4, a closed-source model
- evaluating only on relatively weaker models (pre-previous gen, 7B) may miss some caveats or synergies from more capable LLMs
- using a single training dataset makes it possible that the proposed method is uniquely powerful in this one scenario but not others
The quality of the limitation section did not affect my score.