Beyond Chinchilla-Optimal: Accounting for Inference in Language Model Scaling Laws

Paper · arXiv 2401.00448 · Published December 31, 2023
Inference-Time Scaling

Large language model (LLM) scaling laws are empirical formulas that estimate changes in model quality as a result of increasing parameter count and training data. However, these formulas, including the popular Deepmind Chinchilla scaling laws, neglect to include the cost of inference. We modify the Chinchilla scaling laws to calculate the optimal LLM parameter count and pre-training data size to train and deploy a model of a given quality and inference demand. We conduct our analysis both in terms of a compute budget and real-world costs and find that LLM researchers expecting reasonably large inference demand (~1B requests) should train models smaller and longer than Chinchilla-optimal. Furthermore, we train 47 models of varying sizes and parameter counts to validate our formula and find that model quality continues to improve as we scale tokens per parameter to extreme ranges (up to 10,000). Finally, we ablate the procedure used to fit the Chinchilla scaling law coefficients and find that developing scaling laws only from data collected at typical token/parameter ratios overestimates the impact of additional tokens at these extreme ranges.

Introduction. Large language models (LLMs) have substantial training and inference compute and energy costs (Knight, 2023; Pope et al., 2022). Training computation costs are primarily determined by the size of the model and the amount of data seen during training (Hoffmann et al., 2022). For state-ofthe-art models with tens of billions of parameters trained on trillions of tokens, training costs can easily exceed millions of dollars. Similarly, inference costs depend on the size of the model and the volume of user queries over the lifetime of the model. This volume can be significant; demand for popular models can exceed billions of tokens per day (OpenAI & Pilipiszyn, 2021; Shazeer & Freitas, 2022).

Accounting for both training and inference, how does one minimize the cost required to produce and serve a high quality model?

Recent studies have proposed scaling laws, empirical formulas that estimate how changes in model and training data size impact model quality (Kaplan et al., 2020; Hoffmann et al., 2022). Hoffmann et al. (2022) is perhaps the most influential of these works, finding that to scale language models most efficiently, parameters and tokens should grow approximately linearly. The authors applied this scaling law to train a 70B parameter model (dubbed Chinchilla) that outperformed much larger and more expensive models such as GPT-3. As a result, many subsequent LLMs have been trained following the Chinchilla scaling laws (Dey et al., 2023; Muennighoff et al., 2023).

However, the Chinchilla scaling laws only account for the computational costs of training. By contrast, the Llama 2 family of models were trained on 2 trillion tokens and the Llama 3 family of models were trained on 15 trillion tokens, which is far more data than the Chinchilla scaling laws would deem “optimal” (Touvron et al., 2023a;b; AI@Meta, 2024). Since inference costs are lower for smaller models, the extra training compute required to train a Llama-style model over a Chinchilla-style model of equivalent quality pays off after enough inference requests.

Related work. Prior work has discussed the training-inference compute trade-off (Touvron et al., 2023a;b; Tow et al., 2023; De Vries, 2023; Villalobos & Atkinson, 2023). Touvron et al. (2023a) cites the lower inference cost of smaller models as inspiration for the LLaMA series. De Vries (2023) calculates the compute overhead of training longer than Chinchilla, but does not discuss quantify compute savings from inference. Villalobos & Atkinson (2023) discuss this trade-off in more detail, but show the shift in scaling laws for only a single particular number of inferences.

Other related work includes Muennighoff et al. (2023), which adapts the Chinchilla scaling laws for the dataconstrained regime, where we have to repeat training tokens.

A few studies have also (gently) critiqued the general parametric function fitting approach of Hoffmann et al. (2022). Besiroglu et al. (2024) attempted to replicate the methodology used in Hoffmann et al. (2022) and found that the confidence intervals reported in the original study were implausibly narrow. The broader implication is that confidence intervals are quite wide for parametric function fitting on a small number of data points. This is also a potentially valid critique of our empirical results, as our analysis only includes 47 separate experiments (Chinchilla included more than 400 experiments). In a similar vein, Porian et al. (2024) investigate discrepancies between the scaling laws from Kaplan et al. (2020) and Hoffmann et al. (2022).

