Gemstones: A Model Suite for Multi-Faceted Scaling Laws

We thank the reviewer for their time and effort, and appreciate that they acknowledge the impact that our new convex hull approach and results on width and depth will have on the NeurIPS community.

Figures 4,5 and 6

To the best of our knowledge, Figure 4 is a novel result in a scientific setting. Figures 5 and 6 are demonstrating that our suite of models also extends to benchmark scaling laws. During the rebuttal period we have also conducted experiments fitting scaling laws over the datasets shown in Figure 4 to more effectively demonstrate the impact of using laws fitted on different datasets. Moreover, we extend our current benchmark scaling laws in Figures 4 and 5 by detailing their extrapolation when the 2b parameter models are held out. All of these experiments will be included in the camera ready version even though we cannot directly share them during the rebuttal process. In summary, we find that the formula proposed by Gadre et al. (fitted on all data in Figure 5) to be more robust to extrapolation than the formula proposed by Bhagia et al. (fitted on all data in Figure 6) with the Gadre et al. formula giving near perfect extrapolation to the 2 billion parameter models when fitted on all smaller models. With regards to fitting laws over multiple validation sets, we find that the fitted laws only vary slightly as the current Figure 4 implies.

Convex Hull Method

We agree the convex hull methodology has the potential to be highly impactful and widely used within the community. We hope that future research considers this method alongside the classical approach 1 fitting procedure and highlights how it can be used by practitioners to fit stable scaling laws with fewer data points.

Recommendations

We caution making sweeping recommendations for future practitioners due to the fragility we find to be inherent to scaling laws. Hence, our main recommendation is to be very careful of the pitfalls we highlight, and to make sure that the final configuration for training a larger “production” model mirrors the set of assumptions made when building the smaller models used to fit the scaling law. We also think our other contributions will also be useful to future practitioners: how width and depth impact training speed, how width and depth impact of benchmark accuracy versus validation loss, and how using the convex hull method for approach 1 can drastically reduce variability during fitting.

Learning Rate Schedule

Hägele at al. [1] suggest practitioners should use a warmup-stable-decay scheduler when fitting scaling laws as it allows intermediate checkpoints to be used for fitting and offers similar performance to a cosine decay scheduler. Appendix 1 details how we use a scalable initialisation and learning rate transfer to models of different shapes to be compared fairly in our study.

In response to the reviewer’s valuable feedback, we have added extra references listed to our camera ready draft, and have made significant efforts to incorporate all other the feedback into the manuscript. If your primary concerns have been adequately addressed, and you do not have any other questions, please kindly consider raising your score as a stronger recommendation of acceptance.

[1] Hägele A, Bakouch E, Kosson A, Von Werra L, Jaggi M. Scaling laws and compute-optimal training beyond fixed training durations. (Neurips 2024)

We thank the reviewer for taking time to read our draft and their acknowledgement of the large contribution that our Gemstone model suite will be to the community.

Our Choice of Track

We believe our scientific contributions, including our analysis of the pitfalls of scaling laws and the scientific insights gained from varying the width and depth of the models, make this paper more appropriate for the main track. In addition, we introduce a convex hull fitting approach and analyze the impact of fitting scaling laws across different validation sets. Our results include several novel findings related to width and depth: increasing width decreases loss quicker than depth over time but slower over FLOPs, and increasing depth improves benchmark scores more than width under a fixed compute budget.

Model Sizes Greater Than 2B

We agree with the reviewer that studying larger models would be an interesting extension of our work, but as the compute required scales greatly with model size, we chose to train for longer token horizons rather than increasing model scale. We highlight that several scaling law papers have fitted on models with less than 2 billion parameters and produced results that generalise to much larger parameter counts [1,2,3]. Moreover, scaling laws fitted on small models have become commonplace in large SoTA model releases, predicting the behaviour of much larger models with extraordinary accuracy [4,5,6]. Hence, we believe our paper will be impactful for the NeurIPS community even thought it does not include models beyond 2b parameters.

Non-Gemma Architectures

In order to conduct a rigorous scientific study, we restrict some architectural considerations in order to analyze width and depth with high fidelity. Although many of our architectural choices are based upon Gemma models, we also take inspiration from Llama and Pythia models. Specifically, we take the convention of Llama models where the head dimension is always the embedding dimension divided by the number of heads and use the Pythia tokenizer in all experiments.

