Data Uniformity Improves Training Efficiency and More, with a Convergence Framework Beyond the NTK Regime

TL;DR

This paper shows that selecting more uniformly distributed data improves training efficiency and performance in neural networks, supported by a new convergence framework beyond the NTK regime applicable to various architectures.

Contribution

It introduces a theoretical framework for gradient descent convergence beyond NTK, linking data uniformity to training speed and accuracy, and validates findings with extensive experiments.

Findings

Uniform data selection accelerates training.

Smaller $h_{min}$ slows down gradient descent.

Maximizing pairwise data distance improves performance.

Abstract

Data selection plays a crucial role in data-driven decision-making, including in large language models (LLMs), and is typically task-dependent. Properties such as data quality and diversity have been extensively studied and are known to enhance model performance. However, it remains unclear whether there exist other quantitative and general principles of data selection that can consistently improve performance, especially for complicated tasks. In this paper, we demonstrate that selecting more uniformly distributed data can improve training efficiency while enhancing performance. Specifically, we establish that more uniform (less biased) distribution leads to a larger minimum pairwise distance between data points, denoted by $h_{\min}$, and prove that a smaller $h_{\min}$ can slow down the training dynamics of gradient descent (GD). Moreover, we theoretically show that the approximation error of neural networks decreases as $h_{\min}$ increases. Our analysis introduces a convergence framework for GD beyond the Neural Tangent Kernel (NTK) regime, applicable to a broad class of architectures, including transformers, without requiring Lipschitz smoothness. This framework further provides theoretical justification for the use of residual connection and function composition in deep neural architectures. In the end, we conduct comprehensive experiments for supervised fine-tuning across various settings, including different optimization strategies, model sizes, and training datasets. The results consistently demonstrate that selecting data by maximizing pairwise distance significantly accelerates training and achieves comparable or better performance in LLMs across diverse datasets. Code and Datasets are available at the link: https://github.com/SafeRL-Lab/data-uniformity.