Task-driven Layerwise Additive Activation Intervention (2025)

Abstract: Modern language models (LMs) have significantly advanced generative modeling in natural language processing (NLP). Despite their success, LMs often struggle with adaptation to new contexts in real-time applications. A promising approach to task adaptation is activation intervention, which steers the LMs’ generation process by identifying and manipulating the activations. However, existing interventions rely heavily on heuristic rules or require many prompt inputs to determine effective interventions. In this paper, we propose a layer-wise additive activation intervention framework that optimizes the intervention process, thereby enhancing sample efficiency. We evaluate our framework on various datasets, demonstrating improvements in the accuracy of pretrained LMs and competing intervention baselines.

## introduction transformer-based language models (lms) have revolutionized generative modeling for natural language processing (nlp). this is demonstrated by the impressive performances of lms in various important nlp tasks @xcite @xcite @xcite @xcite . one of such is in-context learning @xcite , where a pretrained lm can perform nlp tasks without fine-tuning their parameters. this is achieved by providing the model with prompts that include demonstrations of the task, allowing it to learn from the examples and make predictions without requiring additional training. despite this, performing icl on lms remains challenging, as lms still struggle to adapt quickly to new context shifts in real-time applications. one possible method for adaptation is activation intervention @xcite @xcite @xcite @xcite , where one uses the activations of the model that are most likely responsible for icl to steer the generation process. however, most of these works either derive the intervention based on a heuristic rule or require a large amount of prompt input. contributions. in this work, we aim to design a principled, optimization-based intervention that delivers competitive results with limited training demonstrations. we propose a layerwise additive activation intervention method for task-driven learning. the intervention is an optimal vector that minimizes the mismatch between the intervened decoding output and the target desired output in the training data. additionally, we impose a joint lasso and group lasso regularization to mitigate overfitting on the sample size and promote the component and head sparsity of the intervention. @xcite @xcite , which can negatively affect the effectiveness of the intervention at later layers due to the representation shifts of the activations generated at earlier layers. to address this issue, we propose to focus the intervention on the same layer, which can be easily formulated as a layerwise optimization problem. the layerwise optimization problem has shown effectiveness in driving the llm-generated content to human alignment @xcite . moreover, our intervention can facilitate task calculus by focusing on the same layer across tasks. by an additive composition of different task-specific interventions, we obtain a new intervention for the corresponding composition of tasks, as we will demonstrate in the numerical experiments. ## related works in-context learning. since its introduction by @xcite , icl in lm has been studied extensively in various directions. for example, @xcite ; @xcite analyzed the role of prompts in improving the icl performance. theoretical analysis of how lms perform icl has been proposed by @xcite ; @xcite @xcite . these works study the internal mechanism -either with regularized linear regression or gradient descent -of the transformer architecture, which is the workhorse behind most current state-of-the-art lms. ## methodologies we have a pre-trained decoder-only transformerbased lm (for example, llama3-8b) that is not yet fine-tuned for the few-shot in-context learning task (icl). the lm has l layers; each layer has h heads of dimension d; overall, the activation vector at each layer has a dimension @xmath0. we use ℓ ∈ {1, . . . , l} as the layer index, and use @xmath1, . . . , h} as the head index. for llama3-8b, we have @xmath2, @xmath3. we consider the layer-wise intervention consisting of finding a task-specific modification vector to be added to the activations of the input's last token so that the lm's output is steered toward our desired direction. to formalize this problem, we consider a task τ dataset consisting of n τ samples. each sample i, @xmath4, . . . , n τ , can be described by a tuple (s iτ , r τ , t iτ ), where s iτ is the input text, r τ is a special token padded to the end of the input, and t iτ is the desired (ground-truth) target output corresponding to the input s iτ . when there is no possible confusion, we will omit the task index τ to avoid cluttered notation. our method aims to find a task-specific ∆ from the training data. then, at inference time with a test input s test , we intervene by adding ∆ to the activations of the last token corresponding to the input (s test , r) to generate ttest . the success of the intervention is measured by the discrepancy in the test set between the generated output ttest and the true desired output t test . the last token's activations at layer ℓ of the input (s i , r) are denoted by a ℓ (s i , r); consequently, the additively-intervened activations become a ℓ (s i , r)+∆. the activations at the last layer (layer l) after the intervention become a l,∆ (s i , r). the decoder will transform a l,∆ (s i , r) into the distribution of the next token for generation. a good intervention vector ∆ should minimize the generation loss averaged over the training dataset simply minimizing (1) leads to overfitting: in general, the number of training samples n is small, while the dimension d of the vector ∆ is much larger (@xmath5. we propose using both lasso regularization and group lasso regularization to combat overfitting. thus, the intervention ∆ solves min here, a natural group assignment is by head, where we decompose ∆ = (∆ 1 , . . . , ∆ h ), where each ∆ @xmath6. the group lasso term penalizes the sum of the 2-norm of headwise interventions ∆ h . we choose group lasso regularization to promote sparsity within heads of activations, as empirical evidence from previous work such as @xcite ; @xcite and @xcite suggests that only a portion of attention heads is responsible for the transformer's ability to generate controllable outputs. the lasso penalty is also added to promote an additional degree of sparsity across all elements of ∆. next, we describe two specific applications of this task-driven intervention. ## numerical experiments we perform benchmarks to demonstrate our algorithm's performance on two tasks: rule understanding and opinion dynamics. all experiments are run on 4× nvidia a5000 gpus. our implementation will be published at https://github. com/hieunt91/layerwiseintervention.git ## conclusions in this paper, we propose and showcase an effective approach using layer-wise additive activation interventions to steer the output of lms. our approach effectively enhances the model performance by optimizing an intervention vector to minimize the mismatch between the intervened decoding output and the desired target output in the training data. additionally, incorporating both lasso and group lasso regularizations addresses overfitting and promotes sparsity in activation heads, ensuring efficient interventions. our evaluations on the rule understanding task and the opinion dynamic task demonstrate that this method significantly improves the performance of pre-trained lms across various tasks, outperforming existing intervention techniques.