PaperlayerView exact revision on arXiv

Source-linked AI summary

Clustering-Based Balanced Sampling and Allocation with Data Parallelism for High-Performance Fine-Tuning

Hyunjin Kim, Youngeun Nam, Jaemin Han, Wonhyeok Choi, Jae-Gil Lee

arXiv:2609.12584v1cs.LGcs.AI

TL;DR

Large, imbalanced instruction-tuning datasets can leave informative samples undercovered and cause redundant computation in data-parallel training. CluSTER clusters samples in gradient space, balances selection and worker allocation, and uses weighted updates to preserve the original distribution. Across multiple instruction-tuning datasets, it reports 69.6% less training time with almost no performance degradation compared with full-dataset tuning.

  • Problem

    Large-batch fine-tuning of long-tailed instruction data can overrepresent dominant groups, under-cover informative samples, and spend computation on redundant training signals.

    Source: §1 · p. 1

  • Method

    CluSTER clusters samples in gradient space, balances cluster sampling across workers and within clusters, and uses importance-weighted updates to preserve the original data distribution.

    Source: §1 · p. 2

  • Results

    69.6% less training time with almost no degradation in model performance across multiple instruction-tuning datasets versus full-dataset instruction tuning.

    Source: §1 · p. 2

  • Takeaways & Limitations

    CluSTER enables faster training with competitive performance under reduced data budgets while improving batch balance across GPUs and reducing selection preprocessing time.

    Source: §6 · p. 9

    Evaluation is limited to single-node multi-GPU training and moderately sized instruction-tuning datasets, without extensive multi-node or substantially larger-corpus experiments.

    Source: Limitations · p. 10

Abstract

from arXiv · show

Instruction-tuning datasets for large language models (LLMs) are often large, redundant, and imbalanced, limiting efficient adaptation. Naive large-batch fine-tuning repeatedly includes overrepresented sample groups while weakly covering underrepresented but informative ones, especially under data parallelism (DP) across multiple GPUs. We propose CluSTER, a Cluster-aware balanced Sampling framework for Training Efficient data Reduction in DP instruction tuning. CluSTER curates a representative reduced dataset through gradient-space clustering and DP-aware balanced allocation, ensuring dual-level coverage across clusters and workers, while preserving the original data distribution by weighted update. As a result, CluSTER reduces redundant computation and improves training stability without compromising model quality. Across multiple instruction-tuning datasets, CluSTER reduces training time by up to 69.6% with almost no accuracy loss compared to prior sampling and data reduction methods. Code is available at https://github.com/kaist-dmlab/CluSTER.

1 Introduction

Large, heterogeneous instruction-tuning datasets can waste large-batch computation on redundant signals and leave informative groups undercovered, especially across data-parallel workers. CluSTER uses gradient-space clustering and balanced sampling to improve coverage while weighted updates preserve the original distribution; it reports faster training with almost no performance degradation.

1 Introduction

  • Large, heterogeneous instruction datasets can overrepresent dominant regions, leaving informative samples undercovered and spending computation on redundant signals.Under data parallelism, dominant groups may be assigned to multiple workers, repeating gradient computations without proportional convergence gains.Source: §1 · p. 1
  • Dual-level coverage requires different sample groups across workers and diverse, complementary sample types within each group.This organizes selected samples for distributed mini-batch optimization rather than ranking samples only by individual informativeness.Source: §1 · pp. 1–2
  • CluSTER clusters samples by gradient similarity, then balances cluster sampling across workers and selects diverse samples within clusters.Clusters are groups of samples providing similar training signals; the approach uses importance-weighted updates to preserve the original data distribution.Source: §1 · p. 2
  • 69.6% less training time with almost no degradation in model performance across multiple instruction-tuning datasets versus full-dataset instruction tuning.The reported gain comes from replacing repeated updates from redundant samples with more complementary updates across workers.Source: §1 · p. 2

2 Related Work

Prior sampling work addresses imbalance, sample utility, or diversity, while data-parallelism research has emphasized system efficiency. CluSTER is positioned as combining sample selection with worker-aware allocation.

