1 Introduction↩︎

Learning effectively from limited labeled data remains a central challenge in high-stakes domains such as AI security, where annotations are often scarce because attacks and threat vectors are constantly evolving and need to be defended against quickly [1]. In this setting, data-efficient text classification is critical for flagging large language model (LLM) inputs [2]–[5] and LLM outputs [6] that can be dangerous.

While LLMs exhibit some data-efficient behavior through in-context learning (ICL) [7], [8] and parameter-efficient fine-tuning (PEFT) [9], [10], their performance often degrades sharply when only a small number of task-specific examples are available. One issue with these methods is that they generally optimize in a large, unconstrained parameter space, which can lead to overfitting.

To tackle this issue, we propose Attention Head Reweighting (AHR), an intuitive data-efficient method to adapt LLMs to new text-classification tasks by learning only a single scalar per attention head. This drastically reduces the number of parameters that need to be adapted (see 1). We hypothesize our method is effective and avoids overfitting by leveraging the functional specialization of attention heads noted in recent works [11]–[14]. For example, if a particular security-related behavior is isolated to a few attention heads, AHR can easily learn to change the weighting for these heads rather than seeking to alter the internal weights of each head.

In our experiments we focus on text classification tasks, as these can often be effectively learned from a few examples, and many important security tasks such as detecting phishing attempts or model jailbreaks are text classification tasks. Classifiers are a critical bottleneck in AI security, which is constantly faced with data scarcity as new threats emerge, for filtering and guarding LLM responses that are interacting with real-world users and attackers.

In our experiments, AHR outperforms baselines across a variety of settings and models, improving by an average of around \(3\%\) points over standard baselines like LoRA [9] when training data is very limited i.e. \(\leq 100\) examples. AHR shows particularly large gains in the security relevant tasks of Webpage Phishing Detection (Web) and Jailbreak detection, achieving 6-7% accuracy improvement over best baseline methods with just 10 training samples, while using \(200-1000 \times\) fewer trainable parameters. We can analyze the changes our method makes to the model, and see that these gains are largely fueled by a combination of upweighting a few task specific heads as well as a couple general in-context learning heads. Finally, to maximally learn from just a few examples, we utilize in-context finetuning (IC-FT), where we update model parameters to maximize the probability of correct answers within a few-shot prompt, which we show improves accuracy by \(\sim 10\%\) points over finetuning on just one sample at a time.

2 Background and related work↩︎

2.0.0.1 Parameter-efficient finetuning

A variety of parameter-efficient finetuning methods have been proposed  to adapt large models while updating only a small subset of parameters [15]–[17]. Prominent examples include LoRA [9], AdaLoRA [10], (IA)³ [18], and related approaches [19]–[21]. These methods substantially reduce adaptation cost, but still introduce thousands to millions of trainable parameters, which can lead to overfitting in data-scarce settings.

Simultaneously, there has been a spectrum of increasingly constrained prompting approaches, ranging from continuous prompt tuning [22], [23] to discrete prompt tuning [24], prompt ensemble construction [25]–[27], and generating a single natural-language prompt [28], [29]. While effective in some settings, prompt-based methods are often very sensitive to minor variations and may struggle with under-specified tasks under severe data scarcity.

2.0.0.2 In-context learning

ICL obviates the need for task-specific parameter adaptation altogether, but can be highly sensitive to minor variations in the provided examples [30] and sometimes unfaithful to those examples [31], [32]. ICL can be improved through directed instruction tuning [33] or by including explanations along with examples [34]. A growing body of work has studied ICL, showing that it can implement linear models [35], [36], discrete functions [37], and more general algorithms [38], [39]. Some works have further argued that ICL implicitly performs optimization steps analogous to gradient descent [40]–[42] and higher-order optimization methods [36], [43].

2.0.0.3 Attention heads

Attention heads have been a primary focus for analysis [44], [45], with evidence that attention heads often latch onto task-specific patterns that are useful for ICL, e.g. induction heads [11], n-gram heads [12], and more easily explainable patterns [46]–[48]. This growing evidence of functional specialization has motivated methods that explicitly manipulate attention. Prior works have shown that modifying attention scores can be useful for post-hoc steering [13], [49], attention steering for transfer learning [50], attention-guided retrieval [51], interpretable modeling [52], boosting attention for instruction following [53], and improving long-context LLM inference [54].

3 Method: Attention Head Reweighting↩︎

