Layer by Layer: Uncovering Where Multi-Task Learning Happens in Instruction-Tuned Large Language Models (2024)

Abstract: Fine-tuning pre-trained large language models (LLMs) on a diverse array of tasks has become a common approach for building models that can solve various natural language processing (NLP) tasks. However, where and to what extent these models retain task-specific knowledge remains largely unexplored. This study investigates the task-specific information encoded in pre-trained LLMs and the effects of instruction tuning on their representations across a diverse set of over 60 NLP tasks. We use a set of matrix analysis tools to examine the differences between the way pre-trained and instruction-tuned LLMs store task-specific information. Our findings reveal that while some tasks are already encoded within the pre-trained LLMs, others greatly benefit from instruction tuning. Additionally, we pinpointed the layers in which the model transitions from high-level general representations to more task-oriented representations. This finding extends our understanding of the governing mechanisms of LLMs and facilitates future research in the fields of parameter-efficient transfer learning and multi-task learning. Our code is available at: https://github.com/zsquaredz/layer_by_layer/

## introduction while pre-trained llms exhibit impressive performance across diverse tasks and demonstrate remarkable generalization capabilities @xcite @xcite @xcite , the representations they learn and the task-specific information encoded during pre-training remain largely opaque and unexplored. recent research has investigated fine-tuning strategies to adapt llms to specific tasks, including supervised fine-tuning on task-specific datasets and instruction tuning @xcite ; chung 1 our code is available at: https://github.com/ zsquaredz/layer_by_layer/ task 1 task t ... transformer layer 1 transformer layer 9 ... transformer layer 10 transformer layer 15 ... transformer layer 16 transformer layer 32 ... shared layers transition layers refinement layers output figure 1 : an illustration of our findings using the llama 2 7b model @xcite as an example. we show that when instruction tuning on t different tasks, the layers are divided into three functional sections: the shared layers (layers 1 to 9) form general representations shared among all tasks, the transition layers (layers 10 to 15) transition the representations into task-specific information, and the refinement layers (layers 16 to 32) continue to refine the representations toward specific tasks. @xcite . while these approaches have shown promising results in tailoring llms for improved task performance, a comprehensive understanding of their impact on the learned representations is still lacking. in this study, we perform a set of analyses to investigate task-specific information encoded in pre-trained llms and the effects of instruction tuning on their representations. the analysis leverages a sub-population analysis technique called @xcite , which provides an alternative to traditional probing methods for analyzing model representations within specific sub-populations of the training data. mossa involves comparing two models: a control model trained on the data relevant to the sub-population of interest (e.g., a particular task), and an experimental model that is identical to the control model but is also trained on additional data from different sources (e.g., multiple tasks). by analyzing the representational differences between these models, we can isolate the task-specific information encoded within the control model for the sub-population of interest. to compare the representations learned by different llm variants, we leverage the center kernel alignment (cka; @xcite metric. cka measures the alignment between representations in a kernel space, providing a robust measure of similarity that is insensitive to scaling and centering. by using mossa and cka, we investigate the following research questions: through an extensive analysis spanning over 60 diverse nlp tasks following the flan framework @xcite , we shed light on the underlying mechanisms that govern the encoding and adaptation of task-specific information within llms under instruction tuning. a key finding of our work is the identification of three functional groups of layers: a) shared layers, in which more general information is learned and shared across tasks; b) transition layers, in which task-specific information is intensified; c) refinement layers, in which the llms continue to refine their representations towards task-specific predictions. our findings contribute to a deeper understanding of the inner workings of llms and hold promising implications for future research in parameter-efficient fine-tuning (peft), multi-task learning (mtl), and model compression, benefiting a wide range of nlp applications. we structure this study as follows: §2 describes our methodology for our analysis, while §3 outlines the experimental setup and tools used to train and analyze our llms. §4 then attempts to answer each of the research questions outlined above by presenting and analyzing our results. finally, in §5, we summarize our key findings and discuss their potential implications. ## methodology we use the mossa framework introduced by @xcite . unlike standard probing methods @xcite , which build a model to predict a downstream task for quantifying encoded information, mossa compares representations from two models: a control model trained on data of interest and an experimental model trained on additional data from different sources. here, the data of interest refers to tasks. probing methods, while useful, can be limited because they rely on different metrics to evaluate performance across various tasks, making it challenging to directly compare the amount of information stored about tasks as diverse as sentiment analysis and translation. mossa, on the other hand, circumvents this issue by comparing the latent representations of models rather than their downstream performance metrics. mossa calculates the similarity between the representations of the control and experimental models, thus representing the information captured from the relevant sub-population of data through their latent representations. by comparing different models to each other, we can learn what information is captured when a subset of the data is used versus the whole dataset. we use matrix analysis to compare representation similarity between the experimental model, such as pre-trained, instruction-tuned, and corresponding single-task control models trained on individual tasks. intuitively, a high similarity between the experimental and control models indicates the experimental model stores task-specific information learned by the control model, which was finetuned solely on data from that task. the similarity is measured using the cka metric, which quantifies the similarity between two representations in a kernel space. formally, let [t ] be an index set of tasks, and let e be the experimental model and c t be the control model for task @xmath0. we assume a set of inputs @xmath1, where each x @xmath2,1 , . . . , x t,n } represents a set of input instructions for task t. for simplicity, we assume that all sets have the same size n, although this is not a strict requirement. 2 for each @xmath3, we apply the experimental model e and the control model c t to the input instruction x t,i to obtain two corresponding representations y t,@xmath4,@xmath5, respectively. here, d is the dimension of the experimental model representations, and d t is the dimension of the control model representations for task t. to obtain the representations y t,i and z t,i , we use the last token representation following previous work @xcite , as llms are decoder-only and the last token captures all input information. these representations can be extracted from any layers of the respective models. by stacking these vectors into two matrices for each task t, we obtain the paired matrices y @xmath6. we calculate the cka value between y t and z t following the procedure: and k @xmath7, respectively, using the same kernel function (e.g., linear, gaussian, or polynomial). foot_1 while other similarity metrics like svcca @xcite exist, they have a limitation due to the constraint of being invariant to invertible linear transformations, which requires the number of data points to be greater than the number of representation dimensions. we use cka as it has shown robust results when the data sample is smaller @xcite , as is sometimes the case for datasets used in our work. our method provides an approach to quantify the task-specific information encoded in the representations of llms. by comparing the experimental model's representations with those of singletask control models, we can gain insights into the extent to which the experimental model captures task-specific knowledge and how this knowledge is distributed across its representations. ## experimental setup data we use the flan 2021 dataset @xcite to fine-tune our llms. the flan dataset is a comprehensive collection of more than 60 nlp datasets, including both language understanding and generation tasks. these datasets are organized into twelve task clusters, where datasets within a given cluster belong to the same task type. to enhance instruction diversity, we follow the approach of @xcite and use ten unique natural language instruction templates for each dataset. these templates provide varying descriptions of the task to be performed. our instruction tuning pipeline combines all datasets and randomly samples from each dataset during training. to mitigate the impact of dataset size imbalances, we limit the number of training examples per task cluster to 50k and use the examples-proportional mixing scheme @xcite with a mixing rate maximum of 3,000 per task. this means that no task receives additional sampling weight for examples in excess of 3,000. we provide further details about the dataset in appendix a. ## experiments and results to shed light on the underlying mechanisms of mtl @xcite in llms, we start by examining what nlp tasks are encoded in the pretrained llm representations, establishing a baseline for comparison with the instruction-tuned model ( §4.1). then, using matrix analysis methods, we contrast the representational properties of the pre-trained and instruction-tuned llms to understand the effects of instruction tuning ( §4.2, 4.3, and 4.4). finally, we evaluate the generalization of our findings to unseen tasks ( §4.5). ## discussion our study offers comprehensive insights into the impact of instruction tuning on the representations learned by llms. previous work has discussed the benefits of instruction tuning @xcite @xcite , but ours is the first to analyze their effects from a representational perspective. our analysis revealed that llms instruction-tuned on multiple tasks learned different representations in the lower layers compared to llms tuned on individual tasks. similar to mtl, such representations can be shared and used across tasks @xcite . our analysis uncovered a key novel finding -we observed clear differences between pre-trained and instruction-tuned models, with the most significant representational transformations occurring in the middle transitional layers. this finding highlights the critical role of middle layers in encoding the specialized task knowledge induced by instruction tuning. similarly, previous studies in multilingual settings have also identified languageneutral transformations in the middle layers of the network @xcite . furthermore, our analysis suggests that in the refinement layers, instruction-tuned models continue to shape representations toward specific tasks but without substantial representational changes with respect to task-specific information. overall, our finding about functionality for different layers in llms generally aligns with previous findings on bert, which have shown that lower layers are more general, while upper layers are known to be more task-specific @xcite . our correlation analysis also revealed insights into the relationship between representations and task complexity. instruction-tuned models exhibited a positive correlation with reading complexity measures in the transition and refinement layers, suggesting better encoding of task-specific information for tasks with more specific vocabularya capability not observed in pre-trained models. notably, instruction tuning enabled models to preserve and enhance task-specific information across a broader range of layers, as evidenced by higher cka similarities compared to control models. our evaluation of unseen tasks further underscored the benefits of instruction tuning for improving generalization, with instruction-tuned models outperforming their pre-trained counterparts in deeper layers responsible for encoding complex task knowledge. this aligns with empirical evidence from @xcite but also highlights how representational changes facilitated by instruction tuning strengthen cross-task transfer capabilities. ## conclusion our study used several analyses to investigate how instruction tuning shapes representations in llms. these analyses revealed that unlike the pre-trained llm (llama 2), the instruction-tuned model (llama 2-sft) retained a high amount of task-specific information for all tasks from the middle layers onward. moreover, we were able to map the layers of llama 2-sft into three groups based on their functionality: shared layers (layers 1-9), transition layers (10-15), and refinement layers (16-32). in addition to expanding our understanding of llms, such mapping can greatly benefit future research in the fields of peft, mtl, and model compression. we also demonstrated that our mapping does not generalize to unseen tasks, revealing that a potential additional reason for the strong generalization capabilities of instruction-tuned models to unseen tasks can be related to their multi-task nature of producing more general representations.