Learn More, Forget Less: A Gradient-Aware Data Selection Approach for LLM

TL;DR

GrADS introduces a gradient-aware data selection mechanism for efficient domain-specific fine-tuning of LLMs. By extracting Embedding and LM Head gradients after a brief one-epoch pre-training and employing a KDE-based PDF over the combined gradient signal, GrADS automatically selects a representative subset that improves domain understanding while substantially mitigating catastrophic forgetting. Across medicine, law, and finance, GrADS achieves strong domain performance with as little as 5–50% of the data and demonstrates robustness across model scales and architectures, including LoRA-tuned and full-parameter setups. The approach offers a cost-effective, scalable solution for domain adaptation in resource-constrained settings, with transferability to larger LLMs and future public release of code.

Abstract

Despite large language models (LLMs) have achieved impressive achievements across numerous tasks, supervised fine-tuning (SFT) remains essential for adapting these models to specialized domains. However, SFT for domain specialization can be resource-intensive and sometimes leads to a deterioration in performance over general capabilities due to catastrophic forgetting (CF). To address these issues, we propose a self-adaptive gradient-aware data selection approach (GrADS) for supervised fine-tuning of LLMs, which identifies effective subsets of training data by analyzing gradients obtained from a preliminary training phase. Specifically, we design self-guided criteria that leverage the magnitude and statistical distribution of gradients to prioritize examples that contribute the most to the model's learning process. This approach enables the acquisition of representative samples that enhance LLMs understanding of domain-specific tasks. Through extensive experimentation with various LLMs across diverse domains such as medicine, law, and finance, GrADS has demonstrated significant efficiency and cost-effectiveness. Remarkably, utilizing merely 5% of the selected GrADS data, LLMs already surpass the performance of those fine-tuned on the entire dataset, and increasing to 50% of the data results in significant improvements! With catastrophic forgetting substantially mitigated simultaneously. We will release our code for GrADS later.

Figures (4)

• Figure 1: Pilot study: From left to right on the x-axis, we sort the CMedQA training data by gradients from largest to smallest, and select 10% by rank at each time, conducting 10 subsets, and predict their responses with untuned LLMs.
• Figure 2: The illustration of the GrADS method.
• Figure 3: Experiments of fine-tuning Qwen1.5-7B, ChatGLM3-6B, Llama3-8B with subsets of different sizes selected from FinQA. Baselines 'base' and 'all' represent performances of the base models without SFT and the models fine-tuned on the entire data.
• Figure 4: Semantic distribution of training instances. The green dots indicate selected instances whereas the red dots indicate dropped instance.