Accelerating Training with Neuron Interaction and Nowcasting Networks

## Abstract

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 n euron i nteraction and no wcasting (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.

## 1 Introduction

Figure: Figure 1: Adam without and with nowcasting using our NiNo vs WNN (Jang et al., ) on a language task that NiNo and WNN have not seen during their training.

Modern deep learning models, such as large language models (Touvron et al., ), are trained using classic adaptive optimizers, such as Adam(^1^11We use Adam throughout the paper, but our discussion and methods are in principle applicable to any optimizers that produce a trajectory of parameters, including SGD with/without momentum, AdamW, Adagrad, etc.) (Kingma & Ba, ; Loshchilov & Hutter, ).
These optimizers update neural network parameters $\bm{\uptheta}\in\mathbb{R}^{n}$ using gradient descent at step $t$ as $\bm{\uptheta}_{t+1}^{i}=\bm{\uptheta}^{i}_{t}-\Delta\bm{\uptheta}_{t}^{i}$, $\forall i\in[1,n]$, where $\Delta\bm{\uptheta}_{t}^{i}$ is the update computed analytically based on the history of parameter values, gradients and the learning rate.
Recently, Jang et al. (); Sinha et al. () showed that parameters $\bm{\uptheta}$ follow a predictable trend so that optimization can be accelerated by nowcasting (predicting near future) parameters using a learnable function $f^{\phi}$: $\hat{\bm{\uptheta}}_{t+k}^{i}=\bm{\uptheta}_{t}^{i}+f^{\phi}(\bm{\uptheta}_{t}
^{i},\bm{\uptheta}_{t-1}^{i},\bm{\uptheta}_{t-2}^{i},...)$, where $k\gg 1$ is the future horizon.
More popular learnable approaches to speed up optimization, such as “learning to optimize” (L2O), are recurrently applied at every step $t$ (Andrychowicz et al., ; Metz et al., ). Compared to L2O, the parameter nowcaster
$f^{\phi}$ is applied very rarely reducing its overhead, e.g. a base optimizer such as Adam is run for 1k steps followed by the prediction step (Fig. 1). Moreover, training such $f^{\phi}$ is simpler than L2O, since a supervised loss can be used instead of a more challenging meta-learning loss with recurrent inner steps.
However, akin to Adam, $f^{\phi}$ from prior works does not directly leverage the structural information of $\bm{\uptheta}$, 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 et al., ; Zhou et al., ; Kofinas et al., ).

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., ; Lim et al., ).
- 2.
We improve Transformer’s neural graphs of ; by more accurately modeling 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 $k$ 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 reducing 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/SamsungSAILMontreal/nino](https://github.com/SamsungSAILMontreal/nino).

## 2 Related Work

Learning to optimize (L2O). The L2O literature has offered many approaches to learn a neural network (optimizer) that optimizes the parameters of other neural nets (optimizees) (Andrychowicz et al., ; Chen et al., ; ; Amos, ). Among them,
Safeguarded L2O (Heaton et al., ; Prémont-Schwarz et al., ) is most related to ours as it switches between an L2O optimizer and SGD/Adam. While Safeguarded L2O alleviates the meta-generalization challenge (Thérien et al., ), training an L2O 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., ; ). Moreover, overheads of L2O add up at each iteration making it more computationally intensive than Adam. In contrast, our approach follows weight nowcaster networks (WNNs) (Jang et al., ), where the nowcaster model is applied very rarely (e.g. once per 1k 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 reducing computational costs of training neural nets (Peebles et al., ; Ashkenazi et al., ; Schürholt et al., ; ; Knyazev et al., ; ; Zhou et al., ; Soro et al., ; Wang et al., ).
Most related to our work are IntrospectionMLP (Sinha et al., ) and WNNs (Jang et al., ) 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 a neural net given its past parameters (Section 3.1). However, their MLPs predict parameter coordinate-wise (for each parameter independently) similar to optimization methods without leveraging the structure of neural networks. Moreover, their MLPs predict parameters only for a predefined future horizon $k$, whereas different tasks and different optimization stages can have different parameter evolution trends (Morchdi et al., ; Guille-Escuret et al., ; Lange et al., ). We address these shortcomings by conditioning the prediction on $k$.

Representation learning of neural network parameters. This area has also developed fast recently (Navon et al., ; Schürholt et al., ; ; Ashkenazi et al., ; Andreis et al., ; Zhou et al., ; ; ). 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-Nielsen, ).
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. (); Lim et al. () proposed neural graphs (graphs of neurons) enabling the usage of graph neural networks (GNNs) (Kipf & Welling, ; Corso et al., ). Remarkably, as 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 a 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 neural graphs of Kofinas et al. (); Lim et al. () do not accurately model neuron permutation symmetry of Transformers (Section 4.1). We address that shortcoming to make better predictions by NiNo.

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., ; Knyazev et al., ) and physical systems (Kipf et al., ; Sanchez-Gonzalez et al., ; ; Zambaldi et al., ). However, developing a strong parameter prediction model requires many specific design choices and careful modeling of neuron permutation symmetry, motivating our work.

## 3 Background

We consider a neural network parameterized by $\bm{\uptheta}_{t}\in\mathbb{R}^{n}$ trained for $t$ steps with Adam (or another optimizer).
Our goal is to accelerate optimization by predicting future parameters $k\gg 1$ steps ahead given its past values using a meta-network $f_{k}^{\phi}$, parameterized by $\phi$: $\hat{\bm{\uptheta}}_{t+k}=\bm{\uptheta}_{t}+f_{k}^{\phi}(\bm{\uptheta}_{t},\bm
{\uptheta}_{t-1},\bm{\uptheta}_{t-2},...)$.

### 3.1 Weight Nowcaster Networks (WNNs)

WNNs (Jang et al., ) model $f_{k}^{\phi}$ as an MLP predicting the delta (update difference) of the $i$-th parameter: $\Delta\hat{\bm{\uptheta}}_{\tau+k}^{i}=f_{k}^{\phi}(\tilde{\bm{\uptheta}}_{
\tau}^{i},\tilde{\bm{\uptheta}}_{\tau-1}^{i},...,\tilde{\bm{\uptheta}}_{\tau_{
c}}^{i})$, where $\tau_{c}=\tau-c+1$ and $c$ is the context length. Here, $\tau,\tau-1,...$ are epoch indices and $\tilde{\bm{\uptheta}}^{i}$ are the parameters scaled based on the range of values in $\bm{\uptheta}^{i}_{\tau},...,\bm{\uptheta}^{i}_{\tau_{c}}$ (see details in Section A.1).
The predicted parameter value is obtained by unscaling the predicted delta:
$\hat{\bm{\uptheta}}_{\tau+k}^{i}=\bm{\uptheta}_{\tau}^{i}+\mathrm{unscale}(
\Delta\hat{\bm{\uptheta}}_{\tau+k}^{i})$.
WNNs are trained in a supervised way by collecting a training dataset of parameter trajectories $\{[\bm{\uptheta}_{1},\bm{\uptheta}_{2},...]\}_{1}^{C}$
obtained with Adam and applying the $l_{1}$-loss:

$$ ${\text{argmin}}_{\phi}||\Delta\tilde{\bm{\uptheta}}_{\tau+k}-\Delta\hat{\bm{ \uptheta}}_{\tau+k}||_{1},$ (1) $$

where
$\Delta\tilde{\bm{\uptheta}}_{\tau+k}$ are the scaled target parameter deltas.
Sinha et al. () proposed the approach of future parameter prediction originally, but without scaling the parameters and without a simple way to choose at which optimization steps to make the prediction. In contrast, WNNs scale the parameters and are always applied periodically, once per every $c$ epochs of a base optimizer (Adam) with $k=c$, which better accelerates optimization. For example, to use a trained WNN $f_{k}^{\phi}$ on a new task, first Adam is run for $c$ epochs ($c=5$ by default) after which $f_{k}^{\phi}$ is applied to predict future (10-th epoch) 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 MLP for parameter $\bm{\uptheta}^{i}$ 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, e.g. flattening of parameters into a vector (Schürholt et al., ), are not general and do not model neuron permutation symmetry (the neurons can be permuted without changing the network output).
Recently, Kofinas et al. (); Lim et al. () proposed
to represent $\bm{\uptheta}$ using a neural graph $\mathcal{G}(\bm{\uptheta})=(\mathbf{V},\mathbf{E})$ with node features $\mathbf{V}\in\mathbb{R}^{|\mathbf{V}|\times d_{\mathbf{V}}}$ and edge features $\mathbf{E}\in\mathbb{R}^{|\mathbf{V}|\times|\mathbf{V}|\times d_{\mathbf{E}}}$, where $d_{\mathbf{V}}$, $d_{\mathbf{E}}$ are the node and edge feature dimensionalities, respectively.
For example, equation LABEL:[Equation]eq:neural_graph_construction shows $\mathbf{E}$ for the $d_{\mathbf{E}}=1$ dimensional edge features of an $L$-layer MLP (we ignore biases for simplicity): $\bm{\uptheta}=\left\{\mathbf{W}^{(1)},\dots,\mathbf{W}^{(L)})\right\}$, where $\mathbf{W}^{(l)}\in\mathbb{R}^{d_{l-1}\times d_{l}}$ are the weights for $l\in[1,L]$.

Neural graphs can theoretically represent the parameters of any neural network taking into account neuron permutation symmetry as the rows and columns of weight matrices correspond to the nodes in the neural graph.
This way, neural graphs impose a strong inductive bias similarly to using convolution for images.
So a model operating on neural graphs, such as a 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 better with fewer samples than, for example, MLPs. However, to fully leverage the neural graphs in practice, we need to accurately construct them, which is not trivial for such networks as Transformers, motivating our approach in this work.

Transformers.
For input $\mathbf{x}\in\mathbb{R}^{N\times d}$, where $N$ is a sequence length, a multi-head self-attention (MSA) layer (Vaswani et al., ) for each head $h\in[1,H]$ projects $\mathbf{x}$ using $\mathbf{W}_{h}^{q},\mathbf{W}_{h}^{k},\mathbf{W}_{h}^{v}$ followed by self-attention and projection using another weight matrix $\mathbf{W}_{h}^{o}$ (see Section A.4 for details):

$$ $\mathbf{A}_{h}=\frac{(\mathbf{x}\mathbf{W}_{h}^{q})(\mathbf{x}\mathbf{W}_{h}^{ k})^{T}}{\sqrt{d}}\quad\mathrm{\text{and}}\quad\mathbf{y}_{h}=\mathrm{softmax} (\mathbf{A}_{h})(\mathbf{x}\mathbf{W}_{h}^{v})\quad\mathrm{\text{and}}\quad \textrm{MSA}(\mathbf{x})=\sum_{h=1}^{H}\mathbf{y}_{h}\mathbf{W}^{o}_{h}.$ (3) $$

To construct neural graphs of MSA, Kofinas et al. (); Lim et al. () use $d_{\mathbf{E}}=3$ dimensional edge features ${\mathbf{e}}_{ij}=\left(\left(\mathbf{W}^{q}\right)_{ij},\left(\mathbf{W}^{k}
\right)_{ij},\left(\mathbf{W}^{v})\right)_{ij}\right)$, assuming that the multi-head case is automatically handled by the neural graph.
However, as shown in Section 4.1, this neural graph does not model the neuron permutation symmetry correctly, therefore we refer to it as a naive neural graph (Fig. 2a,b).

Figure: Figure 2: (a) MSA weights with $d=6$ and $H=2$ heads; (b) its naive (Kofinas et al., ) and (c) our neural graph with color-coded edge/parameter types. (d-e) Neural graph embeddings obtained using a GNN and projected to 2d using tSNE for 1k good and bad permutations of the MSA weights: (d) naive graph, (e) our graph. Accuracy is computed by fitting logistic regression for neural graph embeddings with labels being if the MSA output changes or not as described in Section A.6.
Refer to caption: x1.png

## 4 Methods

In this section, we describe the details of our neural graphs for an 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., ), Llama (Dubey et al., ) (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) restrict neuron permutations across heads and (2) relax permutations of weights $\mathbf{W}^{q},\mathbf{W}^{k}$ (Fig. 2).

