ClusterUCB: Efficient Gradient-Based Data Selection for Targeted Fine-Tuning of LLMs

TL;DR

ClusterUCB introduces an efficient gradient-based data selection method for fine-tuning large language models, utilizing clustering and a modified UCB algorithm to reduce computational costs while maintaining performance.

Contribution

The paper proposes a novel data selection framework combining clustering and a modified UCB algorithm to improve efficiency in gradient-based data influence approximation.

Findings

Achieves comparable fine-tuning results with reduced computational resources.

Effectively balances exploration and exploitation during data selection.

Demonstrates efficiency across various benchmark datasets.

Abstract

Gradient-based data influence approximation has been leveraged to select useful data samples in the supervised fine-tuning of large language models. However, the computation of gradients throughout the fine-tuning process requires too many resources to be feasible in practice. In this paper, we propose an efficient gradient-based data selection framework with clustering and a modified Upper Confidence Bound (UCB) algorithm. Based on the intuition that data samples with similar gradient features will have similar influences, we first perform clustering on the training data pool. Then, we frame the inter-cluster data selection as a constrained computing budget allocation problem and consider it a multi-armed bandit problem. A modified UCB algorithm is leveraged to solve this problem. Specifically, during the iterative sampling process, historical data influence information is recorded to directly estimate the distributions of each cluster, and a cold start is adopted to balance exploration and exploitation. Experimental results on various benchmarks show that our proposed framework, ClusterUCB, can achieve comparable results to the original gradient-based data selection methods while greatly reducing computing consumption.