GRACE: A Dynamic Coreset Selection Framework for Large Language Model Optimization

Tianhao Tang; Haoyang Li; Lei Chen

cs.DBApr 9, 2026

GRACE: A Dynamic Coreset Selection Framework for Large Language Model Optimization

Authors: Tianhao Tang, Haoyang Li, Lei Chen

Organizations: CSE, HKUST, Hong Kong SAR · Computing, PolyU, Hong Kong SAR · DSA, HKUST (GZ)&HKUST, Guangzhou, China

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities in natural language understanding and generation. However, their immense number of parameters and complex transformer-based architectures result in significant resource demands and computational complexity during training, making it challenging to optimize them efficiently on large datasets. To reduce training costs while preserving performance, researchers have investigated coreset selection techniques, which aim to identify small, representative subsets of the entire training dataset to accelerate LLM training. However, existing coreset selection methods fail to adapt to the dynamic nature of LLM training and often struggle with scalability for models of this size. To address these limitations, we propose a graph-guided adaptive and dynamic coreset selection framework for LLMs, namely GRACE. GRACE dynamically constructs and updates coresets by combining representation diversity with gradient-based importance metrics, ensuring both informativeness and efficiency. To mitigate the computational cost of frequent updates, GRACE leverages a $k$ -NN graph-based propagation mechanism and selectively updates scores and embeddings, adapting to evolving training dynamics. Extensive experiments on three benchmarks demonstrate that GRACE significantly improves training efficiency and downstream performance across diverse LLMs and tasks.

Figures & tables

Notation	Description
$X,x_{i},T$	Input sequence, $i$ -th token, sequence length
$\mathbf{X},\mathbf{x}_{1}$	Initial hidden states, $i$ -th token’s embedding
$\mathbf{H}^{(l)}_{i}$	Hidden states of $X_{i}$ after $l$ -th layer
$\bar{\mathbf{H}}^{(l)}_{i}$	Mean hidden states of $X_{i}$ along sequence
$\mathbf{Z}^{(l)},\mathbf{\hat{O}}$	MHSA output at $l$ -th layer, predicted logits
$\hat{P}_{\boldsymbol{\theta}}(X)$	Predicted probabilities of sample $X$

Table 1. Important notations.

Figure 1. Overview of GRACE. Stage 1: Hidden states and importance scores of all training samples are extracted from a warm-up trained model. A

k

-NN graph is constructed from these indicators. Stage 2: Sample a subset of samples from the graph and select the coreset, then train the model with the coreset. Stage 3: At checking steps, if significant changes in importance scores are detected, the graph will be updated. A subset is selected for recomputation and updated scores and embeddings are propagated through the graph structure. The

k

-NN graph is also updated and then reused in the next selection.

Figure 2. Overview of the graph update process. From left to right: we begin with a

k

-NN graph constructed from initial representations and EL2N scores (blue nodes). When adaptive update checking is triggered, a subset of recently trained samples (green) is loaded and we use their history scores as new scores. Another small number of samples (purple) are then selected for recomputation. Combining the history scores and recomputed scores as updated scores, they are then propagated to their neighbors through the graph structure (blue to green), allowing approximate updates without full recomputation.

Type	Methods	Phi-2		Llama2-7B		Qwen2.5-7B
Type	Methods	In-domain	Out-domain	In-domain	Out-domain	In-domain	Out-domain
-	Base (Pretrained Model)	19.66	22.96	3.39	2.56	47.73	58.39
Static	Random	48.35	50.79	22.66	17.57	60.66	63.39
	Middle Perplexity	47.46	47.55	22.90	18.33	61.02	62.07
	Facility Location	46.16	48.89	21.01	15.54	60.66	65.26
	DivIR	47.58	52.16	23.33	21.90	61.34	64.23

Table 2. MathInstruct results with 10% training budget. We report in-domain (GSM8K, MATH, NumGlue) and out-of-domain (SVAMP, DeepMind, SimuLeq) average accuracies (%).

Figure 3. Time Comparison

Figure 4. Experiment results under different budgets on MathInstruct

Model	Phi-2		Llama-2
Model	In	Out	In	Out
G $\cdot$ \I	48.49	52.98	22.78	18.40
G $\cdot$ \R	50.23	54.82	21.67	21.22
G $\cdot$ \Update	48.79	51.54	23.42	23.62
G $\cdot$ \AdaUpdate	50.08	53.03	23.19	22.57
GRACE	50.54	55.27	23.89	24.05

Table 3. Ablation Experiments to test the Effect of Scores. In means in-domain tasks, and Out means out-of-domain tasks.

