Adaptive learning rate decay improves training of ai and vision models
UTA: Universal Training Adaptive Learning Rate Decay
Machine LearningArtificial Intelligence
Summary
Many AI models get better when the speed of learning slows down at the right times, but choosing how and when to slow down is tricky. The authors created a new way called Universal Training Adaptive (UTA) that changes the learning speed based on the current state of training instead of just time passing. This method works well with modern techniques and helps large language and vision models learn more effectively. They tested UTA on several models and showed consistent improvements without extra work for users.
What this means in practice
- •For machine learning engineers: Improve training performance for language models by using adaptive learning rate decay that responds to training progress instead of fixed schedules.
- •For computer vision developers: Enhance image recognition model training efficiency and accuracy on datasets like ImageNet through near-optimal adaptive learning rate control.
Authors
Anda Tang, Junjie Wang, Xuanlin Yang, Cong Fang, Zhouchen Lin
Abstract
Learning rate decay (LRD) is critical for training high-performance models, yet most designs remain empirically driven. Although the Universal Budget Aware (UBA) LRD offers a theoretical framework, its reliance on SGD dynamics limits its applicability to modern Adam-type optimizers and leaves its key `optimization difficulty parameter' opaque. To bridge this gap, we propose the Universal Training Adaptive (UTA) LRD. Unlike UBA's fixed offline trajectory anchored solely to iterations, UTA treats LRD as a closed-loop shape-and-scale problem. By continuous-time formulation, UTA is training-state adaptive with segment-wise optimal shape. This adaptivity links external iterations with internal training states, enabling a near-optimal path in both progression shape and learning rate (LR) scale, a property unattainable by other LRDs, including UBA. Empirically, we demonstrate consistent improvements across language models (LLaMA 60M to 1B) and vision models (ResNet50, ViT-B) on ImageNet with plug-and-play ease and negligible overhead. Extensive ablation studies further confirm UTA's robustness across dimensions, including peak LR, final decay ratio, its parameter $\varphi_0$, component efficacy and generalization to fine-tuning scenarios, underscoring its practical implementability.