(1) Restricting permutations across heads.
We first observe that splitting computations into $H$ parallel heads (equation LABEL:[Equation]eq:msa)
breaks neuron permutation symmetry across the heads, so shuffling neurons across the heads may change the overall output of $\textrm{MSA}(\mathbf{x})$. Consider, for example, an MSA layer with $d=6$ and $H=2$ (Fig. 2a) and the dot product between outputs $\mathbf{x}_{1}^{q}$ and $(\mathbf{x}_{1}^{k})^{T}$ obtained after projecting $\mathbf{x}$ using the weights of the first head $\mathbf{W}_{1}^{q}$ and $\mathbf{W}_{1}^{k}$: $\mathbf{x}_{1}^{q}(\mathbf{x}_{1}^{k})^{T}=[\mathbf{x}^{q}_{1},\mathbf{x}^{q}_
{2},{\color[rgb]{0.78515625,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{
0.78515625,0,0}\mathbf{x}^{q}_{3}}][\mathbf{x}^{k}_{1},\mathbf{x}^{k}_{2},{
\color[rgb]{1,0.359375,0.72265625}\definecolor[named]{pgfstrokecolor}{rgb}{
1,0.359375,0.72265625}\mathbf{x}^{k}_{3}}]^{T}$. The result of this dot product does not change when permuting neurons in both $\mathbf{W}_{1}^{q}$ and $\mathbf{W}_{1}^{k}$ using the same $\pi$, since it results in permuting the outputs: $\pi(\mathbf{x}_{1}^{q})\pi(\mathbf{x}_{1}^{k})^{T}=\mathbf{x}_{1}^{q}(\mathbf{
x}_{1}^{k})^{T}$. However, consider shuffling neurons across head 1 and 2 in $\mathbf{W}^{q}$ and $\mathbf{W}^{k}$, e.g. neurons 3 and 4 (highlighted in different colors in Fig. 2a) resulting in: $[\mathbf{x}^{q}_{1},\mathbf{x}^{q}_{2},{\color[rgb]{0.75,0.56640625,0}
\definecolor[named]{pgfstrokecolor}{rgb}{0.75,0.56640625,0}\mathbf{x}^{q}_{4}}
][\mathbf{x}^{k}_{1},\mathbf{x}^{k}_{2},{\color[rgb]{
0.62890625,0.40234375,0.46875}\definecolor[named]{pgfstrokecolor}{rgb}{
0.62890625,0.40234375,0.46875}\mathbf{x}^{k}_{4}}]^{T}$. Now the dot product no longer equals $\mathbf{x}_{1}^{q}(\mathbf{x}_{1}^{k})^{T}$ unless ${\color[rgb]{0.75,0.56640625,0}\definecolor[named]{pgfstrokecolor}{rgb}{
0.75,0.56640625,0}\mathbf{x}^{q}_{4}}{\color[rgb]{
0.62890625,0.40234375,0.46875}\definecolor[named]{pgfstrokecolor}{rgb}{
0.62890625,0.40234375,0.46875}\mathbf{x}^{k}_{4}}={\color[rgb]{0.78515625,0,0}
\definecolor[named]{pgfstrokecolor}{rgb}{0.78515625,0,0}\mathbf{x}^{q}_{3}}{
\color[rgb]{1,0.359375,0.72265625}\definecolor[named]{pgfstrokecolor}{rgb}{
1,0.359375,0.72265625}\mathbf{x}^{k}_{3}}$. So self-attention matrix $\mathbf{A}_{1}$ for head 1 and $\textrm{MSA}(\mathbf{x})$ can change.

To take into account this restriction, we allow for neuron permutations only within a 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).

##### (2) Relaxing permutations for 𝐖 q , 𝐖 k superscript 𝐖 𝑞 superscript 𝐖 𝑘 \mathbf{W}^{q},\mathbf{W}^{k} bold_W start_POSTSUPERSCRIPT italic_q end_POSTSUPERSCRIPT , bold_W start_POSTSUPERSCRIPT italic_k end_POSTSUPERSCRIPT .

We also observe that neurons in
$\mathbf{W}_{h}^{q}$ and $\mathbf{W}_{h}^{k}$ can be permuted in a different way ($\pi^{\prime}$) from $\mathbf{W}_{h}^{v}$ and $\mathbf{W}_{h}^{o}$ without changing the MSA output:

$$ $(\mathbf{x}\mathbf{W}_{h}^{q})(\mathbf{x}\mathbf{W}_{h}^{k})^{T}=(\mathbf{x} \pi^{\prime}_{c}(\mathbf{W}_{h}^{q}))(\mathbf{x}\pi_{c}^{\prime}(\mathbf{W}_{h }^{k}))^{T}\quad\mathrm{\text{and}}\quad(\mathbf{x}\mathbf{W}_{h}^{v})\mathbf{ W}_{h}^{o}=(\mathbf{x}\pi_{c}(\mathbf{W}_{h}^{v}))\pi_{r}(\mathbf{W}_{h}^{o}),$ (4) $$

where $\pi_{c}$ and $\pi_{r}$ denote permutation $\pi$ of the columns and rows in the matrix respectively.
In the naive graph the neurons in $\mathbf{W}_{h}^{q},\mathbf{W}_{h}^{k},\mathbf{W}_{h}^{v},\mathbf{W}_{h}^{o}$ are always permuted in the same way ($\pi_{c}=\pi^{\prime}_{c}$) making the neural graph unnecessary restrictive (Fig. 2b).
To address this issue, we keep $\mathbf{W}_{h}^{q},\mathbf{W}_{h}^{k},\mathbf{W}_{h}^{v}$ as separate 1-dimensional edge features in a neural graph instead of stacking them as 3-dimensional edge features (Fig. 2c). Since $\mathbf{W}_{h}^{q}$ and $\mathbf{W}_{h}^{k}$ share the input and output neurons while $\mathbf{W}_{h}^{k}$ is transposed in equation LABEL:[Equation]eq:msa, it is important to preserve the edge direction, so that for example ${\mathbf{e}}_{3,4}$ corresponds to the weights of $\mathbf{W}_{h}^{q}$ while ${\mathbf{e}}_{4,3}$ corresponds to the weights of $\mathbf{W}_{h}^{k}$.

Figure: Figure 3: Laplacian Positional Encoding (LPE) of our neural graph’s nodes color-coded by the layer type (left), layer index (middle) and neuron index (right) for a Transformer with 6 layers and 384 hidden units. We show tSNE projections from the 8-dimensional LPE to 2d.
Refer to caption: x4.png

Edge features.
In our neural graph, all the model parameters $\bm{\uptheta}$ and auxiliary connections $\bm{\uptheta}^{\prime}$, such as residual and head connections introduced in Section 4.1, are represented using the edge features $\mathbf{E}\in\mathbb{Z}^{|\mathbf{V}|\times|\mathbf{V}|\times d_{\mathbf{E}}}$ so that $||\mathbf{E}||_{0}=|\bm{\uptheta}|+|\bm{\uptheta}^{\prime}|$ (Fig. 2c).
To differentiate between $\bm{\uptheta}$ and $\bm{\uptheta}^{\prime}$, we associate an integer edge type with each edge: $\mathbf{E}^{\text{type}}\in\mathbb{Z}^{|\mathbf{V}|\times|\mathbf{V}|\times 1}$, so our neural graph $\mathcal{G}(\bm{\uptheta})=(\mathbf{V},\mathbf{E}^{\text{type}},\mathbf{E})$.