Type	Models	Llama2-7b			Phi-2			Qwen2.5-7b
Type	Selection Budget	10%	20%	30%	10%	20%	30%	10%	20%	30%
-	Base (Pretrained Model)	19.70			34.77			23.62
Static	Random	32.57	33.38	33.47	33.63	23.13	20.72	38.10	37.03	36.85
	Middle Perplexity	32.62	32.81	33.10	31.24	23.36	16.77	38.98	38.43	38.07
	Facility Location	32.97	33.02	32.85	35.57	37.85	35.16	39.41	37.19	37.62
	DivIR	31.83	33.36	33.34	36.91	33.93	27.43	38.47	38.47	38.75

Table 4. ROUGE-L results on BioInstruct

Type	Models	Llama2-7b			Phi-2			Qwen2.5-7b
Type	Selection Budget	10%	20%	30%	10%	20%	30%	10%	20%	30%
-	Base (Pretrained Model)	15.92			22.24			16.97
Static	Random	37.38	37.98	38.86	36.06	37.47	37.67	37.93	38.99	39.28
	Middle Perplexity	37.12	38.35	38.55	36.35	37.23	37.68	37.63	39.07	39.19
	Facility Location	37.46	38.52	39.21	36.22	36.76	38.44	38.12	39.28	38.99
	DivIR	37.52	38.54	39.11	37.22	38.78	39.13	39.61	40.14	40.05

Table 5. ROUGE-L results on DialogSum

Coreset Selection

Coreset selection chooses compact subsets that preserve a dataset’s learning, evaluation, or optimization value while reducing cost.

See all 59 papers

GRASS: Gradient-based Adaptive Layer-wise Importance Sampling for Memory-efficient Large Language Model Fine-tuning
Apr 9, 2026Kaiyuan Tian, Yu Tang, Gongqingjian Jiang +5LLM Fine-Tuning Fine-Tuning
National University of Defense Technology · Information Support Force Engineering University
BLADE: Scalable Bi-level Adaptive Data Selection for LLM Training
Jun 17, 2026Jiaxing Wang, Deping Xiang, Jin Xu +9Training Data Selection Bilevel Optimization
University of Oxford
EDGC: Entropy-driven Dynamic Gradient Compression for Efficient LLM Training
Nov 13, 2025Qingao Yi, Jiaang Duan, Jun Zhang +5Communication-Efficient Distributed Training Gradient Compression
University of Shanghai for Science and Technology, Shanghai, China · Shanghai Jiao Tong University, Shanghai, China · Zhejiang University, Hangzhou, China

Apr 9, 2026J/K move · Enter open · S save

GRASS: Gradient-based Adaptive Layer-wise Importance Sampling for Memory-efficient Large Language Model Fine-tuning

Kaiyuan Tian, Yu Tang, Gongqingjian Jiang, Baihui Liu+4

National University of Defense Technology · Information Support Force Engineering University

LLM Fine-Tuning Fine-Tuning Memory-Efficient Fine-Tuning Importance Sampling

Open paper AskPDF arXiv

Full-parameter fine-tuning of large language models is constrained by substantial GPU memory requirements. Low-rank adaptation methods mitigate this challenge by updating only a subset of parameters. However, these approaches often limit model expressiveness and yield lower performance than full-parameter fine-tuning. Layer-wise fine-tuning methods have emerged as an alternative, enabling memory-efficient training through static layer importance sampling strategies. However, these methods overlook variations in layer importance across tasks and training stages, resulting in suboptimal performance on downstream tasks. To address these limitations, we propose GRASS, a gradient-based adaptive layer-wise importance sampling framework. GRASS utilizes mean gradient norms as a task-aware and training-stage-aware metric for estimating layer importance. Furthermore, GRASS adaptively adjusts layer sampling probabilities through an adaptive training strategy. We also introduce a layer-wise optimizer state offloading mechanism that overlaps computation and communication to further reduce memory usage while maintaining comparable training throughput. Extensive experiments across multiple models and benchmarks demonstrate that GRASS consistently outperforms state-of-the-art methods, achieving an average accuracy improvement of up to 4.38 points and reducing memory usage by up to 19.97%.

Abstract

Coreset Selection

Explore similar work

GRASS: Gradient-based Adaptive Layer-wise Importance Sampling for Memory-efficient Large Language Model Fine-tuning

BLADE: Scalable Bi-level Adaptive Data Selection for LLM Training

EDGC: Entropy-driven Dynamic Gradient Compression for Efficient LLM Training