Our paper aims to improve the efficiency of how models can learn from a very small number of training examples by modifying an extremely small number of parameters. We do this via Attention Head Reweighting (AHR), which introduces a single learnable scalar parameter per attention head that multiplies that head’s contributions to the residual stream (see 1). By finetuning these parameters on a small amount of data, we can amplify attention heads that are helpful for the current task, while downweighting heads that harm performance. By not modifying the weights inside individual heads we can avoid overfitting to the small amount of training data.

Notation: We follow the notation of [55], where we divide the \(W_O\) matrix into its components within each attention head. Let \(z_r^l \in \mathbb{R}^{s \times d}\) be the model’s residual stream after layer \(l\), where \(s\) is the number of tokens in the sequence and \(d\) is the model’s representation dimension. A typical transformer block consists of a Multihead-Attention (MHA) layer followed by a multi-layer perceptron (MLP) layer. We can write the update of a layer as: \[z^{l+1}_{Att} = z_r^l + \sum_{h \in H^{l+1}}h(z^l_r) \label{eq:orig95attn}\tag{1}\] \[z_r^{l+1} = z_{Att}^{l+1} + MLP^{l+1}(z_{Att}^{l+1})\]

Here \(h(z) = AzW_VW_O\), \(A = \text{softmax}(\frac{(zW_Q)(zW_K)^T}{\sqrt{d_k}})\). \(W_V \in \mathbb{R}^{d \times d_k}\), \(W_Q \in \mathbb{R}^{d \times d_k}\), \(W_K \in \mathbb{R}^{d \times d_k}\) and \(W_O \in \mathbb{R}^{d_k \times d}\) are the parameters of the attention head, and \(d_k = d/n_h\) where \(n_h = |\mathcal{H}|\) is the number of attention heads per layer.

3.0.0.1 Attention Head Reweighting:

Our method simply edits the update rule of the multi-head attention layer (Eq. 1 ) to:

\[z^{l+1}_{Att} = z_r^l + \sum_{h \in H^{l+1}}(1 + \beta_h) h(z^l_r), \label{eq:modified95attn}\tag{2}\]

where \(\beta_h\) is a head-specific parameter that is initialized at zero. During training, we only train the \(\beta_h\) parameters while keeping all other parameters, including the MLP layers unchanged. At test time, the parameters \(\beta\) can be merged into \(W_O\) for no additional inference cost. During training, we utilize supervised fine-tuning (SFT) on few-shot examples, by minimizing the cross-entropy loss on the answer tokens, as well as a regularization objective to minimize changes to the model:

\[\min_{\beta} \sum_{(x, y) \in \mathcal{D}_{train}} -\log P(y |x; \theta, \beta) + \lambda (\sum_{h} |\beta_h|^p)^{1/p}\]

where \(\mathcal{D}_{train}\) is the few-shot training dataset, \(\lambda\) is a hyperparameter controlling regularization and \(\theta\) are the (frozen) model parameters and \(p \in \{1, 2\}\).

3.1 In-Context Finetuning↩︎

Recently, In-Context Learning (ICL) has risen as an alternative to finetuning as a way for models to learn from few-shot examples [7]. While powerful, in-context learning is limited in how many samples it can learn from by the model’s context window, as well as GPU-memory. Additionally, using many in-context samples incurs a large additional cost at inference time.

To get the best of both worlds, we utilize In-Context finetuning (IC-FT) similar to [56]. Different from [56], we finetune separately on each dataset while their goal is to improve general in-context learning ability. In in-context finetuning, we construct an input string with \(k\) few-shot examples, and do supervised finetuning to maximize the probability of all answer tokens within the context. For each input with \(k\) in-context examples, we minimize the following:

\[\min_{\beta} \sum_{i=1}^k -\log P(y_i | X_i; \theta, \beta) + \lambda (\sum_{h} |\beta_h|^p)^{1/p} \label{eq:ic95ft}\tag{3}\]

where \(X_i = \{x_1, y_1, ..., y_{i-1}, x_i\}\) is the concatenated input with few-shot examples. As this technique is independent of the finetuning method, we utilize it for both our method and the baselines in our experiments (with their trainable parameters instead of \(\beta\) in Eq. 3 ). For text classification, we only maximize the probability of the first token in the answer, as predicting subsequent tokens of the answer is trivial given the first one.

4 Results↩︎

4.1 Experimental setup↩︎

4.1.0.1 Models

We consider four LLMs across a range of sizes in our experiments: GPT-2 XL ([57] [57]; gpt2-xl), LLaMa-3.2 1B and 3B ([58] [58]; meta-llama/Llama-3.2-1B and meta-llama/Llama-3.2-3B), and Qwen3 8B ([59] [59]; Qwen/Qwen3-8B). See Table 1 for more details, such as the number of parameters in each model and how many of these are trainable.

