UTA: Universal Training Adaptive Learning Rate Decay

Anda Tang; Junjie Wang; Xuanlin Yang; Cong Fang; Zhouchen Lin

cs.LGOct 26, 2026

UTA: Universal Training Adaptive Learning Rate Decay

Authors: Anda Tang, Junjie Wang, Xuanlin Yang, Cong Fang, Zhouchen Lin

Organizations: the State Key Lab of General AI School of Intelligence Science and Technology Peking University Beijing, China · the Pazhou Laboratory (Huangpu) · the Institute for Artificial Intelligence

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.

Figures & tables

Panel of Figure 1 — Figure 1: Training curves of LLaMA models, ResNet-50 and ViT-B.

LRD	Model	60M			130M			350M			1B
LRD	Step	3k	6k	final	5k	10k	final	7k	14k	final	8k	16k	final
CS		35.79	29.91	28.03	27.12	22.71	21.76	22.03	18.24	17.62	20.14	16.32	15.41
LD		35.76	30.34	27.99	26.60	22.71	21.43	22.16	18.72	17.68	19.96	16.49	15.10
WSD		34.80	30.03	26.03	27.05	23.71	20.83	22.41	19.66	17.25	20.60	17.94	14.94
UBA		35.48	30.30	27.01	27.18	23.18	21.12	22.24	18.68	17.42	19.78	16.41	14.83
UTA		33.72	29.57	25.80	25.74	22.83	20.16	22.14	19.58	17.17	19.56	16.97	14.53

Table 1: Perplexity (

e^{\text{validation loss}}

) on LLaMA 60M, 130M, 350M and 1B models. Lower is better (↓).

LRD	CS	LD	WSD	UBA	UTA
ResNet-50	78.34 ±0.024	77.90 ±0.028	78.30 ±0.024	78.38 ±0.014	78.73 ±0.029
ViT-B	79.92 ±0.057	79.96 ±0.104	81.14 ±0.048	82.05 ±0.041	82.30 ±0.037

Table 2: ImageNet test top-1 Acc. (%) of different LRDs on ResNet50 and ViT-B for 300 epochs. Higher values indicates better performance (↑). Best results in bold.

Model	LRD	$5{\times}10^{-4}$	$1{\times}10^{-3}$	$3{\times}10^{-3}$	$5{\times}10^{-3}$	$7{\times}10^{-3}$	$1{\times}10^{-2}$	Best from LR Sweep
60M	CS	32.61	28.03	26.91	26.91	26.31	26.57	26.31
	LD	32.83	27.99	26.80	26.77	26.25	26.46	26.25
	WSD	29.40	26.03	26.56	26.34	26.06	26.28	26.03
	UBA	30.65	27.01	26.51	26.42	26.26	26.57	26.26
	UTA	29.28	25.80	26.86	27.00	25.81	27.11	25.80
130M	CS	–	21.76	21.32	20.94	21.27	–	20.94

Table 3: Perplexity (

e^{\text{validation loss}}

) of LLaMA-60M and LLaMA-130M under different peak LRs. Lower is better.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Params	Hidden Dim	Intermediate Dim	Heads	Layers	Training Steps	Data amount (Tokens)
60M	512	1376	8	8	10k	5.2B
130M	768	2048	12	12	15k	7.9B
350M	1024	2736	16	24	20k	10.5B
1 B	2048	5461	24	32	25k	13.1B

Appendix

Table 4: LLaMA architecture hyper-parameters for evaluation. Data amounts are specified in tokens.

Panel of Figure 2 — Figure 2: Perplexity curves of LLaMA models.

LRD	LLaMA Model Size
LRD	60M		130M
	D-10	D-0	D-10	D-0
CS	28.03	28.27	21.76	21.99
LD	27.99	28.21	21.43	21.89
WSD	26.03	26.93	20.83	21.18
UBA	27.01	27.23	21.12	21.00

Appendix

Table 5: Perplexity (

e^{\text{validation loss}}

) on LLaMA models. Lower is better (↓). ‘D-10’ and ‘D-0’ denote final learning rates decaying to 10% and 0% of the peak rate, respectively.

LRD	ResNet-50	ViT-B
CS	78.34 ±0.024	79.92 ±0.057
LD	77.90 ±0.028	79.96 ±0.104
WSD	78.30 ±0.024	81.14 ±0.048
UBA	78.38 ±0.014	82.05 ±0.041
UTA (Ours)	78.73 ±0.029	82.30 ±0.037

Appendix

Table 6: ImageNet test top-1 Acc. (%) of different LRDs on ResNet50 and ViT-B under 300 epochs. Higher values indicates better performance (↑). Best results in bold.

Panel of Figure 3 — Figure 3: Training loss and validation accuracy on ImageNet using ResNet50 and ViT-B.

Network	UTA
Network	$\varphi_{0}$ =0.01	$\varphi_{0}$ =0.05	$\varphi_{0}$ =0.1	$\varphi_{0}$ =0.15	$\varphi_{0}$ =0.2	$\varphi_{0}$ =0.25	$\varphi_{0}$ =0.5	$\varphi_{0}$ =1	$\varphi_{0}$ =2.5	$\varphi_{0}$ =5
LLaMA-60M	26.18	25.94	25.85	25.81	25.80	25.80	25.92	26.04	26.39	26.78

Appendix

Table 7: Cross-

\varphi_{0}

