Model Balancing Helps Low-data Training and Fine-tuning (2024)

Abstract: Recent advances in foundation models have emphasized the need to align pre-trained models with specialized domains using small, curated datasets. Studies on these foundation models underscore the importance of low-data training and fine-tuning. This topic, well-known in natural language processing (NLP), has also gained increasing attention in the emerging field of scientific machine learning (SciML). To address the limitations of low-data training and fine-tuning, we draw inspiration from Heavy-Tailed Self-Regularization (HT-SR) theory, analyzing the shape of empirical spectral densities (ESDs) and revealing an imbalance in training quality across different model layers. To mitigate this issue, we adapt a recently proposed layer-wise learning rate scheduler, TempBalance, which effectively balances training quality across layers and enhances low-data training and fine-tuning for both NLP and SciML tasks. Notably, TempBalance demonstrates increasing performance gains as the amount of available tuning data decreases. Comparative analyses further highlight the effectiveness of TempBalance and its adaptability as an “add-on” method for improving model performance.

## introduction recent surges in foundation models (fms) have stimulated research on aligning pre-trained models with specialized domains using small-sized datasets. this "pre-train and fine-tune" paradigm is prevalent in natural language processing (nlp) tasks @xcite @xcite . it is also gaining popularity in other machine learning (ml) fields, such as scientific machine learning (sciml) @xcite @xcite @xcite . from a practical perspective, the challenge of fine-tuning often lies in curating highquality datasets (possibly with labeled examples) to achieve alignment with the new domain. in sciml, people often use fms for training on different types of partial differential equations (pdes) @xcite @xcite and fine-tuning it on a certain domain when accessible scientific data from that domain is limited. as a concrete example, turbulence simulations at extremely high reynolds numbers are computationally intensive and time-consuming, often leading to only a few available trajectories. therefore, training sciml fms on trajectories with different reynolds numbers and fine-tuning it on trajectories simulated at extremely high ones is beneficial for solving the problem of poor training performance caused by insufficient data volume. using sciml fms, researchers can train these models to generalize across a wider range of downstream tasks, thereby enhancing their applicability and efficiency in diverse scientific scenarios. prior research has shown that strong performance can indeed be achieved by fine-tuning with a few carefully selected examples @xcite , but training with low data can still lead to unstable performance @xcite . therefore, finding fine-tuning algorithms that improve performance in low-data settings, especially few-shot alignment, becomes crucial. in this work, we draw inspiration from heavy-tailed self-regularization (ht-sr) theory @xcite , to improve model performance in low-data regimes. ht-sr theory proposes that well-trained neural network (nn) models exhibit strong correlations in weights, resulting in a heavy-tail (ht) structure in the empirical spectral density (esd, usually represented by a histogram of eigenvalue distribution) of each layers' weight matrix. to quantify the ht structure, we can fit a power law (pl) distribution to the ht part of the esd and extract figure 1 : heavy-tail esd analysis and tempbalance learning rate schedule. to characterize the heavy-tailed structure of esd, we fit a power-law exponent pl_alpha_hill on the tail part of the esds (blue histograms at top left), shown as the red dashed line on the histogram. given the imbalanced layer-wise pl_alpha_hill (bottom left), tempbalance assigns lower learning rate to layers with lower pl_alpha_hill (more heavy-tailed), and assign higher learning rate to layers with higher pl_alpha_hill (less heavy-tailed). tempbalance aims to balance the pl_alpha_hill distribution across layers in low-data regimes (bottom right). using ht-sr theory, we analyze the limitations of model training in low-data regimes by measuring the layer-wise distribution of pl_alpha_hill (discussed in 4.2). our main finding is that when we train with sufficient data, pl_alpha_hill becomes more evenly distributed across layers, resulting in better layer-wise balance; in this case, high performance can be achieved without layerspecific manipulations. however, when we reduce the number of training data samples, test performance decreases, and the standard deviation (std) of pl_alpha_hill across layers tends to increase (see figure 2 ), indicating that pl_alpha_hill is more unevenly distributed when training with fewer data, resulting in worse layer-wise balance. this finding indicates that different layers' training quality becomes more poorly aligned as we reduce training data. therefore, layer-wise balancing is beneficial to balance under-trained layers and over-trained layers in low data regimes. motivated by this observation, we incorporate the variance of pl_alpha_hill across layers with the recently proposed layer-wise learning rate scheduling algorithm tempbalance @xcite , to design a novel method to balance the training quality across layers. to evaluate its empirical performance, we use tempbalance in curated low-data regime in llm fine-tuning and sciml tasks. we compare tempbalance with commonly used baseline methods and demonstrate that tempbalance not only achieves superior performance in low-data training and finetuning, but also can be used as a plug-in method on top of existing optimization methods to achieve even better test performance and stability, such as sam @xcite and adafactor @xcite . furthermore, in our analysis, we reveal that tempbalance successfully balances training quality across all layers during training from the ht-sr point of view. we show that tempbalance balances the training quality of each layer by reducing the std of pl_alpha_hill of all layers. we summarize our contributions as follows foot_0 : simsekli et al. (2020); hodgkinson et al. (2022); wang et al. (2024a) proved generalization bounds dependent on the ht distributions in either model weights or the esds of the weight matrices, which are validated through extensive experiments. motivated by these studies, some efforts have begun to leverage the degree of ht for model training @xcite @xcite , model selection @xcite , and model compression @xcite , as well as to enhance model robustness @xcite . resource-constrained fine-tuning. the pretraining and fine-tuning paradigm has been a primary method for adapting foundation models to downstream tasks for resource-limited users. when adapting very large models, people often resort to the low-rank adaptation method (lora) @xcite , which is also considered in this paper. our primary focus is on low-data fine-tuning, an increasingly studied paradigm where the emphasis is often on careful data selection @xcite . furthermore, when training models in a few-shot fashion, such as in-context learning @xcite @xcite , data selection plays a crucial role in improving model performance. our paper, however, explores layer-balancing schemes to improve model performance. ## methodology in this section, we first revisit ht-sr theory and important ht-sr metrics related to model performance. then, we discuss tempbalance @xcite , which works well on different model architectures based on "shape metrics" from ht-sr theory. ## empirical results in this section, we employ ht metrics to diagnose model performance in data-limited regimes and demonstrate the effectiveness of tempbalance in addressing data limitation in two fields: nlp and sciml. in section 4.1, we describe our experimental setup. in section 4.2, we study the correlation between esd behaviors and model performance with limited training data. then, in section 4.3, we evaluate tempbalance in our experimental setup. in section 4.4, we compare our methods with other optimization baselines. we analyze the experimental results in section 4.6. finally, we perform ablation studies in section 4.7. ## conclusions in this work, we leverage ht-sr theory to diagnose the limitations of low-data training and improve the learning rate scheduling algorithm tempbalance to balance the training quality of different layers in low-data regimes. our extensive experiments demonstrate that tempbalance effectively balances layer-wise training quality and improves performance in nlp fine-tuning and sciml training. our analysis reveals that tempbalance achieves greater performance gain as we train with fewer data. furthermore, the compatibility of tempbalance makes it possible to add tempbalance to existing optimization methods, bringing further performance improvements. we show that ht-sr theory brings useful guidance in low-data training and fine-tuning, and we expect it to be a more generalized toolbox for diagnosing model performance in more training scenarios.