LoFiT: Localized Fine-tuning on LLM Representations

We appreciate your detailed comments and valuable feedback!

My main concern about the paper is that the main evaluations are limited to 3 tasks (TruthfulQA, MQuAKE, CLUTRR). This makes it unclear if LoFiT is generally applicable in place of PEFT methods or if it is only competitive for a specific type of fine-tuning task. [...] It would be best to include more tasks from among PEFT papers and general LLM fine-tuning scenarios.

We appreciate your suggestions. Please see the general response for the additional experiments we conducted based on this!

LoFiT algorithm modifies attention heads and keeps FFN / MLP layers intact. In contrast, several popular model editing algorithms focus on updating FFN / MLP layers alone. [...] When should one focus on editing attention layers or FFN layers. Are these interchangable or is there a type of task that does best with attention editing?

In our preliminary experiments, we found that for the tasks in our paper, fine-tuning on attention heads or on MLP layers led to similar performance for most fine-tuning methods. In addition, attention heads have been shown to serve as important model components for truthfulness, QA, and reasoning from the interpretability literature (e.g. Lieberum et al., 2023), so we ended up updating attentions alone. Theoretically, LoFiT can be adapted to be applied to MLP layers in a similar fashion and we can explore this in future versions, but we consider attention heads as a better component for more granular interpretability analyses.

When selecting the heads, do you need to tune A to convergence to get good scales? How many steps / training tokens are required?

Scaling factors (A) need to be trained to convergence. However, we would like to emphasize that the scaling factors are only learned to select important heads and this is much easier than actually learning the task: in our preliminary experiments, we found that the head selection process can converge in fewer epochs and is less sensitive to random seeds and hyperparameters (as shown in Figure 4 of our paper) compared with the final fine-tuning step.

In your experiments how to choose the total number of heads to be modified? What happens to LoFiT if you choose substantially more fewer heads?

Please see the results in the general response.

When reporting the difference in the number of trainable parameters, how do you count the choice of which heads are modified towards the total number of parameters?

We only counted the bias offsets because the bias offsets are the only learned parameters that will be used during the inference time after fine-tuning is finished. Our primary consideration is the statistical efficiency of the final learned model, which depends on how many parameters are tuned as opposed to the total number of parameters that need to be touched during (two-stage) training. This is also common in model compression literature (e.g. The prune-retrain paradigm [1]). We note that even counting the number of learned scaling factors, LoFiT only optimizes approximately half of the parameters of RED and 3% of LoRA.

Typos in L137 and L294; The work can be improved by describing the limitations of the evaluation methodology in more detail.

Thanks for the precious suggestions! We will revise the typos and include a further discussion of our evaluation methodology in the limitations section of our revision.

References:

[1] Zimmer et al., 2023. PERP: Rethinking the Prune-Retrain Paradigm in the Era of LLMs.

We appreciate your detailed comments and valuable suggestions on our work.

Could benefit from a more thorough comparison with ITI [...] It would be helpful if the authors could explore why LoFiT is so much more effective than ITI.

We think that the main performance gain of LoFiT over ITI comes from two factors. First, our head selection step selects a better set of heads than the probing method of ITI (see the ITI-heads results in Table 2 of our paper). Second, the end-to-end optimization of our bias tuning method yields more stable intervention vectors than the heuristic-based vector extraction of ITI.

Are the bias directions found by LoFiT similar to those found using ITI? Are they very different? Do they have similar magnitudes?

We conduct an analysis on the cosine similarity of learned biases between ITI and LoFiT on TruthfulQA with Llama-2-7B. We used the top 48 heads from ITI and learned offset vectors using ITI and LoFiT. Results can be found in the table below. On the same set of heads, LoFiT learns very different offset vectors from ITI.

Detail of inference-time intervention baselines, including the hyperparameters

We tuned the hyperparameters $\alpha$ for ITI, and the scaling factor $\alpha$ and layer $l$ for RepE on the validation set of each dataset and selected them by the validation accuracy, as suggested by the original papers. We will include the details on the layer $l$ for RepE in our revision of the paper.

Head-selection baseline: Whether one could simply learn a scalar to weight the entire output of each attention head or not

This is an interesting idea! We’ll explore this method in future experiments.

Head-selection baseline: How the method behaves as the number of selected heads $K$ is altered

Please see the general response.

How are the top tokens in Appendix E.1 obtained?

We used Logit Lens [2]: take the hidden state of the language model, multiply by the unembedding matrix, apply softmax to the projected vector, and then decode the logits. We took the learned bias offsets from LoFiT and applied Logit Lens to get the top tokens. Note that the vanilla Logit Lens only applies to the hidden state rather than attention outputs, so we follow [3] to adapt Logit Lens. Details of the adapted method can be found in section 6.1 of [3].

References:

[1] Zimmer et al., 2023. PERP: Rethinking the Prune-Retrain Paradigm in the Era of LLMs.

