Accelerating Training with Neuron Interaction and Nowcasting Networks
Authors: Boris Knyazev
Abhinav Moudgil
Guillaume Lajoie
Eugene Belilovsky
Simon Lacoste-Julien
Venue: ICLR 2025 Poster (2025)
Keywords: accelerated optimization
parameter prediction
graphs
transformers
learning to optimize
Primary area: learning on graphs and other geometries & topologies
TL;DR: We present a meta-model predicting future parameters along the SGD/Adam-optimization trajectories using graphs of neurons and graph neural networks to speed up training of neural networks.

Abstract: Neural network training can be accelerated when a learnable update rule is used in lieu of classic adaptive optimizers (e.g. Adam). However, learnable update rules can be costly and unstable to train and use. Recently, Jang et al. (2023) proposed a simpler approach to accelerate training based on weight nowcaster networks (WNNs). In their approach, Adam is used for most of the optimization steps and periodically, only every few steps, a WNN nowcasts (predicts near future) parameters. We improve WNNs by proposing neuron interaction and nowcasting (NiNo) networks. In contrast to WNNs, NiNo leverages neuron connectivity and graph neural networks to more accurately nowcast parameters. We further show that in some networks, such as Transformers, modeling neuron connectivity accurately is challenging. We address this and other limitations, which allows NiNo to accelerate Adam training by up to 50% in vision and language tasks.

Published as a conference paper at ICLR 2025
ACCELERATING TRAINING WITH NEURON INTERACTION AND NOWCASTING NETWORKS
Boris Knyazev' Abhinav Moudgil Guillaume Lajoie
Eugene Belilovsky2. Simon Lacoste-Julien
Samsung- SAIT Al Lab, Montreal "Concordia University *Université de Montréal "Mila
b knkyazeeessamuung.ccom
ABSTRACT
Neural network training can be accelerated when a learnable update rule is used in lieu of classic adaptive optimizers (e.g. Adam). However, learnable update rules can be costly and unstable to train and use. Recently, Jang et al. (2023) proposed a simpler approach to accelerate training based on weight nowcaster networks (WNNs). In their approach, Adam is used for most of the optimization steps and periodically, only every few steps, WNN nowcasts (predicts near future) parameters. We improve WNNs by proposing neuron interaction and nowcasting (NiNo) networks. In contrast to WNNs, NiNo leverages neuron connectivity and graph neural networks to more accurately nowcast parameters. We further show that in some networks, such as Transformers, modeling neuron connectivity accurately is challenging. We address this and other limitations, which allows NiNo to accelerate Adam training by up to 50% in vision and language tasks.
INTRODUCTION
Modern deep learning models, such as large language models (Touvron et al., 2023), are trained using clas- sic adaptive optimizers, such as Adam' (Kingma Ba, 2015; Loshchilov & Hutter, 2017). These optimizers up- date neural network parameters E R" using gradient descent at step t as 0/+1 = e Ae, Vi e [1, n]. where Ae; is the update computed analytically based on the history of parameter values, gradients and the learning rate. Recently, Jang et al. (2023); Sinha et al. (2017) showed that parameters 0 follow predictable trend so that optimization can be accelerated by noweasting (pre- dicting near future) parameters using a learnable func-
validation perplexity
8-100
WNN [Jlnn et at. 2023)
NNNo
4-10
target perplexity
2-10
Speeduup 48%
2-10
10'
2000
4000
8000 10000 12000 14000 training iteration
6000
Figure 1: Adam without and with nowcasting using our NiNo vs WNN (Jang et al., 2023) on language task that NiNo and WNN have not soen during their training.
tion f 0/+k o + (O)),, 0 2......., where k >> 1 is the future horizon. More popular learnable approaches to speed up optimization, such as "learning to optimize" (L20), are recurrently applied at every step (Andrychowicz et al., 2016; Metz et al., 2022b). Compared to L2O, the parameter nowcaster f is applied very rarely reducing its overhead, e.g. a base optimizer such as Adam is run for Ik steps followed by the prediction step (Fig. 1). Moreover, training such fo is simpler than L20, since a supervised loss can used instead of a more challenging meta-lcarning loss with recurrent inner steps. However, akin Adam, f from prior works does not directly leverage the structural information of O such as connectivity between neurons and layers. This structure has been shown to be critical for many parameter representation tasks, such as property prediction (Navon al., 2023; Zhou et al., 2024b; Kofinas et al., 2024).
We use Adam throughout the paper, hut our discussion and methods are in principle applicable to any opti- mizers that produce trajectory of parameters, including SGD with/without momentum, AdamW, Adagrad, etc.

---

