Source-linked AI summary
Accelerating Training with Neuron Interaction and Nowcasting Networks
Boris Knyazev, Abhinav Moudgil, Guillaume Lajoie, Eugene Belilovsky, Simon Lacoste-Julien
TL;DR
Periodic weight nowcasters can accelerate optimization while Adam handles most steps, but prior predictors do not directly use neuron connectivity, and Transformer connectivity is difficult to model accurately. NiNo uses neural graphs and GNNs to predict future parameters, and the paper reports faster Adam training across vision and language tasks.
Problem
Prior periodic parameter nowcasters do not directly use neuron connectivity, while accurately modeling Transformer connectivity is challenging.
Source: §1 · p. 1, §3.2 · pp. 3–4Method
NiNo uses neural graphs and GNNs to predict future parameters from parameter histories, with Transformer graphs designed to model attention’s neuron permutation symmetry more accurately.
Source: §1 · p. 2Results
Across nine tasks, NiNo consistently outperformed Linefit, WNN, and their improved variants, reducing Adam steps to target performance by 48.9% on average.
Source: §5.2 · p. 8Takeaways & Limitations
The experiments show potential for periodic parameter prediction to reduce neural-network training time, with the reported approach spanning ConvNets and Transformers.
Source: §1 · p. 2, §6 · p. 10NiNo’s peak prediction memory is similar to a Transformer training step with batch size 16–32, and its inference efficiency remains open to improvement.
Source: §5.3 · p. 10
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper builds on periodic parameter nowcasting, which uses a learnable function to predict near-future parameters while a base optimizer runs for most steps. It proposes NiNo to use neural-network structure for more accurate predictions and reports faster Adam training across vision and language tasks.
1 INTRODUCTION
- Figure 1 compares Adam with and without NiNo or WNN on a language task neither nowcaster saw during training.The supplied caption identifies the unseen-task setting but does not state the result of the comparison.Source: Figure 1 · p. 1
- Periodic nowcasting runs a base optimizer such as Adam for many steps before predicting near-future parameters, reducing the frequency of learned-model calls compared with per-step L2O.Prior nowcasters do not directly use parameter structure such as connectivity between neurons and layers.Source: §1 · p. 1
- Up to 50% fewer steps to reach Adam’s target performance is reported for NiNo on ConvNets and Transformers.The paper also reports code and model release.Source: §1 · p. 2
- NiNo represents network parameters with neural graphs and uses graph neural networks to model neuron connectivity when predicting future parameters.Its Transformer graphs are designed to model multi-head self-attention’s neuron permutation symmetry more accurately.Source: §1 · p. 2
2 RELATED WORK
The related work situates NiNo among learned optimization and parameter-prediction approaches. It distinguishes NiNo through graph-based modeling of parameter structure and prediction conditioned on the future horizon.
2 RELATED WORK
- L2O methods learn optimizers for other neural networks, but meta-training and per-iteration overhead can be costly; NiNo follows WNNs in applying its predictor rarely.Rare application is described as making total overhead negligible while still speeding up training.Source: §2 · p. 2
- WNNs and IntrospectionMLP periodically predict future parameter values from past values, but use coordinate-wise MLPs and a predefined future horizon.NiNo instead uses neural graphs and conditions predictions on the horizon.Source: §2 · p. 2
- Neural graphs and GNNs represent neuron permutation symmetry and can process graphs with different sizes and connectivity, supporting one model across architectures and tasks.The cited Transformer neural graphs do not accurately model Transformer permutation symmetry, a shortcoming NiNo addresses.Source: §2 · p. 2
- NiNo’s graph-node interactions for future prediction are related to dynamic interactive models, but parameter prediction requires design choices that account for neuron permutation symmetry.The related work describes this modeling challenge as motivating the paper.Source: §2 · pp. 2–3
3 BACKGROUND
The background introduces periodic prediction of future parameters from an optimization trajectory and neural graphs as a structured alternative to coordinate-wise prediction. It also identifies a weakness in applying existing neural graphs to Transformer attention.
3.1 WEIGHT NOWCASTER NETWORKS (WNNS)
- WNNs use past parameter values to predict future parameter values and are applied periodically during base-optimizer training.The supplied background describes parameter scaling and supervised training from parameter trajectories, then gives an example of applying a trained WNN after Adam runs for several epochs.Source: §3 · p. 3, §3.1 · p. 3
3.2 (NAIVE) NEURAL GRAPH OF TRANSFORMERS
- WNNs predict each parameter with an MLP using its own past values, so their predictions do not directly use neuron connectivity.The background motivates neural graphs as a way to represent network structure and permutation symmetry.Source: §3.2 · p. 3
3 BACKGROUND
- A neural graph represents parameters as node and edge features, using weight-matrix rows and columns to encode neuron relationships and permutation symmetry.The passage notes that accurately constructing such graphs for Transformers is nontrivial.Source: §3.2 · p. 3
3.2 (NAIVE) NEURAL GRAPH OF TRANSFORMERS
- The naive Transformer neural graph assumes the multi-head case is handled automatically, but does not model neuron permutation symmetry correctly.Figure 2 contrasts the naive graph with the paper’s graph and depicts their embeddings for good and bad MSA-weight permutations.Source: §3.2 · p. 4, Figure 2 · p. 4
4 METHODS
NiNo represents parameter histories as neural graphs and uses a GNN to predict future parameter states. Its Transformer graph design encodes attention-specific permutation constraints, and its inference procedure can vary the prediction horizon over optimization.
4.1 NEURAL GRAPH OF TRANSFORMERS
- NiNo’s Transformer graph restricts neuron permutations across attention heads while allowing distinct permutations of Wq and Wk weights where MSA output is preserved.Separate head nodes encode the across-head restriction, while separate directed edge features represent the relevant weights.Source: §4.1 · pp. 4–5
4.2 EDGE AND NODE FEATURES
- The graph encodes parameters and auxiliary connections as typed edge features, while Laplacian positional encodings provide structural node features.For Transformer word-embedding layers, an additional positional feature is used; other layers receive zero for that feature.Source: §4.2 · p. 5
4.3 NEURON INTERACTION AND NOWCASTING NETWORKS (NINO)
- NiNo stacks c past parameter states into one neural graph and uses a GNN to predict future edge features, which are mapped back into parameter deltas.The model predicts K future steps directly; inference decays k from a large horizon toward smaller values as optimization proceeds.Source: §4.2 · p. 5, §4.3 · pp. 5–6
4 METHODS
- The method is designed for rare application during optimization, with the figure describing one NiNo call per 1,000 Adam steps on a new task.Its prediction pipeline uses the improved neural graph and a GNN.Source: §4 · p. 4, Figure 4 · p. 6
5 EXPERIMENTS
Across nine vision and language tasks, NiNo reduces the steps needed to reach target performance and is compared with other acceleration methods. The experiments also examine its design choices, behavior on larger models, parameter trajectories, graph embeddings, and computational costs.
5 EXPERIMENTS
- The evaluation spans nine vision and language tasks, using four in-distribution tasks for meta-training and five out-of-distribution tasks that differ in architecture and/or dataset.Task-specific networks use Adam for vision and AdamW for language, with the meta-model applied every 1,000 steps; performance is measured by steps to a baseline-defined validation target.Source: §5 · pp. 6–7, §5.1 · p. 8
5.1 SETUP
- The setup uses convolutional vision models, GPT2-style language models, and checkpoints from 1,000 in-distribution models for meta-training.The checkpoint collection contains around 1.6 × 10^6 intermediate states; baselines include tuned Adam or AdamW, Linefit variants, WNN+, and pretrained L2O/FM16.Source: §5.1 · p. 7
5.2 RESULTS
- 48.9% average speedup over Adam: NiNo reduces the steps to target performance and outperforms Linefit, WNN, and their improved variants across all nine tasks.The comparison uses relative reductions in steps to reach a target validation performance.Source: Table 2 · p. 8, §5.2 · p. 8
- L2O/FM16 performs best on its in-distribution task but overfits on an unseen task, whereas NiNo/FM16 performs well on unseen-task validation.L2O has overhead at every optimization step; the authors compare its curves only on FM/16 and C10/32 because it misses target validation scores on most tasks.Source: §5.2 · p. 8
- Layerwise parameter scaling and k-decay are NiNo’s most important ablated components; context length c = 5 is a sweet spot.Small context may miss parameter trends, while larger context applies NiNo less often; deeper and wider GNNs help, but depth above 3 or width above 128 makes training a performant model more challenging.Source: §5.2 · pp. 8–9
5.3 ANALYSIS
- On C10/32, NiNo predicts better than WNN at later stages, while its trajectory closely matches 2× faster AdamW despite using Adam as its base optimizer.L2O diverges early in the parameter-space comparison.Source: §5.3 · pp. 9–10
- Across three GPT2-style architectures, NiNo’s graph embeddings remain visually distinct at lower losses, while WNN+ embeddings virtually collapse.The shared embedding space allows the architectures to be projected together; the authors note that different parameters at similar training losses can yield different downstream behaviors.Source: §5.3 · p. 10
- NiNo’s wall-clock overhead is below 0.5% for a 7M Transformer, compared with around 5% for L2O.NiNo’s peak prediction memory is similar to a Transformer training step with batch size 16–32.Source: §5.3 · p. 10
6 CONCLUSION AND FUTURE WORK
The conclusion presents NiNo as a parameter-prediction approach with potential to reduce training time and provide low-dimensional views of training dynamics.
6 CONCLUSION AND FUTURE WORK
- The experiments suggest periodic parameter prediction has substantial potential to reduce neural-network training time.The authors identify scaling to more data and larger models as a future direction, noting their observed scaling trends indicate better speedups in those settings.Source: §6 · p. 10
- NiNo’s low-dimensional encoding of network parameters can be used to analyze diverse models and their training dynamics.Source: §6 · p. 10
- The authors propose investigating scale-up, since observed trends indicate potential for better speedups with more data and larger models.Source: §6 · p. 10
A.1 WNN PARAMETER SCALING DETAILS
The appendix describes WNN parameter scaling as a per-parameter min–max-based transformation with a corresponding unscaling step.
A.1 WNN PARAMETER SCALING DETAILS
- WNN scaling begins by computing minimum and maximum values from the input parameter values.Source: §A.1 · p. 15
- Scaling and unscaling are performed independently for each parameter.Unscaling multiplies the prediction by the parameter’s scale factor.Source: §A.1 · p. 15
- Subtracting the last element from the input sequence, as in Jang et al. (2023), was found redundant.Source: §A.1 · p. 15
A.2 LINEFIT
Linefit predicts future parameters by extrapolating a fitted line based on each parameter’s past values.
A.2 LINEFIT
- Linefit predicts future parameter values by extending a line to a future point [Sinha et al. (2017); Jang et al. (2023)].Source: §A.2 · p. 15
- The baseline fits a separate line for each parameter using only that parameter’s past values, without a training dataset.Source: §A.2 · p. 15
- Despite its simplicity, Linefit was reported to improve on using Adam or SGD alone.Source: §A.2 · p. 15
A.3 LINEFIT+
Linefit+ modifies Linefit’s uniformly weighted fit to emphasize later parameter values in the trajectory.
A.3 LINEFIT+
- Linefit+ penalizes errors on later trajectory values more heavily than earlier ones.Source: §A.3 · p. 15
- The change responds to Linefit’s uniform weighting, unlike momentum’s greater influence from more recent values [Sutskever et al. (2013)].Source: §A.3 · p. 15
- The Linefit+ weighting sequence µ runs from 1/c to 1.Source: §A.3 · p. 15
A.4 NAIVE NEURAL GRAPH OF TRANSFORMERS
A Transformer multi-head self-attention layer uses four d × d weight matrices, split across columns for query, key, and value and across rows for output. Existing neural-graph variants differ in how they represent biases and attention edges, making the multi-head case challenging to use directly.
A.4 NAIVE NEURAL GRAPH OF TRANSFORMERS
- Multi-head self-attention uses Wq, Wk, Wv, and Wo, each a d × d matrix; heads split the first three across columns and Wo across rows.The layer processes the input in parallel across H heads.Source: §A.4 · p. 15
- Lim et al.’s neural graphs add bias and normalization nodes, while their MSA description also stacks Wq, Wk, and Wv edges.The paper builds on this variant because it is easier to implement for layers with complicated neuron permutation symmetry.Source: §A.4 · p. 15
- The authors use Kofinas et al.’s neural graph as the naive MSA baseline because the multi-head case lacks a detailed description and public implementation.The cited work’s stacked Wq, Wk, and Wv edges are part of the implementation challenge described in the passage.Source: §A.4 · p. 15
A.5 GRAPH NEURAL NETWORK
The appendix details the graph-network components, permutation-symmetry tests, meta-training and use of NiNo, and additional language and vision evaluations. It reports a small ImageNet accuracy improvement and describes variable speedups under changed language-model learning rates.
A.5 GRAPH NEURAL NETWORK
- The GNN preserves graph size and is permutation equivariant, making it suitable for predicting parameters on neural graphs with permutable neurons.For the described neural graphs, the input edge features have shape |θ|1 × dE and predictions |θ|1 × 1.Source: §A.5 · p. 16
A.6 NEURON PERMUTATION SYMMETRY
- The permutation experiment indicates that the proposed neural graphs distinguish function-preserving from function-changing Transformer permutations better than naive graphs.It uses 1,000 permutations, about 500 of each type; a logistic-regression classifier and embedding visualization indicate that naive graphs confuse the classes.Source: §A.6 · p. 16
A.7 NINO TRAINING DETAILS
- NiNo achieves lower meta-training loss than WNN+ with a comparable parameter count, but remains severely underfit on all training checkpoints.Reducing the checkpoint models to four yields much lower NiNo loss with 128 hidden units; the authors suggest further scaling could improve fit.Source: §A.7 · p. 18
A.8.1 LEARNING RATE ANALYSIS
- NiNo’s speedup varies but remains significant across tested learning rates on the Llama3-style task, including rates different from checkpoint collection.The reported validation-perplexity standard deviation is generally ≤0.1; the tested setup also uses a new task and an architecture 18 times larger than the checkpoint architecture.Source: §A.8.1 · p. 18, Table 6 · p. 19
A.8.2 IMAGENET
- 50.4% vs 49.6% validation top-1 accuracy (AdamW+NiNo vs AdamW) after 50 ImageNet epochs, with an approximately 1% runtime overhead for NiNo steps.The ViT had 11M parameters; each NiNo step took about 10 seconds, similar to 10 Adam steps, and sampled 5% of edges to avoid GPU out-of-memory issues.Source: §A.8.2 · pp. 18–19
A.9 PSEUDO-CODE
- The checkpoint-collection pipeline samples trajectories from four tasks and produces 1,000 trajectories totaling 1,609,900 checkpoints.The pseudocode describes collecting model-parameter histories after Adam updates.Source: §A.9.1 · p. 20, §A.9 · p. 20
A.9.2 TRAINING NINO
- NiNo training samples five past parameter states to predict the next 40 and minimizes the mean absolute error of the predicted parameter delta.The pseudocode constructs a neural graph from the sampled trajectory and updates NiNo’s parameters with AdamW.Source: §A.9.2 · p. 20
A.9.3 USING NINO
- During use, NiNo wraps the base optimizer: AdamW updates parameters on most steps, while NiNo nowcasts every 1,000 steps.The pseudocode stores parameter history every 200 steps and skips gradient computation for the NiNo forward pass.Source: §A.9.3 · p. 20