We would also like to emphasize how our scaling law fragility results relate to these fixed architectural parameters. Since we find that many factors impact the fit of a scaling law we caution against transferring scaling laws without careful attention to the individual design decisions implicit in that scaling law. That said, we believe that in our setting our scaling laws would generalise as well as the ones shown in [2] and [3] who also fit on models of less than 2b parameters and generalise to much larger parameter counts with high accuracy. We also believe that our key methodological contributions such as the convex hull will transfer to other settings. Moreover, our findings such as the choice of aspect ratio to optimize loss with respect to time and to step count are widely impactful for all practitioners as we will open source all of our data. Finally, we hope this work drives extra consideration in other areas of architecture optimization such as expansion factor in transformer models and number of experts in a mixture of expert models as these are highly related to aspect ratio.

Scaling

We have to politely disagree with the reviewer and think scaling should be discussed in terms of optimal performance for a given compute budget (FLOPs), not model size. We believe the primary motivation of scaling laws is to optimize performance of larger models by training many smaller ones and fitting laws to the data those runs produce. We aim to offer the community a new angle to optimize performance at a given parameter count by leveraging our findings on how width and depth can be varied to achieve better performance. We note that it is common to overtrain models past any prescribed “compute optimal” point (Gadre et al.) and note how practitioners could also use our work to further understand how width and depth affects their FLOPs per training step and optimize over this dimension as well.

We are aware that training models with varying widths and depths does lead to some models being suboptimal for a given FLOP budget and highlight how our convex hull approach is designed to stop this impacting our fitted laws. We also note the impactful additional information gained for our community by identifying which models are slightly suboptimal and our analysis of them throughout our paper. Finally, we highlight how the convex hull approach can also be used by the community in the future to fit high quality scaling laws with less data as shown in Figure 3.

Common Pitfalls of Scaling Laws

In the fragility and common pitfalls section, we are highlighting inherent fragility and pitfalls of all scaling laws. The asymptotic flatness of power law curves makes estimation of their parameters an ill-conditioned problem. Moreover, this is compounded when fitting (interpolating) what are already inherently noisy obervations from real training runs and we believe that our results show how this ill-conditioning can quickly impact predictions in many situations. We want to emphasize that all scaling laws in general should be considered in as close to their original setting and context as possible and practitioners should not expect them to transfer to significantly different settings without non-trivial effort. However, this is not a flaw of our study or methodology and to our knowledge, there is no specific proposed method to compare to that solves this variability problem entirely. With that said, in Figure 3, we do demonstrate how our convex hull method for the approach 1 fitting method (black crosses, blue line) compares to prior fitting approaches (red crosses, red line), and think that the reduction in the noise it yields is extremely promising.

In response to your feedback, we have made significant efforts to improve our camera-ready version. If your primary concerns are at least somewhat addressed, and you have no further questions, please kindly consider raising your score.

[1] Li M, Kudugunta S, Zettlemoyer L. (Mis) Fitting: A Survey of Scaling Laws. (ICLR 2025)

[2] Kaplan J, McCandlish S, Henighan T, Brown TB, Chess B, Child R, Gray S, Radford A, Wu J, Amodei D. Scaling laws for neural language models. (arXiv)

[3] Bhagia A, Liu J, Wettig A, Heineman D, Tafjord O, Jha AH, Soldaini L, Smith NA, Groeneveld D, Koh PW, Dodge J. Establishing task scaling laws via compute-efficient model ladders. (arXiv)

[4] Grattafiori A, Dubey A, Jauhri A, Pandey A, Kadian A, Al-Dahle A, Letman A, Mathur A, Schelten A, Vaughan A, Yang A. The llama 3 herd of models. (arXiv)

[5] Hu S, Tu Y, Han X, He C, Cui G, Long X, Zheng Z, Fang Y, Huang Y, Zhao W, Zhang X. Minicpm: Unveiling the potential of small language models with scalable training strategies. (arXiv)

[6] Bi X, Chen D, Chen G, Chen S, Dai D, Deng C, Ding H, Dong K, Du Q, Fu Z, Gao H. Deepseek llm: Scaling open-source language models with longtermism. (arXiv)

[7] Gadre SY, Smyrnis G, Shankar V, Gururangan S, Wortsman M, Shao R, Mercat J, Fang A, Li J, Keh S, Xin R. Language models scale reliably with over-training and on downstream tasks. (ICLR 2025)

We thank the reviewer for engaging with our rebuttal, taking the time to understand the nuances of our paper and increasing their score.

