GradientSpace: Unsupervised Data Clustering for Improved Instruction Tuning

TL;DR

GradientSpace presents an unsupervised, gradient-space clustering framework for instruction tuning that uses an online SVD-based algorithm on full-dimension LoRA gradients to partition data into gradient-aligned clusters. Each cluster trains a dedicated LoRA expert, with a lightweight encoder-based router selecting the appropriate expert at inference, achieving single-expert routing that outperforms ensembles while reducing latency. The approach reduces gradient variance, accelerates SGD convergence, and yields consistent accuracy gains across math, code, finance, and creative-writing tasks, outperforming state-of-the-art gradient-based clustering and embedding-based baselines. This work demonstrates a scalable, data-driven method to mitigate gradient interference in heterogeneous instruction-tuning data and provides practical routing that scales to real-world deployments.

Abstract

Instruction tuning is one of the key steps required for adapting large language models (LLMs) to a broad spectrum of downstream applications. However, this procedure is difficult because real-world datasets are rarely homogeneous; they consist of a mixture of diverse information, causing gradient interference, where conflicting gradients pull the model in opposing directions, degrading performance. A common strategy to mitigate this issue is to group data based on semantic or embedding similarity. However, this fails to capture how data influences model parameters during learning. While recent works have attempted to cluster gradients directly, they randomly project gradients into lower dimensions to manage memory, which leads to accuracy loss. Moreover, these methods rely on expert ensembles which necessitates multiple inference passes and expensive on-the-fly gradient computations during inference. To address these limitations, we propose GradientSpace, a framework that clusters samples directly in full-dimensional gradient space. We introduce an online SVD-based algorithm that operates on LoRA gradients to identify latent skills without the infeasible cost of storing all sample gradients. Each cluster is used to train a specialized LoRA expert along with a lightweight router trained to select the best expert during inference. We show that routing to a single, appropriate expert outperforms expert ensembles used in prior work, while significantly reducing inference latency. Our experiments across mathematical reasoning, code generation, finance, and creative writing tasks demonstrate that GradientSpace leads to coherent expert specialization and consistent accuracy gains over state-of-the-art clustering methods and finetuning techniques.

Figures (6)

• Figure 1: GradientSpace accuracy compared to other methods on different benchmarks.
• Figure 2: GradientSpace framework. Stage I (LoRA warm-up): train a small LoRA adapter on a warm-up split to obtain meaningful, low-dimensional gradient representations. Stage II (online SVD clustering): estimate K via SVD on a validation gradient matrix, initialize centroids in the dominant gradient subspace, then refine them online with a cluster cache and EMA updates to form gradient-aligned partitions. Stage III (experts + router): finetune one LoRA expert per cluster and use a lightweight encoder-based router to select the most appropriate expert at inference.
• Figure 3: GradientSpace Router.
• Figure 4: Comparison of input embedding similarity and gradient similarity across 500 randomly selected instruction pairs from the MTEB STS benchmark dataset using LLaMA-2-7B.
• Figure 5: Cluster distribution for Data Mix.

Theorems & Definitions (3)

• Lemma 1
• Theorem 2: Variance Reduction
• Theorem 3: Asymptotic Stationarity Improvement