Task-Specific Skill Localization in Fine-tuned Language Models (2302.06600v2)

Abstract: Pre-trained LLMs can be fine-tuned to solve diverse NLP tasks, including in few-shot settings. Thus fine-tuning allows the model to quickly pick up task-specific ``skills,'' but there has been limited study of where these newly-learnt skills reside inside the massive model. This paper introduces the term skill localization for this problem and proposes a solution. Given the downstream task and a model fine-tuned on that task, a simple optimization is used to identify a very small subset of parameters ($\sim0.01$% of model parameters) responsible for ($>95$%) of the model's performance, in the sense that grafting the fine-tuned values for just this tiny subset onto the pre-trained model gives performance almost as well as the fine-tuned model. While reminiscent of recent works on parameter-efficient fine-tuning, the novel aspects here are that: (i) No further re-training is needed on the subset (unlike, say, with lottery tickets). (ii) Notable improvements are seen over vanilla fine-tuning with respect to calibration of predictions in-distribution ($40$-$90$% error reduction) as well as the quality of predictions out-of-distribution (OOD). In models trained on multiple tasks, a stronger notion of skill localization is observed, where the sparse regions corresponding to different tasks are almost disjoint, and their overlap (when it happens) is a proxy for task similarity. Experiments suggest that localization via grafting can assist certain forms of continual learning.

Task-Specific Skill Localization in Fine-tuned LLMs

Overview

The paper "Task-Specific Skill Localization in Fine-tuned LLMs" addresses the phenomenon of skill acquisition in pre-trained LLMs following fine-tuning for NLP tasks. The focus is on identifying the specific subset of parameters responsible for performing well in a fine-tuned model, termed as skill localization. The approach proposed is intriguing, involving an optimization method that identifies these critical parameters, the so-called 'grafting region', which constitutes about 0.01% of the entire model's parameters yet accounts for most of the model's task performance.

Methodology

The core innovation lies in the concept of model grafting, where skill localization is achieved without additional re-training. The grafting involves copying the values of identified sparse parameters from the fine-tuned model to the pre-trained model, essentially localizing the skill. This results in nearly equivalent performance with significant improvements observed in calibration errors (40%-90% reduction) and out-of-distribution predictions without the susceptibility to catastrophic forgetting. The implications extend to multi-task and continual learning, where disjoint parameter subsets emerge for different tasks, indicating a potential method to transfer skills across related tasks.

Numerical Results

Key quantitative results include the findings that sparse graft regions consisting of only 0.01% of model parameters retain over 95% efficiency compared to fully fine-tuned models, with noticeable calibration and generalization benefits. For models optimized across multiple tasks, these regions show minimal overlap, suggesting a clear demarcation of task-specific parameters, a revelation suggesting compositional skill capabilities of models.

Implications and Future Directions

The implications of this research are profound in parameter-efficient model deployment, potentially reducing the computational and storage overhead required for fine-tuning. The observed calibration improvements may have practical applications in safety-critical NLP tasks, where confidence calibration is crucial. This work opens up new avenues in transferring learned capabilities through minimal parameter grafting, offering significant insights into model interpretability and explainability. The identified disjoint task-centric regions provide a new lens to explore modular task learning and potentially robust continual learning frameworks without forgetting learned skills. Future research could explore deeper the underlying mechanisms of skill localization and extend this method to other model architectures and domains beyond NLP to offer generalized solutions for efficient model training and deployment.