Investigating Layer Importance in Large Language Models (2024)

Abstract: Large language models (LLMs) have gained increasing attention due to their prominent ability to understand and process texts. Nevertheless, LLMs largely remain opaque. The lack of understanding of LLMs has obstructed the deployment in safety-critical scenarios and hindered the development of better models. In this study, we advance the understanding of LLM by investigating the significance of individual layers in LLMs. We propose an efficient sampling method to faithfully evaluate the importance of layers using Shapley values, a widely used explanation framework in feature attribution and data valuation. In addition, we conduct layer ablation experiments to assess the performance degradation resulting from the exclusion of specific layers. Our findings reveal the existence of cornerstone layers, wherein certain early layers can exhibit a dominant contribution over others. Removing one cornerstone layer leads to a drastic collapse of the model performance, often reducing it to random guessing. Conversely, removing non-cornerstone layers results in only marginal performance changes. This study identifies cornerstone layers in LLMs and underscores their critical role for future research.

## introduction the rapid advancement of large language models (llms) has revolutionized natural language processing, enabling unprecedented capabilities in text generation, translation, and comprehension taskscitewei2022chain, hu2021lora, rafailov2024direct, ouyang2022training. these models, exemplified by architectures such as gpt-3 @xcite , llama @xcite , and bloom @xcite , rely on transformer-based neural networks with numerous layers @xcite . despite their successes, llms suffer from issues such as hallucinations, biases, and unstable reasoning abilities @xcite @xcite . regardless of the effort to mitigate these issues @xcite @xcite , they remain unsolved nowadays, hindering the deployment of llms in more safety-critical domains. when a neural network makes errors or underperforms, it is valuable to identify the specific part of the model responsible for these issues. therefore, understanding the inner workings of neural networks and recognizing the role of individual components is key to addressing the challenges associated with llms. in this paper, we advance the understanding of llms by investigating the importance of individual layers in llms across multiple tasks. to quantify the contribution of each layer to the overall model performance, we extend the shapley value framework @xcite , originally from cooperative game theory, to layers in llms. we employ an efficient sampling method to estimate layer importance within a practical runtime. to further analyze the impact of the key layers characterized by high shapley values, we perform layer ablation to observe a specific layer's impact on performance. our study reveals that certain early layers in llms, which we term cornerstone layers, play a dominant role in influencing the model's performance. notably, removing one of these cornerstone layers can cause a significant performance drop, reducing the model performance to near random guessing. in contrast, removing other layers typically results in only marginal performance degradation. we hypothesize that these cornerstone layers handle some fundamental tasks in llms and hope this discovery inspires future studies on understanding the role of cornerstone layers. ## related work there is a significant body of research focused on interpreting and understanding large language models (llms). this section provides an overview of some key approaches. ## estimating layer shapley prior works usually calculate shapley values between different input features or different data points. in a nutshell, shapley values are usually applied to data, not models. nevertheless, in this work, we adopt the well-established shapley value framework to measure the contribution of a layer to the model performance. for a model with a sequential structure, we can construct its mathematical form in nested functions: one major drawback of calculating shapley values is the enormous number of required samples. to calculate shapley values for n players precisely, one needs to sample 2 n times, which is computationally not feasible for llms. therefore, we aim for a reasonable estimation of layer shapley with efficient sampling that exploits the structure of llms and achieves orders of magnitude speed-ups, which we explain below. early truncation: as discussed in prior work that estimates shapley values @xcite , the model performance degrades to random guessing for cases where many layers are removed. consequently, the value difference will be neighborhood sampling: besides early truncation, we leverage the sequential structure of the model to perform efficient sampling. each layer primarily influences its immediate subsequent layers and is influenced by its immediate preceding layers. consequently, interactions between distant layers are weaker than those between closely positioned layers. to capture meaningful interactions with fewer samples, we implement neighborhood sampling that only samples subsequent layers for shapley value estimation. formally, a set s of n elements under neighborhood sampling has the form @xmath0, a + 1, . . . , a + (n -1)}, where a is an offset value. ## mechanistic interpretation via layer-wise ablation one limitation of the shapley value is that, on its own, it does not provide a concrete understanding of how the model's performance degrades when layers are removed. aside from using shapley values to quantify layer contribution considering layer-wise interactions, it is also important to understand the functional significance of individual layers. specifically, we perform layer ablation for this endeavor. layer ablation involves selectively removing a target layer from the model and observing the resulting impact on performance across various tasks. this approach helps us isolate and evaluate the unique contribution of that specific layer, independent of others. in conventional llm architectures, skip connections are employed in each layer, as discussed in section 3.1 and equation 1. skip connections, which bypass one or more layers, allow information to be transferred directly from one layer to another non-adjacent layer. hence, it is possible to remove a layer without entirely disrupting the flow of information through the network. without skip connections, the removal of a layer would likely result in a complete breakdown of the model, as the information flow would be interrupted. therefore, we perform layer ablation by removing one layer while keeping the skip connection around the removed layer to maintain the information flow within the model. for a module with skip connection in the form of: removing this module results in an ablated structure in the form of: here, the effect of the module is completely removed, but the information processed by previous layers can still pass to subsequent layers. an illustration of the single-layer ablation is shown in figure 1 . layer ablation complements shapley values by providing a mechanistic perspective on the contribution of each layer. while shapley values offer a mathematical framework to understand the importance of each layer in the context of all possible self-attention ffn self-attention ffn transformer block i transformer block i + 1 decoder-only language models embedding figure 1: illustration of single-layer ablation. a layer is ablated by removing the block while keeping the skip connection across the layer. we choose to ablate layers we used for layer shapley calculation, that are, attention layers and ffn layers. for mixtral 8x7b, we ablate attention layers and moe layers. more details can be found in section 5. ## experiments 6.1 experimental setup we evaluate the models on various datasets to ensure a comprehensive analysis of a wide spectrum of language understanding and reasoning tasks. in our study, we utilize three recent large language models to assess the impact of individual layers: llama3-8b, llama3-70b, and mixtral-8x7b @xcite . llama3-8b contains 8 billion parameters, making it a mid-sized model suitable for a range of nlp tasks. llama3-70b have 70 billion parameters and are significantly larger than llama3-8b. this model is expected to capture more complex language patterns and dependencies. mixtral-8x7b replaces ffn layers with mixture-of-expert (moe) layers, each containing 8 experts. the ensemble approach aims to combine the strengths of multiple models to achieve superior performance. we perform our experiment on 6 datasets ranging from simple to hard tasks. boolq @xcite is a reading comprehension dataset consisting of questions that can be answered with "yes" or "no" based on a given passage. arc-easy and arc-challenge @xcite are part of the ai2 reasoning challenge (arc), which provides multiple-choice questions derived from science exams. piqa @xcite as- openbookqa @xcite comprises questions that require knowledge from elementary science topics, testing the model's ability to combine factual knowledge with reasoning skills. ## conclusion and future work in this study, we investigated the significance and contribution of individual layers in llms using shapley values and layer ablation. our results based on shapley values revealed that certain layers, typically early in the model, exhibit a dominant contribution to the model performance, which we term as cornerstone layers. layer ablation experiments demonstrated that removing a single cornerstone layer can cause the model to collapse and perform random guessing, highlighting their critical role. conversely, removing other non-cornerstone layers resulted in marginal performance changes, indicating redundancy in the model architecture. future works can continue the study on layer importance in groups of layers instead of one single layer. investigating the specific reasons behind the importance of cornerstone layers could provide deeper insights into llm functionality and inspire newer llm structures that promote model transparency, remove redundant parts, and improve inference efficiency.