Model Balancing Helps Low-data Training and Fine-tuning

TL;DR

A recently proposed layer-wise learning rate scheduler, TempBalance, is adapted, which effectively balances training quality across layers and enhances low-data training and fine-tuning for both NLP and SciML tasks.

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.

Figures (8)

• 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).
• Figure 2: Test performance and STD of PL_Alpha_Hill across all layers of RoBERTa-base model trained on MNLI (Accuracy$\uparrow$) and QNLI (Accuracy$\uparrow$) under different subsampling ratios.
• Figure 3: (Main Results on LLM Fine-tuning).TempBalance (TB) achieves better test metric ($\uparrow$) than baseline Full Fine-tuning (FT) on GLUE tasks, especially if training data is small. \ref{['fig:nlp_improve_trend']} compares test performances of baseline FT (Full Fine-tuning) and TempBalance to train RoBERTa-base model on four larger GLUE datasets (color-coded as in \ref{['fig:nlp_trend_in_one']}). \ref{['fig:nlp_trend_in_one']} shows the trend of performance improvement of TempBalance.
• Figure 4: Domain Specific Language Modeling.TempBalance demonstrates significant performance gain when training the RoBERTa-base model on five low-resource domain-specific datasets.
• Figure 5: (Main Results on PDE Learning).TempBalance (TB) achieves lower nRMSE($\downarrow$) than baseline method on CFD tasks, especially if subsampling ratio is small. \ref{['fig:sciml_improve_trend']} compares test performances of baseline trained and TempBalance trained FNO and UNet models on 1D and 2D CFD datasets (color-coded as in \ref{['fig:sciml_trend_in_one']}). \ref{['fig:sciml_trend_in_one']} demonstrates the trend of performance improvement brought by TempBalance.