2.1 Data Sampling

  • Prior sampling methods include random, uniform or stratified, utility- or curriculum-based, difficulty-based, and representation- or gradient-space selection.Some LLM fine-tuning methods, including LESS and S2L, estimate influence or use smaller-model training trajectories, but can require costly preprocessing.Source: §2.1 · p. 2
  • CluSTER differs from example-level selection by using target-model gradient-proxy clusters for joint data-parallel worker allocation and cluster-weighted updates.Figure 2 contrasts random sampling, which can repeat semantic clusters across batches, with uniform sampling, which encourages balanced cluster coverage.Source: §2.1 · p. 2, Figure 2 · p. 3

2.2 Data Parallelism (DP)

  • Data-parallelism research has improved scalability, memory footprint, and communication overhead through architectures such as parameter servers, ZeRO, and FSDP.In data parallelism, each worker processes a different mini-batch and gradients are aggregated into a global update.Source: §2.2 · p. 3

3 Preliminary

The analysis motivates balanced cluster sampling through gradient-variance comparisons, then extends the argument to non-uniform clusters using weighted stratification. In that setting, the estimator remains unbiased but does not guarantee lower total variance.

3 Preliminary

  • The analysis defines a dataset of K equal-sized clusters and decomposes total gradient variability into intra-cluster and inter-cluster components.The inter-cluster component measures dispersion of cluster-wise mean gradients around the population mean.Source: §3 · p. 3
  • The mini-batch gradient estimator has expectation equal to the population mean gradient, with sampling ratio f = B/N.The analysis relates gradient-estimate variance to convergence, with lower variance contributing to faster and more stable convergence.Source: §3 · p. 3
  • The framework overview presents gradient-based clustering, cluster-aware balanced sampling, and weighted gradient update as its three stages.Together, these stages are described as ensuring balanced data allocation.Source: Figure 3 · p. 4
  • In the idealized balanced-cluster setting, uniform sampling reduces gradient variance by removing the inter-cluster variance component when that component is positive.Random sampling can over-sample dominant regions and under-sample minority ones, producing high mini-batch gradient variance.Source: §3 · pp. 3–4
  • For non-uniform clusters, weighted stratified sampling remains unbiased but does not guarantee lower total variance than random sampling.It removes variance due to random fluctuations in cluster composition.Source: §3 · p. 4
  • CluSTER's default allocation uses m_c = B/K and w_c = Kπ_c, corresponding to the weighted stratified estimator for practical non-uniform cluster distributions.Here π_c is the original proportion of cluster c, and the cluster means are computed from m_c sampled items.Source: §3 · p. 4

4 CluSTER Framework

CluSTER reduces redundant data through gradient-based semantic clustering, cluster-aware balanced sampling, and weighted gradient updates. Its design combines worker-level and within-cluster coverage with weights that preserve clusters’ contributions under the original data distribution.

4 CluSTER Framework

  • CluSTER balances selection between clusters to mitigate overfitting while preserving the intrinsic data distribution, using clustering, balanced sampling, and weighted gradient updates.The framework has three stages: gradient-based semantic clustering, cluster-aware balanced sampling, and weighted gradient update.Source: §4 · p. 4

4.1 Gradient-Based Semantic Clustering

  • Gradient-proxy embeddings combine final-layer hidden states with token-level prediction confidence, then K-means groups samples with similar update directions.The number of clusters is set to the number of GPUs, and capacity-constrained K-means limits cluster size using α.Source: §4.1 · pp. 4–5

4.2 Cluster-Aware Balanced Sampling

  • CluSTER balances the sample budget across clusters, downsamples larger clusters to the smallest cluster, and retains farthest-from-centroid samples within each cluster.For selection ratio r, it retains the top r · nmin farthest samples per balanced cluster; peripheral selection is the default because its expected gradient diversity exceeds random and core-centric selection.Source: §4.2 · p. 5