Just because it is so relevant to both our work and this reviewer's comments, we wanted to bring an extremely recent release to the attention of the reviewer (posted during the rebuttal period). The trade off between width and depth has now also been analyzed by Zuo et al. [1] but in a different architecture family. They find strikingly similar trends in their analysis of hybrid state-space and transformer models in Sec 2.3.2 where they show that deeper models achieve lower loss at a fixed token budget but are less efficient to train. While our model architecture is itself a heterogeneous combination of Llama-style attention configuration, Gemma-style MLP layers, and the Pythia suite's vocabulary size, this new study closely corroborates our findings in a completely different architectural setting. Moreover, with what appear to be increased computational resources, they are able to demonstrate that a deeper 1.5b model can even match a 3b or 7b model of the same style, confirming that studies at the sub 2B scale may provide useful insight as to how to optimize performance at larger scales.

We think that these results directly reinforce ours and that their models complement our own model suite. It is promising to see more thorough, scientifically reproducible investigations by our community, and this contemporary work has raised our confidence in the fact that the fully open source models, data, and rigorous analysis we present will serve as a valuable starting point for future work. We will be adding this citation and a short discussion to our camera ready draft.

If this additional evidence and independent corroboration of our results has helped alleviate your remaining concerns, please consider raising your score as a stronger recommendation for acceptance.

---

[1] Zuo J, Velikanov M, Chahed I, Belkada Y, Rhayem DE, Kunsch G, Hacid H, Yous H, Farhat B, Khadraoui I, Farooq M. Falcon-H1: A Family of Hybrid-Head Language Models Redefining Efficiency and Performance. (arXiv)

We thank the reviewer for their thoughtful review of our work and appreciate that they have noted the impact that open sourcing the Gemstone model suite will have on the scientific community we are a part of. We also hope that our release will enable more researchers to participate in nuanced discussions on width and depth, which the reviewer notes have previously been understudied.

Conceptual Contributions

We acknowledge we are not the first to see variability in scaling law analyses, and discuss this in our related work section. However, we do believe that showcasing this variability along both previously studied axes as well as novel ones in a single controlled setting is of value to the community as scaling laws are widely used for planning model training runs. In particular, we think that a large-scale study of width and depth provided by this work along paired with scaling law analyses represents a non-trivial addition to the open source literature. For example, in the Gemma 2 report [1], Table 9 presents a single result on how depth increases benchmark performance more than width, but few details about the experiment are included. In contrast, our work corroborates this in a scientific setting and open source all results for the community to build upon. We also emphasize our other novel, impactful contributions that the reviewer already notes, such as finding width and depth to affect loss differently over time vs. FLOPs, analyzing transfer of scaling laws over different validation datasets, and introducing the convex hull methodology.

The Convex Hull Approach

We provide a mathematical definition below, loosely based on the wikipedia entry for reconstructing functions from epigraphs, and have added this to our camera ready draft:

We can define the set of points we have to fit on as FLOPs/GPU hours $(x)$, loss value $(L)$ pairs as $\mathcal{D} = \\{ (x_i, L_i) \\}_{i=1}^n \subset \mathbb{R}^2$

Let $\text{conv}(\mathcal{D})$ denote the convex hull, the linear interpolation of any two points in $\mathcal{D}$:

\text{conv}(\mathcal{D}) = \left\\{ \sum_{i=1}^{n} \lambda_i (x_i, L_i) \ \middle|\ \lambda_i \geq 0,\ \sum_{i=1}^{n} \lambda_i = 1 \right\\}

Define $\hat{L}(x) = \min(\\{y | (x,y) \in \mathcal{D}\\})$ as the lower convex hull is the graph of this new function: $\\{(x, \hat{L}(x)) | x \in \text{Dom}(\hat{L})\\}$

We think the easiest visualisation of this is in Figure 7 where the red line is the convex hull, the colored crosses are the vertices and the grey lines are all possible points in the dataset. The exact code to generate the convex hull is in our supplementary material, in the function plotters.figure_2_correct_data.get_resource_hull

Rigour

The reviewer has noticed a key merit of our paper in that we present two distinct sets of contributions. One part of our work looks at the established set of design choices one must make when fitting scaling laws and the other considers varying the shapes of the models we train along dimensions of width and depth.

