Layer-Aware Task Arithmetic: Disentangling Task-Specific and Instruction-Following Knowledge (2025)

Abstract: Large language models (LLMs) demonstrate strong task-specific capabilities through fine-tuning, but merging multiple fine-tuned models often leads to degraded performance due to overlapping instruction-following components. Task Arithmetic (TA), which combines task vectors derived from fine-tuning, enables multi-task learning and task forgetting but struggles to isolate task-specific knowledge from general instruction-following behavior. To address this, we propose Layer-Aware Task Arithmetic (LATA), a novel approach that assigns layer-specific weights to task vectors based on their alignment with instruction-following or task-specific components. By amplifying task-relevant layers and attenuating instruction-following layers, LATA improves task learning and forgetting performance while preserving overall model utility. Experiments on multiple benchmarks, including WikiText-2, GSM8K, and HumanEval, demonstrate that LATA outperforms existing methods in both multi-task learning and selective task forgetting, achieving higher task accuracy and alignment with minimal degradation in output quality. Our findings highlight the importance of layer-wise analysis in disentangling task-specific and general-purpose knowledge, offering a robust framework for efficient model merging and editing.

## introduction existing large language models (llms) exhibit robust conversational abilities but often require finetuning for specific tasks. model merging integrates multiple fine-tuned models into a unified multi-task system. a popular merging strategy, task arithmetic (ta) @xcite , manipulates parameter differences (task vectors) from fine-tuning to add or remove task capabilities. fine-tuned models are typically derived from instruction-following llms @xcite , resulting in intertwined instruction-following and task-specific signals within task vectors, which destabilizes their utility ( effectively isolating task-specific segments from instruction-following components is challenging. we find that ta can decompose task vectors into layer-specific vectors, with similarity to instructionfollowing models indicating instruction-dominated (high similarity) versus task-specific (low similarity) layers. in addition, transformer layers naturally differ in captured information: lower layers encode general linguistic features aligned with instruction-following, while higher layers specialize in task nuances during fine-tuning, consistent with recent findings @xcite . we quantify these differences via cosine similarity between layer-specific and instruction vectors. following the above observation, we propose layer-aware task arithmetic (lata), which weights each task-vector layer based on alignment with task-specific capabilities. layers focusing on tasks receive higher weights; instruction-focused layers receive lower weights or are disregarded. lower layers generally reflect pre-training-based instruction-following; upper layers emphasize taskspecific signals. although layers may embed multiple competencies, our similarity analysis highlights segments most distinct from instruction-following baselines. unlike related methods @xcite , lata uniquely addresses both task learning and forgetting. experiments demonstrate that lata maintains model quality, enhances multi-task performance compared to existing methods, and effectively manages task forgetting by selectively removing undesired capabilities with minimal collateral impact. contribution we introduce lata, a method leveraging layer-wise analysis to selectively enhance task-specific segments and mitigate redundant instruction-following components. key novelties include: lata ensures high-quality merged models, improved multi-task performance, and efficient removal of unwanted capabilities. ## related work combining model capabilities without additional training has attracted growing attention. model merging fuses weights of separately fine-tuned models for multi-task learning @xcite , and simple averaging can improve accuracy and robustness @xcite . ties @xcite resets negligible changes to address sign conflicts, reducing performance drops; dare @xcite discards up to 99% of fine-tuning deltas to merge multiple homologous models. most research aims to minimize utility loss of merged llms @xcite @xcite @xcite @xcite an alternative line of research, task arithmetic (ta), views tasks as weight update vectors composed via vector operations. ilharco et al. (2023) define a task vector as the difference between a fine-tuned model and its base, enabling multiple tasks to be learned simultaneously and new tasks to be inferred without retraining. negating a task vector selectively unlearns a specific task with minimal impact on others, implying that model weights shift independently per task. ta has been considered in fine-tuning @xcite and alignment @xcite @xcite contexts. we focus on ta for both task learning and forgetting. existing methods generally merge or edit entire models without distinguishing which layers encode task-specific versus general knowledge. in contrast, lata performs a layer-wise analysis to separate generic utility from task-specific effects, enabling selective amplification or removal of tasks while preserving overall performance. ## background knowledge given θ pre as the weights of a pre-trained llm and θ ft as the parameters of the llm fine-tuned for a target task, ta @xcite proposes the following formula to obtain the task vector τ : in ta, task vectors for different target tasks can be added to a single model, enabling the model to simultaneously perform multiple target tasks. this achieves the effect of task learning: in the task forgetting, the task vector can be used to remove the model's ability for specific tasks: here, τ is the task vector for the task to be removed. ## proposed method here, we present our proposed method, layer-aware task arithmetic (lata). first, we define a base model as a model that does not possess instruction-following capabilities, such as llama-3-8b @xcite . we also define a pre-trained model as a model with instruction-following capabilities, such as llama-3-8b-instruct @xcite . moreover, for multiple target tasks, we obtain models that are fine-tuned from the pre-trained model for each specific task, resulting in models tailored to their respective tasks. we refer to these models as fine-tuned models, which are derived from the pre-trained model through fine-tuning. lata consists of the four steps below. step 1: deriving instruction vector and complex vector we define the instruction vector by subtracting the base model's parameters from the pre-trained model's parameters: this captures the instruction-following capability. we then define the complex vector by subtract-ing the base model's parameters from those of each fine-tuned model: this vector reflects both instruction-following and the target task capability. figure 2 shows how we obtain the instruction and complex vectors. step 2: computing layerwise similarity we split the instruction and complex vectors into layer vectors, with each layer's parameters forming a small vector. thus, the complete task vector is τ = {τ 1 , . . . , τ l }, where l is the number of layers. to isolate target-task elements in the complex vector from instruction-following elements, we compute the cosine similarity between the instruction and complex vectors at each layer: figure 3 illustrates that layers showing higher similarity primarily capture instruction-following capabilities. assigning smaller weights to these layers during ta reduces their impact on the merged model, preserving instruction-following quality. in contrast, layers with lower similarity have less effect on instruction following, so we assign them greater weights to boost target-task performance while maintaining overall utility. step 3: deriving pure vector we obtain the target-task vector τ by subtracting the pre-trained model's parameters from the fine-tuned model: next, we split τ into layer vectors and compute each layer's cosine similarity to the instruction and complex vectors. layers with higher similarity receive smaller weights, and those with lower similarity receive larger weights. the resulting weighted vector is called the pure vector because it preserves the task's core functionality. we propose three approaches to obtain this pure vector τ ′ . 1. linear-drop-by-rank: we rank each layer by its cosine similarity between the complex and instruction vectors, then assign weights from 0 to 1 based on rank: figure 3 : method for identifying important layers: we compute the cosine similarity of each layer vector between the complex vector and the instruction vector. layers with lower similarity are less related to instruction-following and likely enhance the target task, so we strengthen them. conversely, layers with higher similarity align more with instruction-following and have lower task relevance, so we attenuate them to reduce their impact on utility. here, r i is the rank, and higher ranks receive larger weights, indicating greater emphasis on the target task. this reduces weight differences among higherranked layers, better reflecting inter-layer correlations in some architectures. ## evaluation we conduct two experiments. the first is the task learning scenario (see section 3), merging three target tasks (unalignment, math, and code) into a single model via ta's additive operation. the second is the task forgetting scenario (see section 3), using ta's subtractive operation to reduce harmful content and improve alignment @xcite . all evaluations and generations were conducted using zero-shot, greedy search without additional randomness (e.g., using default temperature, top-p, or top-k sampling). as a result, the model always selects the token with the highest probability during generation, leading to identical outputs for the same prompt across different inference runs. we also ran multiple tests to confirm this consistency. throughout the lata experiments, we use cosine similarity to measure the similarity between two vectors, but we also conducted experiments on l2norm. the results are shown in appendix a.4. all experiments were conducted on an nvidia h200 gpu with 141gb of memory and dual in-tel® xeon® platinum 8480c processors (112 cores, 2.00-3.80 ghz). ## conclusion in this work, we introduced a novel approach (lata) to ta, demonstrating its effectiveness in merging and fine-tuning llms across diverse tasks. lata leverages dynamic task representations to achieve improved alignment and utility without compromising model performance. through extensive experiments on benchmark datasets such as wikitext-2, gsm8k, and humaneval, we showed that our approach consistently outperforms existing methods like dare and ties in balancing taskspecific performance and generalization. notably, our framework enables efficient model merging while mitigating interference between tasks, as evidenced by superior results in multi-task scenarios. our findings highlight the potential of ta as a scalable and adaptable solution for optimizing llms in multi-task and cross-lingual settings.