The results presented in Gadre et al. (2024) are particularly relevant to this paper. The authors train 100 models between the sizes of 1.4B and 6.9B parameters and on data with tokens-per-parameter ratios between 20 and 640. Similar to our study, they find reliable scaling laws in these model and data regimes. They also find that downstream task performance is strongly correlated to LLM perplexity.

Method. Our problem setting is the opposite: We assume we are data rich but compute constrained, and seek to minimize computation costs assuming we have enough data to train high-quality models.

In this paper, we modify the Chinchilla scaling laws to account for inference costs by calculating the optimal parameter and training token counts—both in terms of compute (Sec. 2) and dollar costs (Sec. 6)—to train and deploy a model of any given quality and inference demand. Our principled derivation estimates that LLM practitioners expecting significant demand (~109 inference requests) should train models substantially smaller and longer than Chinchilla-optimal. Figure 1 illustrates the benefits of our compute-optimal method for a realistic scenario.

In inference-heavy regimes, our modification predicts that it is optimal to train a model far smaller and longer than Chinchilla predicts, up to thousands of tokens per parameter. Does this hold true in practice? Do transformer models see continued improvements at such extreme cases? Or is there a point beyond which models saturate, and additional tokens provide no further improvement, as De Vries (2023) suggests? To uncover the behavior of transformer models in these cases, we train 47 models ranging from 150M to 6B parameters, on various data budgets from 10 to 10,000 tokens per parameter. We find that model quality continues to improve as we scale token ratios. We do not find evidence of a “saturation point,” beyond which models do not improve, even with additional data.

Lastly, we ablate the parametric curve fitting procedure from the Chinchilla scaling laws. The Chinchilla scaling laws use empirical data from over 400 training runs to determine coefficients that estimate precisely how additional model parameters and training data impact loss. Hoffmann et al. (2022) collected data only from standard token ratio training runs (≤~100 tokens/parameters). Our ablation indicates that when fitting Chinchilla coefficients using only typical token ratio runs, this formula overestimates the impact of additional training data as we move to the extreme ratio regime.

We seek to minimize the computational costs of a model of a given quality and inference demand. We closely follow the methodology in Hoffmann et al. (2022) (henceforth referred to as “the Chinchilla paper”), using pre-training cross-entropy loss as a proxy for quality, and floating-point operations (FLOPs) as our unit of computational cost.

We model our pre-training loss L(N, Dtr) in terms of the number of parameters, N, and pre-training tokens, Dtr, according to the Chinchilla paper’s third scaling law:

The Chinchilla paper derived the parametric loss function in Eq. 1 and fit values for A, B, E, α, and β from the authors’ empirical training results. The best-fit values for these coefficients depend on the exact dataset and model architecture; however, the Chinchilla paper found largely consistent results across the MassiveText, Github (Rae et al., 2022), and C4 (Raffel et al., 2023) datasets, and subsequent work has replicated these scaling laws on other internet corpora and transformer variants (Dey et al., 2023; Besiroglu et al., 2024; Gadre et al., 2024). We use the coefficient values from the Chinchilla paper in our analysis here, and Additionally, we assume that conditioned on pre-training loss, inference demand is independent of model size and token count. In other words, models of equivalent quality but different parameter counts will see the same requests.1 Let TFLOPs(N, D) and IFLOPs(N, D) be the number of FLOPs required to train and run inference, respectively, on a model with N parameters for D tokens. Denote the number of tokens (input + output) of a single inference request i as D(i) inf . Let Dinf = P i D(i) inf be the sum of all tokens over all inference requests.

N ∗and D∗ tr are functions that describe the optimal parameters and pre-training tokens, respectively, that minimize total training and inference compute. The pre-training loss constraint ensures that we minimize compute for a given quality.

Discussion. In Figure 3, we show the major results of our experiments. Our key finding is that loss continues to decrease as we increase tokens per parameter, even to extreme ratios. Although it takes exponentially more tokens to reduce loss at large ratios, loss does not plateau as we scale to 10,000 tokens per parameter for our 150M model. For models larger than 150M, we only test up to a maximum of 1,000 tokens per parameter due to resource constraints, and at this scale we also see no evidence of loss flat-lining. Further experiments are needed to see how loss scales beyond this point.

