LoFiT: Localized Fine-tuning on LLM Representations

Authors: fangcong  yin
xi  ye
greg  durrett
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 23:50:14

Abstract: Recent work in interpretability shows that large language models (LLMs) can be adapted for new tasks in a learning-free way: it is possible to intervene on LLM representations to elicit desired behaviors for alignment. For instance, adding certain bias vectors to the outputs of certain attention heads is reported to boost the truthfulness of models. In this work, we show that localized fine-tuning serves as an effective alternative to such representation intervention methods. We introduce a framework called Localized Fine-Tuning on LLM Representations (LoFiT), which identifies a subset of attention heads that are most important for learning a specific task, then trains offset vectors to add to the model's hidden representations at those selected heads. LoFiT localizes to a sparse set of heads (3%-10%) and learns the offset vectors from limited training data, comparable to the settings used for representation intervention. For truthfulness and reasoning tasks, we find that LoFiT's intervention vectors are more effective for LLM adaptation than vectors from representation intervention methods such as Inference-time Intervention. We also find that the localization step is important: selecting a task-specific set of attention heads can lead to higher performance than intervening on heads selected for a different task. Finally, across 7 tasks we study, LoFiT achieves comparable performance to other parameter-efficient fine-tuning methods such as LoRA, despite modifying 20x-200x fewer parameters than these methods.

