Data-Efficient Adaptation of LLMs via Attention Head Reweighting

AI summary

Overview

Research area: Parameter-efficient fine-tuning (PEFT) of large language models, specifically few-shot text classification, with applications to AI security tasks such as phishing URL detection and jailbreak prompt detection.

Technical level: Intermediate. The method itself is simple, but the paper assumes familiarity with transformer attention heads, the residual stream, LoRA-style adaptation, and in-context learning terminology.

Scope: The paper proposes and evaluates Attention Head Reweighting (AHR), a method that adapts LLMs to new text-classification tasks by training one scalar per attention head, and studies when it helps, why, and what the learned scalars reveal about in-context learning.

What This Paper Is About

Adapting LLMs to new classification tasks usually means updating a large set of parameters, which overfits badly when only a handful of labeled examples exist — a common situation in security domains where new attacks emerge faster than labels can be collected. The authors ask whether adaptation can be reduced to learning a single number per attention head, exploiting the fact that individual heads already specialize in particular patterns. The goal is to match or beat methods like LoRA on few-shot tasks while training a vanishingly small number of parameters and producing an interpretable record of what changed.

Key Contributions

1. Attention Head Reweighting (AHR): a method that introduces one learnable scalar, β_h, per attention head, initialized at zero, which multiplies that head's contribution to the residual stream. Only these scalars are trained; all other model weights, including MLP layers, stay frozen. The scalars can be merged into the W_O matrix at test time for no additional inference cost.

2. Extreme parameter efficiency: AHR trains 1,200 parameters for GPT2-XL, 512 for Llama-3.2-1B, 672 for Llama-3.2-3B, and 1,152 for Qwen3-8B, which the paper reports as 200–1,000 times fewer than baselines and roughly 0.0001% of the model's parameters (the paper also states it trains less than one millionth of the model's parameters in total).

3. In-Context Finetuning (IC-FT) for few-shot adaptation: training on inputs containing k in-context examples and maximizing the probability of all answer tokens within the context, applied to both AHR and the baselines. The paper reports this improves accuracy by about 10 percentage points over finetuning on one sample at a time, and that standard finetuning fails to beat the in-context learning baseline in all settings tested.

4. Interpretability analysis of learned weights: since each learned parameter corresponds to a named head, the authors trace which heads matter. They identify heads modified across all datasets (e.g., Layer 15 Head 14 and Layer 15 Head 3) as general in-context learning heads, and heads modified for one dataset only (e.g., Layer 14 Head 11 for the Web phishing task).

Main Findings

• AHR wins in the low-data regime: Averaged across four models, AHR reaches 74.95% average accuracy at 10 training examples versus 71.22% for AdaLoRA, 68.27% for LoRA, and 67.88% for IA3. The paper reports improving over the best baseline by 2–4% absolute when |D_train| ≤ 30, corresponding to around 10% relative error reduction, despite having over 200 times fewer trainable parameters.

• Security tasks show the largest gains: With just 10 training samples on the Webpage Phishing Detection (Web) task, AHR reaches 79.44% versus 73.60% for AdaLoRA, 66.81% for LoRA, and 60.47% for IA3. On Jailbreak detection at 10 samples, AHR reaches 79.68% versus 72.33% for AdaLoRA, 72.94% for IA3, and 69.59% for LoRA. The paper describes these as 6–7% accuracy improvements over best baselines.

• Advantage narrows as data grows: The authors report that performance equalizes at around 100 samples, and in the averaged table the methods are close at 100 examples (LoRA 83.05%, AdaLoRA 82.70%, AHR 82.68%, IA3 81.27%). They state that with sufficient data (e.g., ≥ 300 examples) other finetuning methods outperform AHR.

• Per-dataset variance exists: The paper notes variance between individual datasets. For example, at 10 training examples on AG News, AHR reports 76.48% while LoRA reports 78.18%; full fine-tuning is omitted from the results because the authors found it performed considerably worse and was significantly more computationally intensive.

• AHR is less destructive to the model: Finetuning Llama-3.2-1B on the Web task, the authors measured performance on SST2, AGNews, and Emotion. With 10 examples, all methods except AHR overfit and harmed performance on the original task and other datasets. With AHR, finetuning on the Web task slightly improved performance on other tasks.

• Only a few heads need modification: On Llama-3.2-1B with the Web dataset, modifying just 1 attention head out of 512 improved accuracy by about 2% with 100 training samples, and modifying 5 heads improved it by around 4%. More training data and more modified heads generally improved results, though low-data settings required some regularization.

• Two heads appear to be general in-context learning heads: L15H14 and L15H3 received large positive weights on every dataset. Manually setting β = −1 for both dropped average accuracy across four datasets from 70.68% to 66.81%, while setting β = 1 increased it to 71.67%.

• One head appears task-specialized: L14H11 is the only head significantly upweighted for the Web (phishing URL) dataset but not for others; its attention pattern on Web inputs focused only on the "Phishing" tokens, while on other datasets it showed no clear pattern.

• Regularization choice barely matters: Comparing L1 and L2 norms on Llama-3.2-3B with SST2, the paper reports very little difference in performance, and uses L2 by default (L1 for the interpretability analysis).

