Paper: NanoLM: An Affordable LLM Study Benchmark via Accurate Loss Prediction Across Scales (ICLR-2024)

Reviewer comment (Evaluation):
The paper focuses on loss prediction as the primary metric for benchmarking and comparison, but it does not explore how well this loss prediction correlates with performance on various downstream tasks. It's not clear if this is a limitation as loss on a pre-training task is not always indicative of performance in practical applications.

Author response:
Great question! The success of Large Language Models (LLMs) mainly relies on their base models, where the key is the loss value in set datasets. For example, GPT-4 [1] aims for a loss value of 1.25. LLMs are trained on tasks like predicting the next or hidden words, and this loss is closely linked to perplexity. The LLaMA study also hints at a connection between pre-training loss and later performance, but it's hard to fully evaluate. In LLAMA’s Figure 1 [2], you can see training loss decreases in models from 7B to 65B. Figure 2 shows better results in question answering and reasoning with more training tokens. Essentially, lower pre-training loss often means better performance later, highlighting the importance of nanoLM's loss prediction.
We appreciate your feedback about the unclear analysis. We revised it in Section 3.1 for better clarity.

Reviewer comment (Experiments):
How robust is nanoLM's loss prediction across different types of transformer architectures, beyond GPT, BERT, and T5 structures? Have you tried more recent open-source models like Mistral or Llama? Are there certain architectures or hyperparameter configurations where nanoLM's predictions are less accurate?

Author response:
We have confidence in our results for two reasons: (a) $\mu$P is theoretically sound, and (b) scaling laws have shown empirical reliability across various scales. We've included additional analysis to assess nanoLM's loss prediction accuracy:
(1) Embedding counts in the model size. Our new experiment shows that our scaling law fits less accurately if embedding weights are excluded, unlike in OpenAI established scaling law [4]. This could be due to $\mu$P's finding that embedding layer learning rates shouldn't be scaled down with widths. Our unified learning rate for all layers at each model size might result in slower learning for embeddings, with matrix-like parameters predominantly influencing training dynamics.
(2) We recommend initializing the word embedding and query parameters to all zeros to avoid Gaussian Processes. Accoding to $\mu$P thereactical proofs, a Gaussian Process with may cause misalignment of landscapes between small and large models.
(3) Based on our current experiments, nanoLM performs well across different transformer-like architectures (decoder-only, encoder-decoder, encoder). However, we encourage users to test nanoLM on other structures as part of future work.
We appreciate your suggestion and these analyses is added in a new appendix A of the paper.

Reviewer comment (Evaluation):
Have you considered extending nanoLM's capabilities to predict other important metrics beyond loss, such as energy efficiency or training time, to provide a more comprehensive view of a model's trade-offs?

Author response:
This is a good catch! As noted in Chichilla's[5] paper, both training time and energy efficiency can be expressed as functions of model size, exhibiting a linear relationship. We will add this clarification in the new appendix, and apologize for any confusion caused. Thank you for highlighting this important aspect.

Reviewer comment (Writing):
To make the paper stronger, the authors should phrase the critical component of the method more clearly.

Author response:
### Q4: The paper need phrase the critical component of the method more clearly.
We agree. We clarified the key elements of our method in Section 3.

Reviewer comment (Novelty):
In terms of method, the paper lacks its own insights into the scaling law. The two main components uP and uScaling are refer to others’ work and not discussed enough.

Author response:
### Q3: The paper lacks its own insights into the scaling law compared with main components.
We believe there are misunderstandings: 
(1) We wish to clarify that nanoLM is an extension of $\mu$scaling [1] (preprint on arXiv). We've contacted the area chair for more details and will update upon receiving feedback. Specifically, Unlike $\mu$Scaling uses 8M-676M GPT models to predict the loss for 1.4B models, nanoLM extends this to 77M-3.4B for predicting loss in 52B models. This scale is 184x larger than $\mu$Scaling, marking a critical step for LLMs. And nanoLM providing extensive datasets ranging from 100B to 2T for comprehensive testing. Additionally, nanoLM supports a variety of architectures, including decoder-only structures (e.g., GPT, LLAMA), encoder-only structures (e.g., BERT), and encoder-decoder structures (e.g., T5), and is compatible with data parallelism. These enhancements are particularly crucial for researchers in the field of LLMs, offering a more robust benchmark for researcher.
(2) The key distinction of nanoLM from prior work is its **avoidance of re-searched on hyperparameter (HP) for large language models (LLMs).**  There are very few related work to ours. Except those mentioned in the paper, we (as well as [4]) noticed a contemporary work [3] that applies $\mu$P in language model pre-training. Their main finding is that $\mu$P results in a constant (thus predictable) deviation from the scaling laws fitted with Standard Parametrization (SP). However, SP scaling laws need HP search on large scales because they have large variances. Besides, they varied some non-$\mu$Transferable HPs for different model scales, while we only vary the widths and show that the $\mu$P process itself results in a model sequence with accurate loss prediction. For these two reasons, the stories are very different.
Thank you for pointing out that this comparative analysis is unclear; we updated the above-mentioned analysis in new Appendix A to clarify this.