Main paper: 
 introduction: 
 A significant body of work has studied how to localize model behavior within pre-trained Transformer language models [7, 32, 30, 23, 45] . One practical benefit of this localization is that we can use lightweight interventions on the localized modules to modify that behavior, such as to change entity knowledge [30] , steer the style of generated text [47, 43] , correct a model's reasoning [12, 38] , or improve factuality [21, 60] . These approaches are unified in recent work on representation intervention [60, 54] , which adds offset vectors into various layer representations of a model to achieve desired behaviors. Computing these vectors from model activations is reported to require less data and compute than fine-tuning approaches [21] .At the same time, a distinct line of work has established the effectiveness of parameter-efficient fine-tuning (PEFT) methods [16, 2, 13, 22] by updating only parts of the pre-trained weights. Very recently, new PEFT methods have been proposed to change models' representations rather than weights [52, 54] , in line with the motivation of representation intervention. However, these methods are typically applied to a network uniformly or treat the choice of modules to fine-tune as a hyperparameter; they do not use any explicit interpretation or localization step.Q c K u + D o G v X i M Y B 6 w W c P s Z J I M m Z 1 Z Z n q F u O Q z v H h Q x K t f 4 8 2 / c Z L s Q R M L G o q q b r q 7 wF i 5 K R K K Y k Z G Z j u R J E a 4 j x 5 I S 1 O O Q i L d d H L 9 C B a 1 0 o V B J H R x B S f q 7 4 k U h V I O Q l 9 3 h k j 1 5 K w 3 F v / z W o k K L t 2 U 8 j h R h O P p o i B h U E V w H A X s U k G w Y g N N E B Z U 3 w p x D w m E l Q 7 M 1 C E 4 s y / P k / p J 2 T k v n 9 2 d F i r X W R x 5 s A 8 O Q A k 4 4 A J U w C 2 o g h r A Y A ik i k M W i o a P J G G U k 6 q i i p F G J A g K f E b q f v 9 m 7 N c f i Z A 0 5 P d q E B E 3 Q A + c d i l G S k t e b q 8 w 9 F Q 7 K b J j f j g y n a M r j 7 W 1 n L d L 9 g T W P H F S k o c U F S / 3 1 e q E O A 4 I V 5 g h K Z u O H S k 3 Q U J R z M j I b M W S R A j 3 0 Q N p a s p R Q K S b T K 4 f W Q W t d K x u K H R x Z U 3 U 3 x M J C q Q c B L 7 u D J D q y V l v L P 7 n N W P V v X Q T y q N Y E Y 6 n i 7 o x s 1 R o j a O w O l Q Q r N h A E 4 Q F 1 b d a u I c E w k o H Z u o Q n N m X 5 0 n t p O S c l 8 7 u T v P l 6 z S O L O z D A R T B g Q s o w y 1 U o A o Y h v A M r / B m P B k v x r v x M W 3 N G O n M L v y B 8 f k D T T u T 2 g = = < / l a t e x i t > 1 + A n l … … 🔥 🔥 🔥 concat & project < l a t e x i t s h a 1 _ b a s e 6 4 = " Z c b l U q e Q f 2 2 S c L V g 3 U U h c o f s d x U = " > A A A B 8 n i c b V D L S g N B E J y N r x h f U Y 9 e B o M Q Q c K u + D o G v X i M Y B 6 w W c P s Z J I M m Z 1 Z Z n q F u O Q z v H h Q x K t f 4 8 2 / c Z L s Q R M L G o q q b r q 7 w l h w A 6 7 7 7 e S W l l d W 1 / L r h Y 3 N r e 2 d 4 u 5 e w 6 h E U 1 a n S i j d C o l h g k t W B w 6 C t W L N S B Q K 1 g y H N x O / + c i 0 4 U r e w y h m Q U T 6 k v c 4 J W A l / 6 k D D 2 l Z n H j H 4 0 6 x 5 F b c K f A i 8 T J S Q h l q n e J X u 6 t o E j E J V B B j f M + N I U i J B k 4 F G x f a i W E x o U P S Z 7 6 l k k T M B O n 0 5 D E + s k o X 9 5 S 2 J Q F P 1 d 8 T K Y m M G U W h 7 Y w I D M y 8 N x H / 8 / w E e l d B y m W c A J N 0 t q i X C A w K T / 7 H X a 4 Z B T G y h F D N 7 a 2 Y D o g m F G x K B R u C N / / y I m m c V r y L y v n d W a l 6 n c W R R w f o E J W R h y 5 R F d 2 i G q o j i h R 6 R q / o z Q H n x X l 3 P m a t O S e b 2 U d / 4 H z + A G R u k K 8 = < / l a t e x i t > z (l,1) t < l a t e x i t s h a 1 _ b a s e 6 4 = " 6 v q 1 0 8 d W M o F 6 y N e C g g d l C f N U y k 0 = " > A A A B 8 n i c b V D L S g N B E J y N r x h f U Y 9 e B o M Q Q c K u + D o G v X i M Y B 6 w W c P s Z J I M m Z 1 Z Z n q F u O Q z v H h Q x K t f 4 8 2 / c Z L s Q R M L G o q q b r q 7 w l h w A 6 7 7 7 e S W l l d W 1 / L r h Y 3 N r e 2 d 4 u 5 e w 6 h E U 1 a n S i j d C o l h g k t W B w 6 C t W L N S B Q K 1 g y H N x O / + c i 0 4 U r e w y h m Q U T 6 k v c 4 J W A l / 6 k D D 2 l Z n M j j c a d Y c i v u F H i R e B k p o Q y 1 T v G r 3 V U 0 i Z g E K o g x v u f G E K R E A 6 e C j Q v t x L C Y 0 C H p M 9 9 S S S J m g n R 6 8 h g f W a W L e 0 r b k o C n 6 u + J l E T G j K L Q d k Y E B m b e m 4 j / e X 4 C v a s g 5 T J O g E k 6 W 9 R L B A a F J / / j L t e M g h h Z Q q j m 9 l Z M B 0 Q T C jE = " > A A A C B H i c b V C 7 T s M w F H X K q 4 R X g L F L R F W p C F Q l i N d Y Y G E s E n 1 I b R o 5 r t N a d Z z I d i q V q A M L v 8 L C A E K s f A Q b f 4 P T Z o C W I 9 2 r o 3 P u l X 2 P F 1 E i p G V 9 a 7 m l 5 Z X V t f y 6 v r G 5 t b 1 j 7 O 4 1 R B h z h O s o p C F v e V B g S h i u S y I p b k U c w 8 C j u O k N b 1 K / O c J c k J D d y 3 G E n Q D 2 G f E J g l J J r l E o P b i y m 5 T p M T u c 6 C X 7 6 M q l X a a P 0 u 4 a R a t i T W E u E j s j R Z C h 5 h p f n V 6 I 4 g A z i S g U o m 1 b k X Q S y C V B F E / 0 T i x w B N E Q 9 n F b U Q Y D L J x k e s T E L C m l Z / o h V 8 W k O V V / b y Q w E G I c e G o y g H I g 5 r 1 U / M 9 r x 9 K / d B L C o l h i h m Y P + T E 1 Z W i m i Z g 9 w j G S d K w I R J y o v 5 p o A D l E U u W m q x D s + Z M X S e O k Y p 9 X z u 5 O i 9 X r L I 4 8 K I A D U A Y 2 u A B V c A t q o A 4 Q e A T P 4 B W 8 a U / a i / a u f c x G c 1 q 2 s w / + Q P v 8 A e V F l l w = < / l a t e x i t > v n l < l a t e x i t s h a 1 _ b a s e 6 4 = " N 8 W T r j 6 i X w K m d m 4 l C G 9 u K G 4 Q j q A = " > A A A C B H i c b V C 7 T s M w F H X K q 4 R X g L G L R V W p C F Q l i N d Y Y G E s E n 1 I b R o 5 r t N a d Z z I d i q V q A M L v 8 L C A E K s f A Q b f 4 P 7 G K D l S P f q 6 J x 7 Z d / j x 4 x K Z d v f R m Z p e W V 1 L b t u b m x u b e 9 Y u 3 s 1 G S U C k y q O W C Q a P p K E U U 6 q i i p G G r E g K P Q Z q f v 9 m 7 F f H x A h a c T v 1 T A m b o i 6 n A Y U I 6 U l z 8 o V H j z V T o v s m B + O z I J z d O W x N j c H u j u e l b d L 9 g R w k T g z k g c z V D z r q 9 W J c B I S r j B D U j Y d O 1 Z u i o S i m J G R 2 U o k i R H u o y 5 p a s p R S K S b T o 4 Y w Y J W O j C I h C 6 u 4 E T 9 v Z G i U M p h 6 O v J E K m e n P f G 4 n 9 e M 1 H B p Z t S H i e K c D x 9 K E g Y V B E c J w I 7 V B C s 2 F A T h A X V f 4 W 4 h w T C S u d m 6 h C c + Z M X S e 2 k 5 J y X z u 5 O 8 + X r W R x Z k A M H o A g c c A H K 4 B Z U Q B VK 9 H + V W M K X b V N 6 o u q B M O X P / e 7 C A k = " > A A A C C 3 i c b V C 7 T s M w F H X K q 4 R X g J E l a h W p C F Q l i N d Y Y G E s E n 1 I b R o 5 r t t a d Z z I d i q V q D s L v 8 L C A E K s / Aj Z g E K o j W T c 9 N w M + I A k 4 F G 9 q t V L O E 0 B v S Y 0 2 D k k R M + 1 m + 1 x A 7 R u n g b q z M k Y B z 9 e d E R i K t B 1 F o O i M C f T 3 p j c T / v G Y K 3 W M / 4 z J J g U n 6 / V A 3 F R h i P A o J d 7 h i F M T A A K G K m 7 9 i 2 i e K U D B R 2 i Y E b 3 L l v 1 D b q 3 i H l Y O r / d L J 2 T i OIn this work, we investigate whether the idea of localization from representation intervention can be useful for fine-tuning LLMs. At the same time, we study whether representation offsets can be more effectively obtained via learning than via representation intervention methods. We propose Localized Fine-Tuning on LLM Representations (LOFIT; Figure 1 ). LOFIT first selects a subset of attention heads to modify for the target task. We compare several methods for this step, but find that it is most effective to fine-tune scaling factors on the model's attention head outputs and select the heads with the largest norm of learned scaling weights. Then, we perform a localized fine-tuning step to learn offset vectors added to these heads' representations, which gives our final model.We compare LOFIT with representation intervention methods on truthfulness and reasoning tasks. We show that LOFIT is substantially more effective than Inference-Time Intervention [21, ITI], even when using heads selected via the ITI localization strategy. Our approach requires learning to fine-tune our offset vectors, but is still effective even on modest amounts of labeled data that methods like ITI and Representation Engineering [60, RepE] used to compute representation offsets.We conduct analysis of which heads are selected and find that localization is important for LOFIT. Even across related tasks about surfacing knowledge from Transformers (e.g., improving truthfulness in TruthfulQA [24] and processing counterfactual knowledge in MQuAKE [58] ), using the set of heads specialized to a particular task improves the final fine-tuning step. Across models at different scales, including Gemma-7B, Llama 2-7B, and Llama 2-13B, localization identifies different subsets of heads, and these subsets of heads are not interchangeable without performance degradation.Finally, we compare LOFIT against existing PEFT methods, specifically LoRA [16] , RED [52] , and ReFT [54] . LOFIT is on par with these across settings for different LLMs, despite using 20x-200x fewer learned parameters. LOFIT also demonstrates better generalizability to out-of-domain settings.The main contributions of this work are the following: (1) We introduce LOFIT, a localized finetuning method that achieves competitive downstream performance on truthfulness and reasoning tasks by modifying the representations of a small number of attention heads. (2) We show the benefits of localization to particular sets of heads across tasks and across models, suggesting that interpretability methods can be combined with the effectiveness of PEFT for strong performance. EQUATIONMultiHead represents the multi-head attention outputs with H attention heads of head dimension d head after a linear projection W O into the residual stream:EQUATIONHere, z(l,i) t ∈ R dhead for i ∈ [1, H]represents the activations output by the ith attention head at layer l. z (l,i) t is essentially the output representation of a single attention head.I = ⟨T, V ⟩. Here, T = {(l 1 , i 1 ) . . . (l K , i K )} is a set of K attention heads to intervene on, where l j denotes the target layer and i j denotes the target head at the target layer. V = {v i l } (l,i)∈T is the set of offset vectors, where v i l ∈ R dhead . At a high-level, localized representation intervention methods involve two steps:1) Localize the set of target attention heads T to intervene. This is typically done by scoring every head (l, i) in the network using a scoring function S : (Z + , Z + ) → R. Specifically, the top-K locations are chosen according to S(l, i).2) Compute offset vectors V for these targeted heads. V can be learned or extracted in a learning free way. During inference time, these vectors will be added to offset the targeted attention activations. That is, z (l,i) t will be overwritten as a linear combination with the offset vectors z(l,i) t ← z (l,i) t + αv i l if (l, i) ∈ T where α is a constant.Instantiations in the literature Several representation intervention methods in literature can be cast in this framework [21, 57, 38] . For example, Inference-time Intervention [21, ITI] localizes T by training a logistic regression classifier to predict the truthfulness of responses using z l,i t=-1 from the last timestep as the input features. It then selects top-K heads where the scoring function S(l, i) is the probe accuracy on the validation set. Finally, ITI extracts the difference in representations of heads in T between truthful and untruthful responses as V through a forward pass, and adds V to the pre-trained representations of heads in T to improve truthfulness.Our framework is also related to another family of representation intervention methods that intervene on MLP layers [47, 45] such as Representation Engineering [60, RepE] , and even PEFT methods as well. For example, RepE extracts the difference vectors between two contrastive prompts as V and adds them to the representations of some MLP layers. BitFit [2] learns V in an end-to-end way and adds offset vectors to all modules that have a bias term. We choose to focus on attention heads, as recent interpretability research indicates that attention heads are semantically interpretable units for many key functionalities of LLMs such as induction [32] and retrieval [53] , but we also compare with stronger PEFT methods in Section 6.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: Our main claims of the effectiveness of LOFIT and the benefits of localization are accurately stated within the scope of truthfulness and reasoning tasks we investigate in Section 5 and Section 6Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