Published as a conference paper at ICLR 2025
In this work, we propose neuron interaction and nowcasting (NiNo) networks making better predictions of future parameters to accelerate training with a base optimizer such as Adam and make the following contributions:
1. We introduce NiNo, which directly uses the structure of neural network parameters by leveraging recently proposed neural graphs (graphs of neurons) (Kofinas et al., 2024; Lim et al., 2024).
2. We improve Transformer's 's neural graphs of Kofinas et al.; Lim et al. by more accurately model- ing the neuron permutation symmetry of multi-head self-attention. Our neural graphs combined with graph neural networks create a strong inductive bias to make rare but accurate predictions.
3. NiNo is conditioned on the future horizon l to allow for prediction in the near and far future without retraining it, facilitating its usage in diverse tasks and at different optimization stages.
4. We demonstrate that NiNo accelerates training with Adam for ConvNets and Transformers reduc- ing the number of steps to achieve the target performance of Adam by up to 50%. We release our source code and models at https: /github. com/ SamsunggATIMontzeal /nino.
2 RELATED WORK
Learning to optimize (I.20). The I.20 literature has offered many approaches to learn neural network (optimizer) that optimizes the parameters of other neural nets (optimizees) (Andrychowiez et al., 2016; Chen et al., 2022b;a; Amos, 2022). Among them, Safeguarded L20 (Heaton et al., 2023; Prémont-Schwarz et al., 2022) is most related to ours as it switches between an L20 optimizer and SGD/Adam. While Safeguarded L200 alleviates the meta-genent reta-generallzation challenge (Thérien et al., 2024), training an L20 optimizer remains costly and unstable due to the use of meta learning methods and its long inner loop unrolls required at each meta-training iteration (Metz et al., 2019; 2022b). Moreover, overheads of L20 add up at each iteration making it more computationally intensive than Adam. In contrast, our approach follows weight nowcaster networks (WNNs) (Jang et al., 2023), where the nowcaster model is applied very rarely (e.g. once per Ik steps of Adam) making the total overhead negligible, yet still speeding up training significantly.
Parameter prediction and generation. This area has been active recently, primarily aiming at reduc- ing computational costs of training neural nets (Peebles et al., 2022; Ashkenazi et al., 2023; Schürholt et al., 2022; 2024; Knyazev et al., 2021b; 2023; Zhou et al., 2024d; Soro et al., 2024; Wang et al., 2024). Most related to work are IntrospectionMLPP (Sinha et al., 2017) and WNNs (Jang et al., 2023) serving the basis of our approach. These methods train simple MLPs to periodically (e.g. every few epochs of Adam) predict future parameter values of neural net given its past parameters (Sec- tion 3.1). However, their MLPs predict parameter coordinate-wise (for each parameter independently) similar optimization methods without leveraging the structure of neural networks. Moreover, their MLPs predict parameters only for predefined future horizon k, whereas different tasks and different optimization stages can have different parameter evolution trends (Morchdi et al., 2022; Guille-Escurer et al., 2024; Lange et al., 2023). We address these shortcomings by conditioning the prediction on k.
Representation learning of neural network parameters. This area has also developed fast re- cently (Navon al., 2023; Schürholt al., 2021; 2024; Ashkenazi et al., 2023; Andreis et al., 2023; Zhou et al., 2024b;c;a). One of the main goals in these works is to model neuron permutation symmetry- the property of neural networks that if the neurons in one layer are permuted in the same way as the neurons in the next layer, the neural network preserves its function (Hecht-Nielsed Accurate modeling of this symmetry allows for better estimation of network properties or encoding implicit neural representations. To model neuron permutation symmetry, Kofinas et al. (2024); Lim et al. (2024) proposed neural graphs (graphs of neurons) enabling the usage of graph neural networks (GNNs) (Kipf& & Welling, 2017: Corso et al., 2020). Remarkably, GNNs can digest graphs of any size and connectivity, the synergy of neural graphs and GNNs enables processing diverse parameters from different architectures and tasks using single GNN. We leverage both the neural graphs and GNNs to learn a single NiNo model that can accelerate optimization in diverse tasks. However, the neu- ral graphs of Kofinas al. (2024); Lim al. (2024) do not accurately model neuron permutation sym- metry of Transformers (Section 1).. We address that shortcoming to make better predictions by NiNo.
1990).
Dynamic models. Our approach of learning from the interaction of nodes in a neural graph to make future predictions is also related to dynamic interactive models, where message passing networks and GNNs are used to learn the interaction within relational temporal networks (Trivedi et al., 2019; Knyazev al., 2021a) and physical systems (Kipf al., 2018: 8; Sanchez- et al., 2018; 2020;
2

---

Published as a conference paper at ICLR 2025
Zambaldi et al., 2019). However, developing strong parameter prediction model requires many specific design choices and careful modeling of neuron permutation symmetry, motivating our work.
BACKGROUND
We consider a neural network parameterized by 0, E I trained for i steps with Adam (or another opti- mizer). Our goal is to accelerate optimization by predicting future parameters k >> 1 steps ahead given its past values using meta-network f% parameterized by Φ: 0 O. + (0,, 0 11-100,220 2.-2,..)
3.1 WEIGHT NOWCASTER NETWORKS (WNNS)
WNNs (Jang al., 2023) model f an MLP predicting the delta (update difference) of the i-th parameter: A ++ = f (θ 0 1 where Te = T c + 1 and c is the context length. Here, T,7 1,... are epoch indices and are the parameters scaled based on the range of values in ....... 0 (see details in Section A. The predicted parameter value is obtained by unscaling the pre-
)².))
dicted delta: 6 ++ = e* + ( ++)). WNNs are trained in a supervised way by collecting a training dataset of parameter trajectories [[0), 0.. obtained with Adam and applying the / llos:: argmin a A - A0,+*|11,
(1)
where A0, are the scaled target parameter deltas. Sinha al. (2017) proposed the approach of future parameter prediction originally, but without scaling the parameters and without a simple way to choose which optimization steps to make the prediction. In contrast, WNNs scale the parameters and are always applied periodically, once per every epochs of a base optimizer (Adam) with = c, which better accelerates optimization. For example, to use a trained WNN fø on new task, first 5 by default) after which S is applied predict future (10-th epoch)
Adam is run for eepochs parameters, after which the procedure repeats (the base optimizer is continued for another 5 epochs followed by parameter prediction) until convergence. So the WNN is applied very rarely, e.g. in the case of 200 Adam steps per epoch, the WNN is applied only once per 1,000 steps (Fig, 1).
3.2 (NAIVE) NEURAL GRAPH OF TRANSFORMERS
WNNs apply an MIP for parameter O² given only its past values. Such an MLP is inherently limited, since future parameter values depend on many factors, including connectivity of neurons. Modeling neuron connectivity in a way that generalizes to arbitrary network structures to make parameter prediction as general as possible is challenging. Simple parameter representations, c.g. flattening of parameters into a vector (Schürholt et al., 2021), are not general and do not model neuron permutation symmetry (the neurons can be permuted without changing the network output). Recently, Kofinas et al. (2024); Lim et al. (2024) proposed to represent 0 using neural graph ((0) = (V, E)) with node features V E IR Vixdy and edge features E E R/Vix[Vixd xdg, where dy, dg are the node and edge feature dimensionalities. respectively. For example, Equation shows E for the dg = dimensional edge features of an I.-layer MLP (we ignore biases for simplicity): 0 = (w())...,,((())
where W!!) E Rd are the weights for E [1, LJ.
Neural graphs can theoretically represent the parameters of any neural network taking into account neuron permuta- tion symmetry as the rows and columns of weight matrices correspond to the nodes in the neural graph. This way, neu- ral graphs impose a strong inductive bias similarly to using convolution for images. So a model operating on neural graphs, such as graph neural network (GNN) and our NiNo model, can be used in diverse tasks and should be able to learn parameter prediction rules that generalize bet- ter with fewer samples than, for example, MLPs. However, to fully leverage the neural graphs practice, we need to accurately construct them, which is not trivial for such networks as Transformers, motivating our approach in this work.
W
W
E
(2)
W
Transformers. For input x E R Nxd, where N sequence length, a multi-head self-attention (MSA) layer (Vaswani al., 2017) for each head la [1, H] projects using W Wk. h. Wy followed
3

---

Published as a conference paper at ICLR 2025
[0]
(c)
good pemuuitoos
tor permuaattoss
Naive graph: accuracy = 57.3%
(b)
graph: accuracy
Figure 2: (a) MSA weights with d = 6 and Hf = heads; (b) its naive (Kofinas et al., 2024) and (c) our neural graph with color-cnded edge/parameter types. (d-e) Neural graph embeddings obtained using GNN and projected to 2d using tSNE for Ik good and bad permutations of the MSA weights: (d) naive graph, (e) graph. Accuracy is computed by fitting logistic regression for neural graph embeddings with labels being the MSA output changes or not as described in Section A.6.
by self-attention and projection using another weight matrix W (see Section A.4 for details):
(x (xWX)xWW)) and Yn softmax(An)) (x and MSA(x) Σ Yb W (3)
H
A
vd
To construct neural graphs of MSA, Kofinas et al. (2024); Lim et (2024) use de = 3 dimensional edge features eij- (WW)).. (W*) < : (W"));)) assuming that the multi-head case automati- cally handled by the neural graph. However, as shown in Section 1,1, this neural graph does not model the neuron permutation symmetry correctly, therefore we refer as a naive neural graph (Fig. 2a.b).
METHODS
In this section, describe the details of our neural graphs for MSA layer of Transformers and the neuron interaction and nowcasting (NiNo) networks operating on such graphs (Fig. 4). Our neural graphs can be constructed for different Transformer architectures, including GPT2 (Radford et al., 2019), Llama (Dubey et al., 2024) (see graph examples in Fig. 9) and others (see our implementation). Following WNNs, NiNo is also applied very rarely during the optimization process (Section 3.1). However, NiNo leverages the representation power of neural graphs and GNNs to speed up optimization more significantly than WNNs, as we show in Section 5.
4.1 NEURAL GRAPH OF TRANSFORMERS
To construct neural graphs that accurately model neuron permutation symmetry in MSA, we (1) re- strict neuron permutations across heads and (2) relax permutations of weights W W (Fig. 2).
(1) Restricting permutations across heads. We first observe that splitting computations into H parallel heads (Equation 3) breaks neuron permutation symmetry across the heads, so shuffling neurons across the heads may change the overall output of MSA(x). Consider, for example, an MSA layer with
and = 2 (Fig. 2a) and the dot product between outputs x and (xx)" obtained after projecting x using the weights of the first head WT and Wt: x (x+) ) [x]. x², x xx][x x J". The result of this dot product does not change when permuting neurons in both W and W using the same n, since results in permuting the outputs: π (xx])=(xt)= = x (x+)²² However, consider shuffling neurons across head 1 and 2 in W and W e.g. neurons and 4 (highlighted in different colors in Fig. 2a) resulting in: [x x3, x[][ x x Now the dot product no longer equals So self-attention matrix A for head I and MSA(x) can change.
x (x+)T unless xxx] = x3x
To take into account this restriction, allow for neuron permutations only within head by adding a separate node for each head that connects the appropriate neurons so that permuting neurons 3 and 4 from our example changes the graph while permuting withing a head does not change it (Fig. 2c).
4

---

Published as a conference paper at ICLR 2025
Figure 3:
Laplacian Positional Encoding (LPE) our neural graph's 's nodes color-coded by the layer type (left), layer index (middle) and neuron index (right) for a Transformer with layers and 384 hidden units. We show (SNE projections from the 8-dimensional LPE to 2d.
(2) Relaxing permutations for W W We also observe that neurons in W and W can be permuted different way (T') from W and Wi without changing the MSA output:
(xWX) = (xT'( (xT) (Wi)) and (xWT) W = (xT,( Wz)) (WR), (4) where n. and - denote permutation of the columns and rows in the matrix respectively. In the naive graph the neurons in W W W W are always permuted in the same way (Te x²)) making the neural graph unnecessary restrictive (Fig. 2b). To address this issue, we keep W separate I-dimensional edge features in a neural graph instead of stacking them as 3-dimensional edge features (Fig. 2c). Since W and W share the input and output neurons while W is transposed in Equation 3, it is important to preserve the edge direction, so that for example e3,4 corresponds to the weights of W while 64,3 corresponds to the weights of W h
h W W as
4.2 EDGE AND NODE FEATURES
Edge features. In our neural graph, all the model parameters 0 and auxiliary connections 0', such as residual and head connections introduced in Section 4.1, are represented using the edge features /VV xdæ so that EElo = |0|+ O'| (Fig. 2c). To differentiate between and 0', we associate an integer edge type with each edge: E Z V so our neural graph ((0) (V,E
Eype EE)
Node features. For the node features we use the Laplacian Positional Encoding (LPE) (Belkin & Niyogi, 2003; Dwivedi et al., 2023) that allow graph neural networks to capture structural information more easily. For example the LPE implicitly embeds nodes of the same layer and same type close to each other even though this information is not explicitly provided (Fig. 3). To compute the LPE, we unweighted edges, transform the graph to undirected one and extract smallest non-trivial eigenvectors, so our node features are Vipe E TRl x| In addition, for Transformer's s word embedding layers, we found it beneficial to leverage positional feature V e Z V| since these layers often have the same size with ordered neurons. For other layers, this feature set to zero.
Sequence of parameters as a neural graph. We consider history of e past parameter vectors following WNNs (Section 3.1): Or:r. [0,, - ..., 0-_] Raxe, where т, 7 C + in Section 3.1. Since each parameter vector in θr:r, represents the same neural network structure, we transform the entire 7rr,, into a single neural graph g(O = (Vlpe, V", Elye, Erire), where Err is obtained by stacking edge features: &T.T. [E,; E,- E..] E [VVVxx[V]xxxxe whereas node features Vip V and edge types E do not change over time.
4.3 NEURON INTERACTION AND NOWCASTING NETWORKS (NINO)
Given an input neural graph ((:r:))) our NiNo model processes using layerwise scaling layer, node and edge embedding layers and GNN layers with a hidden size D. NiNo then predicts edge features for multiple steps in the future ε1.K that mapped back the parameter space [A /20 a using the inverse neural graph construction Δθ g (Ex) (Fig. 4).
Layerwise scaling. Since parameter values can vary in scale significantly across architectures and tasks, it is important to scale them appropriately. Compared to WNNs (Jang et al., 2023) scaling each parameter independently using min-max, we extend layerwise scaling (Schürholt et al., 2022) to sequence of parameters. Specifically, given weights w (!)
Rdd of the I-th layer of the input (8)
7iTe
Er:,. we obtain scaled weights as ((()) (w!)_- - n )o0 where (() a0 are the scalar mean and standard deviation of (l)
Fiter After repeating this step 44 [1, T.J. we obtain scaled E.....
5

---

Published as a conference paper at ICLR 2025
features
g
GNN
Deeet- mull-sko
Noural Graph
GNN layer
Af....
A6..K
pes anc
010 mmo-k3)
Edge typex
Loyerwos soaing
Nedio and sigg laskon erbaadinng
8.15
6.
6.1x
Mir F
Figure 4: Our neuron interaction and nowcasting (NiNo) model. We feed past parameter states (without gradients) input and predict future K states leveraging our improved neural graph (Section 4.1) using GNN. a new optimization task, NiNo is applied rarely over time: only once per steps of Adam (Section 5).
Embedding layers and GNN. We use linear Oipe de and embedding ow Ptyp layers to project node and edge features to the D-dimensional space followed by M GNN layers and the output layer DDM::
v Qlpe (Vlpc) Qw (V"),
V RIV XD
(5)
ED = de (E) + drype (E
E E X (6)
R
TTn m GNN pm (vm-1 E
m 1...M, (7)
E ÓDMS (E
EER VIXII (8)
Our GNN layers are based on Kofinas et al. (2024), but enable better efficiency large models, use a simple mean aggregation. Specifically, given node i,j features v."" 1 mn and edge features
n--1 the m-th GNN layer is formulated as: v! di WOJ IjEN(O) m m)-1) where m mh-1
message computed by aggregating features V v V ,e 0e is MLP; Mii) the neighbors of node i; oj is then obtained using another MLP given v V o (see Section A.5).
Direct multi-step forecasting (DMS). The final layer dDMs outputs values for each edge corre- sponding to future steps from + 1 to 7 + K. This is motivated by direct multi-step forecasting performing well in time series forecasting by avoiding error accumulation effects of autoregressive forecasting (Chevillon, 2007; Zeng al., 2023). To train the model given training dataset of parameters trained with Adam, we use Equation 1, but applied for ks = [1,..., KJ instead fixing k:
argmin ΣK 1 (1A0+++ ,b.+xll1),
where _0,- (-E(( and 0, are the target parameters scaled using μr, oy. The predicted pa- rameters are obtained by unscaling the predicted delta using Mr, o,: 0, = 0, + unscale( - )+
Inference with k-decay. Once model is trained, it can predict future parameters for l: E [1, K]. While k can be treated as a hyperparameter and kept fixed during optimization of the target task, we found that the initial optimization stage using very large k beneficial, whereas in the later optimization stages & should decrease fast, since the parameter values do not change significantly. Therefore, we propose to decay & during optimization as l K(T t)/T)", where T is the maximum number of optimization steps and controls the decay speed (Fig. 12). Although can be tuned, we found that works well in our experiments.
5 EXPERIMENTS
We experiment with nine tasks, each defined by dataset and neural network architecture in the vision or language domains (Table 1). Four of the tasks, the in-distribution tasks, are of a relatively smaller scale and used to train meta-mode modelss (NiNo, WNN and their variants). The other five tasks, the out-of-dist oisiiiiribttiin tasks, differ from the in-distribution tasks in the architecture and/or dataset. We also experiment with larger GPT22based and Llama-based architectures the end of Section 5.2.
Indexing parameters. Since the notion of epoch used for 7 in WNNs (Section 3.1) can be ill-defined in some tasks (e.g. in language tasks the models are often trained only for epoch), in our experiments we found and 7 7+1 indexing parameters at step t and t 200 to be a well performing strategy in all the tasks, so for our default context length 5 we apply the meta-models every 5 200 1, 000 steps.

---

Published as a conference paper at ICLR 2025
Table 1: In-distribution and out-of-distribution tasks.
IN-DISTRIBUTION TASKS
OUT-OF-DISTRIUUTTN TASKS
FM/16
C10/16
LMIB/3-24
LMIB/2-32
FM/32
C10/32
CL00/32
LMIB/3-66 54 WIKI3-61
Final training loss
0.25
0.91
5,94+0.03
5.85
#models
300
300
200
200
*purams
14K
15K
1.2M
16M
56K
57K
63K
3.4M
3.4M
Target (validation) metric
Acc
Acc
Porplexity
Porplexity
Acc
Ace
Acc
Perplexity
Perplexity
Target value
89.5%
66.0%
352
319
90.5%
72.5%
39%
181
147
Adam #steps
8606
8732
23000
23500
8269
8607
8341
23500
13500
NiNo #steps
4582
3775
11500
12000
4395
4323
46-16
12000
7000
avg spgaddd (8)
47.5
Seaee ug (9))
WNN-
45
45.0
30
5129
NiNo
40
40.0
20
@ 37.5
-6
86
32.5
64
120
256
240
2-32
4-64
hidden size of NihloWWNN byyes,
Indistribuutin
transtormer size (aaeershid)
(a)
(b)
Figure 5: Scaling (a.b) and generalization (c) trends for NiNo WNN+.
Training and evaluation pipeline:
1. Meta-training: Train the models with the supervised parameter prediction loss (Equation 1 for WNN and Equation for NiNo) using training checkpoints from the four in-distribution tasks.
2. Usage: Train task-specific neural networks in all the tasks using the Adam (vision) and Adam (language) optimizers with the meta-model applied every 1K steps given
5 past parameters. 3. Evaluation: Define a target validation performance level for each task based on running the baseline Adam optimization (without the meta-model applied) and reporting the relative reduction of the number steps to achieve that performance by other methods.
5.1
SETUP
Vision tasks. We use the FaSHionMNIST (FM), CIFAR-10 (C10) and CIFAR- (C (00 datasets and two convolutional architectures with three layers: with 16, 32 and 32 channels per layer (e.g. task 6M/16) or 32, 64 and 64 channels per layer (e.g. task FM/32). The convolutional architectures are the same as the I.20 experiments Kofinas et al. (2024) to enable fair comparison.
all cases, these tasks are optimized using Adam (Kingma& & Ba, 2015) without weight decay, with constant learning rate of 6e-3 (for FashionMNIST) or 3c-3 (for CIFAR) with a batch size of 128 for T==0kk steps.
Language tasks. We use the LM (Chelba al., 2013) and WikiText 103 (WIKI) (Merity et al., 2016) datasets and train GPT2 style Transformers (Radford et al., 2019) with 3 layers, 24 hidden units and 3 attention heads (denoted as 3-24); 2 layers, 32 units and 2 heads (2-32) or 3 layers, 64 units and 4 heads (3-64). These tasks are optimized for the next token prediction loss with Adam (Loshchilov & Hutter, 2017), weight decay 1e-2, learning rate 2e-4, batch size of 32, sequence length of 1024 for either epoch (for LM or 4 epochs (for WIKI) corresponding to around 24k or 14k steps respectively. We use predefined GPT2 tokenizer in all the cases with fixed vocabulary.
Meta-training dataset. We use FM/16, C10/16, LM and LM BMMIB/24 B/2-32 as the in-distribution tasks training 300, 300, 200 and 200 models in each task respectively (C' 1000 models in total). We save the checkpoints of intermediate steps, having in total around 1.6 x 10t checkpoints. Our dataset is large and diverse (Fig. 10). yet it is relatively cheap be collected as the tasks are small.
Baselines. As reference optimization method we use Adam for vision and AdamW for language tasks with a constant learning rate that we tune independently on each task based on the validation metric. Following WNNs (Jang et al., 2023) we use Linefit as another baseline (Section A.2). We further improve it with a simple scaling term to promote more recent parameters denoted as Linefit+ (see Section A.3). As the strongest baseline, we use WNN with our layerwise scaling and k-decay which we denote as WNN+. Finally, we use the learning to optimize (L20) model from Kofinas et al. (2024), which showed strong results similar vision tasks. We their L20 (NG-GNN) pretrained on FM/16 (denoted as L.2//MM66 as retraining their model on all our dn-iistibution tasks expensive. For fair comparison with 1.20/FM 6, we also trained NiNo FM/16 only (denoted
7

---

Published as a conference paper at ICLR 2025
Table 2: Reduction (%) of the number of steps until a target performance w.r.t. Adam. We observe that on both in-distribution and out-of-distribution tasks our models improve over the base optimizer and other acceleration methods.
IN-DISTRIRUTION TASKS
DUT-OF-DISTRIRUIION TASKS
AVG SPEEDUP
FM/16
C10/16
LMIB/3-24
LMIB/2-32
FM132
C10/32
C102/32
LMIB/3364
WIKU3-64
Linefit
27.0
26.8
32.6
31.9
16.6
26.5
20.7
29.8
33.3
27.3
WNN
29.3
31.4
26.1
23.4
22.2
27.9
25.0
27.7
25.9
26.5
Lineli++
27.3
28.3
32.6
29.8
23.3
27.0
25.0
25.5
WNN+
33.8
45.1
45.7
44.7
31,6
44.0
37.8
38.3
37.0
39.8
NiNo (naive graph)
48.7
53.1
33.7
12.8
42.1
49.3
37.8
23.4
33.3
37.1
NiNo
46.8
56.8
50.0
48.9
46.8
49.8
44.3
48.9
48.1
48,9
Table 3: Ablations and analysis of hyperparameters. Average speedup (%) is reported.
(a) Varying scaling
(b) Varying p in k-decay
(c) Ablating embeddingg (Equation 5)
layerwise (jt, per layer)
p
45.8
LPE pos enc (V*)
48,9
param
41.5
2
48.9
no LPE
47.5
min-max per param scaling
38.8
p = 10
47.7
no word pos ene (V*)
45.4
32.1
p (no k-decay)
40.8
no edge type
45.7
(d) Varying context length
(e) Varying NiNo size
(f) Other hyperparameters (6, y are the batch size and learning rate used to meta-train NiNo)
128
48.9
48.9
M 3, D 32
42.6
C= = 10?, b = 4,7 = 3e:3
48.9
44.0
M D 128
39.7
C 10?, b 4. - 3e:3
45.1
10 34.0
M D
128
44.7
10?. 4, le.3
41.5
C 10?, b 2, 3e-3
43.8
as NiNo/FM16). We train WNN+ and NiNo with M 2, 3] layers and De [16, 32, 64, 128, 256]. where M =3, D = 128 are used unless otherwise stated (see training details in Section A.7).
Evaluation. We follow Dahl et al. (2023) when choosing the approach to compare training methods and set target based on the validation set performance using tuned Adam, which we consider as the main baseline (Table 1). To reduce evaluation uncertainty we train all the tasks multiple times (10 for vision and 3 for language) and use the median number of steps. We report a relative reduction of the number of steps in %, e.g. for task WIKI/3-64 if the method achieves perplexity 147 in 7000 steps (as our NiNo does), the reduction is 48.1% based on Table 1. For language tasks due more expensive evaluation we only evaluate every 500 steps, whereas for vision we evaluate every step.
5.2 RESULTS
Main results. Our NiNo model shows consistently better performance than Linefit and WNN and their improved variants on all the nine tasks (Table 2). On average speed up Adam optimization by 48.9% reducing the number of steps to achieve the target performance roughly by half. NiNo is followed by WNN+ and NiNo with naive graph (Section 4.... While the latter achieves slightly better performance on the FM/16 6 in-distribution task, in the outtof-distribution tasks this model significantly underperforms
NiNo, especially in the language tasks. This performance gap is expected since NiNo represents neural graphs for Transformers more accurately as we further validated qualitatively (Fig. 2d,e) and quantitatively (Section A.6). Other baselines perform poorly, but we highlight that both Linefit basclines compare favorably to WNN indicating that the latter could have learned a very simple rule despite training on a lot of data.
Comparison with L20. L20/FM 16 performs the best in in-distribution on the same task (Fig. 6a). how- ever when applied to an unseen task it overfits to the training set (Fig. 6b). In contrast, our NiNo/FM 16 trained on the same task as 1.20/MM 6 performs well on the validation set of unseen tasks, presumably because Adam is used for most of the steps and for the input NiNo uses the trajectory of parameters of the target task to make the prediction. In addition, L20/FM 6 and L20 in general have overhead at every optimization step making the total computation cost noticeable (Table 7 in Section A.7). It is also more computationally challenging to meta-train it on multiple tasks, therefore we only use pretrained L20/FM 6 from Kofinas et al. (2024). Since L20/FM does not reach the target validation scores in most tasks, we compare it only using the curves in Fig. 6a.b and only on the FM/16 and C10/32 tasks.
Ablations. We ablate NiNo components and hyperparameters and found that our layerwise parameter scaling and k-decay are the most important components (Table 3a.b). Among the hyperparameters,
8

---

Published as a conference paper at ICLR 2025
(a) FM/16
(b) c10/32
(c) WTKI/3-64
04
TRAINING
50
230
.000
1009
200
6000
1093
8930
10093
2000
4X0 acco ax0 'oxee 1289 SIDE Tolliig Ileraton
Tolting Deraton
training Horatio
20-3
val004t00 prestestty
VALIDATION
4-.17
travee grapt)
volratton
%
LM01
FMN
V984
05.0
L20-FM1 NMN
Sarget perplctly
87.
0000
19900
4003
6033
10090
5930 5300 10000 12300 140CE tniinng
Imining
Imaiting
Figure 6: Training loss and validation performance FaShionMNIST, CIFAR-IO and WIKI
we found that context length c= 5 is a sweet spot (Table 3d), because small may not have enough information to capture the parameter trend, while larger e leads to the NiNo applied too rarely (for == 10 it is applied every 10 x 200 = 2, 000 steps). Using deeper (Table 3e) and wider (Fig. 5a) GNNs also important. However we found that depth M > width D > 128 make more challenging to train a performant model. Increasing the number of models (C) helps NiNo achieve better speedups (Table 3f, Fig. 5b). In contrast, WNN+ saturates fast when D or C' is increased.
Evaluating on larger models and Llama3-style archi- tectures. We evaluated the ability of NiNo speed up training of much larger models than it was trained on. Specifically, we trained 4 and 6 layer GPT2-style Trans- formers WIKI with 128 and 384 hidden units having around 7M and 29M parameters respectively (Fig. 5c). The latter is around 18 times larger than the largest in- distribution architecture. Moreover, the WIKI dataset is different from IM B used in meta-training. Despite these challenges, NiNo able speed up training by around 40% and 15% for the 7M and 29M models respectively, outperforming WNN+ We also trained a Llama3-styll Transformer with layers and 384 hidden units (111M I1M parameters) with AdamW W and Adam W+NiNo, which is an even more challenging setup than in the previous experiments because NiNo did observe Llama models during meta- veeai-training (Table 1). In this ex- periment, we use the Llama 3 tokenizer (Dubey et al., 2024) having larger vocabulary than in GPT2, so we use an ablated NiNo without V Besides the vocabulary, Llama models may pose additional challenges for NiNo, c.g. they differ from GPT2 in parameter initialization, MLP structure, layer normalization, positional embeddings, using biases, grouped query attention, etc. (see Fig. 9 for GPT2 and Llama neural graph examples). Nevertheless, even in this setup our NiNo accelerates training reaching the validation perplexity of the Adam W baseline (24.4) in around 30% less steps (Table 4).
Table 4: WIKI validation perplexity ()) on a Llama3-style architecture (reported the mean and standard deviation for 3 runs).
Method
10000 steps
14000 steps
Adam' AdamWoNiNo
26.29-0:055
24.44+0.01
24.36-0.11
22.37=0.13
5.3 ANALYSIS
Parameter space R We show how the parameters evolve dur- ing training on C10/32 by projecting them to 2d using PCA (Fig. 7). The PCA projection matrix is computed using the entire Adam trajectory, so the same projection is applied to all the methods. We show the first 8,000 steps for Adam/AdamW and 4,000 steps for other methods with 200 steps between points. In the beginning of training, both WNN and NiNo make big steps (dashed lines in Fig. 7) along the
Adam
L20/FM16
AdamW WWN
NiNo
1.5
25
1.0
NiNo steps
2.0
1.5
-3.5
-1.0
0.5
WNN loss increase
-.0
-05
0.0
0.5
1.0
1.5
dim
Figure 7: Comparison of optimization methods on the c10/32 task by projecting parameters to 2d using PCA every training step.
9

---

Published as a conference paper at ICLR 2025
4-128
4-128
6-384
6-384
Crwiningg
drl
(a) WNN+
(b) NiNo
Figure 8: Comparing the embedding quality of WNN+ (a) and NiNo (b) during optimization by projecting neural graph embeddings to 2d using PCA every 1,000 steps. We show the training trajectories of three GPT2-style Transformers with 3, 4 and 6 layers respectively.
trajectory when predicting the parameters, while 1.20 diverges early. At later optimization stages, NiNo makes better predictions than WNN as the latter leads to significant loss spikes. Interestingly, NiNo's 's trajectory closely corresponds to a 2 faster AdamW, even though NiNo's base optimizer Adam in this task. This indicates that NiNo may implicitly regularize optimization.
Graph embedding space (E' - RR.. The parameter space analysis used above can be effective on single task. However, for several tasks with different model sizes ((0), finding a common PCA projection matrix (e } required for the parameter space analysis, becomes challenging. Using neural graphs and NiNo enables such an analysis. Specifically, we can first represent the parameters in the graph embedding space using NiNo and then project the embeddings to see how the graph embeddings evolve when training three different GPT2-style Transformers on the language task (Fig. 8). We construct D-dimensional graph embeddings by computing the average edge features (1/lE M Eiy E after the last NiNo (or WNN+) layer by feeding the past parameter states every 1,000 steps as our other experiments. Despite having different Transformer architectures, the graph embedding space is the same, which allows us to project the embeddings to 2d using PCA. The PCA projection matrix is computed for each method (NiNo or WNN+) based on the entire trajectory for all the three architectures to allow for visualization on the same plot. Surprisingly, even WNN+ can encode parameter in a meaningful way - in general the parameters of different architectures and at different optimization stages visually distinct (Fig. 8a). However, WNN+'s embeddings virtually collapse at later training stages for lower loss values (c.g. for 6-384), while NiNo's embeddings visually more distinct for lower loss values (Fig. 8h). The ability of NiNo to distinguish parameters at later optimization stages can be useful in practice for multiple reasons. For example, despite similar training losses, model parameters can be quite different resulting different behaviors on downstream tasks (Liu et al., 2023; Dahl et al., 2023).
2d. We use this approach
Computational costs. NiNo is applied only every 1,000 steps of the base optimizer, so the wallclock time overhead of using our default NiNo is small, e.g. <0.5% for a 7M Transformer compared to around 5% by L20 (Table 7). As for the memory required for a prediction with NiNo, our implemen- tation NiNo's ' peak memory usage is similar to the peak memory usage of a training step of the Trans- former with a batch size of 16-32. We expect that in future work, NiNo's inference efficiency can be improved in multiple ways. For example, we can potentially parallelize NiNo's 's inference on large neu- ral graphs using approaches similar to large scale model inference image generation (Li et al., 2024).
Additional experiments are presented in Section A.8, the pseudo-code is in Section A.9.
CONCLUSION AND FUTURE WORK
We proposed NiNo, a novel approach to accelerate training of neural networks, including Transform- ers. Our work makes step in the relatively underexplored area of periodic parameter prediction, where our experiments show that there a lot of potential for reducing training time. Our experiments also reveal potentially useful byproducts of NiNo, such as low-dimensional encoding of network parameters during training, which can be used to analyze diverse models and their training dynamics. An interesting direction of future work is to investigate the scaling up of our approach, in particular as the scaling trends that we have observed in our experiments indicate better speedups with more data and larger models.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The experiments were in part enabled by computational resources provided by Calcul Quebec and Compute Canada. Simon Lacoste-Julien is a CIFAR Associate Fellow in the Learning in Machines & Brains program as well as a Canada CIFAR Al Chair. Guillaume Lajoie a Canada CIFAR Al chair as well Canada Research Chair in neural computations and interfacing. E.B. and A.M. acknowledge support from Mila-Samsung Research Grant.
REFERENCES
Brandon Amos. Tutorial on amortized optimization for learning to optimize over continuous domains. arXiv preprint arXiv:2202.0066. 2022. 2
Bruno Andreis. Soro Bedionita, and Sung Ju Hwang. Set-based neural network encoding. arXiv
preprint arXiv:2305. /6625, 2023. 2
Marein Andrychowiez, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016. 1,2
Maor Ashkenazi, Zohar Rimon, Ron Vainshtein, Shir Levi, Elad Richardson, Pinchas Mintz, and Eran Treister. NeRN: Learning neural representations for neural networks. The 11th International Conference on Learning Representations (ICLR), 2023. URL https: http://pppeeeeeeeeeeee net/ forum?id=-gffrrfff..
2
Mikhail Belkin and Partha Niyogi. Laplacian eigenmaps for dimensionality reduction and data representation. Neural computation, 15(6);1333-139 2003.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv: 13/2.3005, 2013.
Can Chen, Xi Chen, Chen Ma, Zixuan Liu, and Xue Liu. Gradient-based bi-level optimization for deep learning: A survey. arXiv preprint arXiv:2207. / 11719, 2022a. 2
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Howard Heaton, Jialin Liu, Zhangyang Wang, and Wotao Yin. Learning optimize: A primer and benchmark. Journal of Machine Learning Research, 23-18));1-59, 2022b. 2
Guillaume Chevillon. Direct multi-step estimation and forecasting. Journal of Economic Surveys, 21 (4):746-785, 2007.
Gabricle Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Veličković. Principal Neighbourhood Aggregation for Graph Nets. In Advances in Neural Information Processing Systems 33 (NeurIPS), 2020. 2, 16
George E Dahl, Frank Schneider, Zachary Nado, Naman Agarwal, Chandramouli Shama Sastry, Philipp Hennig, Sourabh Medapati, Runa Eschenhagen, Priya Kasimbeg, Daniel Suo, et al. Bench- marking neural network training algorithms. arXiv preprint arXiv:2306.07177 2023. 8, 10
Abhimanyu Dubey, Abhinay Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, al. The laama 3 herd of models. arXiv preprint arXiv:2407.21788 2024. 4,
Vijay Prakash Dwivedi, Chaitanya K Joshi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. Journal of Machine Learning Research, 24 (43):1-48, 2023. 5
Charles Guille-Eseured Hiroki Naganuma, Kilian Fatras, and Ioannis Mitliagkas. No wrong turns: The simple geometry neural networks optimization paths. In International Conference on Machine Learning (ICMI). PMLR, 2024. 2
11

---

Published as a conference paper at ICLR 2025
Howard Heaton, Xiaohan Chen, Zhangyang Wang. and Wotao Yin. Safeguarded learned convex optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7848-7855, 2023. 2
Robert Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pp. 129-135. Elsevier, 1990. 2
Jinhyeok Jang, Woo-han Yun, Won Hwa Kim, Youngwoo Yoon, Jaehong Kim, Jaeyeon Lee, and ByungOk Han. Learning to boost training by periodic nowcasting near future weights. In International Conference on Machine Learning (ICML), 2023. 1, 2, 5, 7, 15, 18
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In The 4th Interna- tional Conference on Learning Representations (ICLR), 2015. 1,7
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning (ICML), pp. 2697. PMLR, 2018.
Thomas Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In The 5th International Conference Learning Representations (ICLR), 2017. 2
Boris Knyazev, Carolyn Augusta, and Graham W Taylor. Learning temporal attention in dynamic graphs with bilinear interactions. Plos one, 16(3):e0247936, 2021a. 2
Boris Knyazev, Michal Drozdzal, Graham W Taylor, and Adriana Romero Soriano. Parameter prediction for unseen deep architectures. Advances in Neural Information Processing Systems, 34: 39433-29443 2021b. 2
Boris Knyazev, Doha Hwang, and Simon Lacoste-Julien. Can we scale transformers to predict parameters of diverse imagenet models? In International Conference on Machine Learning (ICML), pp. 17243-17259. PMLR, 2023. 2, 18
Miltiadis Kofinas, Boris Knyazev, Yan Zhang, Yunlu Chen, Gertjan J. Burghouts, Efstratios Gavves, Cees G. M. Snock, and David W. Zhang. Graph Neural Networks for Learning Equivariant Repre- sentations of Neural Networks. In 12th International Conference on Learning Representations (ICLR), 2024. 1, 3, 4, 6, 7, 8, 15, 16,
Richard D Lange, Devin Kwok, Jordan Kyle Matelsky, Xinyue Wang, David Rolnick, and Konrad Kording. Deep networks as paths on the manifold of neural representations. In Topological, Algebraic and Geometric Learning Workshops 2023, pp. 102-133. PMLR, 2023.
Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Kai Li, and Song Han. Distrifusion: Distributed parallel inference for high-resolution diffusion models. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp.
7183-7193, 2024. 10
Derek Lim, Haggai Maron, Marc T Law, Jonathan Lorraine, and James Lucas. Graph metanet- works for processing diverse neural architectures. In 12th International Conference on Learning Representations (ICLR), 2024. 2, 3. 15
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same pre-training loss, better downstream: Implicit bias matters for language models. In International Conference on Machine Learning (ICML), pp. 22188 22214. PMLR, 2023. 10
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In The 5th In- ternational Conference Learning Representations (ICLR), 2017. URL https: :tats:://ai.... www.cccchh........oo..t..aeeeeee
1,7, 18
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. 7
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstee Understanding and correcting pathologies in the training of learned optimizers. In International Conference Machine Learning (ICML), pp. 4556-4565. PMLR, 2019. 2
12

---

Published as a conference paper at ICLR 2025
Luke Metz, C Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl-Dickstein. Practical tradeoffs between memory, compute, and performance in learned optimizers. arXiv preprint arXiv:2203.11800 2022a. 18
Luke Metz, James Harrison, C Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint urXiv:2211.09766 2022b. 1,2
Chedi Morchdi, Yi Zhou, Jie Ding, and Bei Wang. On the mysterious optimization geometry of deep neural networks. 2022. 2
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron. Equiv- ariant architectures for learning in deep weight spaces. arXiv preprint arXiv:2301.12780, 2023. 1,
William Peebles, Ilija Radosavovic, Tim Brooks, Alexei A Efros, and Jitendra Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209. /2892, 2022. 2
Isabeau Srannont-Schwarz,Jaaa Jaroslav Vitkû, and Jan Feyereisl. A simple guard for learned optimizers. In International Conference on Machine Learning (ICML), Proceedings of Machine Learning Research. PMLR, 2022. 2
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. 4.7
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115:211-252, 2015, 18
Alvaro Sanchez-Gonzalez Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In International Conference on Machine Learning (ICML), pp. 4470 4479. PMLR, 2018.
Alvaro Sanchez-Gonzzler,, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning (ICML), pp. 8459-8468. PMLR, 2020. 2
Konstantin Schürholt, Dimche Kostadinov, and Damian Borth. Self-supervised representation learning neural network weights for model characteristic prediction. Advances in Neural Information Processing Systems, 34:16481-1649; 2021. 2,3
Konstantin Schürholt, Boris Knyazev, Xavier Giró-i Nieto, and Damian Borth. Hypee-represennations as generative models: Sampling unseen neural network weights. Advances in Neural Information Processing Systems, 35:27906-27920, 2022. 2,5
Konstantin Schürholt, Michael W. Mahoney, and Damian Borth. Towards scalable and versatile weight space learning. In International Conference on Machine Learning (ICML). PMLR, 2024.
Abhishek Sinha, Mausoom Sarkar, Aahitagni Mukherjee, and Balaji Krishnamurthy, Introspection: Accelerating neural network training by learning weight evolution. In The 5th International Conference on Learning Representations (ICLR), 2017. 2, 3, 15
Bedionita Soro, Bruno Andreis, Hayeon Lee, Song Chong, Frank Hutter, and Sung Ju Hwang. Difflusion-based neural network weights generation. arXiv preprint 1rXiv:2402.18153. 2024. 2
Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton. On the importance of initialization and momentum in deep learning. In International Conference on Machine Learning (ICML), pp.
1139-1147. PMLR, 2013. 15
Benjamin Thérien, Charles Charles-Etienne Joseph, Boris Knyazev, Edouard Oyallon, Irina Rish, and Eugene Belilovsky. μLO: Compute-efficin efficieet meta neeta-generalization of learned optimizers. arXiv preprint 00/0234 2024." 2
13

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXfv:2307.09288 2023. 1
Rakshit Trivedi, Mehrdad Farajtabar, Prasenjeet Biswal, and Hongyuan Zha. Dyrep: Learning repre- sentations over dynamic graphs. In The 7th International Conference on Learning Representations
(ICLR), 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. 15
Kai Wang, Zhaopan Xu, Yukun Zhou, Zelin Zang. Trevor Darrell, Zhuang Liu, and Yang You. Neural network diffusion. arXiv preprint arXiv:2402.1314: 2024. 2
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap. Edward Lockhart, et al. Deep reinforcement learning with relational inductive biases. In The 7th International Conference on Learning Representations (ICLR), 2019. 3
Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting? Proceedings of the AAAI conference on artificial intelligence, volume 37, pp. 11121-11128, 2023. 6
Allan Zhou, Chelsea Finn, and James Harrison. Universal neural functionals. arXiv preprint arXiv:2402.05232, 2024a.
Allan Zhou, Kaien Yang, Kaylee Burns, Adriano Cardace, Yiding Jiang, Samuel Sokota, Zico Kolter, and Chelsea Finn. Permutation equivariant neural functionals. Advances neural information processing systems, 36, 2024b. 1,2
Allan Zhou, Kaien Yang. Yiding Jiang, Kaylee Burns, Winnie Xu, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Neural functional transformers. Advances in neural information processing systems, 36, 2024c. 2
Xinyu Zhou, Boris Knyazev, Alexia Martinuu and Jie Fu. Logah: Predicting 774- million-paramesn parameter transformers using graph hypernetworks with 1/100 parameters. arXiv preprint arXiv:2405.16287. 2024d.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 WNN PARAMETER SCALING DETAILS
In WNNs (Jang et al., 2023), parameter scaling is performed by first computing the minimum and maximum values the input parameter values: 8 = mmx (0 0²,)) ) miinmmin (O',.., )) Vie [1, mt]. Each i-th parameters scaling factor s is then used to scale both the inputs: 0 + c). and,
only for training WNNs, the targets: 0 = 0 e). Jang et al. (2023) also subtract the last element from the sequence, however we found this to be redundant. Unscaling is performed
by multiplying the prediction by xx. Scaling and unscaling are performed per parameter (for each parameter independently).
A.2 LINEFIT
Sinha et al. (2017); Jang et al. (2023) introduced "linefit" baseline to predict future parameter values by extrapolating the line to future point:
Vie [1, n]
6 0+ = 2a + b²,
(10)
where a², b are obtained by optimizing the following objective:
argmin 0t hi ||a' x + b² θ O 112.
(11)
where x = [1,2, c| and O T:7. = [e] - e E R. This baseline fits line (a b¹) for each parameter given only the past values of the same parameter without collecting training dataset. Despite its simplicity, this baseline was shown to improve on using Adam/SGD only.
A.3 LINEFI++
Linefit outlined in Section A.2 can be viewed as a form of momentum that promotes the parameters to be consistent with their global trend. However, Linefit weighs all past values uniformly when fitting the line, while in the momentum more recent values have more effect on the weight update (Sutskever al., 2013). Motivated by this observation, we introduce the Linefit+ baseline that penalizes more the errors for later parameter values in the trajectory:
argmin,i vallu(('' x + b¹ - - -
-t....) 122,
(12)
where f [1/c, 2/c,..,
A.4 NAIVE NEURAL GRAPH OF TRANSFORMERS
A mnulti-head self-attention (MSA) layer of Transformers (Vaswani et al., 2017) consists of three weight matrices W V W applied to the input x Nxd and another weight matrix W re- turning the output, where N is a sequence length and all the four weight matrices are d x d. To compute MSA with heads, the weight matrices are typically split into H groups across columns for W W W and rows for 8 (Fig. 2a). The heads are processing the input in parallel, Yhe [1, H] using Equation 3.
Neural graphs defined by Kofinas et al. (2024) and Lim et al. (2024) are conceptually similar, however Lim et al. (2024) model biases (()) and normalization layers as extra nodes connected to the neurons in the 1-th layer (Fig. 2c). Normalization layers are similarly added as separate nodes, hence the node features of Lim et al. (2024) do not include any parameters. We build on this variant, since more easily to implement it for the layers with complicated neuron permutation symmetry, such multi-head self-attention. Compared to Kofinas et al. (2024), Lim et al. (2024) also provide a slightly different description for the MSA neural graph however the W W W edges are also stacked. Given the lack of detailed description for the multi-her head case and implementation currently available publicly, makes it challenging to use their neural graphs directly. Therefore, we use (Kofinas al., 2024) the base (naive) neural graph when constructing MSA layers.
15

---

Published as a conference paper at ICLR 2025
A.5 GRAPH NEURAL NETWORK
A Graph Neural Network (GNN) layer with edge features be formulated using the mean aggrega- tion as follows (Corso et al., 2020): V; v==óh [Vae]] EjEN(e) mi)) where maj is computed using message passing &m given node and edge features vi, Vi, eig; eij Oa is an MLP; (i) are the neighbors of node i. Kofinas al. (2024) modify how myy is computed to better incorporate edge features and introduce an edge update step to make better per edge predictions (Table 5). Stacking such layers form a GNN that (1) does not change the size of the input graph and (2) permutation equivariant. (1) means that for the neural graphs there are |0|' dg input edge features and [0 x 1 predictions, where 10 is the number of trainable parameters in the input model plus auxiliary non-trainable parameters ( residual and head connections). results in the same permutation of the output nodes and corresponding edges. These properties make GNNs suitable model to predict future parameters and to work with neural graphs where nodes (neurons) can be permuted in the ways described in Section A.6.
(2) means that any permutation of input nodes
Table 5: GNN layer comparison.
Step
Typical GNN (Corso el al., 2020)
Neural Graph GNN (Kofinas et al., 2024)
1. Message passing
mij 6m (LvG, V,, 0:yl) maj Peciia (exz) O om ([vs, Vzl) + daain (e:,)
2. Aggregation
V. de (1/W((1))))))) EjEN(O) mis)
3. Edge update
(k+ 1)
(k)
(k+1) (k+1)
([ (2) (6)
(v og 1 .V. vv("1)
A.6 NEURON PERMUTATION SYMMETRY
Permutation symmetry in neural graphs. Let Bood (0) be a permutation of neurons in O that does not change the network function: f(x, 0) f(x, π 7800 (0)) Likewise, let (O) be permutation that changes the function: f(x, 0) / [xx n bad (0)). Denoting = and ş as "isomorphic" and "non- isomorphic" operators graphs respectively. for neural graphs G the following equations satisfied under some assumptions (Kofinas et al., 2024): ( = G (T (e)) and G()) 전े ( Here, by definition, the neural graphs are assumed to be constructed correctly, however, achieving or validating that in practice is not trivial for such networks as Transformers.
Neuron permutation symmetry experiment. To validate how well our neural graphs correspond to true neuron permutations in Transformers, we run a simple neuron permutation experiment. We use single Transformer layer with d = 12 and H permute rows and columns in its MSA weight matrices 0 using a uniformly sampled permutation æ, and label each permutation as = for x good or yx for ب depending if the Transformer output f(x, 0) changes or not:
(1. iff(x, 0) = f(x,x(0)) (0))
y= (0 otherwise.
(13)
We generate 1,000 such permutations with about 500 of good and bad permutations. Then extract a neural graph embedding for each permutation. For correctly constructed neural graphs the embeddings corresponding to yr = should be close to the non-permuted e. To extract embeddings we use randomly initialized graph neural network with 3 layers and 32 hidden units based on Section A.5. We then visualize the embeddings in 2d (Fig. 2d,e) and color-code each embedding with Vr. We also train a logistic regression model tú that takes a graph embedding h as input and predicts if it is good or bad permutation: ú. = (h), and we compute classification accuracy between Dn and Vr. We use all 1,000 samples for training and evaluation i to estimate how easily can the graph embeddings be separated based on Yn (random guess accuracy is around 50%). We repeated the above experiment for the naive and our neural graphs (Fig. 2, d.e). The visualization and classification results indicate that our neural graph construction respects good and bad permutations well, while naive neural graphs confuse them.
16

---

Published as a conference paper at ICLR 2025
edge features introduced in Section 4.1, color-coded by edge type. (top) A layer GPP22based Transformer with el and 2. (bottom) 2 layer Llama3-basee Transformer with 12, intermediate hidden size equal 24, H = and number key value heads cqual for GQA. The number of word embeddings is reduced for visualization purposes.
Figure 9: Graph and adjacency matrix with head restrictions and separate
F/16
C10/16
LM18/1244
LM18,0-32
30
200
COunT
20
15
51
53
53
0.4
0.5
0.8
1.8
1.2
5.00
5.45
6.50
6.55
5.75
5.00
5.85
5,00
5.95
training loss (0.25e0.06)
training loss (0.91=0.10)
training loss (5.04 04.0003)
braining loss (5.65
Figure 10: Histogram of final training losses in the in-distribution (meta-training) tasks. See Table for more details.
17

---

Published as a conference paper at ICLR 2025
8030
training loss
0.10
6000
0.00
5000
lived (WNN)
0.08
WNN--n1120 deamples
4000
0.07
WNN--hid MWWN--id128---1 1000samples
5000
WNN- hi532 4camples
0.00
NNG-hät28 1ob--hit28--aamlla
2000
NiNohhat2a 100tsamples
1030
0.00
0
2530
5200
7503
10030
12500
15003 17500 20000
2000
4000
5003
5030
10000
training step
traning stop
Figure 11: Meta-training curves.
Figure 12: Different k- decay schedules.
A.7 NINO TRAINING DETAILS
For representing convolutional networks with neural graphs, dg ft==/o h x w (Kofinas al., 2024), where h, w are convolution kernel size, so edge features representing e states are |0]' x chw, where |el' the number of trainable and auxiliary non-traine non-trainable (residual, head connections) parameters in model. To make the feature dimensionality the same for different architectures during meta-training, we use zero-pan padd-p.dinn. We train meta-models for 20k training iterations using AdamW, learning rate 3 with cosine decay and weight decay 0.01. We sample batch of 4 checkpoints in each training iteration and use automatic mixed precision. Training of the NiNo and WNN+ meta-modehs completes in under 7 and 6 hours respectively on single NVIDIA RTX8000 with 48GB of memory.
Meta-training behavior can give important highlights about the model and data. Therefore we explore meta- teta-ttaaning for different hidden sizes (128 vs 32) and numbers of dataset samples (checkpoints sampled either from all 1000 or just 4 models). We found that NiNo achieves a lower meta-training loss than WNN+ despite having a comparable number of parameters indicating the important of leveraging neural graphs (Fig. 11). However our NiNo model is still in a severe underfitting regime when trained on all the training checkpoints, since when the number of checkpoint models reduced to just 4 the loss is much lower for NiNo (with 128 hidden units). This result suggests that further scaling up NiNo could help fit all the training data better and potentially achieve stronger results.
A.8 ADDITIONAL. EXPERIMENTS
A.8.1 LEARNING RATE ANALYSIS
When training language models (Section 5.2), we closely followed the setting of WNN (Jang et al., 2023) and used the same learning rate (Ir=2e-4). We used this Ir for all language tasks, including the collection of training checkpoints (LM1B/3-24 and LM LMMB/2231 BB22-3 tasks). Since NiNo was meta-traines on language models trained with Ir=2c-4, we use the same Ir for all language experiments including Llama3-style (Table 4). However, this Ir may be suboptimal for our Llama3-style task, therefore we investigated if our trained NiNo can speed up Adam training with learning rates that are different than those used to collect the checkpoints. We followed the Llama3-style experiment "Evaluating on larger models and Llama3-style architectures" (Section 5.2) with 5 different learning rates run 3 times (Table 6). The standard deviation the validation perplexity is generally &0. 1, indicating that ours is better including the std even in this challenging setup (new Ir compared to the training checkpoints, new and 18 times larger architecture and new task). Speedup is computed as described Section 5, e.g. 14% means that AdamW+Ne Niioo achieved the target performance in 12k steps 14k steps of AdamW. Based on the results, we conclude that the speedup varies but is significant for different learning rates.
A.8.2 IMAGENET
To investigate the ability of NiNo on more challenging vision tasks, we trained a small ViT (1IM parameters) on ImageNet (Russakovsky et al., 2015), with x2x images as in Metz et al. (2022a); Loshchilov & Hutter (2017). The ViT architecture has 6 transformer blocks with 6 heads and 384 hidden units, the patch size is 2. We trained using Adam W for 50 epochs with learning rate 0.003,
cosine scheduler. weight decay 0.3, batch size 3096 and automatic mixed precision using Py PyTorc and the code base from Knyazev al. (2023). This setup fully utilizes a single NVIDIA A100-80GB
18

---

Published as a conference paper at ICLR 2025
Table 6: Learning rate analysis for the Llama3-style architecture and the WIKI dataset, extending the results in Table 4. WIKI validation perplexity (L) is reported.
learning rate
le-4
2e-4
4e-4
6e-4
8e-4
AdamW (10k steps)
31.85
26.29
24.81
24.50
25.03
AdamW + NiNo (10k steps)
26.18
24.36
24.10
23.82
24.49
AdamW (14k steps) target to compute speed
28.12
24.44
23.59
23,30
23.74
AdamW NiNo (14k steps)
24.01
22.37
22.52
22.52
22.98
speedup w.r.t. A (14k steps)
36%
29%
21%
14%
14%
training the model in around 6 hours. The setup has around 21k steps in total, with NiNo applied every 1000 steps as in all our experiments. We multiply the predicted parameter delta by 0.1, which we found to be helpful to avoid divergence. Also, make the NiNo step more efficient and avoid OOM on GPU, sample 5% of the edges fed to NiNo. The NiNo step takes about 10 seconds, which similar
steps of Adam, so the overhead is around 1%. After training for epochs, the validation top-1 accuracies are 49.6% (AdamW)
50.4% (AdamW+NiNo) So NiNo makes a small yet reasonable improvement given that we are applying the NiNo model to a very different architecture on very different task compared to meta-training.
Table 7: Computational cost estimates. Measured on NVIDIA A100-80GB after running for steps with batch size = 32, sequence length= 1024. Note that these numbers are rough estimates as the time/memory vary depending on the GPU/CPU load and internal logie in different Py PyTorch versions. Also note that time and memory can be traded off by processing only fraction of parameters at time. For Adam and L2O, we report per step measurements, since these methods measurements after steps Adam. "Overhead" denotes wall clock time increase vs passing step (Eq. 1-2 in Table 5) of NiNo is computed by sampling 5% of the edges to avoid OOM on GPU.
report per prediction used every slep. For WNN+ and NiNo, Adam. 'The message
METROO HID sIZE D
3-64 (3.4M, GPT2-BASEDi
4-128 (7 (7.4M. GPT2-BASED
61384 (111.5M LLAMA3-BASED)
PEAK MMM
TIME
OVERnEAD
PEAK MMM
TIME
OVERIEAT
PRAK MEM
TIME
OVERuEAn
Adam
27GB
130ms/siep
0%
28GB/step
180mshtep
0%
76GD/sep
909ms/step
0%
WNN+
32
IGiippedict
172ms/predict
0.13%
3GB/prodict
222ms/mmdict
0.12%
42CBipnedict
2.01 scofpocct
0.22%
WNN+
128
4GRipeederr
245ms/predict
0.19%
294me/podier
1.16%
71GRppodller
0.38%
3GB/podicr 323ms/predict
SGit/prediet 49%msppodicc
57GBtpoddkt
5.52sc//rrsdict
NiNo
128
9GB/ppeddict
483mssprediet
0.37%
16GB/peedict
864mm//preicc
0.48%
27GB/predict
13. 13.35scc/rredict
1.47%
120 (mlp)
27GB8kepp
134ms/stop
5.08%
28GBB/vee
190nshtep
3.56%
766i8/sep
91 7mms/ttep
0.88%
19

---

Published as a conference paper at ICLR 2025
A.9 PSEUDO-CODE
We show the pseudo-code for the three main steps in our pipeline:
1. collecting the dataset of checkpoints on set of training tasks;
2. training NiNo given the dataset of checkpoints;
3. evaluating/using the trained NiNo on new tasks.
A.9.1 COLLECT CHECKPOINTS
checkpoints
ckpt_ eaplnerutter
counter of checkpoints
for task in [fm/16, c10/16, 1mlb/3-24 1mlb/2-32]:
for in range (300 if task in [fm/16, c10/16] else 200): model (lask)
initialize trajectory [] # to store the parameters of the model
for t range (1000(10000):
model adam step(moder task) update model parameters using Adam the task if & (4 task in [fm/16, c10/16] else 200) trajectory append (model. parameters save only parameters without gradients ekpt.counier += checkpoints append (trjjector))
11
12
13
14
print (len (checkpoints)) # 1,000 trajectories (models trained)
print expp__cunteer 1.609.990 ,900 total checkpoints
16
print DDDee collecting the 'ataset')
A.9.2 TRAINING NINO
NINo(c=5, gan layers=3, Kw40)
architecture Direct multi-
forecasting as deseribed in 4.3)
opt torch. optim AdaaW( (arameters 1r=3e-3, weight.decay=lo-2) # train NiNo with AdamW
(20000);
trajectory sample(checkpoings
# sample a trajectory of model parameters sample [ [t-5:1] past 5 parameter states
(trjjeetttt))))))
step the trajectory
theta_ing input = trajectory
target trajectory [t:t+40] future parameter states
theta_input. seales (theta.tttatt...ee... scales).
# use layerwisc sealing
theta.target scale.params hheta_delta.target theia_target theta_input[-]] the last value the inpur enbiteettueetraajcett
predict the delta w.r.t.
get model architecture neural graph Graph arapihhachitettuee. edge features@theta input)
architecture
the graph
12
13
theta_delta.pree preta.preed nino geural
14 loss (theta_delta..peed
theta.della.taarett) abs .)mam)))n( mean()
15 lass backward
opt. step # update NiNo's parameters
17
opt gerd
A.9.3 USING NINO
model AutoModdFForauusaLLM from_config (...) some model wrapper around the base optimizer any uplimizer other than Adam should also be possible to use with NiNo NiNa MiNo(base_opt=torch optim model parameters Ir=le-3), nino. nino. pt" pt". modelzmodel.
NiNo is implemented
period=1000.
max.train.stcps:11000
for in (10000)
10
(step + 1) & 1000 1= D:
Usual Adam updates for of the steps (e.g. 0-999) opt. zero-gradd data. targets
12
get data
14
outputs model(data) # forward pass loss F. cross_catropy (outputs , targets) loss lass backward
15
grads for Adam opt step Tadam=Truee # Adam step ar 200 0:
17
if (step opt. nino. iinttt.nptt... append model. parameters ))) collect the history of parameters
19
21
no need to compute the gradients for the NiNo forward pass opt step Falme) every 1000 steps nowcast params using NiNo
22
20


Reviews:
[{"review_id": "iXrP0luI9K", "reviewer": "Reviewer_8oEj", "summary": "The paper proposes to improve training, by leveraging out specific transformer strcture and representing it is as a Neural Graph to model connectivity. The weights of the model would then be encoded as features on the edges of the graph (the nodes corespond to neuron/activation).  Ideally,  isomorphic permutations of such graphs, \nshould result in no-functional change to the model itself, thus capturing the underlying symmetries of the model. In contrast with earlier work, authors proposed a more fine-grained structure that accurately captures multi-headed-self-attention blocks, and thus allows them to improve performance of their meta-learning method. Once such graph is built authors leverage a combination of nodes and edge embedding, methods from GNN which are then translated to model updates..", "questions": "1. How are  the gradients are fed into the NiNo algorihtm? It is part of V^w? \n2. What are the nodes features? \n3. Was there any ablation study done for AdamW in Liama3 style architecture particularly around using different learning rates.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "I think the construction of the graph whose isomorphic permutations of nodes preserves the model functionality,  for transformer layer is neat and interesting.", "weaknesses": "I have two major concerns about the papers:\n1. My main concern is that the experimental section contains only fairly trivial datasets (FashionMNist, Cifar-10), which are very far from anything reasonable these days, and the models authors consider for forecasting is limited to ~1M parameters, and many are 15K params,  which is barely practical for the simplest tasks. I think for image tasks, showing reasonable performance on something like ImageNet is\na must.  On the other hand authors run their training procedures for thousands of steps, which seems an overkill  for simple problems like this (Gradient Descent for Fashion Mnist for instance can converge in ~100 steps).  On the other hand experiments on Liama are not convincing because it is a well  known fact that it is often easy to speed up initial convergence, but the end result can still be significantly worse than slower algorithm. To demonstrate performance gains on these tasks authors should run at least run ablation study with different learning rates of the original model and show that they all converge (hopefully their method still shows the improvement).\n\n2) The actual algorithm descriptions has very little detail what exactly their meta-learning algorithm does, even when reading the  appendix  carefully. Most of the detail is in section 4.2. In particular having at least a basic definition of GNN layers. Given that it is the crux of the paper, i think it should be carefully expanded and explained. I think the graph-construction section, is fairly well written, but it is very much unclear how it is then translated into actual meta-learning algorithms, as best as i could tell authors use some black-box message passing algorithms to come up with the new update, but i think that section would benefit greatly from expansion. \n\n3) Some discussion about optimality of their graph construction would be nice, but it is a minor comment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "zKSB5LtSMt", "reviewer": "Reviewer_Ubaj", "summary": "This paper introduces a novel approach for accelerating the training of neural networks. This technique called NiNo builds upon the concept of Weight Nowcaster Networks (WNNs), which periodically predict future parameter values thus speeding up optimization. However, in contrast to WNNs, NiNo leverages the inherent network structure by incorporating neuron connectivity using graph neural networks. This allows the model to achieve more accurate prediction of future parameters and leads to faster training. The authors also address challenges in modeling neuron connectivity, particularly in Transformers, and demonstrate NiNo's effectiveness in accelerating Adam optimization in various vision and language tasks.", "questions": "1. How could one bridge the gap between current experiments and practically interesting large model sizes (1-100B parameters)? Would generation of multiple future steps (various values of $k$) present a major obstacle? What about stacked feature representations that gather information from $c$ weight instances?\n2. I could have missed this discussion, but what would happen if one decoupled the number of steps between predicting model weights and the size of the history used to make this prediction (both currently chosen to be $c$ if I am not mistaken)? It would seem that there are roughly three time hyper-parameters (at least locally): (a) how frequently do update predictions; (b) how many past states to use for making these predictions (lower would be more advantageous); (c) how far ahead to predict (larger would be more advantageous).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper is sufficiently well written and is fairly accessible.\n2. The proposed approach is sufficiently sound and novel. Even though it can be seen as a combination of two existing techniques (WNNs and an improved GNN model weight representation), this paper still contains a number of non-trivial innovations. For example, among other things, the authors make a number of logical steps improving on a previously published graph topology for multi-headed self-attention.\n3. Experimental results appear to be promising. When fully realized the proposed technique could potentially be quite impactful for training modern large models including large language models. Even 10-20% improvements in training speed could be of large practical significance.", "weaknesses": "1. Some discussions could perhaps be improved upon to be even more clear. For example, while being sufficiently understandable, Section 4.1 could still be clarified further. Figure 2 is also difficult to interpret in its current form. Color coding takes time to digest.\n2. The training method is fairly computationally expensive as the authors collect on the order of $10^6$ checkpoints. To be practical, this initial computational investment should be compensated by the future computational wins from utilizing a more efficient optimizer. Early results seem to paint an optimistic picture and suggest that trained models generalize to much larger underlying models and novel datasets and tasks. However, most current practical LLMs start at around 1B parameters, which still leaves at least 1 to 2 orders of magnitude from the current 111M models the authors experimented with.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "2r2S480LuO", "reviewer": "Reviewer_PGZt", "summary": "The paper introduces NiNo, a method that improves WNN by utilizing graph neural networks (GNNs). Experiments on relatively small datasets demonstrate that NiNo achieves superior performance compared to the WNN baseline.", "questions": "1. I notice that there are some strange performance jumps in Figure 6. Training loss periodically jumps/spikes for NiNo variants. Could the authors clarify the reasons behind these fluctuations?\n2. How is NiNo or WNN applied during inference stage? Specifically, do you forward NiNo k steps and set the network weights to the output of NiNo while retaining the optimizer states? Further details on the inference process would improve clarity.\n3. Could the authors provide details about the computational costs of the meta-training process?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-motivated. It is an improved version of WNN by integrating GNN.\n2. The experiments cover different tasks including language modeling and image classification tasks.", "weaknesses": "1. It is unclear if training multiple models during meta-training is practical for real-world applications, where typically only a limited number of models are trained. \n2. The generalization performance of NiNo should be further tested. The largest test case is 100 M models on small dataset like Wikitext-103. It may not fully represent NiNo's capabilities in broader applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "JmZwI4Q97E", "reviewer": "Reviewer_qcps", "summary": "This paper proposes a novel approach to accelerate the training of neural networks by using Neuron Interaction and Prediction (NiNo) networks to predict future parameter changes, thereby reducing the number of iterations required by traditional optimizers such as Adam. This approach shows significant speed-ups on visual and language tasks while maintaining low memory overhead. The theoretical analysis and experimental results of the paper support the effectiveness of this method.", "questions": "The paper mainly focuses on the performance improvement of specific tasks, but lacks discussion on the generalization ability of the model on unseen tasks. It is recommended to increase the evaluation of the model's generalization ability.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The method proposed in the paper has significantly improved compared to the existing technology in multiple benchmark tests, especially in reducing the number of training iterations.", "weaknesses": "Hyperparameter tuning: The paper mentions the k-decay strategy, but does not discuss the selection and tuning of hyperparameters in detail. It is recommended that the authors provide more analysis on the impact of hyperparameter tuning on model performance.\nComputational resource consumption: Although the paper mentions the memory and time overhead of NiNo, it does not provide a direct comparison with existing methods. It is recommended to add analysis in this regard, especially on resource consumption on large-scale datasets and large models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Poster)