Accelerating Training with Neuron Interaction and Nowcasting Networks

We thank the reviewer for the time and valuable feedback and would like to promptly respond to the concerns and questions.

W1.1 It is recommended that the authors provide more analysis on the impact of hyperparameter tuning on model performance.

We reported the impact of different hyperparameters (and ablations) in Table 3 and Fig. 5a,b. In Table 3e and Fig. 5a, we show the results for different NiNo's hidden sizes and numbers of layers. In Table 3f, we show the results for different batch sizes, learning rates and numbers of training models (also shown in detail on Fig. 5b). These are the main hyperparameters of NiNo. The ablations in Table 3a,b,c,d complement the analysis.

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

In Table 6 in Appendix we compare to the baseline (Adam), WNN+ and L2O. WNN+ has the same memory/overhead as WNN (Jang et al., 2023), since the architecture is identical. In addition to Table 6, we recently also measured the memory/overhead of NiNo (with hidden size=32) on our Llama-3-based architecture (with 111M params).

Overall, NiNo's computational resource consumption is comparable to other approaches, while it consistently outperforms other approaches in terms of speed up.

Q1 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.

In our paper we systematically investigated "the generalization ability of the model on unseen tasks". For example, in Tables 1 and 2, we specifically highlight five Out-of-distribution tasks. All the five out-of-distribution tasks are unseen: the dataset and/or architecture change compared to the training tasks (FM/16, C10/16, LM1B/3-24, LM1B/2-32). Among the five unseen tasks, the tasks C100/32 and Wiki/3-64 are most different w.r.t. the training ones, because both the datasets and the architectures are unseen. Moreover, in Table 4, we evaluate on the Llama3-based Wiki task, which is the most challenging generalization setup (unseen and very different architecture and unseen dataset). We also evaluated a NiNo model trained only on FM-16 (NiNo/FM16) and compared its performance on an unseen task (C10/32) in Fig. 6 (middle). Our NiNo/FM16 outperformed all the baselines (e.g. L2O, WNN).

==== We will add these clarifications in the revision. We thank again the reviewer and will be glad to clarify further questions.

We thank the reviewer for the time and valuable feedback and would like to promptly respond to the concerns and questions.

W1 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.

We assume the reviewer is concerned that training 1k models used for meta-training NiNo is not practical. In this regard, in our paper we presented the results of NiNo with a smaller number of models (in the range 40-1000, see Fig. 5b). For example, NiNo's average speed up is around 45% with just 100 models, making NiNo's training practical and performant. So we believe that for real-world applications, it should be possible to train NiNo on existing models without collecting a specific meta-training dataset. We collected such a dataset to enable a systematic study and remove potential confounding factors.

W2 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.

In our paper (Fig. 5c), we show how NiNo performs when applied to models of different sizes (from 2 to 6 layers). The trend shows that further generalization is unlikely. Nevertheless, the generalization performance is already remarkable in our view because as Fig. 5c highlights (red rectangular), NiNo can be applied to models that are around 18 times larger than those used in meta-training. So to generalize beyond 100M models, one needs to scale up meta-training models and/or further improve NiNo's architecture. This is left for future work.

Q1 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?

When NiNo makes predictions, most of the parameters are generally predicted in the right direction however some parameters may be predicted inaccurately. For the loss spike it can be enough to have just a few parameters to be predicted way off. However, a few training iterations with Adam can easily correct those wrong parameter values. So NiNo makes a good global-level prediction while Adam can fix any local-level mistakes.

Q2 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.

