Balancing Speciality and Versatility: a Coarse to Fine Framework for Supervised Fine-tuning Large Language Model (2024)

Abstract: Aligned Large Language Models (LLMs) showcase remarkable versatility, capable of handling diverse real-world tasks. Meanwhile, aligned LLMs are also expected to exhibit speciality, excelling in specific applications. However, fine-tuning with extra data, a common practice to gain speciality, often leads to catastrophic forgetting (CF) of previously acquired versatility, hindering the model’s performance across diverse tasks. In response to this challenge, we propose CoFiTune, a coarse to fine framework in an attempt to strike the balance between speciality and versatility. At the coarse-grained level, an empirical tree-search algorithm is utilized to pinpoint and update specific modules that are crucial for speciality, while keeping other parameters frozen; at the fine-grained level, a soft-masking mechanism regulates the update to the LLMs, mitigating the CF issue without harming speciality. In an overall evaluation of both speciality and versatility, CoFiTune consistently outperforms baseline methods across diverse tasks and model scales. Compared to the full-parameter SFT, CoFiTune leads to about 14% versatility improvement and marginal speciality loss on a 13B model. Lastly, based on further analysis, we provide a speculative insight into the information forwarding process in LLMs, which helps explain the effectiveness of the proposed method. The code is available at https://github.com/rattlesnakey/CoFiTune.

## introduction aligned llms mainly undergo a two-step procedure: initial pre-training on web-scale text corpora, followed by fine-tuning on diverse instructions to align with human intentions. they exhibit remarkable versatility, showcasing their ability to handle various real-world tasks, such as reasoning, common sense question-answering, and instruction following @xcite @xcite . despite the versatility, aligned llms still fall short in certain tasks or domains, such as mathematics @xcite , finance @xcite , and law @xcite . to bolster performance in these particular tasks or domains, i.e., to gain speciality, a common practice is fine-tuning. however, during the fine-tuning process, the modification of model parameters often leads to catastrophic forgetting (cf), thereby causing a noticeable loss of versatility @xcite . this loss adversely affects the performance of fine-tuned models across various real-world tasks @xcite , propelling several works to investigate and contribute solutions to the cf in llm versatility @xcite . as a relatively new problem, the cf in llms remains under-explored. we categorize the existing studies into regularization-based, weight-based, and architecture-based methods foot_0 . regularizationbased methods @xcite add extra terms into loss function to penalize parameter changes. weight-based methods @xcite introduce weight coefficients for parameters to regulate their updates. architecture-based methods @xcite design and exclusively fine-tune extra modules outside the original model. however, these methods still have limitations: 1) regularization-and weight-based methods involve fine-tuning all parameters, which poses a significant challenge to preventing cf in versatility. particullary, cf deteriorates as the training iteration proceeds @xcite . 2) architecturebased methods only update the inserted parameters, which blocks the learning of speciality, particularly on challenging tasks @xcite . these limitations raise a key question: how can the model gain speciality while preserving the versatility, thereby enhancing its overall abilities? (shown in fig. 1 ). recent research threads highlight two pivotal findings: the presence of redundant parameters in the model @xcite and the distinct role played by different modules and layers @xcite . these suggest that it is highly feasible to enhance speciality by updating specific parameters within a defined layer range while keeping the remainder frozen to maintain versatility. 2 therefore, we propose a coarse to fine framework, i.e., cofitune. at the coarse-grained level, we perform an empirical tree-search algorithm to identify and update specific modules within a defined layer range that are crucial for speciality without significant versatility penalty. simultaneously, the remaining parameters are frozen to further preserve versatility. subsequently, recognizing that not all units 3 in a module are evenly important for the versatility @xcite , we utilize a fine-grained soft-masking mechanism to regulate the backward gradient flow based on their importance values for versatility. this further mitigates the cf issue while not jeopardizing the speciality. to summarize, our contributions are as follows: ## the framework we start with the task formulation (sec. 3.1) and the backbone we used in this work (sec. 3.2). then, we outline our cofitune framework, covering: 1) an empirical tree-search algorithm to identify modules within a defined layer range that balances speciality and versatility effectively (sec. 3.3); 2) a fine-softmask mechanism regulating the backward gradient flow based on units' importance for versatility to further mitigate the cf issue (sec. 3.4). ## experiment in this section, we conduct extensive experiments to answer the following questions: 1) rq1: training which layers and modules can optimize the trade-off between speciality and versatility? 2) rq2: how does our cofitune framework compared to the baseline methods? 3) rq3: can the fine-softmask mechanism further alleviate the cf issue in versatility? 4) rq4: can our cofitune be seamlessly applied to models across different model scales and families? we also conduct further experimental analysis to delve into more insights for this field in sec. 4.4. ## conclusion in this work, we strive to optimize the balance between llm's speciality and versatility. our proposed cofitune framework employs an empirical tree-search exploration algorithm to identify the module within a defined layer range that is crucial for gaining speciality without significantly affecting versatility. additionally, the fine-softmask mechanism is applied to further alleviate cf issue without impairing the speciality. we introduce a chinese cf setting to advance the research in chinese domain. the experimental results demonstrate that our cofitune outperforms all baselines across various tasks and model scales. our in-depth analysis provides valuable insights for future @xcite . this decision stems from the uncertainty surrounding the replay data ratio and strategy in this domain. increasing the number of replay samples during fine-tuning would also incur higher training resource costs. moreover, we believe that the rehearsal-based method can be integrated with our proposed cofitune framework, and we plan to explore this in future work. additionally, the model employed in this study is mainly based on the llama architecture. this choice is driven by the fact that llama family models demonstrate remarkable dominance and performance in the current nlp research area (constituting over 90% on the llm leaderboard 12 ). besides, based on the fact that most current llms share the same architecture as llama, i.e., transformer decoder-only architecture, we are confident that our identified universal optimal range can also be applied to a majority of models within the same architecture (the experiment results of bloomz in rq4 have validated this point). therefore, our cofitune framework can be directly adopted in these models using insight 1, without incurring any tree-search exploration overhead. furthermore, we plan to explore other model architectures (e.g., encoder-decoder) in future research.