4.1.0.2 Datasets

We test our method on a set of six text-classification datasets, including standard tasks as well as cybersecurity relevant applications:

• SST2 [60]. This dataset contains sentences from movie reviews, and the task is to classify their sentiment into one of {Positive, Negative}.

• AG-News [61]. A collection of text excerpts from news articles, where the goal is to classify the excerpt into one of the following categories: {World, Sports, Business, Sci/Tech}

• Emotion [62]. A dataset of tweets, with labels based on the emotion expressed in the tweet. This dataset contains 6 classes {Sadness, Joy, Love, Anger, Fear, Surprise}

• Webpage Phishing Detection (Web) [63], a dataset of web URLs where the goal is to differentiate between Phishing and Legitimate URLs.

• Toxigen [64]. This is a dataset of potentially toxic/hateful comments for detecting hateful speech. We treated an input as toxic if the average human assessed toxicity score was \(>3\) (on a 1-5 scale) and as benign otherwise.

• Jailbreak Detection [65], a dataset of LLM prompts containing Jailbreak attempts as well as standard prompts. The goal is to classify whether a given prompt is a jailbreak attempt. We use the balanced version from https://huggingface.co/datasets/jackhhao/jailbreak-classification.

Most of these datasets have thousands of examples, but in order to study few-shot learning capabilities of our models, we randomly sample a subset of the training data to be used in few-shot learning. We evaluate performance on the full test set. See an example for each dataset in 4.

4.1.0.3 Baselines:

As our main baseline, we compare against LoRA [9]. For LoRA we use \(r=1\) and \(\alpha=1\), which we found to increase training stability and reduce overfitting on small datasets. We additionally compare against (IA)³ [18] and AdaLoRA [10]. For AdaLoRA we use Initial \(r=4\) and target \(r=2\), with \(\alpha=1\). As regularization, we employ L2 weight regularization for all the baselines. For training the baselines, we utilize PEFT package [66]. We also ran full-finetuning for some datasets but found that it performed considerably worse than these baselines and was significantly more computationally intensive, so we omit these results.

4.1.0.4 Training Details:

In our experiments, we discovered that the particular few-shot examples used can have a very large effect on the final accuracy. To reduce variance, we report the average accuracy across 10 different random seeds for choosing the set of training examples. We utilize In-Context Finetuning (Sec 3.1) for training all the baseline as well as our AHR. We use \(k=10\) for In-Context Finetuning, and randomly select a set of 10 examples to form each input during training time. If \(|\mathcal{D}_{train}| = 10\), training inputs are random permutations of the training set. We train all models using the AdamW optimizer [67].

4.1.0.5 Hyperparameter Selection

To stay true to a real data-limited setting, we do not utilize any additional data for selecting hyperparameters. Instead, we select hyperparameters including learning rate from \(\{0.03, 0.003, 0.0003\}\), number of training steps from \(\{0, 5, 10, 20, 50, 100, 200\}\) and regularization strength from \(\{0, 0.1, 0.3, 1.0, 3.0, 10.0, 30.0, 100.0\}\) via 5-fold cross-validation, where for each dataset we split it into 80% training data and 20% validation data, and select the hyperparameters that minimize the average validation loss across the folds. We select hyperparameters separately for each random seed. Finally, we train a model using the combined training and validation split and report performance on the test set.

Our prompts consists of a short instruction describing the task and the class names, as well as the few-shot examples, see exact prompt in Appendix Table 5.

4.2 Experimental results↩︎

2 and Table 2 show the average accuracy after finetuning as a function of training data size \(|\mathcal{D}_{train}|\). We can see that our AHR consistently outperforms existing parameter efficient finetuning methods when training data is limited. On average, our AHR improves over the best baseline by 2-4%(absolute) in low data regimes (\(|\mathcal{D}_{train}| \leq 30\)), which corresponds to around 10% relative error reduction, despite having over \(200\times\) fewer trainable parameters. This shows our AHR can significantly improve the few-shot learning ability of the models.

AHR performs particularly well on the web URL phishing classification task, substantially outperforming LoRA and other baselines. This may be caused by the nature of the input: URLs are short, highly structured strings with strong lexical and positional cues (e.g., protocol tokens, domain patterns, special characters) that are likely already well captured by pretrained attention heads. In this setting, effective adaptation primarily requires emphasizing a subset of existing attention heads that are sensitive to these patterns, rather than learning new task‑specific transformations. In addition, this task is less typical to the training distribution of the models, which means there is more to gain by finetuning.

