Research map / Architectures and efficiency / RNNs and normalization
Function-Space Learning Rates
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a Monte-Carlo, Kronecker-factored estimator of layerwise function-space learning rates and uses it in FLeRM to transfer optimal learning rates across width, depth, init scale and LoRA rank by matching a small base model's function-space updates.
Problem
Parameter-space learning rates do not reflect how much a network's output function changes, and optimal learning rates shift with model width, depth, initialisation scale and LoRA rank, making hyperparameter tuning at scale expensive.
Approach
Measure the RMS change in network outputs caused by each parameter tensor's update using random output projections (phi = sum of Gaussian-weighted outputs) and a single extra backward pass, with a Kronecker-factored covariance to reduce variance; then rescale per-layer parameter learning rates in a scaled model so its function-space learning rates match a cheap base model's.
Key results
- FLeRM substantially reduces or removes shifts in optimal learning rate across model width for ResMLP and transformers relative to standard practice.
- In standard practice the optimal LoRA learning rate for Llama increases by more than an order of magnitude as rank goes from 2 to 32; FLeRM with rank 2 as base largely eliminates this shift.
- Ablation with uniform base function-space learning rates still transfers hyperparameters but gives slightly worse training loss than measured function-space learning rates.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use FLeRM to match layerwise function-space learning rates when scaling width | ResMLP and transformer models · Width scaling | optimal learning rate transfer | Helped |
| Use FLeRM to match layerwise function-space learning rates when scaling depth | ResMLP · Depth scaling | optimal learning rate transfer | Helped |
| Use FLeRM to match layerwise function-space learning rates when scaling depth | PostNorm transformer · Depth scaling | training loss | Helped |
| Use FLeRM to match layerwise function-space learning rates when scaling depth | PreNorm transformer · Depth scaling | optimal learning rate transfer | Helped |
| Use FLeRM to match layerwise function-space learning rates when scaling depth | PreNormPostMod transformer · Depth scaling | optimal learning rate transfer | No clear effect |
| Use FLeRM to match layerwise function-space learning rates across initialization scales | PreNormPostMod transformer · Initialization scale variation | training loss | Helped |
| Use FLeRM to match layerwise function-space learning rates across LoRA ranks | GPT-2 and Llama-3.2-1B · Cold French Law and Mathpile | learning rate transfer | Helped |
| Periodically update FLeRM learning rates instead of fixing them after the first iteration (ablation) | FLeRM · Width, depth, and initialization scale transfer | hyperparameter transfer | No clear effect |
| Replace recorded base function-space learning rates with equal values summing to 1 (ablation) | FLeRM · Width, depth, and initialization scale transfer | training loss | Hurt |
| Set all base function-space learning rates equal across layers at every depth (ablation) | FLeRM with PreNormPostMod transformer · Depth scaling | hyperparameter transfer | Hurt |
| Set all base function-space learning rates equal across layers at every depth (ablation) | FLeRM with PreNormPostMod transformer · Deepest depth-scaled model | training loss | Mixed |
| Sweep the LoRA A learning rate while holding B fixed (ablation) | Standard AdamW · LoRA rank transfer | training loss | No clear effect |
Surprising findings
- The second feedforward weight matrices in self-attention have function-space learning rates an order of magnitude larger than the readout layer.
- With FLeRM, the PostNorm transformer's instability location aligns across depths, whereas standard practice shifts it to lower learning rates as depth grows.
Limitations
- Depth transfer relies on a heuristic of sharing base function-space learning rates across new residual blocks; naive uniform alternatives lose depth transfer for PreNormPostMod.
- LoRA experiments measured base function-space learning rates with a single seed and needed a workaround because B is zero-initialised.
- Evaluated on small-scale models (ResMLP, transformers up to ~814M parameters, GPT-2, Llama-3.2-1B) with limited training.
Experimental setup
Datasets: CIFAR-10 Wikitext-103 Cold French Law Mathpile
Metrics: training loss (averaged over last 200 batches) test loss optimal learning rate shift across scales
Baselines: standard parameterisation (standard practice) mu-P Modula uniform function-space learning rates (ablation)
Builds on: mu-P Adam LoRA Modula
Closest papers by meaning
- Function-space Parameterization of Neural Networks for Sequential Learning (ICLR 2024 · similarity 0.86)
- Functional Scaling Laws in Kernel Regression: Loss Dynamics and Learning Rate Schedules (NeurIPS 2025 · similarity 0.84)
- On the Duality between Gradient Transformations and Adapters (ICML 2025 · similarity 0.84)
- Improving Function Space Flow Matching with Kernel Optimal Transport (NeurIPS 2026 · similarity 0.83)
- Super Consistency of Neural Network Landscapes and Learning Rate Transfer (NeurIPS 2024 · similarity 0.83)
- F-Adapter: Frequency-Adaptive Parameter-Efficient Fine-Tuning in Scientific Machine Learning (NeurIPS 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.