Methodology in Plain English

The authors take a pretrained transformer and, instead of touching the internal weights of any attention head, they add one extra number per head. That number starts at zero and scales how much the head's output contributes to the running representation the model builds up across layers. Training adjusts only these numbers, using cross-entropy loss on the answer tokens of few-shot examples plus a regularization term that penalizes large changes. Because the head internals are frozen, the method has very little room to memorize the training set.

For training data, they build prompts containing an instruction, the class names, and k = 10 in-context examples, and they train the model to predict the answers for all examples in the context. Experiments span four models (GPT2-XL, Llama-3.2-1B, Llama-3.2-3B, Qwen3-8B) and six datasets (SST2, AG-News, Emotion, Web, Toxigen, Jailbreak Detection), with training set sizes of 10, 15, 20, 30, 50, and 100 examples sampled from the training split and evaluation on the full test set. Because the particular few-shot examples matter a lot, they average results over 10 random seeds. To avoid using extra data for tuning, they select learning rate, number of steps, and regularization strength via 5-fold cross-validation on an 80/20 training/validation split, then retrain on the combined split. Baselines are LoRA (r = 1, α = 1), AdaLoRA (initial r = 4, target r = 2, α = 1), and IA3, trained with the PEFT package, all using L2 regularization and the same IC-FT setup.

The case study trains Llama-3.2-1B on the Web dataset with a range of L1 penalties (0 to 0.05 in steps of 0.01) to control how many heads get modified, then inspects which heads change most and visualizes their attention patterns.

Why This Matters

Impact on research: The paper argues that adaptation can be pushed far below the parameter counts used by standard PEFT methods, and it links the resulting weights directly to the growing literature on attention head specialization and induction heads. It also provides evidence that a small number of heads carry general in-context learning ability, which is testable and relevant to work on how ICL is implemented mechanistically.

Real-world applications:

• Phishing URL detection: the Web dataset task of separating phishing from legitimate URLs, where AHR showed its largest gains and where a small number of samples is often all that is available for a new campaign.
• Jailbreak prompt detection: classifying whether a user prompt is an attempt to break model guardrails, another task where the paper reports large gains with 10 training samples.
• Content moderation: the Toxigen task of flagging toxic or hateful comments based on human toxicity scores.
• Classification in general low-label settings: sentiment (SST2), news categorization (AG-News), and emotion classification (Emotion) serve as the standard benchmarks showing the method is not narrowly tailored to security.

Industry relevance: Because the learned scalars can be merged into W_O, AHR adds no inference cost, and the number of stored parameters per adapted task is tiny (hundreds rather than hundreds of thousands). The paper notes the code is fully compatible with the peft package and that a library of fitted AHR models selected at inference time may be practical, an approach already tried with LoRA. For teams guarding LLM inputs and outputs, this points toward fast, cheap, low-overfit task adaptation when labels are scarce.

Future Directions

• Applying AHR to non-classification tasks: the authors state this requires further investigation, and note that initial investigations found AHR more effective on tasks with simple outputs, while tasks needing fine-grained control over outputs may be less suitable, especially if no head originally focuses on that task type.

• Reweighting beyond attention heads: extending the same scalar-reweighting idea to interpretable circuits or Sparse Autoencoder latents, which could increase the number of components available to up- or downweight and make the edits inherently more interpretable. Upweighting particular SAE latents is suggested as an automated way to find task-specific steering vectors.

• Model libraries and on-the-fly adaptation: building a library of fitted AHR models and selecting among them at inference time, or connecting to work on natural-language explanations to adapt parameters from a task description, possibly with a human in the loop.

• Better understanding of head roles: the analysis of head function is limited to text classification, and the authors state the heads likely have other roles on different input types; the relationship between modified heads and mechanisms like induction heads is a hypothesis the paper supports with a manual intervention experiment but does not fully settle.

Target Audience

Researchers and practitioners working on parameter-efficient fine-tuning, few-shot text classification, and LLM interpretability, as well as security-focused machine learning engineers who must build classifiers for evolving threats with very few labeled examples. It is also useful for readers interested in mechanistic accounts of in-context learning, since the learned per-head scalars give a direct, inspectable view of which heads a task relies on. Readers without background in transformer attention will need to consult the background section's notation before the method and case study fully make sense.

Authors’ abstract

Learning effectively from limited data is critical in domains like security where labeled examples are scarce. Large language models (LLMs) have demonstrated some capabilities for data-efficient learning, especially through parameter-efficient adaptation methods, but continue to struggle when faced with few samples for difficult tasks. To meet this challenge, we propose Attention Head Reweighting (AHR), a data-efficient method that adapts 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 learned by making use of the functional specialization of individual attention heads. Experiments on diverse open-source text classification datasets show that AHR can outperform standard baselines like LoRA when learning from limited samples, despite having 200-1000x fewer trainable parameters, as our AHR only modifies ~0.0001% of the model's parameters. In addition, our learned weights are easy to interpret and can be analyzed to better understand the mechanisms and attention heads responsible for in-context learning abilities in LLMs.