limitations: 
 Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Limitations are discussed in a separate paragraph in Section 8.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
lofit: localized representation fine-tuning: 
 We now present our method, Localized Representation Fine-tuning (LOFIT). Similar to the intervention framework as described in Section 2, LOFIT also aims to compute a set of the offset vectors V in the representation space targeted at a localized set of heads in T . Unlike the learning-free alternatives, LOFIT uses a learning phase to select the heads, then also optimizes the offset vectors. As illustrated in Figure 1 , our approach also follows a two-step framework.Step 1. Attention Head Selection: The first step of LOFIT learns to select a set of potentially impactful attention heads for learning the given task.We incorporate a learnable scaling factor A i l ∈ R dhead for any head at layer l ∈ [1, L] and index i ∈ [1, H] of the pre-trained LLM. A i l can be viewed as a vector of scalars to upscale or downscale each element in the activations z (l,i) t of the attention head i at layer l.With the scaling factors, during a forward pass, the activation z(l,i) t is rescaled by z (l,i) t ← (1 + A i l ) ⊙ z (l,i) t. We freeze all pre-trained weights and learn A i l end-to-end with the cross-entropy loss on a small amount of labeled data from the task of interest. During training, A i l is initialized from N (0, σ A ). We regularize the optimization with an L1 normalization with a scaling hyperparameter λ to encourage sparsity for better head selection.Once we have learned A i l , we score each head using the norm of A, i.e., S(i, j) = ∥A i l ∥. A large score S(i, j) indicates a stronger intervention is needed for a particular head. 2 Therefore, we select the set of top-K attention heads as the target locations T . K and σ A are adjustable hyperparameters. 3 Step 2. Bias Tuning: The second step of LOFIT learns the offset vectors V added to the hidden representations of each attention head in T . We freeze all pre-trained weights and add learnable parameters V = {v i l |(l, i) ∈ T } such that during a forward pass, the activation z (l,i) twill have an added offset bias vector: z(l,i) t ← v i l ⊕ z (l,i) t.During training, we learn V with the cross-entropy loss on the same training data. v i l is initialized from N (0, σ v ) and σ v is an adjustable hyperparameter. At inference time, the learned biases V are added to the hidden representations of the target attention heads T in the same way as a forward pass during training.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: As described in Section 4, all datasets we use for experiments are publicly accessible datasets. We will also release the exact data splits of these datasets we use in our code release. Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: Details of computational resources are provided in Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
