Learning rate schedules are typically fixed before training begins, yet when fine-tuning large language models with low-rank adapters (LoRA), the appropriate scale varies substantially with the model, task, and adapter configuration. Practitioners re-tune a schedule for every new problem, yet a tuned schedule still cannot react to the training dynamics it encounters. We instead amortize learning rate selection across LoRA fine-tuning problems. We leave the base optimizer, Muon, untouched, so that its normalized update fixes the direction, and meta-learn only the global step size. Our method MuPACE is a small recurrent policy, meta-trained across a distribution of fine-tuning runs, that conditions on the observed trajectory and emits bounded adjustments to an interpretable analytic schedule. On 64 fine-tuning problems, MuPACE outperforms standard fixed schedules at a strong shared default on 64-78% of problems and is on par with the same schedules with per-task tuned peaks. Meta-trained only on models with up to 16 layers, MuPACE transfers one-shot to unseen tasks on models as deep as 64 layers, such as Qwen3-32B and Qwen3.8-27B, where a single run remains competitive with swept schedules.
Figures & tables
Figure 1: Shape-normalized schedules varying one parameter at a time, default in dark.
Figure 2: Learned scheduler versus defaults on 64 tasks. (a) Win rates and Holm-adjusted sign-test p values. (b) Relative test-loss reduction versus each baseline, where positive favors learned. Dots show tasks, diamonds medians, and bars interquartile ranges.
Figure 3: Learned scheduler versus validation-tuned baselines on 64 tasks. (a) Win rates and Holm-adjusted sign-test p values. (b) Relative test-loss reduction versus each tuned baseline, where positive favors learned. Dots show tasks, diamonds medians, and bars interquartile ranges.
Figure 4
Figure 6: MuPACE versus validation-tuned baselines and Hypergradient Descent on (a) Qwen3-32B and (b) Qwen3.8-27B. Values show baseline minus MuPACE test loss (positive favors MuPACE).
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
Component
Values or sampling range
Base-model deck
37 nanochat checkpoints
Model depth / width
{8,12,16} / 128 – 2048
Dataset
100 Lots-of-LoRAs sub-tasks
LoRA rank r
{8,16,32}
LoRA scaling a
LogUniform[8,32]
Batch size b
{8,16,32}
Appendix
Table 1: Meta-training task distribution. Scale augmentation is disabled at evaluation.
Table 3: Lots-of-LoRAs sub-tasks used for meta-training.
Dataset
Domain
No Robots ( Rajani et al., 2023 )
Instruction following
GSM8K ( Cobbe et al., 2021 )
Mathematical reasoning
MBPP ( Austin et al., 2021 )
Code generation
SAMSum ( Gliwa et al., 2019 )
Dialogue summarization
CNN/DailyMail ( See et al., 2017 )
News summarization
XSum ( Narayan et al., 2018 )
Extreme (single-sentence) summarization
Appendix
Table 4: Out-of-distribution evaluation datasets.
Hyperparameter
Value
Input features / LSTM hidden width
20 / 48
Meta-parameters
19,877
Output head initialization
Zero
Outer optimizer
AdamW, weight decay 10−4
Outer learning rate
10−3 , cosine decay
Meta-gradient clipping
Global norm 1
Appendix
Table 5: Policy architecture and meta-training hyperparameters.
Parameter
Default λi0
Reach ci
Level λ1
10−3
4.5
Decay exponent λ2
2
1.0
Time-warp exponent λ3
1
1.4
Tail floor λ4
0.01
2.3
Warmup time constant λ5
0.03
0.7
Appendix
Table 6: Schedule defaults and adjustment ranges. Each head emits λi=λi0exp(citanhhi) , giving the range [λi0e−ci,λi0eci] .
Figure 7: Expected test performance under validation-based trial selection. Shading shows the interquartile range across task expectations, not confidence intervals.
Figure 8: Illustrative MuPACE train, validation, and test evaluation losses on MBPP and CNN/DailyMail.
Figure 9: MuPACE versus the untrained zero-head policy on in-distribution draws (top) and Qwen3-32B (bottom).
Figure 10: Learning rates applied by the direct LR head on 16 in-distribution tasks.
Figure 11: Direct learning-rate head on the same 64 in-distribution draws. (a) compares against the default schedules, while (b) compares against the swept baselines.
Figure 12: Direct learning-rate head transferred one-shot to Qwen3-32B. Bars show baseline minus learned test loss, where positive values favor the learned scheduler.
Figure 13: Comparison of MuPACE with and without shape normalization on in-distribution draws (top) and Qwen3-32B (bottom). Positive values mean shape normalization performs better.
Figure 14: Depth-scaled fixed schedules versus MuPACE on Qwen3-32B (left) and Qwen3.8-27B (right). Positive differences favor the learned policy. Target-tuned schedules and Hypergradient Descent are shown for context. Depth-scaled fixed schedules use no target-task peak selection.
Figure 15: MuPACE versus a separately meta-trained variant without depth-dependent post-scaling on in-distribution draws (top) and Qwen3-32B datasets (bottom). Positive values favor the full scheduler.
Schedule-Free Learning has shown promise as a practical anytime training method for machine learning, showing success across dozens of standard benchmark problems. However, strong performance for LLM training has only been demonstrated at small scales. We identify a number of fixes necessary to scale up Schedule-Free Learning to larger batch sizes and model sizes, and present a learning-rate-free and schedule-free method (ScheduleFree+) for training large language models which greatly outperforms Warmup-Stable-Decay (WSD) schedules. We also demonstrate that Schedule-Free Learning is most effective for long duration training, and at 1000 tokens per parameter, it outperforms SOTA schedules by 31%. Schedule-Free Learning provides a theoretical foundation for the use of model averaging and checkpoint merging during pretraining.
Low-Rank Adaptation (LoRA) significantly reduces compute and memory costs for finetuning Deep Learning models but is often harder to tune than dense training: when using factor-wise optimizers such as AdamW, it is sensitive to initialization choices, its optimal learning rates transfer poorly across ranks, and it often fails to beat dense baselines. We derive LoRA-Muon by applying the Muon optimizer's spectral steepest-descent rule to the low-rank setting. Along with our split weight-decay rule, our main claim is that LoRA-Muon is a good low-rank proxy for full-rank Muon and Shampoo-family optimizers. Its optimal learning rates transfer across rank, width, depth, and factor-rescaling. In our compute-matched TinyShakespeare study, a rank-2 proxy recovers the dense best tested learning rate, and a rank-32 LoRA-Muon run attains lower mean validation loss than the dense baseline in the seed-averaged sweep. We further show that the Spectron optimizer depends on arbitrary factor scaling, so it would likely be a poor fit when finetuning starts from badly imbalanced factors, and that LoRA-RITE's simplified QR-coordinate core implements the same spectral update. LoRA-Muon computes that update without QR-decomposition and avoids storing second moments, making it more accelerator-friendly and memory-efficient.
Franz Louis Cesista, Katherine Crowson, Cédric Simal +1
Ateneo de Manila University · EleutherAI · NaXys +1
The Muon optimizer shows clear benefits versus alternatives when pretraining neural networks. However, it is used less frequently for parameter-efficient fine-tuning (PEFT). One potential reason is that the most common PEFT method, LoRA, does not naturally combine with Muon since it is not mathematically possible to orthogonalize the weight update given by a low-rank parameterization. In this paper, we address this issue by approximating the solution to a relaxed Muon objective in the low-rank setting via linearization and then least-squares. We provide an efficient implementation that uses matmul operations only, as opposed to more complex linear algebra decomposition routines. Our method, sMuon (small Muon), performs favourably across SFT and a ReLoRA pretraining experiment. While results are model- and eval-dependent, we find overall that using Muon for low-rank fine-tuning provides moderate performance improvements.
Ben Anson, Conor Houghton, Edward Milsom
University of Bristol School of Mathematics · University of Bristol School of Engineering Mathematics and Technology · University of Bath Department of Computer Science