When Parts Are Greater Than Sums: Individual LLM Components Can Outperform Full Models (2024)

Abstract: This paper studies in-context learning by decomposing the output of large language models into the individual contributions of attention heads and MLPs (components). We observe curious components: good-performing ones that individually do well on a classification task, even when the model performs poorly; bad-performing ones that do much worse than chance; and label-biased components that always predict the same label. We find that component accuracies are well-correlated across different demonstration sets and perturbations of prompt templates. Based on our findings, we propose component reweighting, which learns to linearly re-scale the component activations from a few labeled examples. Given 24 labeled examples, our method improves by an average of 6.0% accuracy points over 24-shot ICL across 8 tasks on Llama-2-7B. Overall, this paper both enriches our understanding of ICL and provides a practical method for improvement by examining model internals.

## introduction the rapid progress in large language models (llms) has popularized prompting, which guides llms to perform tasks with instructions or examples. notably, in-context learning (icl; @xcite adapts llms to a new task using only a few labeled examples without parameter updates. however, how llms react to the in-context examples is sometimes unintuitive @xcite . recently, @xcite and @xcite find that even for instruction-tuned @xcite or very large models, adding a space or newline in prompts can greatly affect accuracy. we look into the llm internals to understand what causes the surprising behavior across various icl settings. our work stands in contrast to prior studies, which often treat llms as black boxes and alter either the input @xcite or output @xcite ; ## decomposing the transformer in icl we introduce a new view of in-context learning by decomposing the transformer architecture @xcite . our decomposition is exact-a mathematically equivalent formula for the model's outputs-and enables us to analyze model internals without training additional parameters (unlike, e.g., probing). we first discuss what our new view offers over the standard view of icl, and then walk through the mathematical details. ## characterizing components for icl we conduct in-context learning across 8 classification tasks on 4 llms: llama-2-7b, llama-2-13b, mistral-instruct-7b, and llama-3-8b. icl is sensitive to prompts, so we randomly sample 5 disjoint sets of demonstrations formatted with 3 templates and report the standard deviation across the 15 runs. to avoid majority and recency biases @xcite , each prompt consists of the same number of demonstrations from every class in shuffled order. we use @xmath0. except for §5.1, we refer to @xmath1, 4} without further notice. we sample 2000 examples with balanced labels as the test set for every task. please see a.1 for details about the tasks and templates. ## transferability of components we observe moderate to high component transferability across demonstrations, minimally contrastive templates, and data distributions, whereas there is little transferability across randomly sampled templates. our decomposition uncovers hidden abilities of individual components when the full model performs poorly. ## related work and discussion improving icl. prior work shows that icl performance varies greatly across different choices of demonstrations and templates @xcite . specifically, @xcite and @xcite find no universally better prompt template that can transfer across tasks and models, implying that it is not easy to explain icl through prompt engineering. while several approaches, such as prompt selection @xcite @xcite , prompt ensemble @xcite @xcite , and many-shot icl @xcite , substantially improve accuracy, they treat llms like black boxes without understanding the internals. besides, they greatly increase inference time or require a large set of labeled data, which deviates from true few-shot learning @xcite . in comparison, our paper studies this problem by looking inside the llms. rather than selecting prompts, we select components in a soft, learnable way. our method only requires {12, 24} examples and has negligible computation overhead over 4-shot icl at inference. ## conclusion we introduce a new perspective of icl via decomposing the model output into the sum of individual contributions of components. we then identify three types of component characteristics across 3 llms and 8 classification tasks. our extensive analyses reveal consistency in component accuracy across prompts and suggest the promising direction of improving icl by selecting components. to this end, we propose component reweighting, which learns to scale components differently on few-shot examples. our method achieves the best average accuracy compared to prior methods. we hope this work can deepen our grasp of llms while motivating more methods for practical use. 10288 our component reweighting method requires a small set of labeled data d train to train the component weights w. however, we believe it is not unreasonable to have at least @xmath2. on the other hand, we do not compare with fine-tuning-based baselines, such as lm-bff @xcite , t-few @xcite , because they usually require a larger gpu memory for training and more sophisticated early stopping criteria to prevent overfitting on few-shot examples. another limitation is that we only experiment with classification tasks for ease of evaluation. we leave it for future work to generalize our method to generation tasks by doing decomposition and reweighting at every token during generation. despite similarities in model decomposition, the focus of this paper is not circuits in llms @xcite . thus, we only have limited experiments towards mechanistic understanding of the curious components in §3.3 and a.7. unlike prior work that uses synthetic tasks to testify whether a head attends to certain tokens @xcite , we work on standard nlp benchmark datasets without obviously correct or incorrect tokens to collect answers, making mechanistic interpretation more challenging.