4.3 Weighted Gradient Update

  • Cluster weights proportional to original cluster size restore each cluster’s contribution to the global gradient after balancing changes the data distribution.This is intended to maintain consistency with the original unbalanced distribution while retaining the variance-reduction benefits of balanced sampling.Source: §4.3 · p. 5

4.4 Data Selection Complexity Analysis

  • CluSTER’s lightweight gradient-proxy embeddings yield lower selection-time complexity than IFD, LESS, and S2L by avoiding expensive backpropagation or proxy-model training.The embeddings use final-layer hidden states and token-level probabilities.Source: §4.4 · p. 6

5 Experiments

Experiments evaluate CluSTER across code and medical instruction tuning, GPU and model scales, and analyses of selection efficiency, hyperparameters, weighting, and coverage. Across settings, it maintains competitive performance with reduced data and training time.

5.2.1 Results on Code Generation; 5.2.2 Results on Medical QA

  • At 59% data usage, CluSTER scores 56.5 average versus 57.6 for full-data random; at 31% medical data, it remains comparable to full-data training while cutting time 69.6%.The code comparison reports average scores; the medical comparison reports performance comparable to full-data training.Source: Table 2 · p. 6, Table 3 · p. 6, §5.2.1 · p. 7

5.2.3 Analysis of Selection Strategy Efficiency

  • CluSTER has the shortest selection time among compared methods, requiring one forward pass rather than the additional passes or trajectory collection described for baselines.The comparison covers data-selection, training, and total model-building time.Source: §5.2.3 · p. 7, Table 4 · p. 7

5.3 Analysis of GPU Scalability

  • On eight GPUs, CluSTER remains competitive with full-data training at 41.6% data usage and improves over random sampling at matched data ratios.The passage also reports consistent gains over random sampling at a 31.2% data budget.Source: §5.3 · p. 7, Table 5 · p. 7

5.4 Hyperparameter Sensitivity Analysis

  • CluSTER degrades gradually as selection ratio falls, remains robust across cluster counts, and stays stable across capacity settings, including the default α = 1.5.The reported sensitivity analysis covers selection ratio r, cluster count K, and capacity α.Source: §5.4.1 · p. 8, §5.4.2 · p. 8, §5.4.3 · p. 8

5.5 Analysis of Model Scalability

  • CluSTER remains effective on CodeLlama-Python-13B while using reduced training data, with additional 13B medical results reported in the appendix.The main-text experiment uses Magicoder-OSS-Instruct-75K.Source: §5.5 · p. 8, Table 6 · p. 8

5.6 Effect of Weighting Mechanism

  • Across data ratios, weighting samples by cluster size improves overall scores and robustness, helping preserve each cluster’s contribution during optimization.The comparison evaluates the weighting mechanism across different data ratios.Source: §5.6 · pp. 8–9, Table 7 · p. 8

5.7 Analysis of Intra-Level Coverage

  • Peripheral selection performs best across most data ratios and has the highest gradient diversity, supporting stronger within-cluster coverage than random or core-centric selection.The reported ordering is peripheral, random, then core-centric selection for both performance and gradient diversity.Source: §5.7 · p. 9, Table 8 · p. 9

5.8 Analysis of Inter-Level Coverage

  • CluSTER lowers inter-worker gradient similarity; its clusters align with task categories at 82% annotation agreement, and intra-cluster actual-gradient similarity exceeds inter-cluster similarity (p = 0.0464).The actual-gradient difference is statistically significant under a cluster-size-preserving permutation test.Source: §5.8 · p. 9, Figure 6 · p. 9

6 Conclusion

The paper studies sampling for efficient and stable data-parallel fine-tuning with redundant, imbalanced instruction data. It proposes CluSTER and reports faster training with competitive performance under reduced data budgets.

6 Conclusion

  • CluSTER combines gradient-space clustering, cluster-aware balanced sampling, and weighted updates for data-parallel fine-tuning.The approach aims to reduce redundant computation and stabilize multi-GPU optimization.Source: §6 · p. 9
  • CluSTER improves batch balance across GPUs while preserving the original data distribution, and reduces preprocessing time for data selection.Source: §6 · p. 9
  • The results support structured data sampling as a direction for variance-aware data parallelism, with faster training and competitive performance under reduced data budgets.Source: §6 · p. 9

