LoFiT: Localized Fine-tuning on LLM Representations

Summary

The paper proposes LoFiT, a procedure for *localized* fine-tuning of LLMs. LoFiT chooses a task-specific subset of attention heads by tuning head-wise learnable scales and selecting the heads with largest scales (by absolute value). After that, the algorithm tunes the biases for the chosen attention heads to solve the chosen task. The paper includes experiments on three benchmarks using LLMs from Llama-2 and Gemma families. Authors compare LoFiT against alternative head selection methods and against several PEFT algorithms (LoRA and RED). The paper also analyzes the impact of heads found by LoFiT in different scenarios.

Strengths

1. Authors propose a very simple algorithm for selecting heads that seemingly works well (within the 3 chosen tasks). The fact that it is simple is a significant advantage: such algorithm would be easy to modify or reuse in other circumstances. 2. The paper includes several ablations and sub-analyses that answer the questions that arise when reading it. This is a sign that the experiments are well structured. 3. The paper is generally well written, well organized and easy to read.

Weaknesses

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. If latter is the case, the paper would be substantially stronger if it explained which tasks LoFiT is capable of and where it isn't. If latter is the case, it would be best to include more tasks from among PEFT papers and general LLM fine-tuning scenarios. For instance, the papers you compare against also evaluate on GSM8K, MMLU, MNLI-m, RTE, (super)GLUE for smaller models and more. Naturally, I do not suggest that you evaluate on *all* of the benchmarks, but the paper could be strengthened by either showing that LoFiT generalizes to more tasks or describing (and demonstrating) the types of tasks where it works poorly. This would help readers understand where to use LoFiT (or its components) instead of other intervention / peft methods.

Questions

**Q1:** To the best of my understanding, the LoFiT algorithm modifies attention heads and keeps FFN / MLP layers intact. In contrast, several popular model editing algorithms (e.g. MEMiT [1] or Transformer-Patcher[2]) 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 tasks that do best with attention editing? - [1] https://arxiv.org/abs/2210.07229 - [2] https://arxiv.org/pdf/2301.09785 **Q2:** when selecting the heads, do you need to tune A to convergence to get good scales? If not, how many steps / training tokens are required? **Q3:** in your experiments how to choose the total number number of heads to be modified? What happens to LoFiT if you choose substantially more of fewer heads? **Q4:** 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? > L137 ‘Evaluation is exact match (EM).” Possibly missing a noun (e.g. evaluation *metric* or *criterion*) > L294 Finally, we note that the method we present here requires the ability to take gradients, giving it a similar memory footprint as other PEFT methods and making it only usable on open-weight models. Technically speaking, this statement is false: owners of closed-weight models can use these methods on their models and have, in past, been known to implement some fine-tuning as a service for users (e.g. see openai api finetuning, tunethemodel and others).

Rating

6

Confidence

4

Soundness

2

Presentation

4

Contribution

2

Limitations

Authors have sufficiently addressed the limitations of the proposed method. The work can be improved by describing the limitations of the evaluation methodology in more detail.

Summary

The paper introduces Localized Fine-Tuning (LoFiT) - a two step method that involves (1) localizing attention heads that are important for a given task, and (2) learning an additive intervention for each important attention head. The authors evaluate the method over various tasks, and show that LoFiT outperforms other inference-time intervention methodologies, and is competitive with PEFT methods despite being much more parameter-efficient.

Strengths

- The paper is well-written. - I found the paper easy to read - it is clear and well-organized. - Strong results (Section 5 & 6) - LoFiT outperforms other inference-time intervention methods by a very significant margin (Table 1). - LoFiT is competitive with PEFT methods, despite being more parameter-efficient (Table 3). - Interesting additional investigations (Section 5 & 6) - The authors go beyond mere evaluation, and ask interesting follow up questions. - They show that localization is important by comparing to a baseline of selecting random heads, that the set of important heads are generally task specific, and that LoFiT shows promise in generalizing out of distribution compared to other methodologies.