Reviewer comment (Writing):
Does “different in design” mean each input model is different from not only the widths but also the architecture?

Author response:
### Q1: Ambiguity in Algorithm 1.
* Yes, line 1 "different model design" includes width and architecture variations, as shown in our paper's Figure 2. For example, $M_1$ uses Llama, and $M_2$ GPT. nanoLM allows for loss predictions and comparisons in these models at larger sizes without full-scale training.

Reviewer comment (Writing):
Line 3: “Train above small-width models” reads as the width varied in line 2 should be small. So stating it clearly in Line 2 would help readers to understand the whole algorithm.

Author response:
### Q1: Ambiguity in Algorithm 1.
*  Yes, line 2 "Generate some models varying widths only" mean generating models for each $M_i$ by varying the widths and there is no limit to varying the width according $\mu$P [1] condition. 
*  Thanks for your suggestion, we implemented it for the revised version of Alg 1 Line 2.

Reviewer comment (Experiments):
In Figure 4, the training loss at 7k and 10k of different sizes of models are shown. 7k and 10k iterations are very early stages of 26B and 52B models. What about the loss prediction accuracy for longer training steps (e.g., 20k or more)?

Author response:
### Q2: What about the loss prediction accuracy for longer training steps(eg, 20k or more)?
Thank you for your valuable suggestion. We've now increased the training of our 12-layer models (GPT, LLAMA, BERT, T5) to 20k steps. We included the updated results in Section 4.2 and Figure 3. Additionally, it's worth noting that the theoretical correctness of $\mu$P [1] is independent of the number of training steps.

Reviewer comment (Evaluation):
The author should put the loss prediction values (with multiple fitting and std included) and ground truth into a table and display in the main paper. So the readers have a clear sense of the robustness and performance of the proposed benchmark.

Author response:
### Q5: Loss prediction values and ground truth should put into a table in the main paper. 
Thank you for the suggestion. We incorporated loss prediction values and ground truth in new Table 3 of the main paper.

Reviewer comment (Writing):
Minor: there are some inconsistent notations and grammar errors, please fix them accordingly.

Author response:
### Q6: Typos.
We're sorry for these mistakes. We have rechecked the paper and corrected the typos.

Reviewer comment (Novelty):
The paper's novelty is questioned as the methodology seems to be a synthesis of pre-existing approaches.

Author response:
Thank you for your thoughtful review and for raising concerns about the novelty of our work.  We wish to clarify that nanoLM is an extension of $\mu$scaling [1] (preprint on arXiv). We've contacted the area chair for more details and will update upon receiving feedback. Unlike $\mu$Scaling uses 8M-676M GPT models to predict the loss for 1.4B models, nanoLM extends this to 77M-3.4B for predicting loss in 52B models. This scale is **184x** larger than $\mu$Scaling, marking a critical step for LLMs. Additionally, nanoLM supports various architectures, such as decoder-only (e.g., GPT, Llama), encoder-only (e.g., BERT), and encoder-decoder (e.g., T5), and is data parallelism compatible. These features are vital for LLM researchers, offering a more comprehensive benchmark.

Reviewer comment (Novelty):
The datasets introduced appear to be aggregates of data already available.

Author response:
Thank you for pointing this out. The primary intention behind curating and categorizing datasets in the nanoLM benchmark is to empower researchers to conduct meaningful comparisons between different model architectures and algorithms, especially when operating under computational constraints. We acknowledge that the provision of datasets is not an innovation in itself, but rather a facilitation for researchers.

Reviewer comment (Novelty):
The study does not offer an in-depth comparative analysis or discussion on how their proposed methodologies diverge from OpenAI's established scaling law[1,2] for loss prediction, as detailed in prior technical reports.

