Reproduce “Model Balancing Helps Low-data Training and Fine-tuning” on OpenResearch

About this paper

Zihang Liu, Yuanzhe Hu, Tianyu Pang, Yefan Zhou, Pu Ren, Yaoqing Yang

Researchers from Dartmouth College, UC Berkeley, UC San Diego, and Lawrence Berkeley National Lab developed TempBalance, an algorithm that dynamically adjusts layer-wise learning rates based on Heavy-Tailed Self-Regularization theory. It consistently improved model performance in low-data training and fine-tuning across NLP and SciML tasks, demonstrating a 9.9% accuracy gain on SST2 with 0.02% data and up to 14.47% nRMSE reduction for neural PDE solvers.

Key insights

• Low-data training significantly increases the standard deviation of PL Alpha Hill across model layers, indicating a direct correlation between data scarcity and imbalanced layer-wise training quality.
• Balancing the PL Alpha Hill exponents across layers, through dynamic learning rate adjustments, directly leads to improved overall model performance, especially in data-constrained environments.
• TempBalance is a versatile 'add-on' method, compatible with existing optimizers like AdaFactor, and its effectiveness is inversely proportional to the amount of available training data.

Summary by alphaXiv.