Function-Space Learning Rates
Authors: Edward Milsom
Ben Anson
Laurence Aitchison
Venue: ICML 2025 poster (2025)
Keywords: deep learning
optimisation
function space
hyperparameter transfer
muP
learning rate
kronecker
Primary area: deep_learning->everything_else
TL;DR: We provide a way to measure the step-sizes of a neural network optimiser in function space, rather than the usual parameter space, and use this to enable reusing the optimal learning rate for small models in much wider and deeper models.

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 contrasts with traditional learning rates, which describe 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 computational overhead through a few additional backward passes that can be performed at the start of, or periodically during, training. We demonstrate two key applications: (1) analysing the dynamics of standard neural network optimisers in function space, rather than parameter space, and (2) introducing FLeRM (Function-space Learning Rate Matching), a novel approach to hyperparameter transfer across model scales. FLeRM records function-space learning rates while training 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 hyperparameter transfer across model width, depth, initialisation scale, and LoRA rank in various architectures including MLPs with residual connections and transformers with different layer normalisation schemes.

Fupction-Space Learning Rates
Edward Milsom
Ben Anson
Laurence Aitchison
Abstract
an important question: can we meaningfully quantify and control learning in function space?
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 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 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 FL.RM (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 small, cheap base model, then automatically adjusts parameter-space layerwise learning rates when training larger models to maintain consistent function-space updates. FL&RM gives hyperpa- rameter transfer across model width, depth, initial- isation scale, and LoRA rank various architec- tures including MLPs with residual connections and transformers with different layer normalisa- tion schemes.
We consider the concept of layerwise function-space learn- ing rates which measure the magnitude of change in network output induced by updates to individual parameter tensors. Unfortunately, naive approaches to measuring function- space learning rates would be computationally prohibitive. We solve this problem by developing a Monte-Carle esti- mate to measure function-space learning rates using only a single additional backward pass, which can be performed a handful of times at the start of training (e.g. 40), or periodi- cally during training (e.g., once every 100 steps), resulting in negligible computational overhead. We then consider two immediate applications of function-space learning rates.
First, function-space learning rates provide a novel lens for analysing the behavior of standard neural network op- timisers (e.g. Adam, Kingma, 2014), giving important in- sights into how different parts of the network contribute functional changes during training. Second, function-space learning rates enable new approach to hyperparameter transfer (for previous work see e.g. Yang & Hu, 2022; Borde- lon et al., 2023; Large et al., 2024). As large language model pre-training can cost millions of dollars in compute (Cottier et al., 2024), running extensive hyperparameter sweeps at full scale is impractical. Instead, one might hope opti- mise hyperparameters
smaller models and transfer them to larger ones. However, this is complicated by the fact that optimal learning rates change with model width and depth (Yang & Hu, 2022; Noci et al., 2024; Yang et al., 2023; Bor- delon et al., 2023). We address this challenge with FLeRM (Function-space Learning Rate Matching). which maintains consistent function-space learning rates models are scaled up by automatically adjusting the parameter-spate learning rates, thereby keeping the optimal value of the user-defined learning rate hyperparameter stable.
1. Introduction
The fundamental purpose of neural network training is to learn function that maps inputs to desired outputs. How- ever, we typically understand optimisation methods 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
A key advantage of our approach is its flexibility: our meth- ods for measuring and controlling function-space learning rates work with any network architecture and at any point during training. This contrasts with traditional approaches hyperparameter transfer that often make restrictive assump- tions about architectures or initialisation schemes (Everett et al., 2024; Yang & Hu, 2022; Bordelon et al., 2023; Large et al., 2024). We demonstrate FLeRM's utility across a
University of Bristol. Correspondence to: Edward Mil- som mel.waadd.mlloomeee. bristol.ac.uk>.. Laurence Aitchison <lau- rence.aitchison@ araioooo@@mmmmmmmmmmmm gmail.com>.
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Fupction-Space Learning Rates
range of scenarios, including model width scaling, depth sealing, initialisation scale variation, and even LoRA rank adjustment.
urally exhibits hyperparameter scaling. Rather than study- ing scaling asymptotically, Modula follows metrisation- based approach (for other metrisation works, see e.g. Yang et al., 2024; Bernstein et al., 2021; Bernstein & Newhouse, 2024a;b). Its carefully designed modules allow the com- putation of the network's Lipschitz constant, which can be used to normalise updates to enable hyperparameter scaling, However, one important difficulty with Modula is that requires setting the "mass" of each parameter. This mass can be seen as analagous to the layerwise function-space learning rate in our work, for the first step of the optimiser (i.e. at initialisation). This introduces a large number of new hyperparameters that must be tuned. In contrast, our approach to hyperparameter transfer, FLeRM, does not re- quire the user to specify masses / function-space learning rates for each parameter, because it directly measures the function-space learning rates in a base model, and then uses them in scaled model. We show in Section 4.2.5 that us- ing function-space learning rates measured directly from a base model leads to better training loss than simplistic user- defined function-space learning rates. Furthermore, FLcRM can be applied to any existing neural network in Pytorch, whereas Modula requires the user to rewrite their network architecture using the library's 's modules.
2. Related work
As far as we are aware, this work is unique in proposing methods to measure and set function-space learning rates in arbitrary neural networks far from initialisation, and using this approach to study the dynamics of optimizers.
There is an existing body of literature on hyperparameter transfer (e.g. Yang & Hu, 2022; Bordelon et al., 2023; Large et al., 2024; Yaida, 2022). Broadly speaking, these works analytically derive scaling laws for e.g. the initialisations and parameter-space learning rates, such that the function- space learning rates do not change as e.g. width is increased. This is radically different from our approach to hyperparam- eter scaling. In particular, none of these works provide mechanism to empirically measure the function-space learn- ing rates in an arbitrary neural network. Furthermore, prior works tend to rely on rigid assumptions such as being close to specific random initialisation, which we do not make.
Perhaps the earliest and best-known work hyperparam- eter scaling is μP (Yang& & Hu, 2022). μP derives how to scale random initialisations and learning rates as you in- crease model width, such that the function-space learning rates remain asymptotically constant (i.e. the magnitude of the activations does not blow up to infinity or shrink to zero as the model gets wider). μP has since been extended to depth-scaling (Yang et al., 2023; Bordelon al., 2023) and networks trained with sharpness-aware minimisation (Haas et al., 2024), and is closely related to the mean-fiel field analysis of neural networks grounded in statistical mechanics (Mei et al., 2018; Rotskoff& & Vanden-Eijnden, 2022; Sirignano & Spiliopoulos, 2019; Chizat & Bach, 2018; Geiger et al., 2020: Bordelon & Pehlevan, 2022). Because this approach derives the function-space learning rates analytically, (in contrast to our approach of measuring them), it requires restrictive assumptions, including that the network is wide, and close to a random initialisation. Extending these results to more general cases is made complicated by the fact that they typically rely on heavy-duty mathematical machinery such as Prosor-Programs (Yang, b021b; 2020; 2021a; Yang & Hu, 2022; Yang et al., 2022; 2023) or dynamical mean- field theory (Bordelon & Pehlevan, 2022; 2023; Bordelon et al., 2023; 2024).
Chizat & Netrapalli (2024) quantify feature learning in neural networks as the angle between feature updates and backward passes, enabling analysis of hyperparameter seal- ing laws and development of improved scaling rules for deep networks. By contrast, our work directly measures the funetion-spaac learning rates using autodiff and Monte- Carlo approximation.
Finally. Everett et al. (2024) recently showed empirically that alignment (concerning the size of dot products between activations and updates across different layers) in real mod- els is highly dynamic and complex throughout training. This can make choosing the correct alignment assumptions in μP and mean-field parametrisations (Yang & Hu, 2022; Mei et al., 2018; Bordelon & Pehlevan, 2022), a very difficult task. By contrast, we propose methods that can directly mea- sure the function-space learning rates throughout training, avoiding the need for such assumptions and analysis.
3. Methods
In Section 3.1 we describe how we can empirically measure the layerwise function-space learning rates using a Monte- Carlo estimate. Then, in Section 3.2, we propose using Kronecker factorisation to reduce the variance of our esti- mates. Finally, in Section 3.3, we introduce FLeRM, which modifies the parameter-space learning rates of scaled mod- els their function-space learning rates match small base when the modules are combined, the overall network nat-
The full μP scheme can be complex to apply in practice, because it e.g. requires distinct treatment of the initialisation and learning rates for the embedding weights and output heads. Later work (Large et al., 2024) sought to address this issue, by providing a library (Modula) of modules, such that
model, thereby enabling hyperparameter transfer.
2

---

Function-Space Learning Rates
3.1. Monte-Carlo estimation of the layerwise function-space learning rate
Algorithm Recording (red) setting (FLeRM, blue) function-space learning rates in a training loop.
At the core of our contributions is the estimation of the lay- erwise function-space learning rates, i.c. the magnitude of the change in output logits arising from particular change in the 1hh parameter tensor. We begin by considering the full change in the function output, fnt- Here, n indexes the N datapoints in the minibatch, and k indexes the K output features. We use Def to denote a first-order Taylor ap- proximation of the change in the outputs due particular change, A W in the éth parameter, W E R RNNxXNe-1.
Input: | (hase.)) RMS EMA Z2 [.. EMA ZZT
Base model function-space LRs
EMA ZTZ[(] = 0
for t 1 to T do
f f(X)
Std. forward pass Std. loss
C -- loss(targeess,, f)
9iy + dC/dW,, 1j
Std. backward pass
W < -W'
Save weights before update W t optimiser(no. W .g') Std. update (base LR)
butter
W < optimiser( ( (n²) ( W (or FLeRM LR)
EMA warmup: run below code for few different X'
df dWp, 1/
DeSnk- = Σ AW
(1)
if % 100 = 0 then
Wnk 10,
f' = f(X))
Fresh batch of data for ø
Here, W is the ijth element of W and das is the gradi- aw?;
Φ - JNK E Wak
Compute & (Eq. 3)
ent of the output w.r.t. Wiy. Note that for ease of notation we assume W is a matrix, but W can be a tensor of any rank. We are interested in the layerwise function-space learning rate, i.e. the RMS norm of Aef,
96j f- do/dW?
Backward pass for ø.
AW + 7 (W
- W
LR=1 update
buffer
Zi f gi SW
Compute Zij as in Eq. (8)
Update EMAS
| RMMM En (AeInx)?
(2)
EMA EMA_Z2[(] <- (1-3)EMA 22 Ei Zzy
EMA EMA.ZZT(E)) (1-3))) ZZT 2 3 [ Zik)²
(L2-norm can be used if preferred, the only difference is the term NR) ). Naïvely computing Eq. (2) via Eq. (1) is intractable as it requires NK backward passes, one for each dfak. Instead, we exploit the fact that we only need the a
EMA. EMA.ZTZ(4) (1-
ZTZ (2, Zky)²
Function spacoonssspceeoeo LR (EMA bias correction hidden)
| I + VEMA [ (CTT(EMAAZTZ( /EMA 22[
Set parameter-space learning rates (FLeRM)
magnitude || not the full change Af. Specifically, we use a Monte-Carlo approach. Consider the following scaled random combination of outputs,
rf < To || | (base.t) / /14, (()
RMS
RMS
end if end for Output: ||
VNK EnkWinkfk
Wik (O, 1).
(3)
IRMS Recorded function-space LRs
As in Eq. (1) we can write the change in & arising from a change in the εth parameter,
3.2. More efficient function-space learning rate estimates using Kronecker factorisation
Σ AW de dw?
DeΦ-
(4)
The approach in Sec. 3.1 requires one backward pass per sample, which could still be inefficient we need multiple samples for a good estimate. We remedy this by exploit- ing the structure of é. Noting that i, index the rows and columns of W we rewrite Eq. (4) using Z E R NeX., 1,
Importantly, note that we can compute &eo in a single back- ward pass. To see how AeΦ helps compute || | we substitute the definition of & (Eq. 3) into Eq. (4),
df dwp
A = Σ
Σ AW!
(5)
VNK
Zu = AW deb
rik
A Σ Zij
(8)
aw;
and note that the inner sum is Aefnk (Eq. 1), so simplifying.
DeΦ-
Σ Wik VNR Defnk:
(6)
Hence the function spacc learning rate can be written as
and since Wiek are IID standard Gaussian (Eq. 3), we have
| IRMs = Var Aeol = Σ Cov |ov| Zeyl
(9)
$3.3'''
A ~ N (o, (0,11 (RM))
(7)
Also note that, substituting the definition of á (Eq. 3) into the definition of Ziy (Eq. 8), we see that the Z 's are zero- mean Gaussian, as they are a linear combination of zero-
Hence we can estimate || RM by computing Δgo with multiple samples of Wpk, and estimating the variance.
3

---

Function-Space Learning Rates
We will now show how to compute this efficiently from Z, which itself can be computed with single backwards pass (Eq. 8). Since the Ziz 's are zero-mean Gaussian, and we have assumed Cov [Zij. Z = Ua Vagr, Z is zero-mean Matrix-Normal distributed, and so (Gupta & Nagar, 1999)
Bisses
Irput
Funciion-Sie Space LR
10"
Outout
Höden
ResMLP P
10
E "Zz"" Utr(V)
(13a)
10
I T Z
Vtr(U).
(13b)
2000
1000
6000
8000
10000
Dividing Eq. (13a) by tr(V) and Eq. (13b) by tr tr(U), we obtain U and V and can substitute them into Eq. (12).
Iteration
10a
Blases
Embedding
Transformer (PostNorm)
[Esp E ZZ Jis
EE E Z 33'
Function Spacce LR
10-1
Readcut
Weights
|
(14)
tr( tr(U)
VD Weights
10-2
Weights
FF Weights 2
To obtain the denominator, we take the trace of both sides of Eq. (13a) or Eq. (13b), giving us
tr tr(U)tt(V) tr (E[ZZ"))
(15)
10
Iteratio
E Ei
=E [1Z13]]
Figure 1. Function-space learning rates over time, measured using approach, for the ResMLP model (top) and the Transformer (PostNorm) model (bottom). "QK Weights" refers to Wy or Wk (query and key weight matrices), whilst "VO Weights" refers to Wy or Wa (values and head-connatennttin projection weight ma- trices).
Substituting this into Eq. (14), we obtain
E EΣ [ZZ
1:2)
E EEir
[Z
| RMMS =
(16)
and we note that the sums in the numerator can be computed in quadratic time (not cubic ZZT or ZT suggests), e.g.
mean Gaussian terms Wik (we will use this fact later),
Ew 20 = Zek
AWE df 6j dwr
Zij = Σ
Wak
(10)
= EE, Zix) (E
VNK
nk
E. EE,7it)
(17)
From Eq. (9) we can see that instead of directly estimating the variance of samples of AgØ, we could instead first esti- mate the covariance of the Ziy's 's and sum over them. This opens up the possibility of assuming some covariance struc- ture over the Zzy's 's to reduce the variance of our estimate for Var [o)) at the expense of some bias. One possibly restrictive example is to assume that all the Zij's are IID, in which case Var [ = E; Var [zg]. Instead, we assume a frattoedr-facttoeed covariance matrix (Martens & Grosse, 2015). Specifically, we assume that we have a covariance matrix U NeNN over rows and V TR.NNIx.M XN over columns, giving
Hence to compute the layerwise function-space sppcionn-ppceeeeee learning rate | RMS we only need to estimate 3 scalar expectations.
E
B EI, (E,
(18)
Since usually measure | 2 we estimate these expectations using exponential moving averages (EMAs) to further reduce the variance.
RMS as it changes over time,
This approach can be generalised to tensor-valued param- eters (see Appendix B); the resulting algorithm has very similar computational cost and memory; in particular, we only need to store R+ scalar EMAs for each parameter tensor, where R is the rank of the tensor (e.g. 2 for matrix). In summary (see Algorithm 1, red text), estimating the lay- erwise funetion-space learning rates involves computing a sample Z (Eq. 8), which requires a single forward pass and backward pass using a fresh batch of data to compute aw?
Cov (Zij, Zuyl = Uw Vip-
(11)
Under this assumption, Eq. (9) becomes
dó
(in addition to the usual training step which gives AW(), and updating 3 scalar EMAs (for matrix parameters). We warm up the EMAs by computing several samples of Z at the start of training, and only use one sample of Z every e.g. 100th iteration thereafter.
| = Eyjoy Cov (Zg, 20d
EiyyyUw Vis
(e) (Eny
(12)

---

Function-Space Learning Rates
ResMLP
Transformer (PostNorm)
Transformer (PreNorm)
Transtormer (PraNormPostMod)
0.001
6.5
6.3
6.6
0.000
Waltt
Train Loss
Wish 2
Standard
0.004
Mich
Werh 34
6.3
5.5
5.5
0.002
5.0
5.0
5.0
0.000
10-1
10m2
0..008
6.3
6.5
6.5
0.006
6.0
6.0
6.0
Train Loss
FLaRM
0.004
0.002
5.-
5.0
5.0
0.000
4.5
4.5
4.5
10
10
10
10-1
10
10
10
10
10
10
10
10
10
10"
10
10"
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure FLeRM dramatically improves optimal learning rate transfer across widths. Top: standard practice. Bottom: FLeRM.
3.3. Function-space learning rate matching (FLeRM)
we use the heuristic of sharing the base model's function- space learning rates between the new blocks. For example, (hase) if parameter in the first residual block has |f Il in the base model, and the scaled model has 4x as many blocks, the corresponding parameters in the first 4 blocks of the scaled model will use base function-space learning rate of |
FLcRM uses the machinery developed above to record the function-space learning rates || (hase.t) at iteration RMS
in a small, cheap base model. Then, in the larger, more expensive model, FL.RM uses the ratio between the cur- (t)
rent function-space learning rates |ff] RMS and the base model function-space learning rates | II (base.) to set the parameter-spe space learning rates at time t, such that the function space learning rates match those in the base model (see Alg. 1, blue text, for details). Note that for efficiency reasons, this is usually done at regular intervals once at the start of training, rather than every iteration.
(base)
RMS
4. Experiments
In this section we analyse function-space learning rates for concrete neural networks (Section 4.1), and investigate the use of FL.eM to enable hyperparameter transfer when scaling model width, depth, initialisation scale, and LoRA rank (Section 4.2).
There are two extra details worth discussing. First, re- call that we use EMAs for the three scalars in Eq. (18). These EMA estimates will have considerable bias if the parameter-space learning rates vary (e.g. due to schedul- ing, or FLeRM), as previous updates of the EMA may have used very different learning rate. Hence, in Algorithm (1), we always consider a learning rate of 1 when computing || IRMS- rather than the actual learning rates. Computing the A W implied by a learning rate of ensures that the learning rate seen by the EMA is always consistent. This also applies when recording the base function-space learn- |A I RMM- so the modified layerwise learning rate of becomes
Full details of the models' used in the experiments can be found in Appendix A. The base ResMLP an MLP with hidden layers, each with residual connections, trained for 50 epochs flattened CIFAR-10 10 images (Krizhevsky & Hin- ton, 2009). The base transformer is decoder-only, has two self-attention + feedforward blocks (Vaswani et al., 2017), and is trained on a subset of the Wikitext-103 dataset (Mer- ity et al., 2016). The widest transformer has roughly 814M parameters. We compare 3 different types of Layernorm (Ba al., 2016) in the transformers, with affine transforma- tions disabled: PostNorm is Norm(z + fx)) PreNorm is x+ (Norm())) and PreNormPostMod is x Norm( (())). In the ResMLP, scaling width increases the hidden dimen- sion, whilst in the transformers, we scale the embedding dimension, the feedforward hidden dimension, and the num-
ing rates
To |
hase IRMS
(19)
A RMS
where 7/n is the learning rate of the base model.
Second, if the scaled model has more layers than the base model, it is not possible to "match" the layerwise function- space learning rates one-to-one. Since we scale depth by in- creasing the number of "residual blocks" our experiments,
We provide our code at https:
//github.com com/ edwardm somm function- space- Sopacc--eaandaee Learning- watenng--attspppppe paper
5

---

Function-Space Learning Rates
ResMLP
Transformer (PostNorm)
Transformer (PreNorm)
Transformer (PreNormPeatMad)
7.0
6.00
6.00
0,008
Depth
Depth 2
6.75
5.76
Train Loss
0.006
Depth 1
5.30
5.50
Standard
Dept
0.004
heph, /*
Depth 3ugh
0.002
5.00
5.00
4.75
10-1
10
10-5
7.0
6.00
6.00
0.008
5.75
6.3
Train Loss
0.006
5.50
5:50-
F
6.0
0.004
5.20
0.002
5.5
6.00
5.00
5.0
4.75
4.75
10
10
10
10
10-14
20
10
10"
10
101
10
10"
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure 3. FLeRM improves or maintains optimal learning rate transfer across depth. Top: standard practice, Bottom: (FLcRM).
ber of heads. In all models, depth scaling increases the number of "residual blocks" that form the hidden layers of the model. Both ResMLP and the transformers used the Adam optimiser (Kingma, 2014) with a constant learning rate schedule. In all plots, train loss is averaged over the last 200 batches of training.
We also observe that the different types of layers, such as feedforward weights or QK weights, form very clear "groups" or "bands" in these plots. From this, we can see that the second feedforward weight matrices in self-attention have the strongest influence over the transformer's 's learned function, with function-space learning rates an order of magnitude larger than those of the readout layer. This is a surprising discovery, since one might naively expect the readout layer, whose weights directly project to the output logits, to have the largest effect on the learned function.
4.1. Analysing function-space learning rates
In Figure we measure the function-space learning rates using the techniques presented in Sections 3.1 and 3.2 for the ResMLP and PostNorm transformer models. Plots for the PreNorm and PreNormPostMod transformer can be found in Figure 6, though they are very similar. We use 40 batches of data warm up the EMAs as suggested in Section 3.2, then measure the funetion-space learning rate every 100 iterations. We use an EMA decay rate of B = 0.9.
4.2. FLeRM hyperparameter transfer experiments
We evaluated the effect of FLeRM on hyperparameter trans- fer when scaling model width. depth, parameter initialisa- tion scale, and LoRA rank. We ensure width invariance at initialisation by using Kaiming initialisation (He et al., 2015), and depth invariance at initialisation by introducing
In Figure| 1 we see that in both models, the function-spact learning rates change over time, despite the parameter-space learning rates being fixed. The most obvious pattern is that the function space learning rates fall monotonically for parameters, except the input embedding layer, revealing an implicit scheduling of the function-space learning rates under standard Adam training. Interestingly, in the ResMLP, whilst the function-spi space learning rates initially decay, the hidden layers and input layer eventually start increasing again, whilst the output layer plateaus. In the transformer, the embedding layer's function-space learning rate actually increases over time. One possible explanation is that the noisy initialisation in hidden and output layers effectively scrambles the signal from the input layer, but as these layers are trained, the input embedding can have clearer, stronger effect on the output.
factor of 1 into the residual stream (Hayou et al., 2021; Hayou & Yang. 2023), using FLeRM to achieve invariance throughout training. We first record the function- space learn- ing rates of the base models as in Section 4.1 using 8 random seeds, and then average over these seeds. This process is repeated for every learning rate used in our plots. When running the "scaled" models, i.e. those with altered width, depth, initialisation scale, LoRA rank, use 40 batches of data to warm up the EMAs. and then modify the initial learning rate as specified in Algorithm (1).
VL
We then have a choice. We can either use these learning rates for the rest of training, or we can periodically update them every 100 iterations with single batch of data as in Algorithm (1). We found these approaches to give very sim- ilar results, so we present the results for the fixed learning rates here, and provide plots for the periodically updated
6

---

Function-Space Learning Rates
Standard
FLeRM
The PreNormPostMod transformer is an interesting case. When scaling depth, the optimal learning rate in the standard setting is already depth-invariant, so in this setting, there is arguably no reason to apply FLeRM. Reassuringly, when we do apply FLeRM, the location of the optimal learning rate is preserved. In practice, we will usually be scaling width and depth together, so FLeRM is likely useful overall.
Int.
5.8
0.8
Int. 4-4
Int.
Train Lass
56
Int.
5:4
5.4
Int.
62-
Int. X
5.0
5.0
48
10-3
10-3
10-1 10-3
10-3
10
4.2.3. PARAMETER INITIALISATION SCALE
Learning Rate
Learning Rate
FLeRM's uses not constrained to width and depth scal- ing. In the PreNormPostMoo transformer, which uses resid- ual connections of the form I + Norm(fxz) and, like all 3 transformer variants, uses QK normalisation (layernorms applicd after W and WKX), all learnable functions in the network (with the exception of the input embedding and the readout layer) are followed by a layernomm The network should be invariant to the magnitude of these parameters, in the sense that multiplying such weight matrix by con- stant does not change the output of the network. However, as shown in Figure 4, the loss of networks trained in the standard setting varies wildly with initialisation scale, and the optimal learning rate even shifts. With FLeRM, however, we ensure that the updates are invariant in function-space, and so the loss learning rate curves line up very closely for all initialisation scales, removing the need to tune this hyperparameter.
Figure 4. FLeRM allows us to train initialisation scale invariant networks. Top: standard practice. Bottom: FLeRM.
learning rates in Appendix C.2.
1.2.1. MODEL WIDTH
Figure 2 shows the effect of scaling the width on the opti- mal learning rate for the ResMLP and transformer models. In agreement with e.g. Yang & Hu (2022), we find that in standard practice, there is a significant shift in the optimal learning rate for all models as width increases, but when using FLeRM to normalise the layerwise parameter-space learning rates, this shift is either entirely removed or dra- matically reduced. Additionally, in the transformer models, the use of FLeRM does seem to improve the loss at high widths, compared to standard practice.
4.2.4. LORA ADAPTER RANK
4.2.2. MODEL DEPTH
We also investigated the use of FLcRM in a fine-tuning setting. For this, we trained autoregressively using LoRA adapters (Hu et al., 2021) on two LLMs, GPT-2 (Radford et al., 2019) and Llama-3.2-1B (Dubey et al., 2024). We experimented with 44M token subsets of two datasets: Cold French Law (Harvard Library Innovation Lab, 2024) and Mathpile (Wang et al., 2023). For LoRA experiments, we used Adam as our base optimiser with a constant learning rate schedule, sweeping over adapter rank and the learning rate of each of the two LoRA parameters separately. We measure the base function-space learning rates for single seed and use these with FLeRM in the scaled models. As before, we run 40 EMA warm-up iterations and then nor- malise the learning rates using FLeRM at the first iteration, fixing those learning rates for the rest of training. Figure 5 shows results after sweeping over the learning rate of the B parameter LoRA (AW- = BA),, and in Appendix C.5 we show results for sweeping over the learning rate of A. See Appendix A.3 for further experimental details.
Figure 3 shows the effect of scaling depth on the optimal learning rate for the ResMLP and transformer models. Un- der standard practice, the behaviour of the optimal learning rate we increase depth varies by model. In the ResMLP, there is significant shift in the optimal learning rate, and we observe that FLeRM brings the optimal learning rates much closer together, though for higher depths there is a slight shift towards towards larger values.
The PostNorm transformer is relatively unstable, with the loss shooting up once certain learning rate threshold is passed. In standard practice, the location of this instability shifts to lower learning rates as depth is increased, meaning that deeper models actually have a worse optimal loss. How- ever, with FL,RM, these instabilities all occur around the same place as the base model (either at the same learning rate or the one before it, suggesting the "true" location of the instability could be somewhere between, at learning rate did not evaluate). Thus, in this setting, FLeRM gives dramatic improvements in performance for the deeper models.
We find that in the standard setting, the optimal learning rate increases as we increase the LoRA rank (Figure 5); for example, the optimal learning rate for the Llama model in- creases by more than an order of magnitude on both datasets as we increase the LoRA rank from 2 to 32. There is also a
In standard practice the PreNorm transformer shows a small shift in the optimal learning rate, which is rectified by FLeRM.
7

---

Function-Space Learning Rates
MathPile, GPT2
French, GPT2
MathPile, Llama-3.2-18
French, Llama-3.2-11
3.4
4,00
1.52
3.2
3.75
1.11
Train Loss
3.0
20-
Standardd
3.50
2.5-
1.50
2.6
3.95
1.5
2.4-
16-
2.2
1.49
10"0
8.4
4,00
1.52
2.2
8.2
8.76
1.51
Train Loss
3.0-
2.0
3.50
FLaRM
2.8
1.30
2.6
3.25
1.19
2.4
3.00
1.5
2.2
1.08
10
10
10
10
10
10
10
10
10
10
10
10
10
10"
10
Learning Rate
Learning
Rate
Learning
Rate
Learning Rate
Figure FLeRM improves optimal learning rate transfer when changing LoRA rank. The plots show the behaviour training loss under varying the learning rate of B and LoRA rank for two continual pretraining tasks. Top: standard AdamW optimiser. Bottom: (FLcRM). Some lines end abruptly for larger learning rates, indicating a numerical instability.
corresponding shift learning rate instabilities (i.e. where the learning rate is too high). However, when using FL.RM with r (Rank) equal to 2 as the "base model", the shift is either eliminated or greatly reduced, and the instabilities mostly align. Note that some data points missing when learning rates are too high (for example for 2 and r = 2 in the learning rate ž 10-2 range in the bottom left plot), but this is expected because this instability is inherited from the base model.
hibits depthwise hyperparameter transfer, which matches the findings of Large et al. (2024), who observed that the "im- portance" of the input and output layers must be retained as depth is scaled achieve hyperparameter transfer. Interest- ingly, although hyperparameter transfer is lost, performance in the deepest PreNormPostMod model in this setting is slightly better than the standard setting, suggesting that with a more complex heuristic for matching function-space learn- ing rates to models with new layers, it might be possible to achieve greater performance with FLeRM.
4.2.5. COMPARISON TO NAÎVELY CHOSEN FUNCTION- SPACE LEARNING RATES
5. Conclusion
One might ask whether the matching the function-see space learning rates of a base model is important, or if we could re- place them with something simpler. To test this, we repeated the width, depth, and initialisation scale experiments from carlier, but replaced the base model's recorded function- space learning rates with uniform vectors that sum to 1. The results are shown in Figures 10, 11, and 12 (in Appendix C.3). Whilst hyperparameter transfer is still retained, the training loss is slightly worse than the equivalent experi- ments in Section 4.2, suggesting that the function- spacce learning rates induced by the Adam optimiser give ben- efits to performance. This ablation still shares the base dunction-space learning rates across new hidden blocks as depth is scaled (Section 3.3). Importantly, this means that the function-space learning rates for the embedding and readout layers are held constant even as we scale depth, be- cause these layers are not replicated. We also ran further ablation where all function-space learning rates are com- pletely equal at all depths (Figure 13, Appendix C.4). In this case, the PosMormPostMod transformer no longer ex-
In this paper, we developed an efficient way to estimate the magnitude of changes in function-space caused by an update to a neural network's parameters. We used this method to analyse the dynamics of existing models, and then to mod- ify layerwise parameter-space learning rates (FLeRM) for scaled models so that updates in function-space variant, enabling transfer of the optimal learning rate across width, depth, initialisation scale, and LoRA rank. Such method could very useful for training very large founda- tion models, where scaling laws are currently derived on a case-by-case basis (e.g. see the LLama 3 technical report Dubey et al., 2024). In terms of future work, it was noted in Section 4.2.5 that more sophisticated scheme for matching function-spar space learning rates models with more layers than the base model could further enhance the performance of FLeRM. One possible approach would be to use the meth- ods presented in this paper to study the relationship between funetion-space learning rates in models of various depth. However, this beyond the scope of this work.
scale in-
8

---

Function-Space Learning Rates
Acknowledgements
transport, 2018. URL https:://axiiiieee 1 $805.09545.
Edward Milsom and Ben Anson are funded by the Engi- neering and Physical Sciences Research Council via the COMPASS Centre for Doctoral Training at the University of Bristol. This work carried using the computational facilities of the Advanced Computing Research Centre, Uni- versity of Bristol ht tpttpp// /wwwwwwkkk bris ac uk/ acrc/. We would like to thank Dr. Stewart for GPU compute re- sources.
Chizat, L. and Netrapalli, The feature speed formula: flexible approach to scale hyper-parameters of deep neural networks, 2024. URL https::/arxiv..tt abs 11111.1111.
Cottier, B., Rahman, R., Fattorini, L., Maslej. N., and Owen, D. The rising costs of training frontier ai models, 2024. URL https: / arxiv org/ abs 210....5..
Impact Statement
Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A., Caron, M.. Geirhos, R., Alabdulmohsin, I., Jenatton, R., Beyer, L.., Tschannen, M., Arnab, A., Wang. X., Riquelme, C., Minderer, M., Puigcerver, J., Evei, U., Kumar, M., van Steenkiste, S., Elsayed, G. Mahendran, A., Yu, F., Oliver, A., Huot, F., Bastings, J., Collier, M. P., Gritsenko, A., Birodkar, V. Vasconcelos, C., Tay, Y., Mensink, T., Kolesnikov, A., Pavetić, F. Tran, D., Kipf, T., Lučić, M., Zhai, X., Keysers, D., Harmsen, J., and Houlsby, N. Scaling vi- sion transformers to 22 billion parameters, 2023. URL https: https://arxiee /ttpsv////avvov/a org abs /aabss/233322 2202 05442.
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Ba, J. I.., Kiros, J. R., and Hinton, G. E. Layer normaliza- tion, 2016. URL https: /ttp//rxxii/../////////////.aa wttpa//..rgg/b////e//ea/e 06450.
Bernstein, J. and Newhouse, L. Modular duality deep learning, 2024a. URL https::/arxiv..ot / 1410.21263
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang. A., Fan, A., et al. The llama 3 herd of models. arXiv preprint urXiv:2407.21788 2024.
Bernstein, and Newhouse, L. Old optimizer, new norm: An anthology, 2024b. URL atttss//arxiv..ot
Everett, K. Xiao, L., Wortsman, M., Alemi, A., Novak, R., Liu, P. J., Gur, I., Sohl-Dickstein, J., Kaelbling, L. P., Lee, J., and Pennington, J. Scaling exponents across parameterizations and optimizers, 2024. URL https: wwwxxii.rgg/ab/////.o /ttp://////////..0557777777 05872.
abs/2409 20325.
Bernstein, J., Vahdat, A., Yue, Y., and Liu, M.-Y. On the distance between neural networks and the stability of learning, 2021.
arxiv. org/ abs/
2002. 03432.
Geiger, M., Spigler, S., Jacot, A., and Wyart, M. Disentan- gling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experi- ment, 2020(11) 1): 133301, November 2020. ISSN 1742- 5468. doi: 10. 10.1108//742 5465/aabc22-56///bcdc..aa
Bordelon, B. and Pehlevan, C. Self-consistent dynam- ical field theory of kernel evolution in wide neural networks, 2022. URL https:///x/x org/ abs/ 2205 09653.
URL http: //dx doix www.doi....g...g.oooooooooo. 100 1088 /00111000//72255////00 abbded..
Bordelon, B. and Pehlevan, C. Dynamics of finite width kernel and prediction fluctuations in mean field neural networks, 2023. URL
Gupta, and Nagar, D. Matrix Variate Distribu- tions. Monographs and Surveys in Pure and Ap- plied Mathematics. Taylor & Francis, 1999. ISBN 9781 9781588880462. URL htttps:/bbook: google. co
org/ /bs/
2304. 03408.
Bordelon, B., Noci, L., Li, Li, M. B.. Hanin, B., and Pehle- van, C. Depthwise hyperparameter transfer resid- ual networks: Dynamics and scaling limit, 2023. URL https: //triii////vv..///aa org/ abs/ 22099.9... 6620.
Haas, M., Xu, J., Cevher, V., and Vankadara, L. C. μP2; Ef- fective sharpness aware minimization requires layerwise perturbation scaling, 2024. URL https: /aaxxiv. org/ abs attaaa///...... 2411.00075
Bordelon, B., Chaudhry, H. T., and Pehlevan, C. Infinite limits of multi-head transformer dynamics, 2024. URL htg/abs//// https: / arxii.iii.v org/ abs /20.. 15712.
Harvard Library Innovation Lab, C. P. o. T. R. Cold french law dataset, May 2024. URL
Chizat, L. and Bach, On the global convergence of gradi- ent descent for over-parameterized models using optimal
https://auggingfaae wtw.aa/..oodaaaaeteeees harvaro-lil/coles f-rencch-eea

---

Function-Space Learning Rates
Hayou, S. and Yang, G. Width and depth limits commute in residual networks, 2023. URL https://araa org/ aaaaabb//3222 00453.
Mei, S., Montanari, A., and Nguyen, P.M.. A mean field view of the landscape of two-layer neural net- works. Proceedings of the National Academy of Sci- ences, 115(33), July 2018. ISSN 1091-6490. doi: 10. 10.1073/pmas. (tt00/3/paaa..805577007 180657911.. URL http: //t:://xxxxu /dw. 1806579115.
Hayou, S., Clerico, E., He, B., Deligiannidis, G., Doucet, A., and Rousseau, J. Stable resnet, 2021. URL https: //arxiv. org/ abs/ / 2010. 12859.
org ozg/10.1073/pmm
Merity, S., Xiong. C., Bradbury, J., and Socher, R. Pointer sentinel mixture models, 2016. URL httts:://ar.t org/abs/ 609 07843.
He, K., Zhang. X., Ren, S., and Sun, J. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification, 2015. URL hww.h:tt://ta hwttt:t//aaxiiii org/ htt//bbs/1002. 01852.
Noci, L., Meterez, A., Hofmann, T., and Orvieto, A. Su- per consistency of neural network landscapes and learn- ing rate transfer, 2024. URL https:://axxi...oo
Hoff, P. D. Separable covariance arrays via the tucker prod- uct, with applications to multivariate relational data, 2010. URL https: / www.i....oog............. abs////bb///00 2169.
abs/ 2ess/240..177 174577
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, Language models are unsupervised multitask learners. 2019.
Hu, E. J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, I.,, and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106. 200,0,,,,,, 2021.
Rotskoff, G. and Vanden-Fijnden, Trainability and accuracy of artificial neural networks: An interacting particle system approach. Communications Pure and Applied Mathematics, 75(9):188-11955, July 2022. ISSN 1097-0312. doi: 10. 10.1002//pa.22074, URL http: / //dx doi. oww.ooog///.... 101pp/1....000//p...2201
loffe, S. and Szegedy, C. Batch normalization: Acceler- ating deep network training by reducing internal covari- ate shift, 2015. URL https: //tp///aa/vv///a////a/eaeeaaaal. org/ abs/ 1502. 03167.
Sirignano, J. and Spiliopoulos, K. Mean field analysis of neural networks: A law of large numbers, 2019. URL https: arvvv..vvvvv. www.orr///aeo abs/
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint arXiv: 2014.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit. J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need, 2017. URL https://aaxat abs / /1/./77666600 01.706..333 03762.
Krizhevsky, A. and Hinton, G. Learning multiple layers of features from tiny images. Technical Report 0, University of Toronto, Toronto, Ontario, 2009.
Large T., Liu, Y., Huh, M., Bahng, H., Isola, P., and Bern- stein, J. Scalable optimization in the modular norm, 2024. URL htRp:t/:::/ http://axxvvv/////
Wang, Z., Xia, R., and Liu, P. Generative ai for math: Part i-mathpile: A billion-token-scalle pretraining corpus for math. arXiv preprint arXiv:23/2.17/22) 2023.
2405 14813.
Manceur, A. M. and Dutilleul, P. Maximum likeli- hood estimation for the tensor normal distribution: Al- gorithm, minimum sample size, and empirical bias and dispersion. Journal of Computational and Ap- plied Mathematics, 239:37-49, 2013. ISSN 0377- 0427. doi:
Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. L.,, Gugger, S., Drame, M., Lhoest, Q. and Rush, A. M. Huggingface's 's transformers: State-of-the-an natural language processing, 2020. URL https: /tpps://ariivee http/axxvvvvvv org. abs 910. 03771.
017. URL URL.stppss//w...ceeeeeieeeeeeeee com/
Yaida, S. Meta-principled family of hyperaarameter scaling strategies, 2022. URL 2210 04909.
Mangrulkar, S., Gugger, S., Debut, L., Belkada, Y., Paul, S., and Bossan, B. Peft: State-of-the-art parameter- efficient fine-tuning methods. https://githu: com/hugg www.anacceeccoo ace/p 2022.
Yang, G. Tensor programs ii: Neural tangent kernel for any architecture, 2020. URL htttt:t//txxii...ttt Abs/2006.1454g
Martens, J. and Grosse, R. Optimizing neural networks with kromecker-factored approximate curvature, 2015. URL htttss://arxin woo./org.//...............o 05671.
Yang, G. Tensor programs ii:: Neural matrix laws, 202 a. URL https:/tts//a /0000//////////////////////////0
10

---

Function-Space Learning Rates
Yang, G. Tensor programs i: Wide feedforward recur- rent neural networks of any architecture gaussian processes, 2021b.
org////vvrrroeeeo abs/
1910.12478,
Yang, G. and Hu, E. Feature learning in infinite-width neural networks, 2022. UtL.htts::/arxiiv.o
abs/2011 14522.
Yang, G., Hu, E. J., Babuschkin, I,, Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer, 2022. URL https: / arxiv. htrg/abs/220 03.66.
Yang. G., Yu, D., Zhu, C., and Hayou, S. Tensor pro- grams vi: Feature learning in infinite-depth neural net- works, 2023. URL https: /ttp:///aaivi.ee. org/oa/ggaoo /bs/ 2310. 02244.
Yang, G., Simon, B., and Bernstein, A spectral condition for feature learning, 2024. URL https: //arxiv. org/ abs/2310. 17813.

---

Function-Space Learning Rates
A. Model Details
A.I. ResMLP
The base residual MLP model has 4 residual blocks, each containing a single linear layer. Every hidden layer has dimension 128, which is multiplied by the width multiplier in width-scaling experiments. In residual blocks is multiplied by the depth multiplier. We do not use Layernorms or Batchnorms in the ResMLP (Ba et al., 2016; Toffe & Szegedy, 2015).
experiments, the number of
We initialise all weight matrices using Kaiming / He initialisation (He et al., 2015), that is, IID Gaussian weights with Ean_in variance, multiplied by an activation function-speeiiffc scalar (2 in the case of ReLU, 1 for no activation). Biases are initialised to
To ensure depth invariance at initialisation, we introduce a factor of VT into each layer's weight matrix. At initialisation, this is equivalent to multiplying by V/ in the forward pass as in Bordelon et al. (2023), but does not continue affect the forward passes during training, allowing us to better isolate the effect of FLeRM.
We optimise the model using Adam with default settings (other than the learning rate which we set using the methods detailed in this paper). We train for 50 epochs on the CIFAR-I0 dataset (Krizhevsky & Hinton, 2009).
A.2. Transformer
The base transformer model is decoder-only transformer with self-attention + feedforward layers, query-key normalisatinn (Dehghani al., 2023, applying layernorm the queries and keys before using them in multihead attention), and layernorms with affine transformations disabled. We compare 3 different types of layeroor in our experiments: PostNorm is Norm(r + f(x)), PreNorm is x + [Norm(zz) and PreNormPostMod is + Norm(f(z))
When scaling width, we multiply the number of heads, the embedding dimension dmmodet and the feedforward hidden dimension du by the width multiplier. The base model has dmodel = 128, dr = 512, and 2 attention heads per layer. When scaling depth, we multiply the number of transformer "blocks", consisting of self-attention and feedforward network, as detailed in Vaswani al. (2017), by the depth multiplier.
We initialise all weight matrices using Kaiming / He initialisation (He et al., 2015), that is, IID Gaussian weights with Zen_in variance, multiplied by an octivation-function-ss sspeciiiicssccceeeeeeeeo scalar (2 in the case of ReLU, 1 for no activation). Biases are initialised to 0.
To ensure the model is invariant to depth at initialisation, we multiply the residual branch by A during the forward pass, as in Bordelon et al. (2023). More specifically, PostNorm becomes Norm(x+ (x)), PreNorm becomes x (Norm(rr) and PreMorrPostMood becomes x + J. Nfrm(((xx)) Note that with PostNorm and PreNorm we could have absorbed this factor into the weights of the module f and therefore avoided altering the forward pass computation, because f is at the end of the residual stream. However, in PreNormPostMod the layernorm is at the end of the residual stream, so its initialisation will always have unit size matter how we initialise f. therefore requiring to use the factor vT. during the forward pass. We therefore decided to treat all transformer variants the same in this regard.
We train the transformer on roughly 1 of the Wikitext- 103 dataset (Merity et al., 2016), using a batch size of 20 and a sequence length of 256. We tokenise the dataset using the GPT2 tokeniser from the HuggingFace transformer library (Wolf et al., 2020). We optimise the model using Adam with default settings (other than the learning rate which we set using the methods detailed in this paper). We train for 1 epoch (i.c. we only observe each token once).
A.3. LoRA adapters
For the LoRA experiments in Section 4.2.4, we trained LoRA adapters on continual pretraining tasks. The LoRA adapters were initialized using the gaussian initialization provided by Hugging Face's 's peft library (Mangrulkar et al.. 2022). We experimented with two models, GPT2 and Llama-3.2-11 We added LoRA adapters to the default modules of GPT2, and the q/ ///////poor/ modules of Llama-3.2-1B. We trained for 500 iterations with a batchsize of 8 and sequence length 512. The datasets used were 44M token subsets of Cold French Law (Harvard Library Innovation Lab, 2024) and Mathpile (Wang al., 2023).
12

---

Function-Space Learning Rates
LoRA adapter formed of two parameters B and A, with AW BA. When sweeping learning rates, we vary the learning rates of each parameter (A/B) separately, while keeping the learning rate of the other parameter fixed. When sweeping for B, we used fixed learning rate of 10-4 for A. When sweeping for A, we fixed the learning rate of B as follows: 10 for GPT-2, 10 4 for Llama-3.2-1B Cold French Law, and 5 x 10 for Llama-34 1amma-3.2-1 3 1/MMttppie..
As in the transformer and ResMLP experiments, we wish to warm-up the EMAS of FLeRM with 40 batches of data, and then use FLeRM to normalise the learning rates at the initial iteration, fixing these learning rates for the rest of training. However, we cannot immediately use our method FLeRM when using LoRA adapters. This due to the way LoRA adapters (following the original work of Hu et (2021)) are initialized. Since the parameter B is initialized zero (to ensure that AW at the beginning of training), the gradient of A is 0 for the first iteration. This is problematic because it implies ||A:/ll RMS 0, leading to division by zero in Algorithm 1. As a workaround, we warm-up FLeRM on the fifth iteration, and record / set the learning rates on the 6th iteration. We find that in practice, this still gives adequate hyperparameter transfer. In the LoRA experiments, we measure the base function-space learning rates for a single seed and use these with FLeRM in the scaled models. This is because the base model is a lot more expensive than in the other experimental setups. Only using one seed for the base function-spp -ppcc learning rates did not appear to hinder FL.eM''s ability enabled LoRA-rank transfer.
B. Kronecker factored covariance approximation for tensor-valued parameters
Previous work has extended the matrix-normal distribution to tensors, with one covariance matrix per dimension (e.g. Manceur & Dutilleul, 2013; Hoff, 2010). We can use this to extend the methods Section 3.2 to tensor-valued parameters.
Suppose we have D-dimensional tensor X E R X where we will use Xi.....ip denote the (i1, ... iD) ) element. Extend the vee(.) operation tensors in the obvious way, i.e. form a vector by taking elements (1, 1,... 1) to (m,, ...., 1, 1), then (1,2,--.. 1) the final element (#i,.... ,.D). We say the vector X is tensor-normal distributed if
1)1 to (21,2....5 1, 1) etc. systematically iterating along the dimensions until you reach
vec(X) ~ (V(ec(MM) U(D) x
@ U())
(20)
where M E RoI... is the mean tensor, and U E R _ U(D) E R are covariance matrices for the D different dimensions. In other words, assuming M = 0, the second-order moments covariance of any two elements in the tensor factorises across dimensions:
,....
D II iaja
E(X.....h X jn)
(21)
For distributions, had identities for the second order matrix products | Utr(V) and E[X"X] Vir(U), which we used in Section 3.2 to figure out how to estimate the covariance matrices (or more specifically, the sum over all pairs of elements). Do similar identities hold for the tensor normal distribution? It turns out yes (c.g. Proposition 2.1 of Hoff, 2010).
Consider the following shorthand for "contracting" two tensors over all but one of their dimensions d, kind of like a generalisation of matrix multiplication
(A X B)g
Σ
Aki..k jikal1..kk Br 1...p
(22)
ie.. assuming A and B have conformable shapes, you take "dot products" over all dimensions except d. Then for a
13

---

Function-Space Learning Rates
tensor-normal distributed X have the "second-order moments"
E X X)
Σ
E(X KI..... ka- ka tthee .... Xki.... .ae ha-1skatt...ko
(23a)
Σ
Uid') Barskar
(23b)
k1... Kiat1 ...k
d'
Uld)
II
Σ
Uld Ke
(23c)
d' d'
U:d II
tr(Uld))
(23d)
d'
and have
BX = U II tr(Uld)) d'
(24)
As before, we want to compute the sum over all elements of the covariance matrix. The covariance matrix for the vectorised iensoris is (D)) @ ....00 @ UO), which means it contains precisely all the possible products of D elements, one from each U
Hence we wish to compute
E Lads E(X X X)issa
II d=1
Σ
U
isia
(d)
II
He tr(U (r(())
(25a)
d=1
I Eiaja E(X X X Jasa lo (r((d)))--
(25b)
As before, we can express the denominator in terms of X. Note that
D
D
Il (((())) =r U II (T((O
(26a)
d=2 ir (B X1 X))
(26b)
E(X X1 X)+
(26c)
Σ Σ
E O(X Xikk...kp kkk
(26d)
12...DD
E
Σ
x2 X
(26e)
kg...
E
X
(26f)
where we have used the Frobenius norm to represent the sum of all squared elements of the tensor. So the denominator will
be [[ X [2) D-1
Also similar to the matrix norriil case, the numerator very cheap to compute. Observe that
Σ B(X Xd X)iata = E Σ Σ
Ka kp Xba. iicka
(27a)
fa.Ja
saria by kat kati...k
E
Σ
Σ
Xkg richa Xky..kk io
(27b)
ki...kd Kett..... tarde
Σ
Xen. pidkett
(27c)
Ka. pkasi.nkp
14

---

Function-Space Learning Rates
which in pseudocode can be written as Ex sum (d) square () sum ())
Hence, using pseudocode to make things clearer, our normaliser can be estimated as
El sum (d) square () sum ())
||A,f||RM RMMS R
(28)
() .sum )-
where 2 is our tensor of updates times ahiigradients, and the expectation symbols E tell us what we're taking EMAs of. Remarkably, we can see that we only have track D scalar EMAs, and updating these EMAs only requires simple sum () and square () operations. For numerical stability, it is wise to compute this quotient in log-domain.
C. Extra Plots
C.1. Full function-space learning rate vs time plot
Figure 6 shows the function-space learning rates over time as in Figure 1 in the main text, but also with the Transformer (PreNorm) and Transformer (PreNormPosttood) models. The plots for the three transformer layernorm variants all show very similar behaviour.
C.2. FLeRM with periodic updates to the learning rate
In the FLeRM experiments in the main text (Section 4.2) we only used FLeRM to modify the learning rate during the first step of training, and then used that learning rate for the rest of training. There is nothing to stop us from periodically updating the learning rate throughout training using FLeRM if necessary, e.g. if we think the dynamics of training might be affected by width or depth in a time-dependent manner. Here we present plots similar to Figures 2, 3, and 4 from the main text, but in addition to using FLcRM to compute the learning rate at the start, we also update the learning rate every 100 iterations using a single batch of data and an EMA decay rate 3 = 0.9. The resultant width, depth, and initialisation scale transfer plots are given in Figures 7, 8, and 9. The plots are very similar to those using the static learning rate, suggesting that the effects of increasing width, depth, or initialisation scale can effectively cancelled-out using learning rate computed purely at the start of training. This agrees with prior work like μP (Yang & Hu, 2022) and Module (Large al., 2024), which both achieve hyperparameter transfer in width and depth scaling settings using static learning rates.
C.3. FLeRM with equal base model funetion-space learning rates (but still splitting them in the same way as depth scales
Here we replace the recorded base model function-space learning rates with equal values for each layer, that add up 1. As depth is increased, with still split the function-space learning rates up between the new hidden blocks (described in Section 3.3) as we did in the main experiments. See Figures (10), (11), and (12).
C.4. FLeRM with equal base function-space learning rates that are ALWAYS equally divided across layers, even with scaled depth
This ablation is similar to the equal base function space learning rates ablation in Appendix C.3, but instead of sharing the values across the new hidden layer blocks as described in Section 3.3, we simply set all the "base" function-space learning rates for the deeper model to be equal and sum to 1. Note that a key difference in this approach is that the input embedding and readout layer's base funetion-space learning rate is going to shrink as the model gets deeper, whereas in the ablation in Appendix C.3, the base function-space learning rates for the input embedding and readout layer remained constant, whilst only the hidden layers got diluted when scaling depth. As suggested in previous work (Large et al., 2024), ensuring input embeddings and readout layers retain their "importance" during training in deeper models is critical for achieving depth transfer, and we observe in this ablation (Figure 13, only depth scaling is changed from the other equal function-space learning rate ablation) that we no longer have hyperparameter transfer across depth. Interestingly, however, the performance of the deepest PreNormPostMod model is better in this ablation than in the other FLeRM experiments, actually marginally beating the standard practice model too, suggesting further refine of the scheme for splitting base FSLRs as depth increases could be beneficial. However, performance in this ablation is also far worse on other models like PostNorm at high depths.
15

---

Function-Space Learning Rates
C.5. LoRA
In Section 4.2.4, we showed results for sweeping over the learning rate of the LoRA adapters' parameter, with fixed learning rate for A. Here, we show results for the reverse: sweeping over the learning rate of A with fixed learning rate for B. Figure 14 shows the results. Interestingly, we achieve learning rate transfer with the Standard AdamW optimiser, mainly because varying the learning rate of A does not change the performance of the model (until it becomes too large). This suggests that selecting the learning rate of A is not very important compared to selecting the learning rate of B. We find that with FLcRM, we mostly preserve the learning rate transfer, and we also observe an instability shift to the left.
C.6. Bias and variance of estimator under difference covariance assumptions
In Figure 15 we show that our proposed Kronecker-factored assumption in Section 3.2 greatly reduces the variance of the function-space learning rate estimator whilst avoiding the large bias an oversimplifying assumption such as an IID covariance structure would introduce. Here the "no assumption" estimator is used as the "true" value for computing the biases, as is an unbiased estimator.
C.7. Test loss
In Figure 16 we show the width transfer plot from the main text but with test losses instead of train losses. As expected, the test loss plots for the Transformer models look very similar to the main text, since are only training for epoch and therefore expect test loss look like train loss. However, the ResMLP is trained for multiple epochs on CIFAR-IO, and so complex overfitting patterns occur, which probably explains why the test loss curves do not line up perfectly, though the optima are closer together with FL.RM than with standard practice.
C.8. Cosine Annealing Scheduler
For simplicity used constant LR scheduler in the main text. To verify that method works with LR schedulers, reran the width transfer experiments using cosine annealing scheduler. Since FL.RM modifies the layerwise learning rates, we record the ratio of the current scheduled learning rate with the starting learning rate in the base model, then apply these ratios as the scheduler in the scaled model. This is equivalent to scheduling the "target" function-space learning rate for each layer. The results are shown Figure 17.
C.9. Elementwise Affine Transformations
For simplicity in prototyping, elementwise affine transformations in layernorms were disabled in the main experiments. To verify our method still works with these enabled, we reran the width transfer experiment for the PreNormPostMod transformer, and, as seen Figure 18, found no issues.
C.10. Simultancous Width + Depth Scaling
To test whether our method still works when simultaneously scaling width and depth, we ran the PreNormPostMod transformer experiment again, scaling the dedth+ and found that hyperparameter transfer still holds in this setting. The results are shown in Figure
C.11. Other optimisers
To test whether our method still works when using optimisers with Adam, we repeated the Transformer (PreNormPostModd width transfer experiments for range of different optimisers.
In Figure 20 we used SGD with momentum. In Figure 21 used SignSGD. In Figure 22 used AdamW. In Figure 23, we used Adamax. In Figure 24, we used Adagrad. We can see that FLeRM improves hyperparameter transfer in all cases. In SGD, the curves are very noisy both with and without FLeRM, matching the common practice of not using SGD for transformers. However, even though the curves are less neat than the other optimisers, we still see that FLeRM aligns the curves better.
To demonstrate this a problem with using SGD with transformers and not FLeRM, in Figure 25 we ran the ResMLP width transfer experiment with SGD momentum, where observely that FLeRM very cleanly improves hyperparameter
16

---

Function-Space Learning Rates
transfer.
D. Exploiting known structure for lower variance function-space learning rate estimation
A typical neural network ends in a "readout" layer {W4, b"), mapping from the hidden dimension dmodel to the number of output classes/ vocab size via the transformation
f = - --1 + ba
(29)
where h is the output of the previous layer for the n datapoint.
This layer can be very large if K is large (the vocab size of LLM), and can have large effects on the training dynamics of the neural network (Fig. 1). It is therefore important to ensure its function-space learning rate estimate is as accurate as possible. Luckily, it turns out that by nature of being the final layer, there is extra independence structure in the function-space learning rate estimation procedure we can exploit to lower the variance.
Subbing Eq. 29 into Eq. 3, we have
Σ
wikk by + nkk
Σ
Who 1--1
(30)
VNK
and subsequently subbing Eq. 3 into Eq. 8, considering the function-space learning rate for the readout weight matrix WL E IR x dontat we obtain
2 (w"
AWL
Σ
Wash L-1
(31)
VNK-
Note that this only includes single random variable wnit This means the random variables (W different values of (i.e. between different rows), but dependent for different values for (i.e. between different columns). Hence we can assume independence between the rows of / ((W") ddrastcalll reducing the variance of our estimator.
( Jis independent for
In particular, this corresponds to assuming that U is diagonal in Section 3.2, which results in |
RMS
(E. (Eis Vir) tr(U) (Eor Vis)) and therefore our estimate becomes
(E IZ'
Aw RMS tr(U)
E
> Zz Zli3
(32)
tr(U)
33'
Similarly for the biases b4 R have
2,609 AbD
Whi
(33)
VNK
and so 2(b²) are independent for different i, and so we can assume the covariance matrix is diagonal. we assume a diagonal covariance structure over 2 (b*) (bb²) i.e. = (lag((")) then
K
( 2
|| RMS
Σ DN
E
(34)
1=1
t=1
One might ask at this point how this differs from an IID assumption, i.c Σ = II. With an IID assumption, we would have
|
Eas ΣE' no² = =nE (2(b9) = nE
K
2
(bb)))
2
E
Σ
K
(b)) 2
(35)
RMS
i=1
17

---

Function-Space Learning Rates
which is exactly the same! It turns that, because only care about estimating the sum the variances, assuming that z(b") (bb) are IID just independent are equivalent. Intuitively, consider the difference between making "no assumption" on the covariance matrix versus a diagonal assumption, if we care about the sum over all covariance elements. The diagonal assumption allows us to remove cross-terms from our estimate, since we know their true values are all zero, and hence we remove the variance from the randomness in those cross-terms, giving us a lower variance estimator. But if we consider the remaining positive sum over the variance (diagonal) terms, knowing they're 're all equal (an IID assumption) doesn't tt provide us with any more useful information, because doesn't 't tell us any more about the value of the total sum, and it doesn't 't tell us any more about which values should pay more attention (in fact, it tells already did when only knew that the covariance diagonal). If we had more specific information, such "the last element's s variance forms 90% of the total variance" then we could use more sophisticated weighted averages to get lower variance estimate.
treat them all equally, which is what
In Figure 26, we repeat the width transfer experiments from the main paper but using the tricks we just described, and find that some of the instability / noise in the ResMLP FLeRM plot is now gone.
18

---

Function-Space Learning Rates
10a
Blases
Input
Function Space LR
Oulput
Hidden
ResMLP
10-3
10-
2000
0000
8000
10000
lteration
Blases
Embedding
Transformer (PostNorm)
Function Function-Space Space LR
Readout
10-1
OK Weights
VO Weights
10 2
Weights
Weights
10-3
10-*
500
1000
1500
2000
2500
lteration
10a
Blases
Embedding
Transformer (PreNorm)
Function Space LR
Readout
10 1
Weights
Weights
Weights
10-2
FF Weights 2
10-8
4
500
1000
1500
2000
2500
Iteration
Transformer (PreNoorPPsiMMod
10a
Biases
Embedding
Function Space LR
Readout
10
OK Weights
VO Weights
10-2
Weights
Weights
10-9
500
1000
1500
2000
2500
Iteration
Figure Function space learning rates over time, measured using our approach, for the ResMLP model (row 1) and the transformer models with diffcrent layemorm strategics (rows 2, 3, and 4).
19

---

Function-Space Learning Rates
ResMLP
Transformer (PostNorm)
Transformer (PreNorm)
Transtormer (PreNormPostMod)
0.008
0.000
Train Loss
Standardd
0.004
Wich 22
With
6.3
5.5
5.5
5.0
0.000
10-3
10-1
10-3
10m2
10"1
0..08
6.3
6.5
6.0
0.006
6.0
6.0
Train Loos
FLeRI
0.004
5.5
0.002
50-
5.0
3.0
10
10
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure 7. Normalising function-space learning rates periodically using FL.RM gives very similar width transfer to normalising using FLcRM at the first iteration only (see Figure 2). Top: standard practice. Bottom: our scheme (FLcRM), PreNorm, PostNorm, and PreNormPostMod are different Layemorm configurations deseribed in Section A
ResMLP
Transformer (PostNorm)
Transformer (PreNorm)
Transformer (PreNormPosttMod)
7.0
6.00
6.00
0.005
Depth 2
5.75
6.5
Train Loos
0.006
Depth
5.50
5.50
Peph 35
Depth g
6.0
0.25
0.002
5.5
5.00
4.75
10
1.3
10
10
10
30
10"
10"
10
10
10
10
0.008
5.75
5.75
Train Loss
FLeRM
0.004
5.25
5.95
0.002
5.00
5.
4.75
4.75
10
10
10
10
30
10
10
10
10"
10-2
10
10"
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure Normalising function-space learning rates periodically using FLeRM gives very similar depth transfer to normalising using FLeRM at the first iteration only (see Figure 3). Top: standard practice. Bottom: our scheme (FLeRM).
20

---

Function-Space Learning Rates
Standard
FLeRM
6.0
Inf. 4
Tradomme (PreNortt wowooooroo
Int.
Train Loss
3.6
5.6
Init
Init
3.0
5.0
Learning Rate
Learning Rate
Figure 9. Normalising function-space learning rates periodically using FLRRM gives very similar initialisation scale transfer to normalising using FLeRM the first iteration only (see Figure 4). Top: standard practice. Bottom: our scheme (FLeRM).
ResMLP
Transtormer (PostNorm)
Transtormer (PreNorm)
Transtormer (PreNormFostMod)
6.6
0.006
66.0
6.0
6.0%
Train Loss
5.5
5.5
3.5
0.002
0.000
4.5
4.5
4.5
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10
0.008
6.5
0.006
Train Loss
0.004
0.5
5.5
5.5
5.0
0.000
10a
10-2
10-1
10-2
10-1
10-0
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure 10. Equal base FSLR ablation (width): Using equal base model function-space learning rates results in a degradation in performance of the models, compared the main experiments. Top: standard practice. Bottom: our scheme (FLcRM) with equal base model function-space learning rates.
21

---

Function-Space Learning Rates
ResMLP
Transformer (PostNorm)
Transformer (PraNorm)
Transformer (PreNormPostMMdd)
1.0
6.00
6.00
0.008
Depti
Depth 21
Train Loss
0.000
Depth
6.50
5.50
Stardaad
Dapth 2"
Doph 20
5.25
6.5
0.0012
5.00
5.00
5.0
4.75
10
10
10'22
FLeFM (abletion ccpiiion
0.008
0.70
6.70
Train Loss
0.006
8.0
0.004
6.25
5:25-
5.5
0.002
4.70
4.75
10
10"3
102
10
10"
10a
10
10
10/
10*
10
10-1
10P
10'
Learning
Learning Rate
Learning Rate
Learning
Figure Equal base FSLR ablation (depth): Using equal base model function- peation--spce learning rates results in a degradation in performance in some of the models, compared to the main experiments. Top: standard practice. Bottom: our scheme (FLcRM) with equal base model function-space learning rates.
Standard
FLeRM (ablation, see caption)
Inf.
Traneforma (P PoooomPossttdd
5.8
5.8
Ini.
Inft
Train Loss
5.6
5.6
Inii,
5.2
5.2
Int,
5.0
5.0
10
10"2
100
10"
10²
Learning Rate
Learning Rate
Figure 12. Equal base FSLR ablation (init. scale): Using equal base model function-spant learning rates results in a degradation in performance some of the models, compared to the main experiments. Top: standard practice. Bottom: our scheme (FLeRM) with equal base model function-space learning rates.
22

---

Function-Space Learning Rates
ResMLP
Transformer (PostNorm)
Transformer (PreNorm)
Transformer (PreNormPostModd)
0.00
0.008
Dept
Dopth 2
0.3
5.75
5.75
Train Loss
0.00GG
Depth 20
5.30
0.004
Depth 3"
5.25
Dosth 202"
5.00
5.0
4.75
10
7.0
6.00
6.00
0.008
Train Loss
0.006
5.50
5.50
5.25
5.5
0.001
5.00
5.00
50
4.75
10-3
10-1
107
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure 133. Equal base FSLR ablation where don't carefully retain the importance of the input embedding and readout layers (depth): Just setting the model FSLRs to equal and sum to one, no matter the depth is, means that we no longer have depthwise hyperparameter transfer in the PreNormPostMod model. Top: standard practice. Bottom: scheme (FLeRM) with equal base model function-space lcarning rates that don't 't try to account for depth increase.
MathPile, GPT2
French, GPT2
MathPile, Llama-3.2-18
French, Llama-3.2-1B
3.0
4.00
1.52
Baank
2.2
Rank
Hank
3.70
1.51
Train Loos
2.5
Rank =2*
2.0
puepuress
Bank
3.50
1.50
2.6
3.25
1.09
3.00
1.5
2.4
1.08
101
10
10
10
30
10
10
10
10
10
10
10
101
10
10
2.2
1.51
Train Loss
2.5
FLoRM
1.5
1.40
1.8
2.4
1.48
10-4
10
10-2
10-1
10-4
10
10-2
10-1
10-4
10-5
10-4
10-3
10-2
10-4
10-5
10-
10-3
10-2
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure Behaviour of training log likelihood under varying learning rate LoRA rank two continual pretraining tasks. The top row shows results from standard setup with AdamW, the bottom row shows method, FLeRM.
23

---

Function-Space Learning Rates
Layer wttpss....f..et........
1500
1250
40
1000
Variance
[Bias]
30
750
20
500
10
250
No Assumption IID
KFAC
Assumption
KFAC
Figure Comparison of bias and variance of the function-spant learning rate estimator when using different covariance matrix dependence assumptions in Section 3.1 3.2. Computed over 10,000 batches to estimate the FSLRs at the first step of training for feedforward layer the network. Bias taken
absolute difference from "No assumption", since no assumption (i.e. ignoring Section 3.2 and estimating the quantity directly) is unbiased estimator. We see that IID assumption has extremely low variance (not even visible on the plot), very high bias, whilst the unbiased "no assumption" estimator has very high variance. The KFAC estimator proposed in Section 3.2 has small amount of bias and much smaller variance than the "no assumption" estimator.
ResMLP
Transformer (PostNorm)
Transformer (PraNorm)
Translormer (PrsNormPostMod)
6.5
6.5
6.5
Wich
Wich
1.6
Wich
Wich
Standard
Wich
Teat
1.5
Wich
5.0
5.0
5.0
10
10
10
10-3
10
10
10
10
10
10
10
10
10
10
10
8.0
0.0
0.0
FLoRM
1.3
1.3
4.5
4.5
4.5
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10
20
19
10
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure Width transfer plot for plotting test loss instead train loss.
24

---

Function-Space Learning Rates
ResMLP
Transforeer (PostNorm)
Transtormer (PreNorm)
Transtormer (PreNormPostMod)
0.008
6.6
0.006
660-
6.0-
6.0-
Train Loss
Standard
0.004
5.5
5.5
5.5
0.002
0.0000
4.5
4.5
4.5
10-
10
10
10
10
10-1
10-1
10-1
10
10
10
10
10
10-1
10
10
10
10-22
10
0.008
6.5
6.5
0.006
6.0
Train Loss
FLoRM
0.004
1.5
5.5-
5.5
0.002
50-
5.0
0.000
4.-
10-3
10-3
10-1
10-1
10-3
10-2
10-5
10
10-2
10-1
10-5
10-3
10-1
Learning Rate
Learning Rate
Learning Rate
Learning Rate
Figure 17. Width transfer plot for Transformer (PreNormPosttMd) using a CosineAmnealingL. scheduler.
Stardard
FLetM
1.75
5.7
Teal Looa
4.7
475
30-4
30"
Learning
Learning
Figure 18. Width transfer plot for Transformer (PreNormPostMod) with elementwise affine transformations enabled the I Layeroors..
FLORM
6.25
6.25
6.08
With und Cigah
With Caphh
6.6
1000
5.2
5:0
5.c
Rate
Figure 19. Width Depth simultancous scaling transfer plot for Transformer (PreNormPostMod),
25

---

Function-Space Learning Rates
Tea L004
Leaniig Rase
Figure 20. Width transfer plot for Transformer (PreNormPostMool using SGD instead of Adam. Momentum 0.9.
FLRRM
Wilh
Hith Weth
Tral 1.30
0.20
1.2
Figure 21. Width transfer plot for Transformer (PreNormPosttod) using SignSGD / Signum instead of Adam. Momentum 0.9.
Wet
3.2
Loaa
17
4.3
Pate
Teamiin
Figure 22. Width transfer plot for Transformer (PreNormPPotMMod) using AdamW instead of Adam. Weight decay 0.1.
Wich
Warth
1.2
477
Learing
Figure 23. Width transfer plot for Transformer (PreNormPostMod) using Adamax instead of Adam.
26

---

Function-Space Learning Rates
4.75
$25
Tean Tea
473
Teamiig
Figure 24. Width transfer plot for Transformer (PreNormPostMod) using Adagrad instead of Adam.
16-81
0.000
Trant.00
0.027
Learning Firm
Laatring Tate
Figure 25, Width transfer plot for ResMLP using SGD instead of Adam. Momentum 0.9.
ResMLP
Transtormer (PostNorm)
Transformer (PreNorm)
Transformer (PreNormPostMod)
0.008
0.5
6.5
6.5
Wich
Wicht
Train Loss
Wirh
Standard
0.004
Wich
Witth
0.002
5.0
5.0
5.0
0.000
14
0.008
6.5
6.5
0.006
0.0
6.0
6.0
Train Loss
FLoRM
0.002
5.-
50-
5.0
0.000
4.5
4.5
4.5
10
10
10-4
10
10-5
10
10-1
10-1
10
10
10
10
10-2
10
10"
10
10-2
10a
Learning Rate
Learning
Learning Rate
Learning Rate
Figure 26. Width transfer experiments in main paper, but using output independence tricks given in Section D
27


Reviews:
[{"review_id": "9DVoyhhwUt", "reviewer": "Reviewer_pf7i", "summary": "This paper introduces the concept of function-space learning rates, which measure the magnitude of changes in a neural network's output function in response to updates in parameter space. The authors propose an efficient Monte-Carlo-based method to estimate these function-space learning rates and introduce FLeRM, a technique designed to transfer hyperparameters from smaller \"base\" models to larger ones as updates in function space are scale invariance. The effectiveness of FLeRM is demonstrated empirically through multilayer perceptrons (MLPs) and transformer architectures, enabling the transfer of the optimal learning rate across width, depth, initialisation scale, and LoRA rank.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review with many excellent suggestions, acknowledging that \"3. The paper introduces FLeRM, a robust solution for hyperparameter transfer that can be used with any network architecture and at any point during training. 4. The paper extensively evaluates FLeRM in multiple scenarios, including scaling width and depth in MLPs and transformers, varying initialization scales, and adapting LoRA rank, demonstrating its versatility.\"  \n\nWe have added the additional experiments your requested (see below).\n\n### 1) Other optimisers\n\nWe have run the Transformer (PreNormPostMod) width transfer experiment using SGD, signSGD, AdamW, AdamMax, Adagrad instead of Adam. Please see [Rebuttal Figures 6-10](https://fslricml2025rebuttalsfigures.tiiny.site) (hyperlink to anonymous url). As before, FLeRM aligns the optimal learning rates, and in this case even improves the best loss achieved. There is instability when training the transformer with SGD (SGD is not often used for transformers for this reason). We also tried SGD on ResMLP in [Rebuttal Figure 11](https://fslricml2025rebuttalsfigures.tiiny.site), which worked fine.\n\n### 2) Scaling width and depth simultaneously\n\nWe have run the Transformer (PreNormPostMod) experiment scaling both width and depth simultaneously (up to 8x because of computational constraints) and observed hyperparameter transfer. Please see the [Rebuttal Figure 5](https://fslricml2025rebuttalsfigures.tiiny.site).\n\n### 3) LR scheduling\n\nWe have run the Transformer (PreNormPostMod) width transfer experiment with the CosineAnnealingLR scheduler and observed hyperparameter transfer. Please see the [Rebuttal Figure 3](https://fslricml2025rebuttalsfigures.tiiny.site).\n\nNote that the limited time available means we haven't been able to run the full range of models from the original manuscript for all of these new settings. We will run a systematic sweep for the camera ready. We hope you agree that the results we do have indicate it is unlikely that there will be any surprises in the final figures.\n\n\n### 4/5 Model size\n\nWe have added details on model sizes to the manuscript.  Specifically, the widest model (in the width scaling plot in Figure 2) we considered contains 814M parameters. We agree that this is of course far smaller than the very largest modern LLMs.  But at the same time, it isn't so small (e.g. there is alot of interest at the moment in training ~1B parameter reasoning models). Please also remember that FLeRM forms only one part of our contribution, with our main contribution being the efficient estimate of layerwise function-space learning rates, which has many possible uses, including analysing training dynamics (as shown in Section 4.1) and hyperparameter transfer (FLeRM, section 4.2).\n\n\n> The learning rate update in Algorithm 1 (line 138) seems to contradict the theoretical analysis. I believe the numerator and denominator should be in reverse order.\n\nThanks! Fixed!\n\n> Minor typo in line 261 – missing a closing bracket in x+f(Norm(x)\n\nThanks! Fixed!\n\n### Conclusions\n\nThank you for generously outlining the paper's strengths in your original review. We hope this response (and especially the new experimental results) have addressed your key concerns. If so, we would greatly appreciate it if you would reconsider your score.", "claims_and_evidence": "The claims are generally supported with convincing evidence. The claim that FLeRM can be used for hyperparameter transfer in large-scale LLM training is not convincingly supported, as the tested models are significantly smaller (millions of parameters) than foundational LLMs (hundreds of billions to trillions of parameters).", "methods_and_evaluation_criteria": "The methods are extensively tested on a number of different experimental setups. See the weakness section for additional experiments to support your analysis.", "theoretical_claims": "The theoretical claims of the paper appear correct.", "experimental_designs_or_analyses": "The experimental design appears valid.", "relation_to_broader_scientific_literature": "The paper’s approach in measuring and setting function space learning appears to be unique. I am not aware of any other papers in this area. In terms of hyperparameter transfer, the paper extends the current literature by removing some of the rigid assumptions about initialisation and the need to set a number of hyperparameters in existing methods.", "essential_references_not_discussed": "All essential references to my knowledge are discussed.", "other_strengths_and_weaknesses": "Strengths\n1.\tThe manuscript presents an innovative way to understand neural network training dynamics by shifting the focus from parameter-space learning rates to function-space learning rates.\n2.\tThe proposed Monte Carlo-based method, combined with Kronecker factorization, enables the estimation of function-space learning rates with minimal computational overhead, by requiring only a single additional forward and backward pass.\n3.\tThe paper introduces FLeRM, a robust solution for hyperparameter transfer that can be used with any network architecture and at any point during training.\n4.\tThe paper extensively evaluates FLeRM in multiple scenarios, including scaling width and depth in MLPs and transformers, varying initialization scales, and adapting LoRA rank, demonstrating its versatility.\n\nWeaknesses\n1.\tThe paper evaluates FLeRM primarily with the Adam optimizer. It would be beneficial to compare it with other optimizers such as SGD and AdamW to establish its robustness across different optimization paradigms.\n2.\tWhile FLeRM is shown to be effective for width and depth scaling separately, real-world scaling typically involves increasing both simultaneously. Showing results for this experimental setting will be helpful. \n3.\tThe experiments seem to be limited to constant learning rate schedules. Does FLeRM also work with dynamic LR schedules? \n4.\tWhile Appendix A provides model details, explicitly stating the exact parameter counts and layer configurations across different width and depth settings in the main text would help clarify the extent of scaling in the experiments.\n5.\tDespite strong empirical results, the paper lacks rigorous theoretical analysis or formal guarantees regarding the optimality and convergence properties of function-space learning rates. As such, the authors claim that FLeRM could facilitate hyperparameter transfer in large-scale LLM training is not convincingly supported, as the tested models are significantly smaller (millions of parameters) than foundational LLMs (hundreds of billions to trillions of parameters).", "other_comments_or_suggestions": "1.\tThe learning rate update in Algorithm 1 (line 138) seems to contradict the theoretical analysis. I believe the numerator and denominator should be in reverse order. \n2.\tMinor typo in line 261 – missing a closing bracket in x+f(Norm(x)", "questions_for_authors": "See the weakness section", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oeoVhEtviy", "reviewer": "Reviewer_fHgH", "summary": "The paper defines function space learning rate as the rate of change of a neural network's outputs per training iteration. Then, the method FLeRM is introduced for either estimating the per-layer function space learning rate of a model over training, or setting these learning rates (LRs) to fit an arbitrary schedule. Finally, experiments determine if the method can enable LR transfer between smaller width and/or depth networks and larger networks, as well as LoRA adapter LRs, by first recording and then setting per-layer function space learning rates.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review, stating that \"The paper is really well presented and the experiments are thorough. I am particularly looking forward to the possibility of empirically measuring output dynamics and relating them to various theoretical predictions.\" (We're looking forward to that too!)\n\n\n> one can consider a simpler Monte Carlo estimator... Another possibility... estimate change in output from the derivative of loss with respect to output.\n\nOur method is unique in that, from a single forward+backward pass, it returns $L$ quantities where $L$ is the number of parameter tensors.\nThese $L$ quantities are the change in the function induced by the change in only one specific parameter. In contrast, finite differences based approaches like the one you suggest will only tell us how rapidly the outputs are changing overall, not how much individual layers contribute to this change.\nYou could of course compute how individual layers contribute using finite differences, but that would require $L$ forward passes, where, in each forward pass you perturb only one of the $L$ parameter tensors.\n\n> Could the authors discuss the advantages/disadvantages of considering changes in the function output as opposed to loss?\n\nConsidering the change in output, rather than the loss is the usual approach taken e.g. in muP [2] and Modula [3].  In our preliminary experiments, we did try forcing the loss to change by a specific amount, but we found training to be unstable as the loss approached the minimum. This is perhaps expected: once the loss is at the minimum, it can't go down any further, so trying to force the loss to go down a specific amount further is not sensible, and you might expect it to lead to instability.\n\n> What is the empirical performance of the method if one assumes is diagonal or full rank? ... what is the variance of the proposed estimator? What is the magnitude of the bias introduced by assumptions on the covariance matrix?\n\nWe have included a plot comparing the bias and variance of different covariance assumptions in the appendix. See [Rebuttal Figure 1](https://fslricml2025rebuttalsfigures.tiiny.site) (hyperlink to anonymous url).\n\n> it would be informative to do per-layer learning rate sweeps\n\n> another assumption is to share the learning rate over multiple layers when transferring to a deeper model.\n\nThis would be a super-interesting question for follow-up work!\nOf course, it is not in scope for the present paper due to the very extensive new experiments required, along with the careful thought required to draw conclusions from those experiments.\nThat said, we do hope that our notion of function space learning rates would help design these sweeps efficiently.\n\n> I reviewed appendix A. Appendix A should include the details of the training task for ResMLP.\n\nFixed.\n\n\n> Regarding appendix C.2 - when updating the LR at time t, is it made to match the base model's LR at time 0, or time t?\n\nAt time t! Appendix C.2 is \"change the per-layer learning rates every 100 iterations to make the function-space learning rates match the base model at that time point\", whilst the main paper experiments are \"change the per-layer learning rates at the beginning of training to make the function-space learning rates match the base model at initialisation, then use those learning rates for the rest of training\". Our experiments seem to suggest that, for the purposes for hyperparameter transfer, it is sufficient to \"correct\" the learning rate at initialisation.\n\nOther comments and suggestions: All fixed.  Thanks!\n\n> What is the time complexity of the proposed estimator?\n\nFor the width experiments in the main text we found that FLeRM increased the runtime by around 1.7\\%, which is negligible relative to the benefits achievable by accurate hyperparameter transfer.\n\n> Why ResMLP instead of a convolutional ResNet?, Also, why is Layernorm/Batchnorm not used in the ResMLPs? \n\nThe ResMLP serves as an extremely simple setting, and a similar architecture is used in previous work of depthwise hyperparameter transfer [1]. \n\n> why disable affine transformations in the transformer Layernorms?\n\nWe disabled affine transformations early on in the project to reduce complexity in prototyping. We have rerun the Transformer (PreResPostMod) width transfer experiment with affine transformations enabled and found no problems. See [Rebuttal Figure 4](https://fslricml2025rebuttalsfigures.tiiny.site).\n\n[1] Yang, G., Yu, D., Zhu, C., and Hayou, S. Tensor programs vi: Feature learning in infinite-depth neural networks, 2023.\n\n[2] Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi,\nD., Ryder, N., Pachocki, J., Chen, W., and Gao, J. Tensor\nprograms v: Tuning large neural networks via zero-shot\nhyperparameter transfer, 2022\n\n[3] Bernstein, J. and Newhouse, L. Modular duality in deep\nlearning, 2024a", "claims_and_evidence": "Claim: FLeRM efficiently estimates function-space LRs. True - although there are some caveats (see Methods below).\n\nClaim: FleRM enables LR transfer between networks of different width, depth and LoRA adapters of different rank. True - although evidence could be strengthened (see Experimental Design and Questions below).", "methods_and_evaluation_criteria": "The method is elegant but could also be overcomplicating things. The authors say equation (1) is intractable, which is reasonable for large models/datasets. However, one can consider a simpler Monte Carlo estimator of (1) that is just the change in output between two iterations, computed over some minibatch. One could even schedule the same minibatch close together in training so as to get this information \"for free\" during training. Another (speculative) possibility is to take the change in loss between successive iterations, and estimate change in output from the derivative of loss with respect to output. In any case, the authors should justify why their method is preferred over some other estimator (via variance of the estimator, time complexity, etc.) - see Questions (below) for more on this.\n\nSome other assumptions which are reasonable, but could use some empirical support are:\n1. What is the empirical performance of the method if one assumes $Cov[Z_{ij}, Z_{i'j'}]$ is diagonal or full rank? While the latter is probably impractical, the former is an even simpler assumption than the covariance being factorizable, and is also more closely related to existing work (e.g. Adam) - thus it would be helpful to have a comparison of the proposed method versus this simpler approach.\n2. Learning rate sweeps are for global learning rates, but this method already sets learning rates for individual layers. Thus, although costly, it would be informative to do per-layer learning rate sweeps (perhaps via random or grid search) for at least some small settings, as well as per-layer learning rate transfers. Transferring LRs found via per-layer sweeps on small models might even be a way to cheaply improve performance on deeper models.\n3. Related to point 2, another assumption is to share the learning rate over multiple layers when transferring to a deeper model. Although appendix C.4 conducts the ablation of dividing learning rates equally across layers, if point 2 is addressed by finding optimal per-layer learning rates, seeing if those rates are uniform over blocks would strengthen this assumption.", "theoretical_claims": "The derivation of the estimator in section 3.1-3.2 appears sound. Some properties of the estimator are not addressed (see Questions below).", "experimental_designs_or_analyses": "The experiments and analyses are sound, albeit not sufficiently large-scale to demonstrate transfer in cutting-edge models (although it is fair to say this would be out of scope). Some of the model architecture choices are a bit strange - see Questions below.", "relation_to_broader_scientific_literature": "The empirical-first approach to measure the relationship between weight changes and function space changes is complementary to existing theoretical-first approaches. Not only is this method useful for its stated purpose (hyperparameter selection and transfer), but as demonstrated in figure 1, it could also be a way to generate empirical evidence on training dynamics. This would be a useful tool in literature that looks at how outputs evolve due to changes in weight space (e.g. neural tangent kernel literature, Lipschitz-based complexity bounds).", "essential_references_not_discussed": "I am not aware of any missing references.", "other_strengths_and_weaknesses": "Strengths: as discussed above, the method is complementary to existing theoretically oriented work around learning rates and parameter scales. It is also useful both as a tool for LR selection, and as a tool for analyzing training dynamics. The paper is really well presented and the experiments are thorough. I am particularly looking forward to the possibility of empirically measuring output dynamics and relating them to various theoretical predictions.\n\nWeaknesses: as discussed above, the method might be needlessly complicated. The experiments also do not take full advantage of the method's potential and so I am unsure how much significance the results have. If the results cannot extend beyond the slightly non-standard settings explored by the paper, and beyond rescaling per-layer learning rates by constants, then the impact is somewhat limited. If the results do generalize to more settings and more complex learning rate schedules, then the work is very significant.\n\n## Update after rebuttal\n\nThe authors have answered all of my questions, and I stand by my review that this paper has significant contributions and should be accepted.", "other_comments_or_suggestions": "Some comments on notation in section 3:\n- $d$ should be $\\partial$ in equations 1, 4, 5, 8.\n- it would be helpful to give the dimensions of the matrices $Z_{ij}$, $U$, and $V$.\n- $||\\Delta_l f||^{base}_{RMS}||$ should be defined\n\nOther presentation issues:\n- line 301: \"as suggested *in* Section 3.2\"\n- Figure 1 and 2 are too far removed from the discussions. Also, the grid lines could be stronger and the plotted lines thinner (it is hard to tell which direction the trends are in due to the large number of layers being plotted).\n- Algorithm 1 is also separated by several pages from its discussion\n- Some claims are made about minor improvements in performance - adding a table of test loss would make this easier to evaluate than comparing between figures.", "questions_for_authors": "What is the variance of the proposed estimator? What is the magnitude of the bias introduced by assumptions on the covariance matrix? What is the time complexity of the proposed estimator? Either derivations or empirical evaluations (e.g. comparing the method to the naive approach of estimating equation 1) are welcome.\n\nCould the authors discuss the advantages/disadvantages of considering changes in the function output as opposed to loss? I can see some advantages (e.g. function output is not sensitive to choice of loss), but for the sake of argument, why not look at $|| \\Delta_l L ||_2^2$ instead, where $L$ is the loss function?\n\nWhy ResMLP instead of a convolutional ResNet?, Also, why is Layernorm/Batchnorm not used in the ResMLPs? Similarly, why disable affine transformations in the transformer Layernorms?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dFMeacJqlv", "reviewer": "Reviewer_JVrs", "summary": "This paper provides a novel method to transfer learning rates across model sizes. The approach is very flexible as it leverages monte-carlo estimation of the changes in model outputs under a proposed change in one of the weight matrices. The authors show that their approach can enable consistent optimal learning rates while training models of different widths and depths as well as LORA rank during fine tuning. They also show the ability of their algorithm to transfer across initialization scale. The flexibility of the approach makes it especially attractive.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review noting that we introduce \"a novel and flexible approach that would be easy and cheap to implement for the practitioner.\"\n\n> One potential drawback is the need to train a side by side base model...\n\nThis is unavoidable whenever doing any form of hyperparameter transfer, as hyperparameter transfer by definition involves tuning the hyperparameters on a small base model, then using some strategy (e.g. muP) to transfer those hyperparameters to the scaled model.\n\n> potentially many scalars to track, but this is still pretty cheap.\n\nYes, our Kronecker scheme only needs $LD$ scalars if we have $L$ tensors with $D$ dimensions each, which is insignificant compared to the costs of training the network itself.\n\n> In equation 16 it should be $i,i'$ in the sum not $ii$. There are also two equal signs in 16.\n\n> Blue line in algorithm should be flipped...\n\nThanks! Fixed!\n\n> The authors mention the limitations of scaling theory by citing Everett et al...\n\nThanks, this is very interesting! We have updated this section to emphasize that even in theory in the infinite-width limit, alignment is potentially very complex, with values lower than $1$ that can change over time. Do let us know if there are any references we should cite on this point. In any case, our point here is merely that deriving the \"correct\" alignment is very complicated, which suggests that our more \"empirically led\" approach is a useful alternative.  For the purposes of that point, it doesn't really matter whether the complexity can in-principle be captured in the infinite-width limit, or only emerges empirically.\n\n> Do the authors have a sense of how the quality of their estimate degrades with either (A) the length of steps / training time between learning rate adjustments (the value set to 100 in pseudocode) or (B) the number of monte carlo samples?\n\nIn the main paper's hyperparameter transfer experiments, we only adjust the learning rate at initialisation, and keep it constant for the rest of training. In Appendix C.2, we did it every 100 steps (as in the algorithm). We found very little difference between these two settings, suggesting that hyperparameter transfer mainly depends on correcting \"constant\" differences in function-space learning rates, rather than varying wildly throughout training. We did compare the bias and variance of the estimator using different covariance assumptions: See [Rebuttal Figure 1](https://fslricml2025rebuttalsfigures.tiiny.site) (hyperlink to anonymous url).\n\n> Depth scaling\n\nThis is super interesting.\nCertainly, we agree with [1] that you want changes in the attention patterns to be constant as you scale depth, and that isn't at all trivial to achieve (as the function-space learning rate for $W_K$ and $W_Q$ is basically the change in the output of attention, multiplied by the init for $W_V$ and $W_O$).\nAs such, we're pretty sure we agree that with FLeRM, changes in the attention would vanish in the infinite depth limit using the $1/sqrt(L)$ init, but not using the $1/L$ init.\nWe'll have to think about the right way to handle this in the FLeRM setting, to be robust to e.g.\\ different choices of normalization and changes over time during training, but we're confident that FLeRM has enough flexibility to handle it correctly.\nOne interesting approach would be to apply FLeRM to e.g. the self-attention layer with randomized $W_V$, rather than just to the overall network output.\nThat would allow you to isolate the change in the attention patterns, and ensure they remained constant as you scale depth.\nBut we will definitely have to think harder about it!\n\n[1] https://arxiv.org/abs/2405.15712 \n\n> Do the authors have a sense of why their experiments do not always show improved performance with respect to model size?\n\nThis is a very interesting phenomenon. We speculate that as models get larger in the \"standard\" setting, there are shifts in the relative size of the function space learning rates for different parameters, and that sometimes these changes are actually beneficial to performance.\nFLeRM, by fixing the function-space learning rates to those in the base model, might eliminate some of these beneficial changes.\nWe're super-excited to pursue follow-up work which uses function-space learning rates to investigate some of these phenomena in-depth.\n\n> Do the authors think their approach is mathematically similar to controlling the scale of the instantaneous NTK?\n\nThere likely is a connection, and this would be an interesting direction to explore in the future.", "claims_and_evidence": "The theoretical results are supported by proofs and derivations and the entire approach is supported by a large array of experiments.", "methods_and_evaluation_criteria": "Yes, the experiments seem reasonable. They looked at ResNets on CIFAR, transformer models on Wikitext 103, and finetuning language models on Cold French Law and Mathpile.", "theoretical_claims": "The proofs and derivations are correct and straightforward from my reading.", "experimental_designs_or_analyses": "Yes, the experimental design is reasonable in my opinion.", "relation_to_broader_scientific_literature": "This paper is studying an important problem in the science of scaling neural networks (how to make their optimization profiles consistent across model sizes). It introduces a novel and flexible approach that would be easy and cheap to implement for the practitioner.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper provides an interpretable algorithm and several useful experiments. One potential drawback is the need to train a side by side base model to keep track of the rescaled learning rates for each layer. In addition, if the base model has many layers there are potentially many scalars to track, but this is still pretty cheap.", "other_comments_or_suggestions": "1. In equation 16 it should be $i, i'$ in the sum not $i i$. There are also two equal signs in 16.\n2. Blue line in algorithm should be flipped. I think it should be $|\\Delta f_{base}| / |\\Delta f|$ instead of $|\\Delta f | / |\\Delta f_{base}|$5. \n3. The authors mention the limitations of scaling theory by citing Everett et al and the fact that weight alignment to input features in hidden layers are dynamic over time. I would like to point out that **even in the infinite limit** we would not expect perfect alignment $A=1$ which would require the feature vectors to become singular vectors of the weight matrix and there would be complex dynamics for $A(t)$. Concretely, with SGD a weight matrix would have the form $W(t) = W(0) + \\frac{\\eta}{N} \\sum_{t} g_t \\phi(h_t)^\\top$ where $g_i$ and $\\phi_j$ are vectors with $O(1)$ entries. If I pass the last vector $\\phi(h_T)$ through this matrix I get $W(0) \\phi(h_T) + \\eta \\sum_{t} g_t  \\left( \\frac{1}{N} \\phi(h_t) \\cdot \\phi(h_T) \\right)$. In general, $\\phi_T$ is not a singular vector of $W(t)$ so the alignment will be lower than one and also changing over time in a way that depends on the correlation structure of the $\\phi$'s. Thus in my understanding their experiments do not really invalidate the scaling theory to the extent they claim.", "questions_for_authors": "1. Do the authors have a sense of how the quality of their estimate degrades with either (A) the length of steps / training time between learning rate adjustments (the value set to 100 in pseudocode) or (B) the number of monte carlo samples? Since only a small number of scalars need to be estimated to control the norm, it seems likely that very few samples are needed. However, I wonder if longer delays between updating  could really impact performances of models which have been scaled up significantly compared to the base model. Are there any experiments with this?\n2. On depth scaling to a model with $L$ blocks, the authors rescale the base function LR $|\\Delta f_{base}| \\to \\frac{1}{L}$ and introduce branch scale factors $1/\\sqrt{L}$. I believe that this would lead to $\\eta_L = \\eta_0$ for SGD and $\\eta_L = \\eta_0 / \\sqrt{L}$ for Adam. Both of these match the scaling theory for $1/\\sqrt{L}$ branch scaling. However, there are other possible depth scalings. For example, if one adopts a $1/L$ branch scale, the residual blocks do not linearize in the limit and there is within-block feature learning as $L \\to \\infty$ (see section 3.4 here https://arxiv.org/abs/2405.15712).  \n3. Do the authors have a sense of why their experiments do not always show improved performance with respect to model size? Deeper networks seem to be worse at their optimal learning rates in some settings. \n4. Do the authors think their approach is mathematically similar to controlling the scale of the instantaneous NTK? \n$$df \\sim \\sum_{\\ell} \\eta_\\ell \\frac{\\partial f}{\\partial W_{\\ell} } \\cdot \\frac{\\partial f}{\\partial W_{\\ell}}$$", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RtMVmY4cnr", "reviewer": "Reviewer_ALy5", "summary": "This paper introduces FLeRM, an optimization algorithm that trains a smaller model, records function-space learning rates, and uses those recorded learning rate in the training process of a larger model.\n\nThe authors perform experiments on various datasets such as CIFAR-10 and Wikitext-103. They use residual MLP and Transformer architectures for their experiments. They show that using the proposed method reduces the effect of different learning rates on the training loss. They also, explore the effect of increasing the depth and width of neural networks and the learning rate adjustments necessary when such changes are made.\n\n## Update after rebuttal\n\nI wish to thank the authors for the rebuttal and the reviewers for their thoughtful comments. \n\nThe issue of dataset size and pertinence to modern deep learning landscape is an important one. I understand that during the rebuttal timeframe it is not feasible to run additional experiments. Regardless, it is an issue that needs to be addressed. \n\nI have decided to keep my original rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive and thoughtful review!\n\n> ...the proposed algorithm needs to record and store the model parameters at every iteration. This would require double the amount of memory to maintain this buffer... \n\nWe agree that if we needed to store an extra copy of the weights at every step, that would be bad.  We definitely don't do that! Specifically, the extra copy of the weights is only necessary at at the time steps where you use FLeRM to compute the layerwise learning rates.  And we do this very infrequently.  For example, in the FLeRM experiments in the main paper, we apply FLeRM at the start and then fix those layerwise learning rates for the rest of training.  After that, training requires only the same time and space as standard training: the only difference is the layerwise learning rates.  \n\nMoreover, note that an extra copy of the weights wouldn't increase peak memory usage by nearly as much as a factor of 2. \nThat's because Adam already stores not only the weights themselves, but also the average gradients and squared gradients, and because backprop requires storing a large number of intermediate activations.\n\nWe agree that the peak memory usage during a FLeRM step is higher due to this extra copy of the weights, but as it happens so infrequently, there are a number of strategies to mitigate the issue, including using smaller batch sizes for a FLeRM step, or shifting some quantities to CPU memory.\n\n> ...computation requirements of your method versus vanilla approaches...\n\nFor the width experiments in the main text we found that FLeRM increased the runtime by around 1.7\\%, which is negligible relative to the benefits achievable by accurate hyperparameter transfer.\n\n> these smaller datasets are less relevant...\n\nWe agree that it would be ideal to be able to use larger datasets.  However, bear in mind that our datasets are common across previous hyperparameter transfer works, e.g. MuP [1], and that for each panel in the hyperparameter plots, we must train the network dozens of times (for every learning rate, for every width / depth etc.) and so very large scale pretraining tasks are not practical, even given access to a reasonable amount of compute.\n\n> ...validation/test loss is required to see if their method improves generalizability (All figures, including the appendix, focus on the training loss)... Why did you decide not to focus on validation/test set loss, instead focus on the training loss?\n\nWe plotted the test loss for the Transformer (PreNormPostMod) width transfer experiment in [Rebuttal Figure 2](https://fslricml2025rebuttalsfigures.tiiny.site) (hyperlink to anonymous url), which showed exactly the same patterns as the train loss. (Using the train-loss is common in the hyperparameter transfer literature [1,2,3] because in the LLM-pretraining setting where you usually train on each datapoint only once, the expected train and validation loss turn out to be equivalent [4]).\n\n> In various cases, the best training loss achieved by vanilla methods is better (e.g., ResMLP).\n\nSometimes the loss is better with FLeRM.  This is most clear in:\n* Figure 2 (main paper): PreNorm, PreNormPostMod\n* Figure 3 (main paper): PostNorm, PreNorm\n* most of the [new figures we have done for the Rebuttal](https://fslricml2025rebuttalsfigures.tiiny.site).\n\nWe agree though that sometimes the loss is better without FLeRM.  This is most clear in:\nFigure 3 (main paper): ResMLP, PreNormPostMod\n\nIf anything, we would argue that when FLeRM seems to better in more relevant settings (especially PreNorm Transformers). Importantly though, we did not motivate our work as improving performance so we did not perform the large-scale experiments necessary to definitively establish performance improvements.\nNonetheless, this is definitely a super-exciting direction for future work.\n\n> If your method is applied to a larger, noisier dataset, such as ImageNet, what potential issues do you expect to see with your approach?\n\nHyperparameter scaling work e.g. [1,2,3] has not, to our knowledge, found qualitative differences in scaling across different datasets, so we believe we are unlikely to see any such differences here either.\n\n[1] Yang, G., Hu, E. J., Babuschkin, I., Sidor, S., Liu, X., Farhi, D., Ryder, N., Pachocki, J., Chen, W., and Gao, J. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer, 2022\n\n[2] Yang, G., Yu, D., Zhu, C., and Hayou, S. Tensor pro- grams vi: Feature learning in infinite-depth neural net- works, 2023.\n\n[3] Bordelon, B., Noci, L., Li, M. B., Hanin, B., and Pehle- van, C. Depthwise hyperparameter transfer in resid- ual networks: Dynamics and scaling limit, 2023\n\n[4] Aitchison L. Why you don't overfit, and don't need Bayes if you only train for one epoch. arXiv preprint arXiv:2411.14478. 2024 Nov 19.", "claims_and_evidence": "There is a serious issue with the proposed algorithm that poses a major limitation on the usability of the algorithm in real-life modern optimization scenarios. The main claim of this paper is the use of a smaller network to optimizer a larger one. This would reduce the amount of computation required for the larger network. However, according to Algorithm 1, the proposed algorithm needs to record and store the model parameters, $W^l$, at every iteration $t$.\nThis would require double the amount of memory to maintain this buffer. This becomes problematic when training large models, such as LLMs that already require multiple GPUs for training.", "methods_and_evaluation_criteria": "The datasets used (e.g., CIFAR-10 and Wikitext-103) are good starting points for the paper. However, in the modern deep learning landscape, these smaller datasets are less relevant. Datasets such as ImageNet have become the bare minimum to effectively examine the efficacy of new optimization approaches.", "theoretical_claims": "Yes, all of them.", "experimental_designs_or_analyses": "1. Throughout the paper, the authors present and compare training loss for their method and the vanilla approach. While training loss can be an indicator of learning, validation/test loss is required to see if their method improves generalizability (All figures, including the appendix, focus on the training loss). Generalizability is the main focus of optimization. It is possible that lower training loss leads to higher validation loss.\n\n2. In Figure 3, the authors show the values of the training loss for vanilla and the proposed optimization scheme. In various cases, the best training loss achieved by vanilla methods is better (e.g., ResMLP). There is no advantage of less variance to learning rates when the final training loss is worse using the proposed method.", "relation_to_broader_scientific_literature": "Hyperparameter transfer has been focus of research for a long time. This paper proposes using learning rates found in a smaller network and apply them to a larger network to improve invariance to hyperparameters.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "- Without results on validation/test set data, the findings of the paper have limited significance for the academic and practical use cases. Please refer to the comments in **Experimental Designs Or Analyses** for more details.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. Do you have comparative results on the memory and computation requirements of your method versus vanilla approaches that would put into perspective the extent of the limitations of your approach?\n\n2. Why did you decide not to focus on validation/test set loss, instead focus on the training loss?\n\n3. If your method is applied to a larger, noisier dataset, such as ImageNet, what potential issues do you expect to see with your approach?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces a method for estimating the sensitivity of the output of a neural network with respect to the parameters. The method estimates this sensitivity using a monte carlo estimate, together with an assumed kronecker structure of the covariance over elements of the layer. The resulting sensitivity now depends on three scalars per layer (or group of layers), which are estimated in an online fashion using an exponential moving average. The overall idea is new, interesting, and the experiments also validate the approach. Since the reviewers were also positive, I recommend the paper be accepted. But I do see one thing I would think would improve the paper, and that is
- [Baseline comparison] The main application for your technique is hyperparameter transfer across model scales. Yet you do not compare to any baseline for this task, such as the muP or modula. I understand that comparing to modula can be difficult, since this would require re-writing your models in modula. But could you not at least show how your method compares to muP in practice? This would really help to get the attention of practitioners. 
  
I have also some smaller recommendations:

- [Misleading title for theoreticians] For theoretical subcomunities of ICML, the title could be misleading. Function space learning is sometimes used to refer to learning the model over the space of functions, perhaps by lifting to the space of measures. It would typically make theoretical readers think that you are considering optimization analysis in infinite-dimensional function spaces, such as Functional gradient descent (like boosting or kernel methods). I would recommend renaming the paper to more precisely target your audience, for instance
"Learning the sensitivity of a function online, with an application to hyperparameter transfer"
Or something that conveys the same meaning. This is just a recommendation.

- [Relationship between sensitivity and learning rates] I would recommend expanding on why you use the relative change in the sensitivities of the big model to the base models to set the learning rate in (19). This is somewhat intuitive: more sensitive models should have smaller learning rates. But since this is key to how you apply your technique, it would be best to clarify further.

- Because you initialize your momentum buffers at zero EMA Z2[ℓ], EMA ZZT[ℓ], EMA ZTZ[ℓ] = 0, the resulting EMA estimates are not discrete expectations over the past samples. You might glean some minor improvements by correcting for this bias. For this you would need to either use the "bias correction" trick of Adam [Adam], or initialize these buffers at the first sample, instead of zero [Section 2.3, Schaipp2024]. Either one of these approaches will avoid the issue of your momentum buffer starting too small. 



[Adam] Kingma, D. P., & Ba, J. (2015). Adam: A Method for Stochastic Optimization. In Proceedings of the 3rd International Conference on Learning Representations (ICLR).
 
[Schaipp2024]  Fabian Schaipp, Ruben Ohana, Michael Eickenberg, Aaron Defazio, Robert M. Gower, MoMo: Momentum Models for Adaptive Learning Rates, ICLR 2024.
Decision: Accept (poster)