ReFT: Representation Finetuning for Language Models

Thanks so much for raising these great questions and providing helpful feedback!

RepE baseline.

We agree that gradient-free methods such as activation addition or RepE could be effective in steering models for tasks such as style transfer [4]. On the other hand, we argue that it could be hard for these methods to steer models to achieve general finetuning goals (e.g., steer a model to become a binary sentiment classifier or steer a model to answer multiple choice questions). These methods usually rely on statistical methods (e.g., PCA) to propose a set of static steering vectors based on a few-shot training dataset (i.e., usually with a handful of training examples). Mathematically, these methods have the same expressivity as BitFit [5] (only learning the bias term of LMs) yet BitFit is learned with SGD on large datasets. Meanwhile BitFit, although with much lower trainable parameter count (comparable to ReFT), usually underperforms compared to other PEFTs (see Table 4 on pg. 9). However, we agree that learning-free methods such as RepE have other unique advantages:

• it could be very effective when the steering objective is more generic (e.g., style transfer for a chat-model) without heavy shifts in model’s behavior;

• it could be very useful for LM personalization applications with very limited training source and do not require high precision.

Need to include the baseline generations.

Thanks for raising this issue. We agree. For all examples in Appendix I, we will add in the original mode outputs. We include one example from GSM8K in our attached rebuttal pdf.

The best expression of ReFT.

Indeed, there is no clear winner for math reasoning benchmarks. Here are some additional pointers in terms of choosing different intervention functions:

• DiReFT removes the orthogonality constraint which improves the memory footprint and trades compute efficiency (i.e., DiReFT trains faster) for a slight drop in performance.

• LoReFT generally converges quicker and is less sensitive to learning rate and other hyperparameters partly due to its orthogonality constraint.

• The orthogonality constraint offers composability and interpretability. Given the constraint, subspaces are orthogonal to each other (i.e., changing in one subspace should not affect others). As a result, we think this gives a nice property for us to compose LoReFT together at inference-time. We explore this a bit in Appendix G.1..

Hyperparameter selection with ReFT.

We address this question in our general responses by providing additional details on the hyperparameter searching process of ReFT!

PEFT vs ReFT vs other inference-time interventions.

Thanks for the question! It would be best for users to benchmark these methods for a specific domain. If we were allowed to guess, here are some pointers:

• Inference-time intervention or activation addition does not require training. If the use case is not mission critical, one could use these methods for quick turnaround and showcase.

• In general, since ReFT allows gradient descent, it should be more effective than non-training methods since it actively searches causal pathways to steer the model. ReFT also works with quick adaptation (e.g., n-shot training where n <= 10) as shown in G.2..

• The composability of ReFT could be better (e.g., combining a set of directional steering together such as changing the tone, Language and length of the generating text.) compared with PEFT and ITI. We include a preliminary exploration of composability in Appendix G.2..

Again, we feel like this is an open-ended question. And introducing ReFT definitely pushes the community to think about the differences among these methods.

The classes of models evaluated is limited.

We addressed this question in our general responses by applying ReFT to other models and tasks!

Unclear presentation.

Thanks for all these suggestions! We will address the following items in camera-ready.

Line 102: What does "the hidden representation created at row and column" refer to?

Sorry about the confusion here. Row and column map to the layer and position of the intervening residual stream. We will remove these two redundant notations $i$ and $k$, and rewrite the current sentence as: “Let b be the hidden representation created at a specific residual stream located at a specific position and layer., …”

Variable names were used across sections and typos.

We will make the following changes:

• We will revisit our notations in Sec. 3 and keep a consistent format in our next revision. Please feel free to raise additional suggestions, and we will incorporate them.

• We will replace our notation for $b$ in Eqn. 1 and related parts with $h_b$ to represent the base representation.

• We will replace $m$ in Sec. 3.2 with $k$ to avoid overloading our symbols.

• We also noticed this typo ("LLaMA-7B/DiReFT/AQuA: 221.3") after we submitted our draft. The correct entry should be 21.3. We will update the number, and check existing ones in our next revision. Our result tables (i.e., not the hyperparameter tables) are semi-automatically generated, but human error is still possible)

Other suggestions on writings.

We will make the following changes:

• We will motivate how we come up with the current LoReFT formula better. We did cut a bit of text on this from our earlier draft due to length concerns. We will motivate LoReFT better from the interpretability works by bridging Sec 3.2 and Sec 3.2 together better in the next revision.

• We will clarify the use of dropout in Eqn. (2) and (3).

[4] Zou et. al., 2023, “Representation Engineering: A Top-Down Approach to AI Transparency”

[5] Zaken et. al., 2022, “BitFit: Simple Parameter-efficient Fine-tuning for Transformer-based Masked Language-models”

Thanks for your constructive feedback and questions!

LoRA with fewer parameters has been tried.