The "set the network weights to the output of NiNo while retaining the optimizer states" is correct (we also tried resetting the states but that didn't change the results significantly). Overall the inference process is the following. NiNo is applied every 1k steps of Adam (see Fig. 1). For example, given 10k training steps in some task, NiNo is applied at steps 1k, 2k, ..., 9k, 10k. NiNo's forward pass is done for a specific k value, where k is decayed during the optimization process (Fig. 12 in Appendix). For example, for 10k step the k is decayed as: k=40 at 1k, 33 at 2k, ..., 1 at 10k.

Pseudo code of using NiNo (see supplementary material for full code):

model = AutoModelForCausalLM.from_config(...)  # some model
# NiNo is implemented as a wrapper around the base optimizer
# any optimizer other than Adam should also be possible to use with NiNo
opt = NiNo(base_opt=torch.optim.AdamW(model.parameters(), lr=1e-3),
           ckpt='checkpoints/nino.pt',
           model=model,
           period=1000,
           max_train_steps=10000)
for step in range(10000):
    if opt.need_grads:  # True for Adam steps (e.g. 0-999)
        opt.zero_grad() 
        data, targets = ...  # get data
        outputs = model(data)  # forward pass
        loss = F.cross_entropy(outputs, targets)  # loss
        loss.backward()  # grads for Adam        
    opt.step()  # Adam step or, every 1000 steps, nowcast params using NiNo 
    ...

Q3 Could the authors provide details about the computational costs of the meta-training process?

Training of NiNo completes in under 7 hours on a single NVIDIA RTX8000 with 48GB of memory (see Section A.7 in Appendix). It can be further sped up using multiple GPUs as in standard training pipelines. There are many other potential approaches to make training more efficient (e.g. subgraph/subnetwork sampling common in the GNN works, which can be interesting to investigate in future work). Collecting the meta-training dataset is relatively cheap (can be done in a few hours) and the dataset takes around 200GB (raw memmap data, which we will release upon publication). But as pointed out above, it would be intriguing to try already existing datasets of models in future work.

==== We will clarify these concerns in the revision. We thank again the reviewer and will be glad to clarify further questions.

We thank the reviewer for the time and valuable feedback and would like to respond to the concerns and questions.

W1 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.

We just updated the pdf to make Figure 2 more clear (increased the contrast and augmented the notation with different styles). We will be glad to further improve it according to feedback.

To clarify Section 4.1, below we are rephrasing the two problems of naive neural graphs (Kofinas et al., 2024) that we are addressing:

1. in naive neural graphs it’s possible to have neuron/node permutations that do not change the graph (i.e. it remains isomorphic to the original graph), but change the transformer output. Such an example would be swapping neuron 3 (belongs to head 1) with neuron 4 (belongs to head 2). We address this by allowing the permutations only within the head.
2. in naive neural graphs it’s also possible to have neuron permutations that do not change the transformer output, but change the graph. Such an example would be permuting neurons in $W^q$ and $W^k$ while not permuting neurons in $W^v$. We address this by assigning separate edges for $W^q,W^k$ and $W^v$ instead of having 3-dimensional (for q,k,v) edges, so in our case the neurons in $W^q, W^k$ can be shuffled independently from $W^v$ according to the underlying function of MSA layers.

As we show in our synthetic experiments in 4.1 (illustrated in Figure 2d,e and explained in Appendix A.6), our modifications indeed better correspond to the correct neuron symmetry in MSA layers. That is in our neural graphs, when the graph does not change, then the transformer output does not change either, and vice versa: if the graph changes, the transformer output changes as well.

We will be glad to further clarify this or other pieces.

W2 The training method is fairly computationally expensive as the authors collect on the order of $10^6$ checkpoints.

While the total number of checkpoints is $10^6$, the total number of models we trained is only $10^3$. Moreover, in our paper we also presented the results of NiNo with a smaller number of models (in the range 40-1000, see Fig. 5b). For example, NiNo's average speed up is around 45% with just 100 models, making NiNo's training practical and performant. So we believe that for real-world applications, it should be possible to train NiNo on existing models without collecting a specific meta-training dataset. We collected such a dataset to enable a systematic study and remove potential confounding factors. We are answering to "1B parameters" below in Q1 (see Part 2 of the response).

We thank the reviewer for the time and valuable feedback and provide our response below.

High-level response: Please note that with academic-level GPU resources it is often very challenging to scale up the experiments of very novel methods to industry-level standards (>=1B parameters). Our work has made significant and non-trivial contributions in terms of novelty compared to the very recent works, WNN (Jang et al., ICML 2023) and neural graphs (Kofinas et al., ICLR 2024). Regarding the scale, Kofinas et al. applied neural graphs to the models with only ~0.1M parameters, so we have significantly (by 3 orders of magnitude) scaled up the neural graph approach when applied it to a Llama3-style transformer with 111M params (see our Table 4). This kind of scale is also common in optimization literature [1, 2, 3], and scaling up is often studied in follow up works. At the same time, we thank again the reviewer for valuable feedback that helps to improve our paper, which we updated. Please see the detailed response below.

W1.1 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.

For Llama3-style experiments the NiNo is applied to a 111M parameter model, also the results in Figure 5c are for GPT2-based models with up to 30M params, so it is much larger than ~1M parameters. Our experiments are in-line with recent literature on experimental optimization procedures, e.g. [1, 2]. We thus believe it is reasonable to leave further scaling up to future work as it requires a separate study. For example, one key direction to investigate could be developing more scalable GNN layers.

W1.2 For image tasks, showing reasonable performance on something like ImageNet is a must.

We trained a small ViT (11M parameters) on ImageNet (with 32$\times$32 images as in [2, 3], for 50 epochs given the short response time). The validation top-1 accuracies are 49.6% (Adam) vs 50.4% (Adam+NiNo), which we believe is a reasonable improvement on ImageNet, especially given that we are applying the NiNo model to a very different architecture on a very different task compared to meta-training. We will add this experiment to the revision.

W1.3 Authors should 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).

We attempted to stick closely to the setting of the recent work WNN. Therefore we used the same learning rate (lr=2e-4) as in WNN. We used this lr for all language tasks, including the collection of training checkpoints (lm1b/3-24 and lm1b/2-32 tasks). Since NiNo was meta-trained on models trained with lr=2e-4, we use the same lr for all experiments including Llama3-based. As in WNN we limit our analysis to lr used in checkpoints, but using a different lr than those from checkpoints at inference time is an interesting question. We have now done a study for different learning rates (without retraining NiNo on these new lr checkpoints):

We ran each experiment 3 times, the standard deviation is generally <=0.1, indicating that ours is better including the std even in this challenging setup (new lr compared to the training checkpoints, new and 18 times larger architecture and new task). Speed up is computed as described in our paper (e.g. 14% here means that AdamW + NiNo achieved the target performance in 12k steps vs 14k steps of AdamW). As reviewer Ubaj noted, even 10-20% improvements in training speed could be of large practical significance. We will add this analysis to the revision.

References:

• [1] Learning to Generalize Provably in Learning to Optimize, AISTATS 2023
• [2] Practical tradeoffs between memory, compute, and performance in learned optimizers, CoLLA 2022
• [3] Decoupled Weight Decay Regularization, ICLR 2019