We remark that many impactful scaling law studies train fewer distinct model shapes than we do, such as Porian et al. [4] who train 16 distinct model shapes and even one of the largest studies, Chinchilla [2], only considers 50 distinct model shapes. To fit a scaling law, we need data points which consist of a compute budget (often made up of a model parameter count and number of training tokens) and loss values. Some prior work such as Chinchilla [2], Kaplan [3] and Porian [4] choose to do this by running many independent training runs as they use a cosine learning rate scheduler. To reduce cost and increase efficiency, we follow Hägele et al. [5] who propose practitioners should use a warmup-stable-decay learning rate scheduler so that all intermediate checkpoints in the stable regime can be used during fitting, this means we can utilise all 4000 of our checkpoints in different ways when fitting scaling laws much more effectively and train for much larger token counts. Our 4000 checkpoints can be split into 3 categories {main, cooldown, lr ablation}; each can use all points to fit a scaling law. We also highlight that in Table 1 we see fragility over each of these subsets, inferring this is a repeatable pattern. We also refer to our related works section to highlight some fragility in scaling laws has been seen in prior work as it is inherent to all scaling laws due to the asymptotic flatness of power law curves; we corroborate and extend these results leveraging controlled experimental training of the Gemstone suite.

Cooldown

10% here is referring to token counts—there is no data repetition during training run in any of our experiments. For example, a model trained for 40b tokens and then cooled down, would be cooled down on 4b new tokens. However, it is worth noting that the 4b additional tokens the model would see during cooldown are actually the next 4b tokens that the model would have trained on had the learning rate remained constant. This design allows for a controlled comparison where two “hot” and “cooled down” checkpoints can be compared such the only difference is whether or not the most recent 10% of tokens were trained on at a constant lr, or while the lr was linearly decaying.

Reducing Variance During Fitting

We would like to highlight the reduced variance and better fit found by our convex hull approach as demonstrated in Figure 3 and utilized throughout the paper including the experiments behind Tables 1 and 2. In addition, as shown in Figure 19, we also check that our grid search size and delta in the Huber loss are in a low variance regime when fitting Approach 3 laws to limit the impact of fitting noise on our results. During the rebuttal period, we have also conducted a leave-one-out analysis over the 22 models and added this to the camera ready draft appendix so that the reader can visually contextualize fitting variability caused by model selection.

As we have made significant efforts to incorporate all feedback into our camera-ready version, including addressing concerns over readability of figures, if you feel that your questions have been sufficiently addressed, please kindly consider raising your score.

[1] Team G. Gemma 2: Improving open language models at a practical size. (arXiv)

[2] Hoffmann J, Borgeaud S, Mensch A, Buchatskaya E, Cai T, Rutherford E, Casas DD, Hendricks LA, Welbl J, Clark A, Hennigan T. Training compute-optimal large language models. (arXiv)

[3] Kaplan J, McCandlish S, Henighan T, Brown TB, Chess B, Child R, Gray S, Radford A, Wu J, Amodei D. Scaling laws for neural language models. (arXiv)

[4] Porian T, Wortsman M, Jitsev J, Schmidt L, Carmon Y. Resolving discrepancies in compute-optimal scaling of language models. (NeurIPS 2024)

[5] Hägele A, Bakouch E, Kosson A, Von Werra L, Jaggi M. Scaling laws and compute-optimal training beyond fixed training durations. (Neurips 2024)

We thank the reviewer for their time and insightful comments on our draft. We appreciate their acknowledgement of the fact that the Gemstone model suite itself represents a significant contribution, and are also encouraged by the fact that they find our detailed analysis of the fragility of scaling laws and our findings on the trade off of width and depth when training transformer language models interesting.

Larger Scale Models

We too would like to scale to larger models but unfortunately the compute required increases greatly with model size and so training models larger than 2 billion parameters was simply not possible within our compute budget. We still feel that our work will have a significant impact on the scaling laws community as others have shown laws fitted on models with up to 2b parameters can be highly informative scaling laws [1,2,3]. Moreover, in the Gemma 2 report [3], Table 9 already corroborates our prediction that depth improves benchmark performance more than width; our work demonstrates this trend in an open, scientific and reproducible manner for the NeurIPS community. Hence, we believe our work can still be very impactful for the NeurIPS community despite the fact that the models we release are smaller than those that industry labs often train.

Colab Notebook

We agree that in addition to the model suite and fitting datasets, an interactive colab notebook could greatly increase the accessibility of our data for the community and thereby elevate the impact of the work. We are happy to invest effort into this while preparing the camera ready manuscript for publication.