experimental setup: 
 We evaluate LOFIT on question answering (QA), multi-hop reasoning, and counterfactual reasoning tasks, which are common settings for evaluating interpretability-motivated methods [21, 60] . We focus on a relatively low data condition: for each dataset, we sample 500 training points or fewer, to be consistent with the common low-data setup of representation intervention methods.TruthfulQA [24] is a QA dataset with questions where humans are likely to give false answers because of common misconceptions. Representation interventions such as ITI [21] have shown success in eliciting truthful responses from LLMs without tuning. We follow the setup in [21] to split TruthfulQA into train/dev/test sets into 326/82/407 questions and use two-fold cross validation. We report MC1 and MC2 accuracy in the results; details of these metrics can be found in Appendix A.1.CLUTRR [41] is a deductive reasoning dataset requiring multi-hop reasoning over family relationships. We use the subset of 2-hop questions and randomly split the data into train/dev/test sets of 300/450/450 QA pairs. The evaluation metric is exact match (EM).MQuAKE [58] is a knowledge editing benchmark for evaluating the propagation of edited knowledge to related facts. We convert the subset of 2-hop knowledge propagation into an in-context knowledge editing setup by prepending "Imagine that <Edited Knowledge>" to the question of related facts, following [6] ; we use this setup to evaluate if our method can learn simple reasoning over counterfactuals. Data is randomly split into train/dev/test sets of 134/95/864 QA pairs. The evaluation metric is exact match (EM).Representation Intervention Baselines Our primary aim in this work is to compare LOFIT with other representation intervention techniques. In Section 6, we will also compare it against other 2 One could also tune the offset vectors and select attention heads with the largest norm of the offset vectors. In our experiments in Section 5, we find this is less effective than tuning the scaling factors for localization. We hypothesize that this is because scaling only amplifies or suppresses the existing pre-trained representations as opposed to directly editing them, which potentially prevents overfitting in the localization step. 3 Appropriate values for K and σA can be selected based on the preferred sparsity level for a specific downstream task. Further discussion can be found in Appendix D.1 We use Llama 2-7B [46] , Llama 2-13B, and Gemma-7B [44] for experiments.Unless explicitly stated, we use the pre-trained base version of each model. When training with LOFIT, we learn the same scalars and biases for every token position of the input sequence. When performing inference with LOFIT, we use greedy decoding and we add the bias terms at the targeted attention heads to every decoding step. Prompt templates can be found in Appendix B. Configuration details and hyperparameters for LOFIT and the baselines can be found in Appendices C and D.For CLUTRR and MQuAKE, we use cross-entropy loss with gold responses for fine-tuning. For TruthfulQA, we use direct preference optimization [37] by pairing the gold truthful responses and untruthful responses as preference data for fine-tuning, as SFT has been shown ineffective in [21] .For all experiments in Section 5, we select 3% attention heads for each model for LOFIT and ITI: precisely, K = 16 for Gemma-7B, K = 32 for Llama 2-7B, and K = 48 for Llama 2-13B. 4 We use the top-1 layer for layer-based baselines, including RepE. Such granularity has been shown to be the minimum effective level of representation interventions [21, 60] . 