[2] Nostalgebraist 2020. Interpreting GPT: the logit lens.

[3] Yu et al., 2023. Characterizing Mechanisms for Factual Recall in Language Models.

Thanks for your thoughtful comments and feedback! Please find our answers to the question as well as clarification to some misunderstandings in the review.

The comparison with representation steering may be somewhat unfair, as LoFiT requires labeled data and an explicit training stage, while representation steering methods (e.g., RePE) do not.

We think there are some misunderstandings here, which we would like to clarify.

First, note that representation steering methods actually do require labeled data. For example, ITI needs at least 300 labeled examples to work well for TruthfulQA (see Figure 6(a) of ITI). RepE needs to “use labels to identify the layer and direction” and to stabilize the intervention for QA tasks (quoted from Section 4.1 of RepE). We experiment with LoFiT in the same low-data setting as these papers, only using 100 - 300 labeled examples to ensure a fair comparison. We did not use any additional data beyond the baselines.

Furthermore, representation steering methods also require a training stage. ITI trains a linear probe for each attention head to select the set of heads to intervene on. RepE trains a linear model to extract a reading vector for intervention. As a fine-tuning method, LoFiT requires training computation for gradient updates, which is a larger amount of compute than for ITI/RepE. However, both stages of LoFiT complete within 10 mins for our datasets using a single A6000, which is worthwhile given the substantial performance gain compared to ITI and RePE.

Therefore, we view the comparison of LoFiT to representation steering as a fair comparison. We will clarify this point in any future version of the paper.

The technical contribution is minor—the localization step is the main difference from RED, which is not very significant. Moreover, the overall idea of fine-tuning transforms for representations is shared with RED and ReFT.

We see localization as a valuable and significant contribution for two reasons. First, localization of fine-tuning to specific modules in the network has generally been underexplored. Even in work that considers localization (e.g. ReFT), the layer to modify is only considered as a hyperparameter to tune with heuristics. Our localization method provides a principled, targeted way to reduce parameter count for fine-tuning, which can alleviate overfitting (see Table 5).

Second, we think that localization itself provides insights into the research questions from the interpretability community of understanding how language models learn new tasks. Our work builds from a line of work like ROME and ITI, which have two goals: learn how to modify networks and understand which parts of a network are important for task learning. We see LoFiT as a new tool in that toolbox to shed light on how LMs work.

Although LoFiT employs a two-step process, the learned scaling factors from the first step are discarded after attention head selection. It is unclear why these parameters are not used.

The goal of the first step is to select important heads to fine-tune. This is a different goal than final fine-tuning, and is actually a much easier task. In our preliminary experiments, we found that the head selection process can converge in fewer epochs and is less sensitive to random seeds and hyperparameters (see Figure 4) compared with the final fine-tuning step. In addition, jointly optimizing the scalars and vectors (as RED does) is more susceptible to overfitting and forgetting: RED has worse generalization than LoFiT in Table 5. Thus, we separate the two processes for better usage of training compute and better fine-tuning results.

How many training examples $n$ are needed for LoFiT to be successful? A study on data efficiency could further highlight LoFiT's strengths.

We analyze the data efficiency of LoFiT on CLUTRR and MQuAKE with Llama-2-7B. Results can be found in Figure 2 of the PDF. In the low data settings ($n \leq 100$), LoFiT is better than LoRA and RED, showing that LoFiT is very data efficient. For $n \in \{300,500,1000\}$, LoFiT is still comparable to LoRA and RED with fewer parameters.

If the authors could run ReFT experiments (since ReFT is in the related works and seems to be an important baseline), what were the results?

We note that ReFT was released within two months of the NeurIPS submission deadline and should be considered contemporary to our work given the NeurIPS policy. Regardless, we included additional ReFT results for the three tasks in Table 2 of the PDF. With fewer parameters, LoFiT beats ReFT in almost all settings, including the additional datasets we evaluated; see general response.

How does LoFiT perform on tasks other than QA and reasoning tasks? Typical representation engineering methods are only used for alignment problems—can LoFiT provide broader adaptation?

Please see the general response for results on some additional datasets. We believe that the set of datasets covers a broad range of use cases, including open-ended generation (TruthfulQA), QA, math, and commonsense reasoning. Specifically for alignment, TruthfulQA is a commonly used dataset for alignment research and we explicitly evaluate the open-ended answers for it as a generation task (see section 6 of our paper).

In Table 3, do the parameter counts for LoFiT represent the learned scaling factors plus bias offsets, or just the bias offsets?

Just the offsets, because the bias offsets are the only learned parameters that will be used during the inference time after fine-tuning is finished. Our main consideration is the statistical efficiency of the final learned model, which depends on how many parameters are tuned as opposed to the total number of parameters that need to be touched during (two-stage) training. This is also common in model compression literature (e.g. prune-retrain [1]). We note that even counting the learned scaling factors, LoFiT only optimizes half of the parameters of RED and 3% of LoRA.