We want to clarify that the baseline numbers we have in our tables are the best performance after hyperparameter tuning done by the original LLM-Adaptor [3] paper. For instance, this adaptor benchmarking paper searches the best adaptor location (out of ranks {Attention only, MLP only, both}) as well as best ranks (out of ranks {4, 8, 16, 32}) for LoRA. With applying LoRA only to the Attention module and rank=4, LoRA can get much lower parameter count while sacrificing performance (e.g., see Figure 3 on pg. 5 of their paper). We will clarify this in our next revision.

ReFT with long-context tasks.

Thanks for the suggestions! We took up this suggestion in our general responses by applying ReFT to a long-context summarization task.

ReFT with LMs other than LLaMAs.

We addressed this in our general responses by applying ReFT to two other LM types: Mistral and Phi-3. We additionally tried to close the gap in the math benchmark by enhancing LoReFT.

Memory footprint of ReFT.

If we understand the question correctly, the memory footprint should be largely bounded by the number of training parameters. Thus, one can use the Params (%) column in our result tables (e.g., Table 1 on pg. 6) to rank the memory footprint for various methods. One thing to note is that LoReFT’s orthogonal constraint does require more memory due to the orthogonalization process compared with DiReFT. However, this might not be the dominating factor.

Moreover, compared against LoRA, ReFT does require some inference-time overhead since ReFT is an intervention-based method which cannot merge its learned weights into the original model. Nevertheless, since we constrain ourselves to intervene on the prompt tokens, the overhead is limited (fractional inference time increase is less than 1% in various settings). We provided a detailed analysis in Appendix H on pg. 38, and we hope to highlight this in our next version when more space is allowed in the main text.

Prompt prefix is shorter than the number of intervention positions.

Thanks for bringing this up, and this is a great technical question! We introduced the concept of intervention position padding in our ReFT Python library. In short, it will pad the prompt with a single padding token, and we will perform dummy interventions on this token if needed. The attention mask and loss calculation will bypass this token to make sure other tokens are not affected.

[3] Hu et. al., 2023, "LLM-Adapters: An Adapter Family for Parameter-Efficient Fine-Tuning of Large Language Models"

Thanks for raising these methodological points!

Throughout the paper, we report published numbers for other approaches rather than running our own evaluations of those approaches. Our assumption is that the other authors are the most expert in how to get their approaches to work best, and so this provides us with the most ambitious comparison points. We have carefully verified that we are using the same protocols for evaluation and use of the training data.

For the hyperparameter searching comparisons, our method for choosing hyperparameters seems stricter than the norm. In particular, we use GSM8K (also in math reasoning domain) to select our hyperparameters and apply these to both the Math and Commonsense Reasoning benchmarks, to avoid any implicit model selection based on test runs. The other authors use the Math10K train/test split to do model selection.

For the Math and Commonsense Reasoning benchmarks, both LoRA and ReFT are applied to all layers. For each layer, ReFT intervenes on the residual stream (which is weightless and therefore cannot be targeted by LoRA), while LoRA is applied to multiple weight matrices such as the Q/K/V projections. Although the LLM-Adaptor paper did not attempt to tune ranks on the Commonsense Reasoning benchmark, the DoRA paper (a newer variant of LoRA) [6] tried to halve the rank of DoRA (i.e., reducing the parameter count by 50%) and found that performance consistently dropped across all LLaMA models, as we reported in the paper (see DoRA (half) in our Table 1).

We report ReFT evaluations with much smaller parameter counts than the other methods, which would seem to put us at a disadvantage rather than an advantage. We could double-check this by increasing the ReFT parameter count to match the LoRA numbers. We would worry about lowering the LoRA counts and running our own experiments, for the reason we noted above (LoRA advocates might argue for different settings than we would choose).

[6] Liu et. al., 2024, “DoRA: Weight-Decomposed Low-Rank Adaptation”

Additional clarification: regarding the evaluation setup, we would like to clarify that we directly copied the publicly available codebase for the LLM-Adaptor paper to ensure a fair comparison (e.g., same datasets, evaluation metrics, decoding strategies, etc..).

Thanks for the response!

Further, from their paper, it seems they tune the LoRA rank parameter on math reasoning datasets, which, as your results suggest, have potentially different behavior than other datasets

To clarify what I mean by this: you find that your methods reasonably underperform e.g., LoRA, on arithmetic reasoning tasks (Table 2), which, as an example, could arise because LoRA has more trainable parameters. If this is true, then when tuning on an arithmetic reasoning task (GSM8K, Math10k, or otherwise), your hyperparameter selection might favor higher LoRA ranks because it adds additional trainable parameters. The most striking claim of your paper is that your method outperforms LoRA specifically on commonsense reasoning tasks with an order of magnitude fewer parameters. It's possible, then, that one could tune LoRA specifically to perform well on these commonsense reasoning tasks by using the same tricks that you use to tune your method (e.g., applying it to specific layers).

I also did read your global response and see that with some additional tweaking you match LoRA on math reasoning datasets, which is interesting—what I'm arguing is simply that it's important to put the same effort into tuning LoRA as your method. For example, as I proposed above, adding an hyperparameter for LoRA (analogous to the one for your method) to specify which layers it is applied to.