Weaknesses

- Could benefit from more thorough comparison with ITI - The methodology is very similar to ITI, and as such the difference in performance (Table 1) is surprising. It would be helpful if the authors could explore why LoFiT is so much more effective than ITI. - One possible way to explore this would be to investigate the learned bias vectors directly. Are the bias directions found by LoFiT similar to those found using ITI? Are they very different? Do they have similar magnitudes? - Lacks detail in inference-time intervention baselines - Appendix C.2 and D.2 provide some limited details, but it seems important to give more detail, to convince a reader that these baseline methodologies were evaluated properly. Dataset construction is very important for contrastive pair methods, as are the hyperparameters $\alpha$ and the layer $l$ of the intervention (for RepE). Values for $\alpha$ are given, but it would be good to give more detail as to how these values were selected. - Head-selection baseline - I am curious to know if learning $A_{l}^{i} \in \mathbb{R}^{d_{head}}$ is necessary, or whether one could simply learn a scalar $A_{l}^{i} \in \mathbb{R}$ to weight the entire output of head $(l, i)$. This would seem to simplify the method considerably, requiring optimization over only one parameter per head in the first phase. - If the simplified version does not perform as well, then I think the method as presented would be better justified. - I am also curious to know how the method behaves as $K$ is altered. Is it increasingly effective as $K$ increases? Does the performance difference plateau once $K$ is increased past some value? Concretely, a figure which has $K$ on the x-axis, and performance on the y-axis would be informative.

Questions

- See the weaknesses section.

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

- The authors acknowledge the following limitations: - The paper only focuses on English-language evaluation, with short contexts. - The paper only explores 3 particular benchmarks. I think this one is particularly salient as a limitation. - The paper only evaluates models up to 13B parameters, and results may not extend to larger models.

Summary

This paper introduces a lightweight fine-tuning method that trains bias offsets for only a subset of attention heads, achieving significantly lighter adaptation compared to methods that fine-tune all layers, with minimal performance loss. The proposed method involves two steps. First, attention heads to fine-tune are selected using a scoring scheme; in this paper, the norm of learnable scaling factors is used for scoring. Second, offset vectors for the selected layers are trained. Experimental results demonstrate that LoFiT outperforms representation intervention methods by a large margin and shows performance comparable to parameter-efficient methods such as LoRA, but with far fewer parameters.

Strengths

- The concept of LoFiT is interesting as it represents a middle ground between representation steering and fine-tuning. The proposed procedure effectively addresses the main challenges in representation steering: 1) selecting layers and attention heads, and 2) determining the steering direction. LoFiT introduces a novel two-step procedure to address these problems using labeled data. - LoFiT is efficient and delivers performance comparable to LoRA and other parameter-efficient fine-tuning (PEFT) methods. - The experiments are comprehensive, and the analysis of attention head transfer and localization is insightful. - The paper is clearly written and easy to follow.

Weaknesses

- 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. - 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. - Although LoFiT employs a two-step process, the learned parameters (scaling factors) from the first step are discarded after attention head selection. It is unclear why the learned parameters are not used—is it mainly to differentiate this work from RED?

Questions

- Typically, how many training examples are needed for LoFiT to be successful? Considering the number of parameters, the data efficiency of LoFiT should be superior to other PEFT methods. A study on data efficiency could further highlight LoFiT's strengths. - If the authors had the opportunity to run ReFT experiments in a PEFT setting (since ReFT is mentioned in the related works and seems to be an important baseline), what were the results? - In Table 3, do the parameter counts for LoFiT represent the learned scaling factors plus bias offsets, or just the bias offsets? - How are the top tokens in Appendix E.1 obtained? - 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?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have adequately addressed the limitations.

On Author Response

I thank the authors for answering my questions and appreciate the additional evaluations. I have no further questions. With that in mind, I have raised my score by a notch.