ICML 2025

Function-Space Learning Rates
Edward Milsom 1 Ben Anson 1 Laurence Aitchison 1
Abstract
We consider layerwise function-space learning
rates, which measure the magnitude of the change
in a neural network’s output function in response
to an update to a parameter tensor. This con-
trasts with traditional learning rates, which de-
scribe the magnitude of changes in parameter
space. We develop efficient methods to measure
and set function-space learning rates in arbitrary
neural networks, requiring only minimal compu-
tational overhead through a few additional back-
ward passes that can be performed at the start
of, or periodically during, training. We demon-
strate two key applications: (1) analysing the dy-
namics of standard neural network optimisers in
function space, rather than parameter space, and
(2) introducing FLeRM (Function-space Learn-
ing Rate Matching), a novel approach to hyper-
parameter transfer across model scales. FLeRM
records function-space learning rates while train-
ing a small, cheap base model, then automatically
adjusts parameter-space layerwise learning rates
when training larger models to maintain consistent
function-space updates. FLeRM gives hyperpa-
rameter transfer across model width, depth, initial-
isation scale, and LoRA rank in various architec-
tures including MLPs with residual connections
and transformers with different layer normalisa-
tion schemes.
1. Introduction
The fundamental purpose of neural network training is to
learn a function that maps inputs to desired outputs. How-
ever, we typically understand optimisation methods as act-
ing in parameter space, e.g. traditional learning rates tell us
how much the parameters change during each step rather
than the functional impact of those changes. This raises
1University of Bristol.
Correspondence to: Edward Mil-
som <edward.milsom@bristol.ac.uk>, Laurence Aitchison <lau-
rence.aitchison@gmail.com>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, ...