Author response:
Thank you for pointing out the inadequacies in this part of our analysis. We add comparisons and experments with OpenAI's scaling law from the following perspectives:
(1) GPT-4 [3] technical report shows that some behaviors of large models can be accurately predicted before the training starts. However, GPT-4 has not disclosed the techniques used for its loss prediction, so we are unable to make a comparison with it.
(2) General conditions for scaling laws? Previous scaling laws [4] directly search for HPs on each scale, and the optimal HPs do not satisfy $\mu$P[2] function. This indicates that $\mu$P is a sufficient but not necessary condition for scaling laws, and scaling law itself may represent a higher level of universality.
(3) The key distinction of nanoLM from prior work is its **avoidance of repeated hyperparameter (HP) searches for large language models (LLMs)**. There are only a few works similar to ours. Beyond those cited in our paper, we, along with [6], recognize a concurrent study [5] applying $\mu$P in language model pre-training. Their main discovery is that $\mu$P leads to a predictable deviation from scaling laws established with Standard Parametrization (SP). However, SP scaling laws demand HP searches at larger scales due to high variances. Additionally, they modified some non-$\mu$Transferable HPs across different model sizes. In contrast, we vary only the widths and demonstrate that $\mu$P alone can produce a model sequence with accurate loss prediction.  For these two reasons, the stories are very different.
(4) Embedding counts as model size. We add experiment shows that our scaling law fits less accurately if embedding weights are excluded, unlike OpenAI established scaling law [4]. This could be due to $\mu$P's finding that embedding layer learning rates shouldn't be scaled down with widths. While [4] unified learning rate for all layers at each model size might result in slower learning for embeddings, with matrix-like parameters dominate the training dynamics. 
Thank you for pointing out that this comparative analysis is unclear; we will update the above-mentioned experiment and analysis in new Appendix to clarify this.

Reviewer comment (Experiments):
Authors show experiments that indicate that nanoLM can predict the loss of extremely large-sized models by fitting the losses of several smaller models. It is not clear how to calibrate the error they report. It would be good to compare the results the authors show with other sources of data for predicting loss such as the results of [https://arxiv.org/pdf/2304.01373.pdf]. In the same vein, it would be good to compare against other work in the field e.g. [https://arxiv.org/pdf/2309.14322.pdf].

Author response:
### Q2: More evidence to explain why the predictions errors are meaningfully small.
Thank you for raising this important question and for providing related papers [3][4] that aid in our explanation. Paper [3] primarily explores the relationship between learning rate and training stability, without focusing on loss prediction, a key aspect of our work. In contrast, nanoLM conducts the $\mu$HPs search on smaller models to ensure entry into the loss basin, facilitating accurate predictions. Additionally, the largest model in [3] is only 1.2B, making direct comparisons with our larger models impractical. We have conducted additional experiments to further calibrate the prediction errors of nanoLM, as detailed below:
(1) Based on Table 8 of the Cerebras-GPT [5] and the scaling law's power function $L=ax^b+c$, we conducted an error comparison analysis of the predicted losses for Cerebras-GPT[5] , Pythia[4], and nanoLM. Cerebras-GPT fitted the 13B model's loss using 111M-6.7B models, Pythia for 12B with 7M-6.9B, and nanoLM for 52B from 77M-3.4B, with their respective errors being **0.025, 0.019, and 0.022**. When fitting losses for models down to 10B, the errors are **0.034, 0.049, and 0.018**, respectively. Additionally, we calculated the covariances of the fitted coefficients {a, b, c}, finding that **nanoLM's covariances are significantly lower** than those of Cerebras-GPT and Pythia. These experimental results  demonstrate the (a) $\mu$P infinite neural network is theoretically correct and (b) scaling laws are empirically reliable on any scale. (c) The loss prediction of nanoLM is more stable and reliable. Loss prediction validation is costly, and we have conducted as many experiments as possible within our computational power limits (nanoLM has reached 52B for this purpose, while Pythia and Cerebras-GPT have only reached 13B). Further experimentation would be prohibitively expensive in terms of computational costs.
(2) **Small models are more vulnerable and can be solved by average fitting**. We predict loss of small model (6-layer models). After grid search for the best $\mu$HPs being inside loss basin. nanoLM works perfectly for this $\mu$HPs. We then explored other $\mu$HPs around it and found that the scaling laws have larger deviations than 12-layer models. This is potentially because small models are more vulnerable to slight misalignment of loss landscapes across $\mu$-Transfer. However, we can easily balance-off this deviation by fitting scaling laws with the average results across all these HPs near the loss basin. This can be very practical because we observe that larger widths have low variance in training loss w.r.t different $\mu$HPs. 
We presented these analyses in the updated Appendix A.

Reviewer comment (Experiments):
The experiments are limited to English language modeling. Have the authors considered testing the approach for other languages and tasks? It would strengthen the claims of applicability.

Author response:
### Q1: Experiments are limited to English language modeling. Plans for suporting/expanding nanoLM are not discussed. 
Thank you for your suggestions. $\mu$P's [2] thereactical proof shows nanoLM suits most transformer-like architectures. We've included Llama[1] and other language in nanoLM's benchmark for loss prediction, demonstrating its robustness; These updates are detailed in the revised Section 4.2 and Figure 3. In the future, we plan to integrate additional popular TF-like architectures and tasks (eg., vision) into the nanoLM and encourage users to test, aiding in the continuous evolution of the nanoLM benchmark.