results: effectiveness of localization: 
 We first discuss the overall performance of LOFIT. We find that learning-based localized intervention can be much more effective than learning-free alternatives, even with limited training data. As shown in Table 1 , fine-tuning the representations of specific attention heads with LOFIT outperforms the baselines by a large margin across all settings.The rest of this section will focus on our primary question on the effectiveness of localization: is localizing attention heads important for learning downstream tasks in the LLM representation space? If so, how task-specific are these sets of attention heads?To validate the effectiveness of our localization method, we compare it with other head selection methods by tuning the biases V for a set of heads T selected by the following baseline methods. Random sampling: We randomly sample K heads from the uniform distribution. 5Probing for layers: Along with other mechanistic interpretability works [30, 40] , RepE focuses on layers as the subject matter for localizing functionality in LLMs. We examine if localizing important layers is better than important attention heads. Given a prompt in training data, we concatenate the gold response with it and a sampled incorrect response with it to create a pair of contrastive responses. We extract the pre-trained representations of all attention heads at each layer at the last token of both responses through a forward pass and concatenate them. We then train a logistic regression classifier for each layer to predict the correctness of responses using the concatenated representations as the input features, i.e., z l t=-1 = concat(z(l,1) t=-1 , ..., z (l,i) t=-1 ) for i ∈ [1, H].We define the scoring function over heads as a scoring function over layers S( * , l), which is the probe accuracy on the validation set.Bias-based selection: Prune-then-retrain is a common paradigm in neural network sparsification literature [59, 36] by re-training models on a sparse set of fine-tuned weights. We adapt this paradigm as the bias-based head selection baseline: we fine-tune the biases for the hidden representations of all attention heads v i l , and select the top-K attention heads where the scoring function S(i, l) = ∥v i l ∥. ITI head selection: ITI shows that training a linear probe using the hidden representations of each attention head can help identify important heads for truthful generation. We select the top-K heads based on ITI head selection method as described in Section 2.Results Table 2 shows that selecting attention heads based on LOFIT or probing results in consistently better downstream accuracy than random sampling and intervening on an entire layer. This indicates the effectiveness of localized interventions at the attention head level for learning the given tasks. Moreover, fine-tuning the representations of LOFIT heads outperforms all other head localization methods in most settings. This shows that LOFIT is a strong localization method for finding some optimal sets of attention heads for learning the interventions.While localized interventions outperform ones without localization for the same task, we are also interested in the localization across tasks: is the localized set of attention heads task-specific, or is it generally helpful for learning any task in the representation space? To answer this question, we run a specificity experiment by changing the sets of heads to tune, specifically using the selected heads from a task different from the one we tune the bias. We compare the accuracy of LOFIT using different-task heads with the same-task accuracy and we use randomly sampled heads as a baseline. Figure 2 shows that localized interventions by LOFIT are task-specific for TruthfulQA and MQuAKE: tuning biases on same-task heads is consistently better than tuning on different-task heads by a significant margin across all models. For TruthfulQA, using different-task heads can lead to as large as 10% absolute performance degradation and can even be worse than random head selection (e.g., for Llama 2-13B), suggesting that the selected heads might have very specific functions. On the contrary, CLUTRR does not require a task-specific set of heads to achieve good performance, possibly because LLMs can be easily adapted without task-specific localization for this relatively easy task. In addition, examples in Appendix E show that some offset biases learned by LOFIT might also carry task-specific concepts.We further examine where the task-specific heads reside in LLMs: do the localized sets overlap, or tend to select heads from similar layers? We illustrate the distribution of LOFIT heads over layers in Figure 3 . Across all models, task-specific distributions peak in some contiguous sets of layers within the same model rather than being widely spread across layers or concentrated in a single layer. For the two Llama 2 models, peaks of different tasks are qualitatively very different from each other. We report Jaccard similarities of the selected head sets in Appendix G; head sets for different tasks are only mildly overlapping. These findings demonstrate that there is not a single set of heads best for fine-tuning across all tasks.  
comparison with peft methods: 
 We also compare LOFIT with existing PEFT methods, LoRA [16] , RED [52] , and ReFT [54] . LoRA learns offsets to weight parameters (rather than representations, as we do) via a product of two low-rank matrices. RED fine-tunes scaling factors and biases on the hidden representations of LLMs, similar to combining Step 1 and Step 2 of LOFIT. ReFT learns to edit the hidden representations with a linear projection in a lower-dimensional subspace. Unlike LOFIT, RED and ReFT involve no localization. Details for ReFT can be found in Appendix C. We also include a half-parameter version of RED (RED (Half)) as an additional baseline where only the layers in the second half of the network are tuned.Our focus is on datasets where models can be steered to have the right behavior, e.g., by interpretability-motivated methods from Section 4. However, for completeness, we also include a representative selection of datasets from [17] for PEFT methods that cover commonsense reasoning, open-book/closed-book QA, and mathematical reasoning: SIQA [39] , ARC-Challenge [5] , BoolQ [4] , and SVAMP [35] . We stick to the low-data setting by sampling 500 training examples or fewer; details can be found in Appendix C. For all the following experiments, we select 10% attention heads for each model: K = 48 for Gemma-7B, K = 96 for Llama 2-7B, and K = 160 for Llama 2-13B. For n ≥ 300, LOFIT is still comparable to LoRA and RED with fewer parameters.updates outweighs the benefit of localization. Nevertheless, these results show that localization can enable further targeting of parameter updates without impacting task accuracy. 6 Data efficiency We further compare LOFIT against PEFT methods using different numbers of training examples n. We analyze the data efficiency of LOFIT on CLUTRR and MQuAKE with Llama 2-7B. Figure 4 shows that in the extremely low data settings (n ≤ 100), LOFIT performs better than LoRA and RED, showing that LOFIT is very data efficient. For 300 ≤ n ≤ 1000, LOFIT is still comparable to LoRA and RED with fewer parameters.Open-ended Generation While being fine-tuned for discriminative tasks in the above experiments, LOFIT also shows good performance on the open-ended generation task of TruthfulQA. We fine-tune with LOFIT and other methods on TruthfulQA with the same setup in Section 4and evaluate its open-ended generation on Truth-fulQA test questions with GPT-4. We prompt GPT-4 to evaluate the informativeness (Info) and truthfulness (True) of model responses given the question and the gold labels. 7 Table 4 shows that LOFIT leads to truthful and informative responses that are comparable to PEFT methods and are better than ITI. Example outputs can be found in Appendix H.Benefiting from the small number of parameter updates, LOFIT can potentially generalize well to out-ofdomain tasks after fine-tuning. We show a case study on TruthfulQA: we first fine-tune Llama 2-7B-chat on TruthfulQA and then evaluate 0shot on three out-of-domain question-answering benchmarks: TriviaQA [18] , Natural Questions [20, NQ], and MMLU [14] . We follow the same evaluation scheme in [21] to convert TriviaQA and Natural Questions into multiple-choice questions and report accuracy. 8 Table 5 shows that LOFIT suffers less from overfitting on TruthfulQA as its 6 Our two-step procedure requires fitting additional parameters in the first stage. However, even accounting for these parameters, which are not used in the final model, we still modify fewer parameters than PEFT methods: half the parameters of RED and 3% of LoRA even if counting both sets. 7 Details of GPT-4 evaluation prompt can be found in Appendix A.2 8 For TriviaQA and NQ, we use the two evaluation splits curated by [21] that are publicly accessible in the honest_llama repository. We refer readers to [21] for details. For MMLU, we use the evaluation scheme of open-source package lm-evaluation-harness [10] . performance on TriviaQA and MMLU does not drop from the non-fine-tuned base model, and it even improves on NQ while OOD performance degradation can be observed in ITI and PEFT baselines. 