More Nuanced Discussion Points

We appreciate the reviewer pointing out a potential connection between our work and Allen-Zhu's "Physics of LLMs." We have updated our camera ready draft to include a discussion these nuanced links focusing on how design decisions in “clean room” scientific settings can transfer to larger scale and more realistic experiments.

As we have made significant efforts to incorporate all feedback into our camera-ready version, if you do not have any other questions we can address, please kindly consider raising your score as an even stronger recommendation of acceptance.

[1] Li M, Kudugunta S, Zettlemoyer L. (Mis) Fitting: A Survey of Scaling Laws. (ICLR 2025)

[2] Kaplan J, McCandlish S, Henighan T, Brown TB, Chess B, Child R, Gray S, Radford A, Wu J, Amodei D. Scaling laws for neural language models. (arXiv)

[3] Bhagia A, Liu J, Wettig A, Heineman D, Tafjord O, Jha AH, Soldaini L, Smith NA, Groeneveld D, Koh PW, Dodge J. Establishing task scaling laws via compute-efficient model ladders. (arXiv)

[4] Team G. Gemma 2: Improving open language models at a practical size. (arXiv)

We thank the reviewer for their thoughtful comments and their acknowledgement of the multi-faceted contributions that the Gemstone model suite will make to the scientific community.

Larger Scale Models

We too would like to scale to larger models but unfortunately the compute required increases greatly with model size. Training models larger than 2 billion parameters was simply not possible within our compute budget. Kaplan et al. [1], one of the most widely known scaling laws, also only scale to 1.5 billion non-embedding parameter models, so we feel our work can still have a significant impact on the scaling laws community at this scale. Moreover, Bhagia et al. [2] only use models of up to 2b parameters to predict the performance of 6b and 14b parameter models. Finally, in the Gemma 2 report [3], Table 9, already corroborates our prediction that depth improves benchmark performance more than width; our work demonstrates this trend in an open, scientific and reproducible manner for the NeurIPS community. Hence, we believe our paper will be very impactful to the community.

During the rebuttal period, we have updated our manuscript to include benchmark scaling laws where we fit to models less than 2b parameters and check the predictions against the actual data for our 2b parameter models observing high agreement. However, the rebuttal process constraints mean that we cannot share the new plots directly, so we will simply summarize the results here and then include them in the camera ready copy. We find that the formula proposed by Gadre et al. (fitted on all data in Figure 5) to be more robust to extrapolation than the formula proposed by Bhagia et al. (fitted on all data in Figure 6). The Gadre et al. formula yields near perfect extrapolation to the 2 billion parameter models when fitted on all models with less than 2b parameters.

Using Intermediate Checkpoints

Hägele at al. [4] suggest practitioners should use a warmup-stable-decay scheduler when fitting scaling laws as it allows intermediate checkpoints to be used for fitting scaling laws and offers similar performance to a cosine decay scheduler. The laws we present are fit based on checkpoints taken every 10b training tokens, but we also fit laws on checkpoints every 2b tokens and found these to be the same; hence for efficiency we fit on checkpoints taken every 10b tokens. We have included this ablation fitting on checkpoints every 2b tokens in our camera ready appendix, showing little to no difference to the laws fitted on checkpoints every 10b tokens.

What would we do next?

Given additional time and computational resources, we would first look at the impact of the expansion factor of the transformer, as this is most closely related to the aspect ratio and it would be interesting to compare and contrast to our current findings.

We thank the reviewer again for their effort in engaging with our work so far. We believe that your feedback has been valuable in preparing a more camera-ready version of our manuscript. If you find our responses satisfactory and there are no further questions we can address, please kindly consider moving your score in the direction of an even stronger recommendation of acceptance.

[1] Kaplan J, McCandlish S, Henighan T, Brown TB, Chess B, Child R, Gray S, Radford A, Wu J, Amodei D. Scaling laws for neural language models. (arXiv)

[2] Bhagia A, Liu J, Wettig A, Heineman D, Tafjord O, Jha AH, Soldaini L, Smith NA, Groeneveld D, Koh PW, Dodge J. Establishing task scaling laws via compute-efficient model ladders. (arXiv)

[3] Team G. Gemma 2: Improving open language models at a practical size. (arXiv)

[4] Hägele A, Bakouch E, Kosson A, Von Werra L, Jaggi M. Scaling laws and compute-optimal training beyond fixed training durations. (Neurips 2024)