With that said I think you have a good paper and I am voting for acceptance, the reason my score isn't higher is because (1) I am not fully convinced that tuning LoRA (e.g., by applying it to specific layers and tuning the rank a bit) could not perform similarly to your method, and (2) because this method might not serve as a good drop-in replacement for LoRA because it underperforms on important datasets (arithmetic reasoning) without extra manual effort.

Thanks for assessing our paper to be a significant contribution, and for your question!

ReFT with LMs other than LLaMAs.

We addressed this question in our general responses by exploring other model types such as Mistral and Phi! As shown by our initial results, both LoReFT and DiReFT work for other types of LMs.

Hyperparameter selection process of ReFT.

We took up this suggestion in our general responses by providing more details on the current hyperparameter searching process.

Practically, when to choose ReFT over LoRA?

Thanks for raising the question! We wish to include more of these insights in the main text with additional space allowed in the next revision. For now, here are some practical guide of using ReFT:

• LoReFT works better when the base LM is strong: in our experiments, we usually find ReFT scales with the quality of the base LM for harder generation tasks. For instance, the gap between ReFT and LoRA is much smaller when applying ReFT to larger LLaMA models for our math reasoning benchmarks as shown in Table 2.

• LoReFT (or ReFT) is composable by its nature: LoReFT localizes task-dependent interventions into orthogonal subspaces: you can partition the subspaces of a single LoReFT intervention for different tasks. Specifically, you can train different subspaces for different tasks, and compose them together to learn a combined skill. We showed some initial results in Appendix G.2.. Although LoRA weights can be squished together, ReFT is much more interpretable. Additionally, the number of intervenable representations are abundant in LLMs. As a result, it becomes much more feasible to overload and stack interventions together in a zero-shot fashion.

• Practically, ReFT could be a better solution in a multi-tenant finetuned model serving service: imagine a case where we are serving thousands of finetuned models: for a batch user queries, we want to call different finetuned models. We cannot serve thousands of SFT LMs without enormous costs. One alternative is to have a single base LM and thousands LoRA weights cached in the memory. For this approach, you have to hot-switch between LoRA weights in memory. For ReFT, it potentially becomes much easier since ReFT only needs to intervene on the prompt representations once for the batch, and pass the intervened prompt KV cache to the inference engine without inference-time overhead. We only realized this after we submit the paper, and we hope to discuss this further in our next revision.

To improve ReFT's performance on tasks such as GSM8K.

We addressed this in our general responses by improving LoReFT’s performance on math with additional interventions on decoding steps!

Thank you for appreciating our work and raising interesting questions!

Generalization of hyperparameters.

Yes, we try to challenge the generalizability of ReFT by testing whether a set of hyperparameters for one task transfers to another as we do hyperparameter search on separate dataset splits. Moreover, we indeed tried the exact same setting you are proposing here for our instruction-tuning benchmarks in Sec. 4.4. We select the hyperparameters on LLaMA-1 7B and test those settings with Llama-2 7B without additional hill-climbing (see L248 where we mention this). Although this is not the full picture, our findings do show that selected hyperparameters transfer well across models as well. We will highlight this in our next revision.

ReFT with the SFT + RLHF alignment paradigm.

We want to clarify that our focus is not to establish a new SoTA on model alignment compared to SFT+RLHF or DPO. Rather, we seek to offer comprehensive comparisons with other PEFTs.

While ReFT saves parameters and maintains better performance compared with LoRA in our instruction-tuning experiments, our results do suggest other important application of ReFT that we wish to discuss in-depth if space is allowed:

• Much quicker iteration on the alignment pipeline. For instance, if we want to evaluate the alignment dataset quality of Alapca-15K and ultrafeedback, we could use ReFT to finetune a model and compare the performance instead of SFT. This potentially allows much quicker iterations on the data pipeline especially when datasets and models are large and constantly evolving.

• Towards understanding the mechanism of instruction-tuning. The fact that a base LM can be finetuned to follow instructions with extremely lightweight interventions that worth no more than 0.0019% - 0.0039% of the original parameter counts is surprising. Our finding can shed lights on uncovering the training dynamics of instruction-tuning.

Clarifications on the prefix and suffix tokens.

In the current draft, we experiment with a simple intervention strategy: intervening only on the leading tokens (“prefix” or the first $n$ tokens) and the trailing tokens (“suffix” or the last $n$ tokens) of the input prompt (i.e., there is no intervention being applied to output tokens). For classification tasks, we only intervene on the prefix tokens.

The intuition of this simple strategy is that interventions on the prefix tokens change the “information read-out” of all following tokens working as anchors (since attention processes them differently now), and the interventions on the suffix tokens steer the generations.

ReFT with DPO and the usage of the preference dataset Ultrafeedback.

Indeed, DPO (or any arbitrary loss function) can be integrated with a single model loaded into the memory (taking the advantage of ReFT by only training the interventions). Additionally, we have integrated DPO trainers in our library which will be open-sourced.

And yes, you are right - we are under-utilizing these preference dataset at this point. We only do SFT with the best rated responses without using the contrastive signals.