Node features.
For the node features we use the Laplacian Positional Encoding (LPE) (Belkin & Niyogi, ; Dwivedi et al., ) 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 use unweighted edges, transform the graph to an undirected one and extract 8 smallest non-trivial
eigenvectors, so our node features are $\mathbf{V}^{\text{lpe}}\in\mathbb{R}^{|\mathbf{V}|\times 8}$.
In addition, for Transformer’s word embedding layers, we found it beneficial to leverage a positional feature $\mathbf{V}^{\text{w}}\in\mathbb{Z}^{|\mathbf{V}|\times 1}$, since these layers often have the same size with ordered neurons. For other layers, this feature is set to zero.

Sequence of parameters as a neural graph.
We consider a history of $c$ past parameter vectors following WNNs (Section 3.1): $\bm{\Theta}_{\tau:\tau_{c}}=[\bm{\uptheta}_{\tau},\bm{\uptheta}_{\tau-1},...,
\bm{\uptheta}_{\tau_{c}}]\in\mathbb{R}^{n\times c}$, where $\tau_{c}=\tau-c+1$ as in Section 3.1.
Since each parameter vector in $\bm{\Theta}_{\tau:\tau_{c}}$ represents the same neural network structure, we transform the entire $\bm{\Theta}_{\tau:\tau_{c}}$ into a single neural graph $\bm{\mathcal{G}}(\bm{\Theta}_{\tau:\tau_{c}})=(\mathbf{V}^{\text{lpe}},\mathbf
{V}^{\text{w}},\mathbf{E}^{\text{type}},\bm{\mathcal{E}}_{\tau:\tau_{c}})$, where
$\bm{\mathcal{E}}_{\tau:\tau_{c}}$ is obtained by stacking edge features: $\bm{\mathcal{E}}_{\tau:\tau_{c}}=[\mathbf{E}_{\tau},\mathbf{E}_{\tau-1},...,
\mathbf{E}_{\tau_{c}}]\in\mathbb{R}^{|\mathbf{V}|\times|\mathbf{V}|\times d_{
\mathbf{E}}\times c}$, whereas node features $\mathbf{V}^{\text{lpe}},\mathbf{V}^{\text{w}}$ and edge types $\mathbf{E}^{\text{type}}$ do not change over time.

### 4.2 Neuron Interaction and Nowcasting Networks (NiNo)

Given an input neural graph $\bm{\mathcal{G}}(\bm{\Theta}_{\tau:\tau_{c}})$,
our NiNo model processes it 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 $\hat{\bm{\mathcal{E}}}_{1:K}$ that are mapped back to the parameter space $[\Delta\hat{\bm{\uptheta}}_{\tau+1},...,\Delta\hat{\bm{\uptheta}}_{\tau+K}]$ using the inverse neural graph construction $\Delta\hat{\bm{\uptheta}}_{\tau+k}=\mathcal{G}^{-1}(\hat{\bm{\mathcal{E}}}_{k})$ (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., ) scaling each parameter independently using min-max, we extend a layerwise scaling (Schürholt et al., ) to a sequence of parameters.
Specifically, given weights $\mathbf{W}^{(l)}_{\tau:\tau_{c}}\in\mathbb{R}^{d_{l-1}\times d_{l}\times c}$ of the $l$-th layer of the input $\bm{\mathcal{E}}_{\tau:\tau_{c}}$, we obtain scaled weights as $\tilde{\mathbf{W}}^{(l)}_{\tau:\tau_{c}}=(\mathbf{W}^{(l)}_{\tau:\tau_{c}}-\mu
^{(l)}_{\tau})/\sigma^{(l)}_{\tau}$, where $\mu^{(l)}_{\tau},\sigma^{(l)}_{\tau}$ are the scalar mean and standard deviation of $\mathbf{W}^{(l)}_{\tau:\tau_{c}}$. After repeating this step $\forall l\in[1,L]$, we obtain scaled $\tilde{\bm{\mathcal{E}}}_{\tau:\tau_{c}}$.

Embedding layers and GNN.
We use linear $\phi_{\text{lpe}},\phi_{e}$ and embedding $\phi_{w},\phi_{\text{type}}$ layers to project node and edge features to the $D$-dimensional space followed by $M$ GNN layers and the output layer $\phi_{\text{DMS}}$:

$$ $\displaystyle\mathbf{V}^{0}=\phi_{\text{lpe}}(\mathbf{V}^{\text{lpe}})+\phi_{ \text{w}}(\mathbf{V}^{\text{w}}),\quad$ $\displaystyle\mathbf{V}^{0}\in\mathbb{R}^{|\mathbf{V}|\times D},$ (5) $\displaystyle\bm{\mathcal{E}}^{0}=\phi_{e}(\tilde{\bm{\mathcal{E}}})+\phi_{ \text{type}}(\mathbf{E}^{\text{type}}),\quad$ $\displaystyle\bm{\mathcal{E}}^{0}\in\mathbb{R}^{|\mathbf{V}|\times|\mathbf{V}| \times D},$ (6) $\displaystyle\mathbf{V}^{m},\bm{\mathcal{E}}^{m}=\text{GNN}_{\phi_{m}}(\mathbf {V}^{m-1},\bm{\mathcal{E}}^{m-1}),\quad$ $\displaystyle m=1...M,$ (7) $\displaystyle\hat{\bm{\mathcal{E}}}=\phi_{\text{DMS}}({\bm{\mathcal{E}}^{M}}),\quad$ $\displaystyle\hat{\bm{\mathcal{E}}}\in\mathbb{R}^{|\mathbf{V}|\times|\mathbf{V }|\times K}.$ (8) $$

Our GNN layers are based on Kofinas et al. (), but to enable better efficiency and generalization when working with large models, we use a simple mean aggregation (see Section A.5).

Figure: Figure 4: Our neuron interaction and nowcasting (NiNo) model. We feed $c$ past parameter states as input and predict future $K$ states leveraging our improved neural graph structure (Section 4.1) using a graph neural network. For a new optimization task, NiNo is applied rarely over time: only once per 1k steps of Adam (Section 5).
Refer to caption: x7.png

Direct multi-step forecasting (DMS).
The final layer $\phi_{\text{DMS}}$ outputs $K$ values for each edge corresponding to future steps from $\tau+1$ to $\tau+K$.
This is motivated by direct multi-step forecasting performing well in time series forecasting by avoiding error accumulation effects of autoregressive forecasting (Chevillon, ; Zeng et al., ).
To train the model given a training dataset of parameters trained with Adam, we use equation LABEL:[Equation]eq:loss-wnn, but applied for $k=[1,...,K]$ instead of fixing $k$:

$$ $\displaystyle{\text{argmin}}_{\phi}\frac{1}{K}\sum\nolimits_{k=1}^{K}(||\Delta \tilde{\bm{\uptheta}}_{\tau+k}-\Delta\hat{\bm{\uptheta}}_{\tau+k}||_{1}),$ (9) $$