In contrast, most PEFT methods such as LoRA actually harm performance compared to the In-Context Learning baseline when \(|\mathcal{D}_{train}|\) is small as shown in Fig 2 (b). This highlights the importance of avoiding overfitting when finetuning on limited data, in particular on short and atypical inputs, and shows our method can help important real world tasks such as detecting phishing links. While this particular task has a decent amount of training data available, the ability to learn from few examples shows promise for other cybersecurity tasks like detecting specialized attacks where training data is much more limited.

4.3 Ablations↩︎

IC-FT: 3 (a) compares test accuracy while using In-Context Finetuning described in Section 3.1 against standard finetuning without in-context samples. Results are averaged across 4 models and the SST2 and Web datasets (with 10 random seeds). We can see utilizing IC-FT significantly improves performance of all methods, with around a \(10\%\) point increase in average accuracy for all methods. Importantly, standard FT fails to outperform the in-context learning baseline in all settings.

Regularization Comparison: In Table 3 we compare different regularization strategies for AHR on SST2 LLama-3.2-3B. We can see this makes very little difference in performance. As a result, we use L2-norm for our results unless otherwise specified, though we use L1-norm for our analysis in 5.

5 Case Study: Finetuning Llama-3.2-1B on Web↩︎

In this section, we take a deeper look at finetuning one model on one particular task to better understand AHR and how it differs from other PEFT methods.

5.1 Number of Attention Heads Modified↩︎

First, to understand whether AHR needs to modify many attention heads, we train models with varying L1 regularization penalties {0, 0.005, 0.01, 0.02, 0.03, 0.04, 0.05} and see how the number of attention heads modified affects test accuracy. The results with Llama-3.2-1B trained on the Web dataset, average of 10 random seeds are shown in 3 (b). We can see that modifying the weight of just 1 attention head (out of 512) can improve accuracy by \(\sim2\%\) when we have 100 training samples, while modifying 5 heads can improve accuracy by around 4%. Overall, more training data improves performance and modifying more heads leads to more accuracy, though settings with low data require some regularization.

5.2 Effect on Overall Model Performance↩︎

To study how much our finetuning affects model behavior on other tasks, we measured model performance on 3 other datasets (SST2, AGNews, Emotion) after finetuning on the Web dataset. In 4 we measure the performance on the Web dataset (blue) as well as average over our other datasets (red). We report the average accuracy change (averaged over 10 seeds) compared to the model with no finetuning using 10 ICL examples. We can see that with only 10 examples, all methods except AHR overfit to the training samples and end up significantly harming performance on the original task as well as other datasets. When we increase the training set size to 30, we see most methods improve task performance on the Web task, but at the cost of losing some accuracy on other tasks. Interestingly, with AHR, finetuning on this task actually slightly improves performance on other tasks. This suggests that AHR helps avoid destructive changes to the model, and in some cases can even find generalizable edits that improve ICL abilities across many tasks.

5.3 Which heads are modified?↩︎

Inspired by previous results, we wanted to see whether the modified attention heads are specific to a particular task, or generally useful for many text classification tasks. To study this, we trained AHR models with varying L1 penalties \(\{0.01, 0.02, 0.03, 0.04, 0.05\}\) on 4 of our tasks (SST2, AGNews, Emotion and Web). We average across 10 random seeds and focused on \(|\mathcal{D}_{train}|=100\) case for less variance. We then measured which heads had the largest average change across different settings, and plot a heatmap of most changed heads in 5.

This reveals a few interesting things. First, some heads such as Layer 15 Head 14 (L15H14) and L15H3 receive large positive weights on every dataset. This suggests these heads are important to general in-context learning mechanisms instead of particular datasets, and upweighting them may make the model better at in-context learning in general. This could happen if, for example, they play a role similar to induction heads [11]. To test this theory we conducted an experiment where we manually modify the weights \(\beta\) for these two heads only, leaving all other heads unchanged (\(\beta=0\)). The results are shown in Appendix 7.2 and 8, and support our hypothesis. Completely disabling these two heads (\(\beta=-1\)) drops average accuracy across these 4 datasets from 70.68% to 66.81%, while doubling their impact (\(\beta=1\)) increases average accuracy to 71.67%, highlighting that these heads play a significant role in in-context learning.

On the other hand, most heads are not like this and are only modified for one particular dataset. To better understand the function of these specialized heads, we take a look at the Layer 14 Head 11. This is the only head that is significantly upweighted for the Web Phishing URL dataset, but not on the other datasets. 6 shows the attention pattern of this head when processing the final token of the input on the Web dataset. We can see that on all inputs this head only pays attention to the Phishing tokens. In comparison, in 7 we plot its attention pattern on the other datasets and find it very noisy with no clear pattern. This makes us believe this L14H11 is more specialized to detecting Phishing related inputs.

