Data-scarce Behavior Editing of Language Models (2025)

Abstract: Large Language Models trained on web-scale text acquire language generation abilities that can solve a wide range of tasks, particularly when task knowledge is refined into the generative prior using in-context examples. However, spurious features learned from noisy data hinder their generalizability. Supervised fine-tuning can enhance task specificity but may lead to data inefficiency. Prior studies indicate that (i) noisy neural circuitries coexist with generalizable ones within LLMs, and (ii) finetuning typically enhances (or suppresses) existing abilities without introducing newer ones. Building upon these, we propose TaRot, a novel method for task adaptation. TaRot intervenes in the neural circuitries using learnable rotation matrices that are optimized using Bayesian optimization, on labelled samples in the order of standard few-shot prompting examples. Experiments on multiple classification and generation tasks using LLMs of varying sizes reveal the efficacy of TaRot, improving upon both zero- as well as few-shot performance, with average improvements (across models and tasks) of 15.6% and 14%, respectively.

## introduction large language models (llms) acquire the ability to associate different language concepts presented in a sequential context by optimizing the prediction probability of the next token given a context. despite its apparent simplicity, when scaled across web-sized text corpora, such a learning strategy introduces the ability to solve a wide range of tasks presented in natural language. however, the web contains almost everything humankind has written, and therefore, it introduces spurious token associations that are irrelevant or even counter-productive to the model to become generalized task-solvers. ## related work our work is primarily relevant to two broad areas of existing literature: adaptation of pretrained language models to downstream tasks, and mechanistic understanding and intervention techniques. task adaptation of pretrained language models. the pretrain-finetune regime for adapting language models to downstream tasks dates back to the early approaches like bert @xcite pretrain a language model (lm) on large unstructured text corpora using self-supervised objective, followed by supervised fine-tuning on task-specific, relatively smaller datasets. despite the apparent simplicity, the pitfalls of this regime have been pointed out in terms of distribution shift @xcite . with the development of large-scale, autoregressive transformer-based language models and their ability to learn from in-context examples @xcite , a definitive shift has happened in the more recent past. current practices of using these models for downstream tasks primarily rely on designing suitable prompt templates and labeled example retrieval for in-context learning (icl) @xcite @xcite ; traditional techniques of finetuning have taken a back seat due to the computational cost and catastrophic forgetting introduced by small-scale task-specific data that hurts the pretrained abilities @xcite . instead, finetuning to follow task instructions, aka instructiontuning @xcite , has gained popularity. instruction-tuning has been shown to introduce zero-shot task adaptation abilities in llms @xcite . additionally, different methods of alignment tuning have been proposed with the primary goal being aligning the generative distribution of the language models with human values and preferences @xcite . despite the popularity of instruction and alignment tuning, their ability to alter fundamental information processing has been put in question in recent literature. @xcite investigated the effects of fine-tuning in toy models trained with formal languages as well as precompiled ones; their findings suggest that supervised fine-tuning does not introduce any new ability into pretrained models but only reinforces (or suppresses) existing ones. similar concerns have been raised upon investigating entity tracking in the neural representation space @xcite . @xcite identified multiple limitations of instruction tuning, including the inability to introduce new knowledge and deterioration of performance due to overreliance on pattern matching. (c) tarot learns a parametrized rotation operator r θ that rotates h to the direction of the desired token (red to blue). the intervention results in a change in the forward pass in (a) that outputs the desired token t ′ 6 . ## experiment setup training setting. @xcite previously found that token associations corresponding to pretrained knowledge primarily resides in the initial half of the model. since the rotational intervention designed in equations 6 and ?? are primarily targeted towards undesired token associations acquired through pretraining, we restrict l to the initial half only. therefore, the total number of parameters to optimise becomes dl 4 . since we want to optimise the rotation matrix for a particular task, only a small subset of training samples is required, i.e, 6 ≤ d @xmath0. models. six different instruction-tuned models with varying size are used for all experiments: qwen2-1.5b-instruct @xcite , phi-3-mini-4k-instruct @xcite (2.8 billion parameter), mistral-7b-instruct-v0.1 @xcite , meta-llama-3-8b-instruct @xcite , qwen2.5-14b-instruct @xcite , and qwen2.5-32b-instruct @xcite ; we refer to these models as qwen2-1.5b, phi-3-mini, mistral-7b, llama-3-8b, qwen2.5-14b, and qwen2.5-32b, respectively. tasks. we experiment with four different classification (i.e., single token generation) tasks and two nlg tasks. classification tasks used are as follows: ag news @xcite , entailed polarity @xcite , navigate @xcite , and winowhy (big-bench authors, 2023). the generation tasks used include imdb positive review @xcite evaluation metrics. for nlg tasks, imdb and detoxify, two different types of reward models are used. to calculate the fluency of the generated text, gpt4 @xcite is used as an oracle. for both the tasks, the average of fluency and the score from the reward models are reported. further details are present in appendix a.3 5 results tables 1 and 2 summarize the performance of various methods across classification tasks in zeroand 6-shot settings, respectively. eigen pruning is included only in zero-shot comparisons, per its original design. in-context examples vs. tarot. unlike eigen pruning (or even, traditional fine-tuning), tarot is optimized with a mixture of m-shot inference to avoid zero-shot bias. consequently, we can observe the improvement over the base model achieved via tarot while provided with in-context examples, except with mistral-7b on ag news and navigate (c.f. importance of rotation over rescaling attention heads. comparing tarot with the rotation ## conclusion in this work, we proposed tarot, a novel, gradientfree, mechanistic intervention method for editing language models. tarot builds on observations from implicit gradient descent bias of causal attention and applies parametrized rotation on the attention output to minimize the effects of undesired memorizations, doing away with effort-intensive localization steps and task-specificity of prior intervention techniques. using bayesian optimization of the rotational parameters, tarot renders as dataefficient as in-context learning; yet, across a variety of tasks and language models of different sizes and families, robust improvement is observed. in a nut-shell, tarot can pave the path for general-purpose model editing methods in the future beyond supervised fine-tuning.