where $\Delta\hat{\bm{\uptheta}}_{\tau+k}=\mathcal{G}^{-1}(\hat{\bm{\mathcal{E}}}_{k})$ and $\tilde{\bm{\uptheta}}_{\tau+k}$ are the target parameters scaled using $\mu_{\tau},\sigma_{\tau}$. The predicted parameters are obtained by unscaling the predicted delta using $\mu_{\tau},\sigma_{\tau}$: $\hat{\bm{\uptheta}}_{\tau+k}=\bm{\uptheta}_{\tau}+\mathrm{unscale}(\Delta\hat{
\bm{\uptheta}}_{\tau+k})$.

Inference with $k$-decay.
Once our model is trained, it can predict future parameters for $k\in[1,K]$. While $k$ can be treated as a hyperparameter and kept fixed during optimization of the target task, we found that in the initial optimization stage using very large $k$ is beneficial, whereas in the later optimization stages $k$ should decrease fast, since the parameter values do not change significantly.
Therefore, we propose to decay $k$ during optimization as $k\approx K((T-t)/T)^{p}$, where $T$ is the maximum number of optimization steps and $p$ controls the decay speed (Fig. 12). Although $p$ can be tuned, we found that $p=2$ works well in our experiments.

## 5 Experiments

We experiment with nine tasks, each defined by a dataset and a 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 our meta-models (NiNo, WNN and their variants). The other five tasks, the out-of-distribution tasks, differ from the in-distribution tasks in the architecture and/or dataset. We also experiment with larger GPT2-based and Llama-based architectures at the end of Section 5.2.

Indexing parameters.
Since the notion of epoch used for $\tau$ in WNNs (Section 3.1) can be ill-defined in some tasks (e.g. in language tasks the models are often trained only for 1 epoch), in our experiments we found $\tau$ and $\tau+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 $c=5$ we apply the meta-models every $5\times 200=1,000$ steps.

Training and evaluation pipeline:

- 1.
Meta-training: Train the meta-models with the supervised parameter prediction loss (equation LABEL:[Equation]eq:loss-wnn for WNN and equation LABEL:[Equation]eq:loss 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 AdamW (language) optimizers with the meta-model applied every 1k steps given $c=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.

**Table 1: In-distribution and out-of-distribution tasks.**
|  | In-distribution tasks | Out-of-distribution tasks |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | FM/16 | C10/16 | LM1B/3-24 | LM1B/2-32 | FM/32 | C10/32 | C100/32 | LM1B/3-64 | Wiki/3-64 |
| Final training loss | 0.25$\pm$0.06 | 0.91$\pm$0.1 | 5.94$\pm$0.03 | 5.85$\pm$0.04 | $-$ | $-$ | $-$ | $-$ | $-$ |
| #models | 300 | 300 | 200 | 200 | $-$ | $-$ | $-$ | $-$ | $-$ |
| #params | 14K | 15K | 1.2M | 1.6M | 56K | 57K | 63K | 3.4M | 3.4M |
| Target (validation) metric | Acc | Acc | Perplexity | Perplexity | Acc | Acc | 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 | 4646 | 12000 | 7000 |

Figure: Figure 5: Scaling (a,b) and generalization (c) trends for NiNo vs WNN+.
Refer to caption: x8.png

### 5.1 Setup

Vision tasks. We use the FashionMNIST (fm), CIFAR-10 (c10) and CIFAR-100 (c100) datasets and two convolutional architectures with three layers: with 16, 32 and 32 channels per layer (e.g. task fm/16) or 32, 64 and 64 channels per layer (e.g. task fm/32).
The convolutional architectures are the same as in the L2O experiments of Kofinas et al. () to enable fair comparison. In all cases, these tasks are optimized using Adam (Kingma & Ba, ) without weight decay, with a constant learning rate of 6e-3 (for FashionMNIST) or 3e-3 (for CIFAR) with a batch size of 128 for $T$=10k steps.

Language tasks. We use the lm1b (Chelba et al., ) and WikiText103 (wiki) (Merity et al., ) datasets and train GPT2 style Transformers (Radford et al., ) 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 AdamW (Loshchilov & Hutter, ), weight decay 1e-2, learning rate 2e-4, batch size of 32, sequence length of 1024 for either 1 epoch (for lm1b) or 4 epochs (for wiki) corresponding to around 24k or 14k steps respectively. We use a predefined GPT2 tokenizer in all the cases with a fixed vocabulary.

Meta-training dataset. We use fm/16, c10/16, lm1b/3-24 and lm1b/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\times 10^{6}$ checkpoints. Our dataset is large and diverse (Fig. 10),
yet it is relatively cheap to be collected as the tasks are small.

Baselines.
As a 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., ) 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 (L2O) model from Kofinas et al. (), which showed strong results on similar vision tasks. We use their L2O (NG-GNN) pretrained on fm/16 (denoted as L2O/fm16) as retraining their model on all our in-distribution tasks is expensive. For a fair comparison with L2O/fm16, we also trained NiNo on fm/16 only (denoted as NiNo/fm16).
We train WNN+ and NiNo with $M\in[1,2,3]$ layers and $D\in[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. () when choosing the approach to compare training methods and set our 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 to more expensive evaluation we only evaluate every 500 steps, whereas for vision we evaluate every step.

### 5.2 Results

**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-distribution tasks | Out-of-distribution tasks | Avg speed up |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | FM/16 | C10/16 | LM1B/3-24 | LM1B/2-32 | FM/32 | C10/32 | C100/32 | LM1B/3-64 | Wiki/3-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 |
| Linefit+ | 27.3 | 28.3 | 32.6 | 29.8 | 23.3 | 27.0 | 25.0 | 25.5 | 33.3 | 28.0 |
| 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 speed up (%) is reported.**
| layerwise ($\mu,\sigma$ per layer) | 48.9 |
| --- | --- |
| $\mu,\sigma$ per param | 41.5 |
| min-max per param | 38.8 |
| no scaling | 32.1 |

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 we 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 a naive graph (Section 4.1). While the latter achieves slightly better performance on the fm/16 in-distribution task, in the out-of-distribution tasks this model significantly underperforms to 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 baselines compare favorably to WNN indicating that the latter could have learned a very simple rule despite training on a lot of data.

Comparison with L2O. L2O/fm16 performs the best in-distribution on the same task (Fig. 6a), however when applied to an unseen task it overfits to the training set (Fig. 6b). In contrast, our NiNo/fm16 trained on the same task as L2O/fm16 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, L2O/fm16 and L2O in general have an overhead at every optimization step making the total computation cost noticeable (Table 6 in Section A.7). It is also more computationally challenging to meta-train it on multiple tasks, therefore we only use a pretrained L2O/fm16 from Kofinas et al. (). Since L2O/fm16 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, we found that context length $c=5$ is a sweet spot (Table 3d), because small $c$ may not have enough information to capture the parameter trend, while larger $c$ leads to the NiNo applied too rarely (for $c=10$ it is applied every $10\times 200=2,000$ steps). Using deeper (Table 3e) and wider (Fig. 5a) GNNs is also important. However we found that depth $M>3$ or width $D>128$ make it more challenging to train a performant model. Increasing the number of models ($C$) helps NiNo achieve better speed ups (Table 3f, Fig. 5b). In contrast, WNN+ saturates fast when $D$ or $C$ is increased.

**Table 4: Wiki validation perplexity ($\downarrow$) on a Llama3-style architecture (reported as the mean and standard deviation for 3 runs).**
| Method | 10000 steps | 14000 steps |
| --- | --- | --- |
| AdamW | 26.29$\pm$0.05 | 24.44$\pm$0.01 |
| AdamW+NiNo | 24.36$\pm$0.11 | 22.37$\pm$0.13 |

Evaluating on larger models and Llama3-style architectures. We evaluated the ability of NiNo to speed up training of much larger models than it was trained on. Specifically, we trained 4 and 6 layer GPT2-style Transformers on 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 lm1b used in meta-training. Despite these challenges, NiNo was able to speed up training by around 40% and 15% for the 7M and 29M models respectively, outperforming WNN+.
We also trained a Llama3-style Transformer with 6 layers and 384 hidden units (111M parameters) with AdamW and AdamW+NiNo, which is an even more challenging setup than in the previous experiments because NiNo did not observe Llama models during its meta-training (Table 1). In this experiment, we use the Llama 3 tokenizer (Dubey et al., ) having a larger vocabulary than in GPT2, so we use an ablated NiNo without $\mathbf{V}^{\text{w}}$.
Besides the vocabulary, Llama models may pose additional challenges for NiNo, e.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 AdamW baseline (24.4) in around 30% less steps (Table 4).

Figure: Figure 6: Training loss and validation performance on FashionMNIST, CIFAR-10 and wiki.
Refer to caption: x11.png

### 5.3 Analysis

Figure: Figure 7: Comparison of optimization methods on the c10/32 task by projecting parameters to 2d using PCA at every training step.

Parameter space ($\bm{\uptheta}\rightarrow\mathbb{R}^{2}$).
We show how the parameters evolve during 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/AdamW trajectory when predicting the parameters, while L2O diverges early. At later optimization stages, NiNo makes better predictions than WNN as the latter leads to significant loss spikes. Interestingly, NiNo’s trajectory closely corresponds to a $2\times$ faster AdamW, even though NiNo’s base optimizer is Adam in this task. This indicates that NiNo may implicitly regularize optimization.

Graph embedding space (${\bm{\mathcal{E}}^{M}}\rightarrow\mathbb{R}^{2}$).
The parameter space analysis used above can be effective on a single task. However, for several tasks with different model sizes ($|\bm{\uptheta}|$), finding a common PCA projection matrix ($\bm{\uptheta}\rightarrow\mathbb{R}^{2}$), 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 2d.
We use this approach 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/|\bm{\mathcal{E}}^{M}|\sum_{ij}{\bm{\mathcal{E}}_{ij}^{M}}$) after the last NiNo (or WNN+) layer by feeding the 5 past parameter states every 1,000 steps as in 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 are visually distinct (Fig. 8a). However, WNN+’s embeddings virtually collapse at later training stages for lower loss values (e.g. for 6-384), while NiNo’s embeddings are visually more distinct for lower loss values (Fig. 8b). 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 in different behaviors on downstream tasks (Liu et al., ; Dahl et al., ).