ablation studies: Perplexity of LLaMA-60M using UTA under different

\varphi_{0}

Network	UTA
Network	$\varphi_{0}$ =0.15	$\varphi_{0}$ =0.2	$\varphi_{0}$ =0.25	$\varphi_{0}$ =0.5	$\varphi_{0}$ =2.5
LLaMA-130M	20.63	20.65	20.66	20.67	20.99

Appendix

Table 8: Cross-

\varphi_{0}

ablation studies: Perplexity of LLaMA-130M using UTA under different

\varphi_{0}

	Ours-Static	Ours-Full	Improvement
CIFAR100-ResNet34(↑)	76.86	77.81	0.95
ImageNet-ResNet50(↑)	78.23	78.55	0.32
LLaMA60M(↓)	27.38	26.78	0.60
LLaMA130M(↓)	21.02	20.15	0.87

Appendix

Table 9: Ablation Study: Impact of Dynamic State-Aware Mechanism.

Method	60M	130M	350M
D-Adaptation	25.80	20.65	17.64
Prodigy	27.37	21.26	17.97
UTA (Ours)	25.80	20.16	17.17

Appendix

Table 10: Perplexity (

e^{\text{validation loss}}

) comparison between automatic LR methods (D-Adaptation, Prodigy) and our UTA on LLaMA models. Lower is better (

\downarrow

	UBA( $\varphi=0.03$ )	UTA( $\varphi_{0}=0.15$ )	Improvement
ImageNet-ResNet50(↑)	78.47	78.64	0.17
LLaMA130M(↓)	21.08	20.15	0.93
LLaMA1B(↓)	14.81	14.53	0.28

Appendix

Table 11: Curvature alignment between UTA and UBA.

Panel of Figure 4 — Figure 4: Curvature alignment between UTA and UBA.

Network	Dataset	CS	LD	WSD	UBA	UTA
ViT-B	CIFAR100	91.93	92.45	92.62	92.58	92.81
ViT-B	OxfordPets	92.59	92.67	92.86	93.05	93.43

Appendix

Table 12: UTA for LoRA with ViT-B.

Figure 5: The learning rate curve in UTA LRD across training steps.

Learning Rate Scheduling

Learning rate scheduling studies how to vary optimizer step sizes during training to improve stability, convergence, efficiency, and generalization.

See all 60 papers

Learning Rate Engineering: From Coarse Single Parameter to Layered Evolution
Apr 30, 2026Ming-Hong Yao, Di Wang, Jian Cui +5Deep Learning Optimization Fine-Tuning
School of AI, Jilin International Sciences University, Changchun, China.
One LR Doesn't Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMs
May 21, 2026Di He, Songjun Tu, Keyu Wang +2Learning Rate Scheduling Efficient Language Model Training
Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences · Peng Cheng Laboratory · University of Chinese Academy of Sciences +5
One Rate Is Not Enough: Adaptive Anisotropic Learning Rates for LoRA Fine-Tuning
Sep 5, 2026Huiyi Wang, Daijiao Liu, Lina Yao +1LLM Fine-Tuning Low-Rank Adaptation
University of New South Wales

Apr 30, 2026J/K move · Enter open · S save

Learning Rate Engineering: From Coarse Single Parameter to Layered Evolution

Ming-Hong Yao, Di Wang, Jian Cui, Jin-Yan Chen+4

School of AI, Jilin International Sciences University, Changchun, China.

Deep Learning Optimization Fine-Tuning Learning Rate Scheduling Adaptive Gradient Methods Layer-Wise Learning Rate Adaptation Adaptive Learning Rate

Open paper AskPDF arXiv

Learning rate scheduling has evolved from the single global fixed rate of early SGD to sophisticated layer-wise adaptive strategies. We systematize this evolution into five generations: (Gen1) global fixed learning rates, (Gen2) global scheduling, (Gen3) parameter-level adaptation, (Gen4) layer-level differentiation, and (Gen5) joint layer-time scheduling. We trace the fundamental motivation behind each transition, showing how the shift from one-size-fits-all to tailoring by layer and time addresses the impossible trinity of transfer learning: lower layers require small updates to preserve general knowledge while higher layers need large updates to adapt to new tasks. Building on this taxonomy, we propose Discriminative Adaptive Layer Scaling (DALS), a unified framework that integrates phase-adaptive cosine scheduling, depth-aware Grokfast gradient filtering, and LARS-style trust ratios into a single coherent optimizer. We benchmark 18 strategies including three DALS variants across all five generations on five datasets: synthetic, CIFAR-10 (from scratch), RTE, TREC-6, and IMDb (fine-tuning). On synthetic, DALS achieves the best accuracy at 98.0%, while DALS-Fast reaches 90% in just 3 epochs. The cross-dataset analysis reveals striking regime-dependent patterns -- no single strategy wins across all regimes. Critically, STLR+Discriminative, the ULMFiT champion, catastrophically fails on from-scratch tasks (43.6% on TREC-6 from scratch vs. 96.8% with RAdam), confirming that directional decay biases are harmful without pretrained features. DALS avoids either extreme, achieving the best synthetic result while maintaining competitive fine-tuning performance.

Abstract

Learning Rate Scheduling

Explore similar work

Learning Rate Engineering: From Coarse Single Parameter to Layered Evolution

One LR Doesn't Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMs

One Rate Is Not Enough: Adaptive Anisotropic Learning Rates for LoRA Fine-Tuning