6 Discussion↩︎

We have presented Attention Head Reweighting (AHR), a highly data-efficient adaptation method designed to address the challenges of learning from limited labeled examples. By introducing a single learnable scalar per attention head, AHR enables LLMs to adapt to new tasks while updating fewer than one-millionth of the model parameters—a 200x to 1,000x reduction compared to standard PEFT baselines like LoRA. Our experimental results on diverse text-classification benchmarks demonstrate the efficacy of this approach:

• Few-shot Learning Performance: AHR consistently outperforms existing methods in extremely data-scarce settings (\(|\mathcal{D}_{train}|\le30\)), achieving a 2-4% average absolute improvement over baseline methods like LoRA, with large improvements in security relevant tasks like detecting Phishing URLs and Jailbreak attempts.

• Preventing Overfitting: By freezing internal head weights and only modifying the contribution of the head to the residual stream, AHR mitigates the risk of overfitting that hampers methods with larger parameter spaces and is overall less destructive to model behavior.

Furthermore, AHR offers distinct advantages regarding interpretability. Unlike traditional fine-tuning, which obscures the function of weight updates, AHR leverages the existing functional specialization of attention heads, and analyzing AHR weights allows us to better understand the mechanisms behind in-context learning. Ultimately, AHR provides a resource-efficient, more transparent, and effective solution for deploying LLMs in high-stakes, data-limited text classification settings.

6.1 Limitations↩︎

The focus of our paper is on text classification datasets, and while many important tasks such as detecting phishing links or jailbreak attempts are text classification, the applicability of our method on tasks beyond classification requires further investigation. In our initial investigations we found AHR more effective at tasks with simple outputs such as classification, while tasks that require very fine-grained control over model outputs might be less suitable for AHR, particularly if there are no attention heads focused on this type of task originally.

Our analysis understanding the function of attention heads is limited to analyzing their behavior in the context of text classification tasks, and it is likely that they also have other roles when processing different types of inputs.

Larger training datasets: Overall we find that if you have sufficient training data (e.g. \(\geq 300\) examples) other finetuning methods outperform AHR. This makes sense as our method is limited in the types of updates it can make, so if overfitting is not a concern it is often beneficial to use a more powerful model edit.

Acknowledgement↩︎

T. Oikarinen and T.-W. Weng are partially supported by National Science Foundation under Grant No. 2313105, 2430539, Hellman Fellowship, Intel Rising Star Faculty Award. The authors would like to thank anonymous reviewers for valuable feedback to improve the manuscript.

7 Appendix↩︎

7.1 Visualizing Attention Patterns↩︎

7.2 Manually Changing ICL Related Heads↩︎

Figure 8 shows the effect of manually changing the weighting of just 2 ICL relevant heads of Llama-3.2-1B in terms of average accuracy across our 4 datasets. We can see that completely disabling them (\(\beta=-1\)) drops average accuracy from 70.68% to 66.81%, while doubling their impact (\(\beta=1\)) increases average accuracy to 71.67%.

7.3 Regularization Comparison↩︎

7.4 Future work↩︎

As the extra parameter count required to store the AHR change is very small compared to the original model, attempts to create a library of fitted AHR-models and select from them at inference time may be successful, as has been done with LoRA [68].

Interpretability guided finetuning: A possible extension of our method is to apply the same general approach to reweighting different model subcomponents beyond attention heads, like interpretable circuits [69] or Sparse Autoencoder (SAE) latents [70]. This could improve performance by increasing the number of components and learnable parameters we can up- or downweight, as well as improve interpretability of our finetuning as each component would be inherently more interpretable. Models could even be trained with this kind of steering in mind, like weight-sparse transforms with interpretable circuits [71]. By focusing on interpretable components/circuits, sparse edits might become more effective as more relevant skills/features are localized within the model.

Using our finetuning method to upweight particular SAE latents could also function as an automated way to find steering vectors from SAEs for particular tasks similar to [72].

Future work could even connect with a line of work on natural-language explanations [73], [74] to adapt parameters on the fly based on task specifications (similar to approaches that have been tried for LoRA [75]) or to involve a human-in-the-loop when specifying and adapting to a new task [76], [77].

7.5 Prompt and Data examples↩︎

8 Detailed Results↩︎

References↩︎

---

1. Code to use AHR(fully compatible with the peft package) and to reproduce our experiments will be available on Github at github.com/tuomaso/attention-head-reweighting↩︎