Figure: 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.
Refer to caption: x17.png

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 L2O (Table 6). As for the memory required for a prediction with NiNo, in our implementation NiNo’s peak memory usage is similar to the peak memory usage of a training step of the Transformer 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 inference on large neural graphs using approaches similar to large scale model inference in image generation (Li et al., ).

## 6 Conclusion and Future Work

We proposed NiNo, a novel approach to accelerate training of neural networks, including Transformers. Our work makes a step in the relatively underexplored area of *periodic parameter prediction*, where our experiments show that there is 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 speed ups with more data and larger models.

#### 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 AI Chair. Guillaume Lajoie is a Canada CIFAR AI chair as well a Canada Research Chair in neural computations and interfacing.

## Appendix A Appendix

### A.1 WNN Parameter Scaling Details

In WNNs (Jang et al., ), parameter scaling is performed by first computing the minimum and maximum values in the input parameter values: $s^{i}_{\tau}=\mathrm{max}(\bm{\uptheta}^{i}_{\tau},...,\bm{\uptheta}^{i}_{\tau
_{c}})-\mathrm{min}(\bm{\uptheta}^{i}_{\tau},...,\bm{\uptheta}^{i}_{\tau_{c}})
,\forall i\in[1,n]$.
Each $i$-th parameters scaling factor $s^{i}_{\tau}$ is then used to scale both the inputs: $\tilde{\bm{\uptheta}}^{i}_{\tau}={\bm{\uptheta}^{i}_{\tau}}/({s^{i}_{\tau}}+{
\epsilon})$, and, only for training WNNs, the targets: $\tilde{\bm{\uptheta}}^{i}_{\tau+k}={\bm{\uptheta}^{i}_{\tau+k}}/({s^{i}_{\tau}
}+{\epsilon})$. Jang et al. () also subtract the last element from the sequence, however we found this to be redundant.
Unscaling is performed by multiplying the prediction by $s^{i}_{\tau}$.
Scaling and unscaling are performed per parameter (for each parameter independently).

### A.2 Linefit

Sinha et al. (); Jang et al. () introduced a “linefit” baseline to predict future parameter values by extrapolating the line to a future point:

$$ $\displaystyle\forall i\in[1,n]:$ $\displaystyle\hat{\bm{\uptheta}}^{i}_{\tau+k}=2a^{i}c+b^{i},$ (10) $$

where $a^{i},b^{i}$ are obtained by optimizing the following objective:

$$ $\displaystyle\text{argmin}_{a^{i},b^{i}}||a^{i}\bm{x}+b^{i}-\bm{\Theta}^{i}_{ \tau:\tau_{c}}||_{2},$ (11) $$