related work: 
 Recent interpretability work has shown that certain semantics can be encoded in identifiable circuits or modules of transformer language models in human-interpretable ways, including entity knowledge [30] , factual associations [11] , logical reasoning [55] , and arithmetic reasoning [42] . Moreover, mechanistic interpretability methods, especially causal intervention, have revealed that such semantics can be attributed to specific attention heads [23, 21, 28] , MLP layers [40, 30] , neurons [49] , or subnetworks [1] . Inspired by these findings, a line of work has explored intervening on pre-trained language models by manipulating and editing hidden representations at some specific locations in transformers with lightweight or no tuning to perform specific tasks or for controllable generation, including alignment [60] , style transfer [43, 47] , reasoning [12] , truthfulness [21] , and knowledge editing [30] . We follow this vein of work, but we do not use neuron-level localization methods.Separately from interpretability, parameter-efficient fine-tuning (PEFT) methods have been broadly used to selectively change a small fraction of pre-trained model parameters to learn specific downstream tasks. State-of-the-art PEFT methods [15, 16, 54, 52, 19] can learn to adjust less than 1% of the pre-trained parameters and match or even outperform vanilla full fine-tuning methods on various benchmarks. Concurrent to our work, ReFT [54] proposes to integrate interventions on representations into PEFT as a strong alternative fine-tuning method. However, ReFT does not use a localization step derived from interpretability methods and instead treats layers to tune as hyperparameters. We believe our approach, focusing on a subset of the network, has benefits for continual learning and can support techniques like model merging. 