Limitations

The evaluation demonstrates the sampling strategy in multi-GPU data-parallel training, while leaving larger-scale settings unexplored.

Limitations

  • Evaluation is limited to single-node multi-GPU training; multi-node environments with larger batches and higher communication overhead are not extensively explored.The experiments scale up to eight GPUs and a batch size of 512.Source: Limitations · p. 10
  • Computational constraints limit experiments to moderately sized instruction-tuning datasets, excluding corpora with hundreds of thousands or millions of samples.Source: Limitations · p. 10
  • The authors identify significantly larger batch regimes and datasets as directions for future investigation.Source: Limitations · p. 10

Ethics Statement

The paper documents research-purpose use of existing artifacts, licensing, and data handling; the supplied appendix passages also outline notation and sampling setups for variance analysis.

Ethics Statement

  • The experiments use public datasets and benchmarks for research, do not collect new user data, and do not redistribute original datasets or model weights.The authors rely on dataset creators’ preprocessing and release procedures for personal or offensive content.Source: Ethics Statement · p. 10
  • The authors reviewed artifact licenses, including the Llama 2 Community License for pretrained models and specified licenses for datasets and benchmarks.The supplied passages name MIT, CC BY-NC-SA 4.0, Apache-2.0, and CC BY 4.0 among those licenses.Source: Ethics Statement · p. 10

A.1 Notations

  • The theoretical notation treats the dataset as N samples in K equal-sized clusters and defines sample gradients, global and cluster means, and intra-cluster variance.Table 9 summarizes the notation.Source: §A.1 · p. 13, Table 9 · p. 13

A.2 Intra- and Inter-Cluster Decomposition

  • The decomposition passage states that total sum of squares splits into intra-cluster and inter-cluster parts.Source: §A.2 · p. 13

A.3 Proof of Proposition 1

  • The random-sampling proof sets up selection of B distinct samples from N and uses indicators to represent whether each sample is selected.The passages give the indicator’s expectation, variance, and covariance and note an approximation for large N.Source: §A.3 · p. 13

A.4 Proof of Proposition 2

  • The uniform-sampling proof draws B/K samples per cluster and forms a weighted gradient estimate across clusters.The supplied passage states that sampling across clusters is independent.Source: §A.4 · p. 13

A.5 Proof of Proposition 3

The proof compares the variance of stratified and random gradient estimators, finding stratified variance no greater and identifying when the inequality is strict.

A.5 Proof of Proposition 3

  • The comparison is between the random-sampling and uniform-sampling estimator variances.The proof passage describes taking the difference between the two expressions.Source: §A.5 · p. 14
  • Stratified sampling has variance no greater than random sampling, with strict improvement when inter-cluster variance is positive.The inequality becomes equality if and only if inter-cluster variance is zero.Source: §A.5 · p. 14
  • Uniform stratification requires prebalanced clusters and equal sampling weights, which can distort the original distribution and increase redundancy in heterogeneous datasets.The authors describe this setup as difficult to apply to real-world heterogeneous data.Source: §A.5 · p. 14

B Gradient Diversity Analysis under Distance-Based Selection

The appendix explains gradient diversity and weighted cluster sampling, validates centroid preservation after pruning, and outlines CluSTER’s algorithm and experimental settings.

B Gradient Diversity Analysis under Distance-Based Selection

  • Gradient diversity decreases as pairwise gradient inner products increase, so less coherent gradient directions correspond to greater mini-batch diversity.The analysis models within-cluster gradients as a cluster mean plus isotropic residual noise assumed small relative to the mean.Source: B Gradient Diversity Analysis under… · p. 14

C.2 Empirical Validation of the Centroid Preservation Assumption

  • Selected-subset gradient centroids have cosine similarity consistently above 0.999 to full-cluster centroids, supporting the assumed preservation of cluster gradient direction.The authors attribute the small centroid shift to redundancy near cluster centroids and selected samples spanning cluster structure.Source: Table 10 · p. 15, §C.2 · p. 16

