Evrim Tan

Pretraining LLMs is expensive, but scaling laws allow us to accurately predict the performance of larger training runs from much cheaper ones. Recent research has drastically improved our practical understanding of scaling laws for LLMs... “We fit scaling laws that extrapolate in both the number of model parameters and the ratio of training tokens to parameters.” - from [1] Some background. Prior work has discovered scaling laws [2, 3] for language models, which show that model performance improves as we i) make the model larger and ii) increase the amount of data over which the model is pretrained. However, there are two main limitations of this work that make scaling laws less practical: 1. These works always study the compute-optimal regime [2], but language models are rarely trained in a compute-optimal manner in practice. Usually, we train smaller models over a larger number of tokens (i.e., overtraining) relative to the compute-optimal approach because smaller models are cheaper to use at inference time. 2. Performance is measured in terms of next-token prediction loss when measuring scaling laws, but we performance on downstream tasks is what actually matters in most applications. With this in mind, we might wonder: Do scaling laws still hold for more practical settings? In [1], authors aim to answer this question by training language models with sizes ranging from 0.011B to 6.9B parameters over a variable number of tokens samples from RebPajama, C4, or RefinedWeb. From these models, we can measure scaling laws for overtrained models, as well as study the correlation between model loss and aggregate performance on downstream tasks. What is a power law? To understand scaling laws, we first need to understand the concept of a power law. In its most basic form, a power law is expressed as: y = lambda * x^{-alpha} A power law simply tells us that two variables (x and y) have a relationship—changing x produces a scale-invariant change in y. The shape of this relationship is controlled by the values of constants lambda and alpha. For LLMs, the quantity y is typically the next-token prediction loss on a validation set, while x can be several different quantities of interest. For example, we can study the validation loss with respect to compute used during training, number of model parameters, or amount of pretraining data. By plotting the power laws with respect to each of these quantities, we get a familiar picture; see below. The model’s loss decreases smoothly—according to a power law—as each of these quantities are increased. Why is this useful? As we know, training a language model is an expensive process. The benefit of power laws is that they can help us to mitigate this expense. Namely, we can: - Train several smaller models. - Fit a power law to their performance. - Use this power law to extrapolate and predict the performance of larger (and more expensive) training runs. Using this approach, we can predict the performance of our LLM in advance and determine the optimal training settings without wasting tons of compute on tuning! “To reduce the cost of finding successful training recipes, researchers first evaluate ideas with small experiments and then extrapolate ... to larger scales.” - from [1] If we can extrapolate LLM performance reliably, we can iterate quickly at smaller scales and select the best settings for large-scale training runs. In the literature, this strategy has been used to predict compute-optimal settings by solving the formulation shown below, telling us the optimal model and dataset size to use that achieves the best possible test loss for a fixed compute budget. Overtrained LLMs. Authors in [1] begin by generalizing the compute-optimal training regime to the overtrained regime, where we train smaller models on more tokens. Although this approach is not compute-optimal, it is often adopted in practice due to the simple fact that the compute-optimal training formulation shown above completely ignores the inference cost of the resulting LLM. In most cases, we will happily pay a larger training cost to reduce the size of the model—assuming the model still achieves competitive performance—that we have to deploy. “While loss should be higher than in the compute-optimal allocation for a given training budget, the resulting models have fewer parameters and thus are cheaper at inference.” - from [1] To study overtraining, we can first introduce a ratio called the token multiplier that can be computed as M = D / N. Compute-optimal settings typically have a token multiplier of M = 20. We see in [1] that scaling laws hold for different values of M ranging from 20 to 640; see below. From this analysis, we learn that overtrained models obey similar power laws as those trained in a compute-optimal manner. Going further, we see in [1] that the value of M does not change the slope of the power law! Rather, only the y-intercept is changed, creating the group of parallel lines shown in the figure above. From these scaling laws, we can easily predict the performance of larger, overtrained models. For example, authors in [1] predict the validation loss of 1.4B and 6.9B parameter LLMs trained on 900B and 138B tokens, respectively, from a group of smaller-scale experiments that use 300X less compute. Downstream task performance. Taking the same group of LLMs used to discover scaling laws based upon the validation loss, we can plot their average top-1 error over a set of downstream evaluation tasks4. When the authors in [1] create these plots, we see that top-1 error in downstream tasks exponentially decays with respect to the model’s validation loss on the C4 dataset. This analysis reveals a clear trend in downstream task performance that authors in [1] use to propose a scaling law for top-1 error. This strategy is used to predict the average top-1 error over downstream tasks for 1.4B and 6.9B parameter LLMs using 20X less compute. Empirical evaluation. The first step of experiments in [1] is to train smaller LLMs with different values of M. We can then use scaling laws to extrapolate model performance to larger values of N and M. The main results of these experiments are shown in the figure above, where we see that the performance of larger-scale training runs—in terms of both validation loss and top-1 error on downstream tasks—can be reliably predicted. In particular, we see 0.7% relative error in predicting validation loss and 3.6% relative error in predicting top-1 error. As such, we clearly see in [1] that scaling laws are useful for predicting the performance of large-scale training runs in practice. -------- Bibliography [1] Gadre, Samir Yitzhak, et al. "Language models scale reliably with over-training and on downstream tasks." arXiv preprint arXiv:2403.08540 (2024). [2] Hoffmann, Jordan, et al. "Training compute-optimal large language models." arXiv preprint arXiv:2203.15556 (2022). [3] Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).

Recent AI research has revealed that reinforcement learning (RL)—more specifically, reinforcement learning from human feedback (RLHF)—is a key component of training state-of-the-art language models. But, what is reinforcement learning and why is it so useful for training LLMs? What is RL? Reinforcement learning (RL) is just another way of training a machine learning model. While supervised learning trains a model (using gradient descent) based upon labeled data, RL is based upon trial and error. Similarly to how humans learn, our agent (a neural net) will produce output, receive feedback, then learn from this feedback as it interacts with an external environment. RL Framework. Problems in RL are commonly formulated as a Markov Decision Process (MDP). MDPs have states, actions, and rewards, as well as transition/policy functions. At each time step, our policy takes the current state as input, predicts an action, receives a reward, and proceeds to the next state. In this way, our agent can iteratively interact with an environment and learn a policy that maximizes total rewards. Simple Example. Let’s consider training a neural network to navigate a 2 X 3 grid. The agent receives a reward of +10 for reaching the final state and a reward of -10 for visiting a red square. Our environment is the grid and our state is given by the current position (encoded as a one-hot vector). We can implement our policy as a feed-forward neural net that takes this state as input and predicts a distribution over actions (up, down, left, right). Using RL for LLMs. The above framework can be easily adjusted to the language modeling setup. Our language model is the policy and the state is the current textual sequence. Given this state as input, the language model outputs an action—the next token—to produce the next state—the textual sequence with an added token. After generating a full sequence, we can rate the quality of the language model’s output, either with a human or a reward model, and use RL to learn from this. Why use RL? Typically, the environment in RL is non-differentiable, meaning that the rewards given to our agent are a black box. But, RL allows us to learn from signals that are non-differentiable! Put simply, this means that we can learn from arbitrary feedback on a neural network’s output! Such a property makes RL powerful for LLMs, as we can learn from a variety of subjective sources of feedback that may not be compatible with supervised learning.