where $\bm{x}=[1,2,...,c]$ and $\bm{\Theta}^{i}_{\tau:\tau_{c}}=[\bm{\uptheta}^{i}_{\tau},\bm{\uptheta}^{i}_{
\tau-1},...,\bm{\uptheta}^{i}_{\tau_{c}}]\in\mathbb{R}^{c}$.
This baseline fits a line ($a^{i},b^{i}$) for each parameter $i$ given only the past values of the same parameter without collecting a training dataset. Despite its simplicity, this baseline was shown to improve on using Adam/SGD only.

### A.3 Linefit+

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 et al., ). Motivated by this observation, we introduce the Linefit+ baseline that penalizes more the errors for later parameter values in the trajectory:

$$ $\displaystyle\text{argmin}_{a^{i},b^{i}}||\bm{\mu}(a^{i}\bm{x}+b^{i}-\bm{ \Theta}^{i}_{\tau:\tau_{c}})||_{2},$ (12) $\displaystyle\text{where}\ \bm{\mu}=[1/c,2/c,...,1].$ $$

### A.4 Naive Neural Graph of Transformers

A multi-head self-attention (MSA) layer of Transformers (Vaswani et al., ) consists of three weight matrices $\mathbf{W}^{q},\mathbf{W}^{k},\mathbf{W}^{v}$ applied to the input $\mathbf{x}\in\mathbb{R}^{N\times d}$ and another weight matrix $\mathbf{W}^{o}$ returning the output, where $N$ is a sequence length and all the four weight matrices are ${d\times d}$. To compute MSA with $H$ heads, the weight matrices are typically split into $H$ groups across columns for $\mathbf{W}^{q},\mathbf{W}^{k},\mathbf{W}^{v}$ and rows for $\mathbf{W}^{o}$ (Fig. 2a).
The heads are processing the input in parallel, $\forall h\in[1,H]$ using equation LABEL:[Equation]eq:msa.

Neural graphs defined by Kofinas et al. () and Lim et al. () are conceptually similar, however Lim et al. () model biases ${\mathbf{b}}^{(l)}$ and normalization layers as extra nodes connected to the neurons in the $l$-th layer (Fig. 2c).
Normalization layers are similarly added as separate nodes, hence the node features of Lim et al. () do not include any parameters.
We build on this variant, since it is more easily to implement it for the layers with complicated neuron permutation symmetry, such as multi-head self-attention.
Compared to Kofinas et al. (), Lim et al. () also provide a slightly different description for the MSA neural graph , however the $\mathbf{W}^{q},\mathbf{W}^{k},\mathbf{W}^{v}$ edges are also stacked. Given the lack of detailed description for the multi-head case and no implementation currently available publicly, it makes it challenging to use their neural graphs directly. Therefore, we use (Kofinas et al., ) as the base (naive) neural graph when constructing MSA layers.

### A.5 Graph Neural Network

A Graph Neural Network (GNN) layer with edge features can be formulated using the mean aggregation as follows (Corso et al., ):
${\mathbf{v}}_{i}=\phi_{a}\left(\frac{1}{|\mathcal{N}(i)|}\sum_{j\in\mathcal{N}
(i)}{\mathbf{m}}_{ij}\right)$,
where ${\mathbf{m}}_{ij}$ is computed using message passing $\phi_{m}$ given node and edge features ${\mathbf{v}}_{i},{\mathbf{v}}_{j},{\mathbf{e}}_{ij}$; $\phi_{a}$ is an MLP; $\mathcal{N}(i)$ are the neighbors of node $i$. Kofinas et al. () modify how ${\mathbf{m}}_{ij}$ 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) is permutation equivariant. (1) means that for the neural graphs there are $|\bm{\uptheta}|^{\prime}\times d_{\mathbf{E}}$ input edge features and $|\bm{\uptheta}|^{\prime}\times 1$ predictions, where $|\bm{\uptheta}|^{\prime}$ is the number of trainable parameters in the input model plus auxiliary non-trainable parameters (e.g. residual and head connections).
(2) means that any permutation of input nodes results in the same permutation of the output nodes and corresponding edges. These properties make GNNs a 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.

**Table 5: GNN layer comparison.**
| Step | Typical GNN (Corso et al., ) | Neural Graph GNN (Kofinas et al., ) |
| --- | --- | --- |
| 1. Message passing | ${\mathbf{m}}_{ij}=\phi_{m}\left(\left[{\mathbf{v}}_{i},{\mathbf{v}}_{j},{<br>\mathbf{e}}_{ij}\right]\right)$ | ${\mathbf{m}}_{ij}=\phi_{\texttt{scale}}\left({\mathbf{e}}_{ij}\right)\odot\phi<br>_{m}\left(\left[{\mathbf{v}}_{i},{\mathbf{v}}_{j}\right]\right)+\phi_{\texttt{<br>shift}}\left({\mathbf{e}}_{ij}\right)$ |
| 2. Aggregation | ${\mathbf{v}}_{i}=\phi_{a}\left(1/{|\mathcal{N}(i)|}\sum_{j\in\mathcal{N}(i)}{<br>\mathbf{m}}_{ij}\right)$ |  |
| 3. Edge update | ${\mathbf{e}}_{ij}^{(k+1)}={\mathbf{e}}_{ij}^{(k)}$ | ${\mathbf{e}}_{ij}^{(k+1)}=\phi_{e}^{(k+1)}\left(\left[{\mathbf{v}}_{i}^{(k)},{<br>\mathbf{e}}_{ij}^{(k)},{\mathbf{v}}_{j}^{(k)}\right]\right)$ |

### A.6 Neuron Permutation Symmetry

Permutation symmetry in neural graphs.
Let $\pi^{\text{good}}(\bm{\uptheta})$ be a permutation of neurons in $\bm{\uptheta}$ that does not change the network function: $f(\mathbf{x},\bm{\uptheta})=f(\mathbf{x},\pi^{\text{good}}(\bm{\uptheta}))$. Likewise, let $\pi^{\text{bad}}(\bm{\uptheta})$ be a permutation that changes the function: $f(\mathbf{x},\bm{\uptheta})\neq f(\mathbf{x},\pi^{\text{bad}}(\bm{\uptheta}))$.
Denoting $\cong$ and $\ncong$ as “isomorphic” and “non-isomorphic” operators on graphs respectively, for neural graphs $\mathcal{G}$ the following equations are satisfied under some assumptions (Kofinas et al., ): $\mathcal{G}(\bm{\uptheta})\cong\mathcal{G}(\pi^{\text{good}}(\bm{\uptheta}))$ and $\mathcal{G}(\bm{\uptheta})\ncong\mathcal{G}(\pi^{\text{bad}}(\bm{\uptheta}))$. 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 a single Transformer layer with $d=12$ and $H=4$, permute rows and columns in its MSA weight matrices $\bm{\uptheta}$ using a uniformly sampled permutation $\pi$, and label each permutation as $y_{\pi}=1$ for $\pi^{\text{good}}$ or $y_{\pi}=0$ for $\pi^{\text{bad}}$ depending if the Transformer output $f(\mathbf{x},\bm{\uptheta})$ changes or not:

$$ $y_{\pi}=\begin{cases}1,&\text{if}f(\mathbf{x},\bm{\uptheta})=f(\mathbf{x},\pi( \bm{\uptheta}))\\ 0,&\text{otherwise}.\end{cases}$ (13) $$

We generate 1,000 such permutations with about 500 of good and bad permutations. Then we extract a neural graph embedding for each permutation.
For correctly constructed neural graphs the embeddings corresponding to $y_{\pi}=1$ should be close to the non-permuted $\bm{\uptheta}$.
To extract embeddings we use a 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 $y_{\pi}$. We also train a logistic regression model $\psi$ that takes a graph embedding $\mathbf{h}$ as input and predicts if it is a good or bad permutation: $\hat{y}_{\pi}=\psi(\mathbf{h})$, and we compute classification accuracy between $\hat{y}_{\pi}$ and $y_{\pi}$. We use all 1,000 samples for training and evaluation $\psi$ to estimate how easily can the graph embeddings be separated based on $y_{\pi}$ (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.

Figure: Figure 9: Graph and adjacency matrix with head restrictions and separate q, k, v edge features introduced in Section 4.1, color-coded by edge type. (top) A 2 layer GPT2-based Transformer with $d=6$ and $H=2$. (bottom) A 2 layer Llama3-based Transformer with $d=12$, intermediate hidden size equal 24, $H=6$ and number of key value heads equal 2 for GQA. The number of word embeddings is reduced for visualization purposes.
Refer to caption: x19.png

Figure: Figure 10: Histogram of final training losses in the in-distribution (meta-training) tasks. See Table 1 for more details.
Refer to caption: x23.png

Figure: Figure 11: Meta-training curves.
Refer to caption: x24.png

### A.7 NiNo Training Details

For representing convolutional networks with neural graphs, $d_{\mathbf{E}}=h\times w$ (Kofinas et al., ), where $h,w$ are convolution kernel size, so edge features representing $c$ states are $|\bm{\uptheta}|^{\prime}\times chw$, where $|\bm{\uptheta}|^{\prime}$ is the number of trainable and auxiliary non-trainable (residual, head connections) parameters in a model.
To make the feature dimensionality the same for different architectures during meta-training, we use zero-padding.
We train meta-models for 20k training iterations using AdamW, learning rate 3e-3 with cosine decay and weight decay 0.01. We sample a batch of 4 checkpoints in each training iteration and use automatic mixed precision. Training of the NiNo and WNN+ meta-models completes in under 7 and 6 hours respectively on a single NVIDIA RTX8000 with 48GB of memory.

Meta-training behavior can give important highlights about the model and data. Therefore we explore meta-training 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 is 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.

**Table 6: Computational cost estimates. Measured on NVIDIA A100-80GB after running for 1k steps with batch size $b=32$, sequence length$=1024$. Note that these numbers are rough estimates as the time/memory vary depending on the GPU/CPU load and internal logic in different PyTorch versions. Also note that time and memory can be traded off by processing only a fraction of parameters at a time. For Adam and L2O, we report per step measurements, since these methods are used at every step. For WNN+ and NiNo, we report per prediction measurements after 1k steps of Adam. “Overhead” denotes wall clock time increase vs Adam. ^*The message passing step of NiNo is computed on CPU to avoid OOM on GPU, resulting in much longer prediction time.**
| Method | hid size $D$ | 3-64 (3.4M, GPT2-based) | 4-128 (7.4M, GPT2-based) | 6-384 (111.5M, Llama3-based) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  |  | Peak Mem | Time | Overhead | Peak Mem | Time | Overhead | Peak Mem | Time | Overhead |
| Adam | $-$ | 27GB/step | 130ms/step | 0% | 28GB/step | 180ms/step | 0% | 76GB/step | 909ms/step | 0% |
| WNN+ | 32 | 1GB/predict | 172ms/predict | 0.13% | 3GB/predict | 222ms/predict | 0.12% | 42GB/predict | 2.01sec/predict | 0.22% |
| WNN+ | 128 | 4GB/predict | 245ms/predict | 0.19% | 9GB/predict | 294ms/predict | 0.16% | 71GB/predict | 3.44sec/predict | 0.38% |
| NiNo | 32 | 3GB/predict | 323ms/predict | 0.25% | 5GB/predict | 498ms/predict | 0.28% | 57GB/predict | 5.52sec/predict | 0.61% |
| NiNo | 128 | 9GB/predict | 483ms/predict | 0.37% | 16GB/predict | 864ms/predict | 0.48% | 71GB/predict^* | 510sec/predict^* | 56.11%^* |
| L2O (mlp) | 32 | 27GB/step | 134ms/step | 3.08% | 28GB/step | 190ms/step | 5.56% | 76GB/step | 917ms/step | 0.88% |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: 1. How are  the gradients are fed into the NiNo algorihtm? It is part of V^w? 
2. What are the nodes features? 
3. Was there any ablation study done for AdamW in Liama3 style architecture particularly around using different learning rates.
soundness: 2
strengths: I think the construction of the graph whose isomorphic permutations of nodes preserves the model functionality,  for transformer layer is neat and interesting.
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, 
should 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..
weaknesses: I have two major concerns about the papers:
1. 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
a 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).

2) 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. 

3) Some discussion about optimality of their graph construction would be nice, but it is a minor comment.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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?
2. 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).
soundness: 3
strengths: 1. The paper is sufficiently well written and is fairly accessible.
2. 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.
3. 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.
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.
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.
2. 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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
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?
2. 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.
3. Could the authors provide details about the computational costs of the meta-training process?
soundness: 3
strengths: 1. The paper is well-motivated. It is an improved version of WNN by integrating GNN.
2. The experiments cover different tasks including language modeling and image classification tasks.
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.
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. 
2. 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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
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.
soundness: 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.
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.
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.
Computational 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.

Meta-review: This paper proposes a method to accelerate the training of neural networks called NiNo, which uses neuron connectivity to better predict the near-future parameters, via a graph meta-network, during the optimisation process. Overall, the reviewers found the construction of the graph network to be interesting, and that the conducted experiments appear to be promising. The main weakness is the scalability of the approach, and how effective the approach would be beneficial to large-scale models, as the validation focuses on relatively simple datasets. The AC also agrees with this point.

Decision: Accept (Poster)