conclusion: 
 In this work, we introduce LOFIT, a two-step localized fine-tuning method for LLMs that selects a subset of attention heads and learns task-specific offset vectors to be added to the hidden representations of the targeted attention heads. We show the strong downstream performance of LOFIT on tasks involving truthfulness and reasoning, outperforming representation intervention methods (ITI and RepE) and matching strong PEFT methods (LoRA) with fewer learned parameters. We also show that LOFIT is effective at localizing task-specific attention heads for learning downstream tasks, showing that interpretability insights have a part to play in improving LLM fine-tuning processes.Limitations: Our experiments in this work focus on English-language truthfulness and reasoning datasets with short contexts and responses; different model behaviors and localization might be observed for tasks with long context or long-form generation [53, 25] . We selected these tasks following past interpretability work and to enable straightforward evaluation. In addition, we only evaluate autoregressive Transformer language models with up to 13B parameters; different behavior could be observed at the largest scales, although we note that past interpretability work has shown localization on larger models (≥ 70B) as well [23] .not know the answer to the question.  
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This work does not introduce any new datasets or new foundation models that can potentially cause risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We include license information in Appendix I.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes]Justification: All details of reproducing the data splits we use for experiments are included in Section 4 and Appendix B.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: This work involves no human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: This work involves no human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: We use the following prompt templates for fine-tuning and evaluating LOFIT and baselines.We follow the prompt strategy in ITI [21] : at the fine-tuning steps, we simply concatenate the question with the gold response or the incorrect response in the preference pair of the training and validation data as:Prompt B.1: TruthfulQA Prompt: "Q: {Question} A: {Response} "For evaluation, we prepend the standard "QA prompt", which can be found in the original implementation of TruthfulQA [24] and are later adopted by others [21, 46, 44] , to the aforementioned prompt as the standard way of evaluating models on TruthfulQA in literature.Each example in MQuAKE comes with a piece of edited knowledge and a multi-hop reasoning question that uses the edited knowledge. We used the following prompt for MQuAKE.Prompt: "Q: Imagine that {edited_knowledge} . {question} A:"Each example in CLUTRR comes with a few-sentence story that describes relations among fictional characters and a pair of characters whose relationship needs to be inferred from the story and answered by the model. We use the following prompt for CLUTRR. In preliminary experiments, we found that LLMs sometimes refuse to answer and indicate that relationships in the provided story are incorrect based on names and relations of known people in the real world, so we include further clarifying instructions in the prompt in addition to basic formatting.Prompt: "Read the following story about a family. {Story} Assume the relations described in the story are all true.Based on relations between the fictional characters in the story (assume the relations are all true) and your commonsense knowledge about family relationship, how is {Character2} related to {Character1} ? Answer: {Character2} is {Character1} 's"We use the prompt from [17] for SIQA. We use the prompt from [17] for ARC-c.Prompt: "Please choose the correct answer to the question. In the prompt, we instruct the models to generate an equation for the math word problem, and we only evaluate the correctness of the final answer derived from the equation rather than the entire equation. We fine-tune LOFIT and baselines using a single NVIDIA-RTX A6000 GPU with 48G memory. We use the huggingface implementation of Transformers [51] in PyTorch for all fine-tuning, and the TRL [50] implementation of direct preference optimization [37] for fine-tuning on TruthfulQA. We use AdamW optimizer for fine-tuning [26] with ϵ = 1e-8 an a weight decay of factor 0.01. For both fine-tuning and inference, we use full precision for Llama 2-7B and bfloat16 mixed-precision for Llama 2-13B and Gemma-7B to fit on a single GPU.For ITI and RepE, we used the original paper implementations. Note that RepE proposes three different representation engineering methods for different tasks. We select the method that works best for each task in the original paper: specifically, we used RepE with contrast vector for MQuAKE and CLUTRR, and RepE with reading vector for TruthfulQA. We refer readers to [60] for further details.For PEFT baselines, we used the huggingface PEFT library [27] implementation of LoRA. We used our replication of RED as the official implementation was not available at the time of writing this paper. For ReFT, we use the official implementation of the most performant variant of ReFT, called LoReFT, from [54] . LoReFT edits the representations of a model in a lower-dimensional subspace by learning a linear projection from the residual hidden states to the subspace.Although our main focus is to evaluate methods on datasets that are commonly used to benchmark interpretability-motivated methods, including representation intervention methods and LOFIT, we include additional results of LOFIT on common benchmarks for PEFT methods in Section 6. To be consistent with our low-data setting in Section 4, we only sampled 100 to 350 training examples rather than using the entire training data. We use the single-task learning setting from [17] where each fine-tuned model only learns to do one task. Details of these datasets are below.SIQA [39] SIQA is a commonsense QA dataset that evaluates the model's understanding of social scenarios. We SVAMP [35] SVAMP is a math word problem dataset. We split the dataset into train/dev/test splits of 350/350/300 examples, and we trained the models using the gold equation and the final answer as the label (rather than just using the final answer).We converted SIQA, ARC-c, and BoolQ into a unified multiple-choice QA format using prompts from [17] . Following [17] , we evaluate the models based on whether they can correctly generate the option (rather than using the log-likelihood scoring of the sequence of each option). The prompts can be found in Appendix B.For all experiments on TruthfulQA, MQuAKE, CLUTRR, and SVAMP, we fine-tuned for 5 epochs with a batch size of 8 for all methods except ReFT; we will discuss the specific hyperparameters for ReFT in later sections. For all experiments on SIQA, and ARC-c, we fine-tuned for 3 epochs with a batch size of 8 to prevent memorization of world knowledge. For BoolQ, we fine-tuned for 3 epochs with a batch size of 4 for Llama 2-7B and of 2 for Llama 2-13B to fit the long passage context into a single GPU. We used the same implicit reward hyperparameter of direct preference optimization β = 0.5 for TruthfulQA experiments. Method-specific hyperparameters can be found in the following subsections.Hyperparameters of LOFIT used in each experiment are summarized in Table 6 . Because LOFIT for different models involves different numbers of learned parameters, we did not use the same set of hyperparameters throughout; instead, we tuned the following hyperparameters with grid search. Specifically, we found that a small L1 regularization term is good for enforcing the learning of a sparse set of effective attention heads, which is also suggested by model sparsification literature [34] . We also found that when using fewer heads, a larger learning rate is needed to stabilize training for LOFIT.In addition to hyperparameters listed in Table 6 , we used the following hyperparameters uniformly across all experiments for LOFIT: σ A = σ v = 0.001 for the initialization of LOFIT. We found that this set of initialization is robust to random seeds.ITI In preliminary studies, we tuned the scaling factor α of the offset vectors over a range of values suggested in the original paper [21] . We tuned α on a validation set of each dataset and we found that α = 15 worked robustly across all settings.RepE In preliminary studies, we tuned the scaling factor α of the offset vectors on a validation set of each dataset and found that α = 5 worked best across settings.LoRA In preliminary studies, we experimented with different configurations of LoRA, including applying LoRA weights to MLP layers and changing the rank and α. We found that the following configuration strikes the best balance across settings between overfitting with too many parameters and under-tuning with too low rank: we fine-tuned the Q projection and V projection matrices of all attention heads, used rank = 8, α = 8, and applied a dropout rate of 0.1 to prevent overfitting. We performed a hyperparameter sweep for the learning rate on the validation set of each dataset and the optimal configurations for each model and dataset can be found in Table 7 .RED Suggested by the original RED paper [52] and confirmed in our preliminary studies, finetuning all attention heads with RED is better than other alternative RED configurations, including tuning MLP layers and tuning all modules, across settings. We performed a hyperparameter sweep for the learning rate on the validation set of each dataset and the optimal configurations for each model and dataset can be found in Table 7 .-3 8e-3 CLUTRR 5e-3 5e-4 1e-2 1e-2 Llama 2-7B TruthfulQA 5e-4 5e-3 1e-2 5e-3 MQuAKE 1e-3 5e-3 1e-2 5e-3 CLUTRR 5e-3 5e-4 1e-2 5e-3 Llama 2-13B TruthfulQA 1e-3 1e-3 2e-2 5e-3 MQuAKE 1e-3 1e-3 8e-3 5e-3 CLUTRR 1e-3 1e-3 1e-2 8e-3 Llama 2-7B SIQA 5e-3 5e-4 - 1e-2 ARC-c 1e-3 1e-3 - 5e-3 BoolQ 5e-3 5e-4 - 8e-3 SVAMP 1e-3 1e-3 - 1e-2 Llama 2-13B SIQA 1e-3 1e-3 - 5e-3 ARC-c 1e-3 1e-3 - 5e-3 BoolQ 1e-3 1e-3 - 5e-3 SVAMP 1e-3 1e-3 - 1e-2ReFT The ReFT paper [54] indicates that ReFT has the following important hyperparameters to tune: layers to apply interventions, the rank of the subspace, token positions in the input where the interventions are applied, and learning rate. In addition, in our preliminary experiments, we found that ReFT is very sensitive to hyperparameter choices and batch size also has an impact on the ReFT performance. Therefore, we performed grid search for the aforementioned hyperparameters to select For the bias tuning stage, the number of heads K has an impact on LOFIT performance. We conduct an analysis on the percentage of attention heads used for LOFIT bias tuning versus the accuracy on MQuAKE and CLUTRR. Results in Figure 5 show that the performance plateaus when K reaches 10% -20% of the total number of attention heads and continues to increase as K gets larger before it reaches the above threshold. This is likely because the number of learned parameters is too small to be expressive when K is smaller than 10% of attention heads (<10K parameters). Based on these observations, we use 3% for extremely parameter-efficient scenarios and 10% for the best balance between parameter counts and performance.Accompanying the qualitative evaluation of the distribution of LOFIT heads for different tasks, we also quantitatively examine the similarity of LOFIT heads with the following metrics:Jaccard similarity For the same model and a fixed number of selected heads K, we compute the overlap between the two sets of LOFIT heads selected from a pair of tasks (i, j) with Jaccard similarity, namely:EQUATIONEarth Mover Distance of Head Distributions Adjacent attention heads might share similar functionality in LLMs [29] , and Jaccard might fail to capture such similarity. Therefore, we compare the similarity of head distributions over layers. For a pair of LOFIT sets, we normalize the number of heads selected in each layer in each set into a probability distribution over layers, and compute Earth Mover Distance between the two distributions.Results are summarized in Figure 6 . Along the diagonals, Jaccard similarities show that for the same task, LOFIT is able to robustly localize approximately the same set of heads when trained with different random seeds. Across all models, pairs of LOFIT sets for different tasks share less than a third of attention heads. In addition, Earth Mover Distances are relatively large between differentNeurIPS Paper ChecklistJustification: We checked every aspect in the Code of Ethics and confirmed that we follow the guidelines.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There are no specific social impacts of this work separate from the general impacts of work on large language models, interpretability, and parameter-efficient tuning methods.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).Our code is available at https://github.com/fc2869/lo-fit.We conduct an analysis on the percentage of attention heads used for LOFIT bias tuning versus the accuracy on MQuAKE and CLUTRR to determine the best K and results can be found in Appendix F.We considered other random selection schemes, including sampling the same number of heads from each layer and biased sampling towards more important layers, but they worked worse than uniform random sampling. 