Existing literature claims there exists a critical model size, below which we cannot match a Chinchilla-optimal model’s quality (De Vries, 2023). Since our experiments do not see any plateauing of loss as we scale token ratios, we find no evidence to support the critical size hypothesis, although further testing is needed at extreme scales—it is possible that beyond 10,000 tokens/parameter, behavior changes.

We also show results for downstream metrics. Fig. 3b shows that the Gauntlet Average improves as we increase tokens per parameter — again, we find no evidence of a “saturation” point, beyond which additional tokens do not result in better performance. In fact, as Fig. 3c shows, as loss decreases, smaller decreases in loss lead to larger improvements in downstream accuracy.

In Fig. 3c, we plot the Gauntlet Average as a function of loss. Loss and Gauntlet Average are tightly correlated, showing that improvements in loss are excellent predictors for improvements in general model quality. Our results show that practitioners interested in predicting downstream metrics as a function of model parameters and token counts can make use of existing scaling laws to accurately understand how their downstream metrics change at scale. best fit are nearly parallel, indicating that models learn with similar efficiency. This result is striking, as it indicates that loss decreases at similar rates for each additional FLOP regardless of whether training occurs in the standard Chinchilla regime or at extreme ratios (500 tokens/param). Our results also confirm that training at low (< 20) token-perparameter ratios is less computationally efficient, in line with the results from the Chinchilla paper.

Our results are shown in Table 1 and visualized in Figure 5. First, we see that our parameteric curves follow a consistent trend—as we use more extreme data in our fitting procedure, our scaling curves become flatter. This trend suggests that if we only use data from typical token ratios to determine our scaling law coefficients, we will overestimate the impact of additional training data as we move towards the longdata-ratio regime. Taken together with our results from Sec.

4, our experiments show that although models continue to learn at extreme training durations, they do so more slowly than scaling laws predict.

We note that our experimentally-derived coefficients differ from those found in the Chinchilla paper. This is expected, as both the data and the model architectures we used to train all checkpoints are quite different from the original work (Hoffmann et al., 2022). Besiroglu et al. (2024) noted in a replication study that the coefficients reported in Hoffmann et al. (2022) were rounded in ways that leads to significant bias in scaling law predictions.

These results have significant impact as the field continues to train longer and longer. For example, LLaMA 7B was trained on 1 trillion tokens (Touvron et al., 2023a), Llama 2 7B was trained on 2 trillion tokens (Touvron et al., 2023b), and Llama 3 8B was trained on 15 trillion tokens (AI@Meta, 2024). How much of the quality difference between these models is due to simply training on more data?

Conclusion. In this work, we modify the Chinchilla scaling laws to account for both the computational and real-world costs of inference. As inference demand approaches pre-training data size, the additional cost pushes the optimal tokens-toparameters ratio towards smaller and longer-trained models.

We experimentally validate the hypothesis that very small models, trained on enough data, can match larger Chinchillastyle ones. Both in terms of loss and downstream metrics, model quality improves as tokens per parameter increase. Further work is needed to show if this scales beyond 10,000 tokens per parameter, or at larger model sizes. Still, our results show that practitioners operating in inference-heavy regimes, or with limited deployment memory, can scale training duration considerably longer than current literature suggests and see quality improvements.

Finally, we show evidence that the Chinchilla parametric coefficient fitting procedure overestimates the reduction from additional training data when applied to extremely dataheavy training runs. More work is needed to develop scaling laws that apply precisely at a wide range of ratios.

Limitations. Further work is needed to fully characterize scaling laws Accounting for Inference in Scaling Laws Eq. 6 is a simplified model of real-world costs: we leave aside latency requirements and assume MFU and cost per FLOP do not depend on model size, configuration, or sequence length. Still, our approximation is flexible enough to account for heterogeneous hardware utilization and costs.

Lines of inquiry this paper opens 24

Research framings built by reading the notes related to this paper — the questions it feeds into.

What fundamentally drives AI inference pricing and cost decline patterns? How does model scale affect fine-tuning and downstream task performance? Why does separating planning from execution enhance reasoning performance? Can structural constraints replace embedding depth in collaborative filtering? How does scaling test-time compute compare to increasing model size? Does parallel reasoning outperform sequential approaches in reasoning efficiency? How do diffusion language models compare to autoregressive architectures? Can AI systems autonomously pursue unstated objectives or circumvent design constraints?