Adaptive optimizers like AdamW apply uniform hyperparameters across all parameter groups, ignoring heterogeneous optimization dynamics across layers and modules. We address this limitation by proposing MetaAdamW - a new optimizer that integrates a self-attention mechanism to dynamically modulate per-group learning rates and weight decay. The modulation factors are produced by a lightweight Transformer encoder that operates on statistical features (gradient norms, momentum norms, correlations) extracted from each parameter group. To train the attention module, we introduce a meta-learning objective that combines gradient alignment, loss decrease, and generalization gap. A key novel contribution is the extension of homoscedastic uncertainty weighting (HUW) with task-specific priorities that directly scale the regularization terms - enabling domain knowledge to guide automatic loss balancing. Extensive experiments on five diverse tasks-time series forecasting (ETT), language modeling (WikiText-2), machine translation (Multi30k), image classification (CIFAR-10), and sentiment analysis (IMDB) - demonstrate that MetaAdamW consistently outperforms the standard AdamW baseline in terms of validation loss, accuracy, or perplexity. Depending on the task, MetaAdamW either reduces overall training time (by up to 17.11%) or improves performance (by up to 11.08%) while introducing only moderate overhead; in some cases, it can also mitigate issues of insufficient convergence caused by premature early stopping. Ablation studies validate the effectiveness of each component, including feature versions, grouping strategies, and the proposed priority-injected uncertainty weighting.
Figures & tables
Parameter
Value
Global learning rate η
5×10−4
Weight decay λ
1×10−2
AdamW betas
(0.9,0.999)
α range
1.0
β range
1.0
Warmup epochs (before meta)
1
Table 1: MetaAdamW Common Hyperparameters
Task
AdamW
MetaAdamW
Improvement
Time Overhead
Metric
Value
Metric
Value
ETTh1 (forecast)
MSE
0.006147
MSE
0.005885
+4.26%
-7.20%
WikiText-2 (LM)
PPL
120.47
PPL
115.51
+4.12%
-17.11%
Multi30k (MT)
PPL
2.0297
PPL
1.9690
+2.99%
+27.35%
CIFAR-10 (classif.)
Acc
87.02%
Acc
88.05%
+1.18%
+27.58%
IMDB (sentiment)
Acc
74.53%
Acc
82.79%
+11.08%
+172.53%
Table 2: Main Results: Best Validation Metrics and Total Training Time
Figure 1: ETTh1 Validation loss over epochs.
Figure 2: WikiText-2 Validation perplexity over epochs.
Figure 3: Multi30k Validation perplexity over epochs.
Figure 4: CIFAR-10 Validation accuracy over epochs.
We introduce SODA, a generalization of Optimistic Dual Averaging, which provides a common perspective on state-of-the-art optimizers like Muon, Lion, AdEMAMix and NAdam, showing that they can all be viewed as optimistic instances of this framework. Based on this framing, we propose a practical SODA wrapper for any base optimizer that eliminates weight decay tuning through a theoretically-grounded 1/k decay schedule. Empirical results across various scales and training horizons show that SODA consistently improves performance without any additional hyperparameter tuning.
Thomas Pethick, Wanyun Xie, Roman Machacek +1
Independent Researcher · EPFL (LIONS) · University of Bern
Learned optimization aims to improve upon hand-designed optimizers (e.g., Adam and Muon) by meta-learning small neural network optimizers over a distribution of tasks. While recent work has greatly advanced the architectural design and inductive biases of learned optimizers (LOs), their meta-training remains biased toward short-unroll learning on particular tasks, resulting in redundant computation and leaving LOs often unable to compete with hand-designed optimizers. We introduce Efficient Long-hOrizon (ELO) learning, an efficient meta-training algorithm that (1) reallocates wasted meta-training compute to longer failure regimes, achieving efficient long-horizon learning, and (2) enforces decoupled progressive expert supervision, providing stable meta-learning signals that additionally improve the generalization of LOs. Our empirical study evaluates ELO for meta-training both element-wise and matrix-based LOs. Across downstream language modeling (GPT-2 and GPT-3 on FineWeb) and image classification (ViT-B/16, ResNet-50 on ImageNet-1K) tasks, ELO substantially improves the long-unroll performance and out-of-distribution generalization of the base LOs. In particular, ELO-Celo2 outperforms well-tuned AdamW and Muon on GPT-2 (30M and 124M) language pretraining. Notably, all ELO baselines require less than 7 H100 GPU-hours for meta-training.
Xiaolong Huang, Benjamin Thérien, James Harrison +1
Mila - Quebec AI Institute · Concordia University · Universit´e de Montr´eal +1
Large language models (LLMs) contain functionally distinct modules with heterogeneous gradient statistics. Although Adam(W) adapts updates at the parameter level, selecting relative learning rates across modules often relies on costly manual tuning. We study this problem through gradient signal-to-noise ratio (SNR). Our analysis identifies structural sources of module-wise SNR differences and relates SNR to the attenuation of Adam's normalized mean-gradient component. Based on these findings, we propose Module-wise Learning Rate Scaling via SNR (\name), a lightweight method that replaces manual multiplier search with one-shot calibration. MoLS estimates module-level SNRs after an initial adaptation period and uses inverse-square-root scaling to determine relative learning-rate multipliers. The multipliers are introduced smoothly during warmup and then held fixed while the global learning-rate schedule continues. Our theoretical results quantify the effect of calibration and its sensitivity to estimation error and relative SNR drift. Extensive experiments across LLM pre-training and fine-tuning show that MoLS reduces perplexity and improves downstream accuracy with minimal computational overhead. It achieves performance comparable to carefully tuned module-specific learning rates and integrates with memory-efficient optimizers and parameter-efficient fine-tuning.
Ziqing Wen, Zhouyang Liu, Jiahuan Wang +5
College of Computer Science and Technology, National University of Defense Technology, Changsha, Hunan, China