Decision: The paper presents LoFiT, a localized fine-tuning method that effectively balances parameter efficiency and performance by selectively tuning attention heads in large language models (LLMs). Reviewers appreciate the simplicity and novelty of the two-step approach, which selects important attention heads using a scoring scheme and trains bias offsets for those heads. LoFiT demonstrates strong results across multiple benchmarks, outperforming representation intervention methods and competing closely with parameter-efficient fine-tuning (PEFT) methods like LoRA, despite using significantly fewer parameters (reviewers 1, 2). The paper is well-organized, clear, and supplemented by comprehensive ablation studies that validate the approach's effectiveness (reviewers 2, 3). However, concerns were raised about the scope of evaluations, as the main results are limited to three tasks, and further analysis of LoFiT's generalizability across more diverse tasks is suggested (reviewer 3). Additionally, a deeper comparison with inference-time intervention baselines and a more detailed examination of LoFiT's data efficiency would strengthen the contribution (reviewers 1, 2). Despite these minor limitations, the paper is deemed technically sound, innovative, and well-presented, supporting its acceptance.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: questions: - See the weaknesses section.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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).
limitations: 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.

Comments: title: On Author Response
Rebuttal of the Authors (Continued)
comment: 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.
We appreciate your thoughtful suggestions. We will include the analysis in any future version of the paper.
Thank you for your suggestions! We will include the study on data efficiency in any future version of the paper.
Thank you for your thoughtful and thorough response. I believe that incorporating the study on data efficiency would significantly strengthen this paper. I've raised my score to 5
I have read the rebuttal and the general response.

I suggest including "Effects of the number of heads $K$ on performance" analysis in the paper, at least as an appendix section. I also suggest making the plots more granular in the 0-20% region (it looks like a very small number of heads are needed to saturate MQuAKE performance).

I thank the authors for their diligence. I will keep my score the same.
> 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.