Pas de taille unique : Un calendrier de taux d'apprentissage unifié pour l'entraînement à itérations budgétisées

Stepsize anything: A unified learning rate schedule for budgeted-iteration training

May 30, 2025

Auteurs: Anda Tang, Yiming Dong, Yutao Zeng, zhou Xun, Zhouchen Lin

cs.AI

Résumé

L'augmentation des coûts de calcul et les ressources limitées soulignent le besoin crucial d'un entraînement avec budget d'itérations, visant à atteindre un apprentissage optimal dans des budgets d'itérations prédéfinis. Bien que les plans de taux d'apprentissage régissent fondamentalement la performance des différents réseaux et tâches, en particulier dans les scénarios avec budget d'itérations, leur conception reste largement heuristique, manquant de fondements théoriques. De plus, le plan de taux d'apprentissage optimal nécessite une sélection extensive par essais et erreurs, rendant le processus d'entraînement inefficace. Dans ce travail, nous proposons le plan Unified Budget-Aware (UBA), un plan de taux d'apprentissage théoriquement fondé qui surpasse systématiquement les plans couramment utilisés parmi diverses architectures et tâches sous différents budgets d'entraînement contraints. D'abord, nous comblons l'écart en construisant un nouveau cadre d'optimisation prenant en compte le budget d'entraînement, qui tient explicitement compte de la robustesse aux variations de courbure du paysage. De ce cadre, nous dérivons le plan UBA, contrôlé par un seul hyper-paramètre varphi qui offre un compromis entre flexibilité et simplicité, éliminant le besoin d'optimisation numérique par réseau. De plus, nous établissons une connexion théorique entre varphi et le nombre de conditionnement, ajoutant une interprétation et une justification à notre approche. Par ailleurs, nous prouvons la convergence pour différentes valeurs de varphi. Nous fournissons des directives pratiques pour sa sélection via une analyse théorique et des résultats empiriques. Les résultats expérimentaux étendus montrent que UBA surpasse systématiquement les plans couramment utilisés dans diverses tâches de vision et de langage, couvrant des architectures de réseaux (par exemple, ResNet, OLMo) et des échelles, sous différents budgets d'itérations d'entraînement.

English

The expanding computational costs and limited resources underscore the critical need for budgeted-iteration training, which aims to achieve optimal learning within predetermined iteration budgets.While learning rate schedules fundamentally govern the performance of different networks and tasks, particularly in budgeted-iteration scenarios, their design remains largely heuristic, lacking theoretical foundations.In addition, the optimal learning rate schedule requires extensive trial-and-error selection, making the training process inefficient.In this work, we propose the Unified Budget-Aware (UBA) schedule, a theoretically grounded learning rate schedule that consistently outperforms commonly-used schedules among diverse architectures and tasks under different constrained training budgets.First, we bridge the gap by constructing a novel training budget-aware optimization framework, which explicitly accounts for the robustness to landscape curvature variations.From this framework, we derive the UBA schedule, controlled by a single hyper-parameter varphi that provides a trade-off between flexibility and simplicity, eliminating the need for per-network numerical optimization. Moreover, we establish a theoretical connection between varphi and the condition number, adding interpretation and justification to our approach. Besides, we prove the convergence for different values of varphi.We offer practical guidelines for its selection via theoretical analysis and empirical results.xtensive experimental results show that UBA consistently surpasses the commonly-used schedules across diverse vision and language tasks, spanning network architectures (e.g., ResNet, OLMo) and scales, under different training-iteration budgets.