C.1 Derivation of the Weighting Coefficient

  • Balanced sampling draws equal numbers from clusters, then weights their contributions to compensate for cluster-size differences and preserve the original data distribution.The derivation assumes expected gradient direction within each cluster is approximately preserved after selection.Source: §C.1 · p. 15, §C.2 · p. 15

C.2 Empirical Validation of the Centroid Preservation Assumption

  • The centroid validation compares gradient centroids from each full cluster and its pruned subset to test the direction-preservation assumption.The results support the weighting derivation’s approximation of each cluster’s original gradient contribution.Source: §C.2 · pp. 15–16

D Algorithm of CluSTER

  • CluSTER’s algorithm builds gradient-proxy embeddings, clusters them with capacity constraints, and selects the farthest-from-centroid samples up to the smallest cluster size.The selected subsets are then pruned to the target ratio and assigned to GPUs for data-parallel training.Source: §C.2 · p. 16

E.1 Implementation Details

  • Across experiments, each configuration ran with three random seeds; for Magicoder-OSS-Instruct-75K, training used four NVIDIA H100 GPUs for 2 epochs with PyTorch DDP and global batch size 512.Source: §E.1 · pp. 16–17

E.2 Dataset Details

The appendix reports dataset and benchmark statistics, then presents an additional dataset comparison and ablations of weighting, selection metric, and batch size.

E.2 Dataset Details

  • Dataset statistics cover training-example counts, effective data usage after pruning and sampling, evaluation benchmarks, and repeated-run settings.The experiments use official benchmark splits when available.Source: §E.2 · p. 17

F.1 Experiment Result on Additional Dataset

  • CluSTER consistently outperforms random sampling on Evol-Instruct-Code-80K across reported data ratios, while training time falls with reduced data usage.The table reports Pass@1 for CodeLlama-7B; performance degradation remains limited as pruning increases.Source: Table 12 · p. 17, §F.1 · p. 17

F.2 Effect of Weight Scaling

  • Removing cluster weighting causes performance drops across all data ratios, while applying weighting yields better results on every reported evaluation metric.The authors describe weighting as important for stable training and preserving informative samples during selection.Source: §F.2 · p. 18, Table 13 · p. 18

F.3 Effect of Data Selection Metric

  • Peripheral selection performs best among peripheral, random, and core-centric within-cluster selection strategies on Evol-Instruct-Code-80K.The passage attributes its advantage to selecting boundary samples after cluster-level diversity is ensured.Source: §F.3 · p. 18, Table 14 · p. 18

F.4 Effect of Batch Size

  • CluSTER remains stable across tested global batch sizes and consistently outperforms random sampling with the selection ratio fixed at 1.00.Source: §F.4 · p. 18

F.5 Additional 13B-Scale Evaluation

The additional 13B-scale evaluation compares CluSTER with random sampling and full-data training, and reports manual evidence that gradient-based clusters align with task categories. Representative cluster examples span several programming and reasoning task types.

F.5 Additional 13B-Scale Evaluation

  • At an effective data ratio of 65.9%, CluSTER scores 52.6 on average, versus random sampling (51.7), and approaches full-data training (53.1).The comparison evaluates Llama-2-13B on MedInstruct-52K.Source: Table 16 · p. 18, §F.5 · p. 18

F.6.1 Statistical Analysis of Cluster Semantics.

  • 328 of 400 inspected samples (82.0%) match their cluster’s dominant task category, indicating meaningful task-level semantic alignment.The evaluation sampled 400 instances and manually assessed cluster assignments against dominant task categories.Source: §F.6.1 · p. 18, Table 17 · p. 19

F.6.2 Representative Examples by Cluster

  • The four clusters cover algorithms and data structures, SQL and front-end interaction, mathematical pattern matching, and basic programming grammar.Examples were grouped using k-means with k = 4 on embeddings described as capturing semantic and uncertainty information.Source: §F.6.2 · p. 19
Loading 2609.12584v1…