MLE, Log-Likelihood and Gradient Descent

So we'd like to figure out how to find the values of the many parameters in our model, i.e. of all the elements of our token embeddings and all the weights and biases of our neural network.

Recall that our goal in optimizing these parameter values is to have the initial next-token probabilities our model produces for any given context mimic the empirical ones found in the training data as well as possible, with the phrase "as well as possible" being used very purposefully here because unlike an n-gram model, which simply introduces a new probability for any new n-gram encountered in the training data, our new neural-network-based model has a fixed size and is by no means guaranteed to be able to match all the empirical probabilities even if its parameters have been optimized.

We have no idea how to perform such an optimization in one fell swoop, so we will likely have to use a more incremental approach, in which we improve the parameter values bit by bit.

This leaves us with two subgoals:

1. We need some quantitative way to evaluate how well we are mimicking the data for a given set of parameter values, which we can use as a guidepost in our improvement journey. In more technical terms, we need to define an Objective Function whose value we are trying to optimize.

2. We need to figure out how to perform that optimization. Meaning once we have defined our objective function, we need a Learning Algorithm for how to adjust our parameters to optimize its value.

Maximum-likelihood estimation (MLE)

Let's start with finding a suitable objective function. It's not obvious how to go about this and the task seems intimidating, so let's apply our proven strategy of starting with a simplified toy problem and go from there.

Coins, dice, and maximizing probability

The core of our problem is that we have a paremeterized model outputting probabilities and would like to adjust it to mimic the probability distributions found in a bunch of data.

Perhaps the simplest parameterized probabilistic model out there is that of a Bernoulli distribution modelling a weighted coin. This model has no inputs (or, if you prefer to think of it that way, always has one and the same input) and its only parameter is the probability that the coin lands on heads:

Let's try using this as our toy model. Say we were handed a bunch of "training data", showing that an unknown coin had been tossed 10,000 times, resulting in 4,327 heads and 5,673 tails.

How would we adjust the paremeter $p$ for our model's output to mimic this data?

Well, we would simply set $p = 0.4327$, since this is the value that will make the probabilities $p_{\text{model}}(\text{heads})$ and $p_{\text{model}}(\text{tails})$ of our model equal the empirically observed probabilities $p_{\text{data}}(\text{heads}) = 4327/10000$ and $p_{\text{data}}(\text{heads}) = 5673/10000$ from the training data.

Unfortunately, this was too easy to be helpful — the problem here is that we essentially recreated the n-gram scenario where we directly paramaterized the probabilities themselves, so we didn't naturally come across any sort of objective function that we might be able to use for our optimization simply because we were able to directly jump to the ideal end result of perfectly matched probabilities without even thinking about it.

But as mentioned before, it is not a given that this result is even achievable — a parameterized model might simply be unable to perfectly match the empirical probabilities exactly, even when its parameters have been optimized. This scenario seems like it might force us to come up with something new, so let's find a different toy model for which we can't perfectly match the probabilities and see where our instincts lead us then.

I'd say the second-simplest "probability thing" after a coin is a die, so let's use the probability distribution associated with the rolling of a die and parameterize it in such a way that we do not have full control over all the individual probabilities. A simple way of doing this is to make the number of sides $n$ of the die a parameter, so $n = 6$ for example would correspond to the probability distribution associated with a regular, six-sided die. Here is what that the resulting model looks like:

Now let's once again assume we were handed a bunch of "training data", this time taking the form of a sequence of random numbers generated by the rolling of some unknown die, and we would like to model this data as best we can with our model. Say the sequence looks like this:

In this scenario, what would be the optimal value of $n$ that best models the training data?

Three different dice with $n\!=\!6, 8, 10$.

Three different dice with $n\!=\!6, 8, 10$.

I'd say it's clearly $n = 8$. This is because the highest number in the training-data sequence is 8, so if the data were indeed generated by a die, that die would have to have at least 8 sides. Furthermore, while the sequence could have been generated by repeatedly rolling a die with more than 8 sides, it would simply be less probable to result from a 10- or 12-sided die than it would be to result from an 8-sided one, so assuming a value higher than 8 would, in a sense, be tantamount to assuming additional coincidences beyond the ones reflected in the data.

Thus, instinctively, $n=8$ is the most likely, and in that sense optimal, parameter value for our model simply because it is the value that maximizes the probability with which our model would generate the training data if we were to use it in a generative fashion.

In statistics, this type of "most likely" parameter value is known as the Maximum Likelihood Estimate and the process of finding it is known as Maximum Likelihood Estimation, both of which we will abbreviate as MLE in the following.

This seems interesting. What if we tried to apply the MLE approach to a case where it is possible to perfectly match the probabilities? Well, let's revisit the coin toss example and see what would happen there.

If $p_{\text{model}}(\text{heads}) = p$ and $p_{\text{model}}(\text{tails}) = 1 - p$, the probability that our simple weighted-coin model would generate 4,327 heads and 6,673 tails upon 10,000 draws is simply

Taking the derivative with respect to $p$, setting said derivative equal to 0 and solving, we find that this is maximal exactly for $p = 0.4327$, which is also the parameter value for which $p_{\text{model}}(\text{heads}) = p_{\text{data}}(\text{heads})$ and $p_{\text{model}}(\text{tails}) = p_{\text{data}}(\text{tails})$.

In other words, maximizing the probability that our model would generate the training data upon 10,000 draws from its output distribution yields the exact parameter value for which the model probabilities match the empirical probabilities reflected in the training data!

Thus, in the toy problem where exact probability matching was achievable, maximizing the probability with which our model would generate the training data got us there, and in the toy problem where such matching was not achievable, it got us as close as possible.

MLE and its perks

Combined with its general intuitiveness, this crushing success of the MLE approach in our primitive toy problems alone would warrant choosing it as our method of choice for now, but it should be noted for the record that there are more fundamental statistics arguments that can be invoked to further strengthen the MLE case.

We will avoid going down the rabbit hole of examining these arguments all too closely, but for the type of situation our toy problems represent, where the dataset consists of a sequence of independent samples from the same underlying probability distribution over a set of discrete outcomes, it can be shown that the maximum likelihood estimate has several very desirable properties so long as some basic regularity conditions are met:

1. Matching empirical probabilities: Generally, the probabily distribution that makes such a training dataset the most probable is the empirical probability distribution it reflects. So if there is some set of finite parameter values for which the model probabilities match the training-data probabilities exactly, finding the MLE will get us there. This is why finding the MLE resulted in model probabilities which perfectly matched the training-data probabilities in the coin toss example.

2. KL minimization: If the model cannot match the training-data probabilities exactly, finding the MLE corresponds to minimizing a quantity called the Kullback–Leibler (KL) Divergence between the model probability distribution and the empirical distribution. Explaining the KL divergence in detail would lead us too far astray here, but it is essentially a proper mathematical quantity measuring how "far off", in an informational sense, a probability distribution is from a second reference distribution. So saying "finding the MLE minimizes the KL divergence of the model distribution with respect to the data distribution" is a more rigorous way of saying "finding the MLE gets you as close as possible to mimicking the training data", which is exactly what we observed in our die-rolling toy problem.

3. Asymptotic consistency: If the model is capable of fully describing the unknown, data-generating probability distribution, it can be shown that the MLE converges towards the true parameter values (i.e. the ones where the model distribution exactly matches the data-generating one) in the limit of an infinitely large dataset. In statistics speak, the MLE is a consistent estimator of the true parameter values. This is basically just a restatement of the empirical-probability matching for the infinite-data limit where those empirical probabilities approach the actual underlying probabilities.

4. Asymptotic efficiency: In the same scenario, it can be shown that the MLE is asymptotically efficient. This means that as the dataset gets larger and larger, the variance we would observe in the MLE of the parameter values if we took a bunch of datasets of the given size and checked how much the MLE fluctuates converges towards being as small as theoretically possible. This is very desirable because we ultimately only get one huge training dataset, and so knowing that the MLE varies (close to) as little as possible from dataset to dataset means that it is a good tool to estimate the parameters in practice, not just in some hypothetical "it would get us somewhere reasonable if we could average all our estimates over a million datasets" sense.

In contrast to our toy models, our language model does not represent a simple probability distribution from which tokens are drawn independently. Rather, our language model is a function representing the conditional probability distributions associated with all possible contexts, and if we use the model in a generative fashion, each individual token is drawn from a different probability distribution, namely the one associated with the context preceding it. Generally, the degree to which the above results hold for this type of model depends on the nature of the training data and the training process and things get messy pretty quickly.

But even if they do not hold exactly, seeking to find the MLE of our parameter values seems very natural, works like a charm in simple toy problems, and has a lot going for it conceptually, so let's roll with it and officially make maximum-likelihood estimation our approach for optimizing our model's parameter values.

Applying MLE to our model

Unlike our initial, vague goal of making the model probabilities "match the empirical ones as well as possible", our MLE goal of finding the parameter values which maximize the probability of the training data under our model immediately suggests a natural objective function we can use as a guidepost throughout optimization: the training-data probability itself.

If we simply consider the entirety of our training data one giant stream of tokens and ignore the very beginning of it where we have too few tokens to feed a full context to our model†Ignoring this tiny fraction of our training data shouldn't matter at all, but also recall that figuring out how to deal with variable context lengths is already on our to-do list for later anyway.[note], we can simply apply the standard multiplication rule for the probabilities of successive events and write this training-data probability as

where $T$ is the total number of tokens in the training data, $t_i$ is the $i^\text{th}$ token in the training data (not to be confused with the $i^\text{th}$ token in the vocabulary, which we sometimes use the same notation for), and $L$ is the length of our context window, so that $p_{\text{model}}(t_i \,|\, t_{i-L},\, t_{i-L+1},\,\dots,\, t_{i-1})$ is simply the probability that our model would generate the $i^\text{th}$ token in the training data next if fed the previous $L$ training-data tokens as a context.

Of course, considering the entirety of the training data one continuous stream of text makes limited sense. It will necessarily contain lots of document/website/text boundaries across which we would not want the model to train next-token prediction since predicting a token in text B based on a context from text A makes no sense if the texts aren't closely related.

So to allow for this while also simplifying our expression, let's simply say we identified a large number $N$ of full-length contexts $c_i$ in our training data whose subsequent tokens $t_{c_i}$ we wish to predict.

The probability with which our model would generate this (very large) subset $\mathcal{D}$ of the training data can then be written as

Let's tweak this notation a bit: In the field of machine learning, the parameters of a model are usually thought of as the elements of a parameter vector $\vct{\theta}$, and the model's probability distribution is commonly referred to as $p_{\theta}$, simply because that is very compact while explicitly reminding us that the distribution depends on the parameter values:

Another convention in the field is that the quantity to be optimized is commonly referred to as the Loss Function or simply the Loss, rather than the "objective function". Even though it does depend on the data, this loss function tends to be written simply as $L(\vct{\theta})$ rather than as $L_{\theta}(\mathcal{D})$ or $L(\vct{\theta}; \mathcal{D})$, simply because people are lazy, the dataset is usually fixed, and it's the dependence on the parameters that we care about for optimization.

As the name suggests, the loss is always defined in such a way that the goal is to minimize it, not maximize it. This is just a convention and does not change anything conceptually, its just one of many ways to talk about what we are doing:

To satisfy this convention, we can simply define our loss function to be the opposite of the above training-data probability:

So far, so good! We have now used the idea of maximum-likelihood estimation to define a conceptually meaningful loss function. If we can figure out how to minimize its value, our model will mimic the training data as well as possible.

Negative log-likelihood (NLL) loss

Unfortunately, there is one problem with this loss function: Its value cannot be calculated in practice.

To see this, say we are at the very beginning of training, when no learning has taken place yet and the parameters simply have whatever values we chose to initialize them with. At this point, the probabilities our model outputs are essentially random, so a very rough estimate for the probability it assigns to any next-token continuation from the training data would simply be $p_\theta(t_{c_i} \,|\, c_i) \approx \frac{1}{V}$, where $V$ is the number of tokens in our vocabulary.

Say we have about $V \approx 50000$ tokens in our vocabulary, then even if our dataset only consisted of 1000 context-plus-next-token pairings (never mind the trillions we would ultimately like it to consist of), this would yield

which is absurdly tiny. For reference, the smallest positive number we can store in a standard 32-bit float format is $2^{-149}\approx 1.4 \cdot 10^{-45}$ and the smallest positive number we can store in a standard 64-bit float format is $2^{-1074}\approx 4.9 \cdot 10^{-324}$.

In other words, we simply cannot calculate the actual probability with which the model would generate the full dataset. Fortunately, this does not have to be a death-blow to our hard work so far. Remember that we do not care about the value of the probability itself, we only care about finding the point in parameter space where it is maximized. In other words, we can apply any sort of scaling or transformation to our loss function that leaves the location of its maximum unchanged, and minimizing that new loss function would still correspond to performing MLE!

Perhaps the most intuitive way to change our loss function such that it doesn't result in ludicrously tiny numbers anymore would be to simply scale it up. In order for the product of many, many numbers not to blow up or shrink dramatically, their geometric mean has to be close to 1. So a simple approach towards making the magnitude of the loss more palatable would be to try to calculate or estimate the geometric mean $\bar{p_{\theta}}$ of all the next-token probabilities at our current point in parameter space and to divide every single next-token probability vaue $p_\theta(t_{c_i} \,\big|\, c_i)$ in our loss calculation by $\bar{p_{\theta}}$.

Unfortunately, this is not a good solution. If we manage to actually have our model learn and increase the training-data probability, the required normalization would quickly change and we would suddenly start having the opposite problem, where the value of the loss starts becoming unconceivably large because we are essentially multiplying a giant bunch of numbers whose geometric mean is larger than 1. And if we start redoing our normalization regularly and messing with our loss values, we are starting to defy the very idea of having a loss function across parameter space whose value tells us how well we are doing and things get messy fast.

But wait! If we are already considering using the geometric mean of all the next-token probabilities for normalization, can't we perhaps just use that geometric mean itself as our loss? After all, if our loss value is too tiny because we are multiplying a giant bunch of numbers, the simplest way to counteract that would be to take a giant root:

Conceptually, this would actually work! It would make the magnitude of our loss value way more palatable and, crucially, would not change the location of the loss minimum in parameter space. We know this because $\sqrt[N]{x}$ is an increasing function, meaning it will always result in larger outputs for larger inputs. Thus, the maximum value of $p_\theta(\mathcal{D})$ is guaranteed to be mapped to the maximum value of $\sqrt[\scriptstyle{N}]{ p_\theta(\mathcal{D}) }$, meaning the minimum value of $- p_\theta(\mathcal{D})$ is guaranteed to be mapped to the minimum value of $-\sqrt[\scriptstyle{N}]{ p_\theta(\mathcal{D}) }$.

Unfortunately, this is still a mess to calculate though. If we tried to calculate it as written, by first calculating $p_\theta(\mathcal{D})$ and then taking the $N$-th root, we would literally have the same problem as before in that we can't calculate $p_\theta(\mathcal{D})$, and even if we try to do in a slightly smarter recursive fashion, it is hard to get around the taking of gigantic roots, which also is a huge pain and quickly becomes prohibitive.

All in all, multiplying a huge number of values simply seems destined to produce either rapidly vanishing or rapidly exploding results, which require gigantic roots to be taken unless an extremely fine balance is struck and the geometric mean of all the numbers happens to be extremely close to 1 in the first place.

This brings about the question if we can somehow change our loss such that it is no longer a product of a giant bunch of numbers without impacting the location in parameter space where it is minimized (remember: we still want to do MLE!).

The standard way to turn a product of things into a not-product of things is to apply a logarithm. This is just high-school math:

This simple observation suggests that we should consider using the logarithm of our training-data probability as our loss, rather than the training-data probability itself:

Crucially, this too leaves the position of the loss minimum untouched, for the same reason taking the $N$-th root did — $\log(x)$ is an increasing function, so whatever parameter vector $\vct{\theta}$ minimizes $-p_\theta(\mathcal{D})$ will also minimize $-\log(p_\theta(\mathcal{D}))$.

Converting the product of all the next-token probabilities to a sum of their logarithms immediately makes the resulting loss values more palatable. For example, if we revisit our earlier estimate and assume that all the next_token probabilities are roughly equal at the very beginning of training, meaning $p_\theta(t_{c_i} \,|\, c_i) \approx \frac{1}{50000}$, then $\log(p_\theta(t_{c_i} \,|\, c_i)) \approx -10.8$ (assuming we are using the natural logarithm, as is custom in math/physics). So even with a trillion (i.e. $10^{12}$) tokens the overall loss value would only be about $L(\vct{\theta}) \approx 10^{13}$.

This is a storable number, but it's still annoyingly large to think about or work with, so let's simply divide it by the number of tokens $N$:

This is called (Per-Token) Negative-Log-Likelihood Loss, and is often simply referred to as the (Per-Token) NLL or NLL Loss for short.

If we were married to ultimately calculating the geometric-mean loss $L_\text{GM}(\vct{\theta})$, calculating the NLL loss first would actually be the smart way to do it, because we can simply exponentiate the latter to get the former:

However, this exponentiation only adds another computational step without adding anything of value to our endeavour, and working with sums is much more convenient than working with exponents of sums, so let's just use the NLL loss directly.

Gradient descent

So far so good. We now have a clearly defined loss function whose value we are trying to minimize. Let's figure out how to do that.

We cannot find the minimum analytically

Knowing calculus, our first instinct might be to attempt what we did when finding the MLE of the coin-toss data in the previous section, that is, to take the derivative of our loss, set it equal to 0, and solve the resulting equation(s) to find the location of the minimum in one fell swoop.

Unfortunately, this is hopeless in the case of our actual model, as we ultimately want it to have billions of parameters, and we have no clue how to solve the resulting system of billions of differential equations analytically. In other words, we have arrived back at the same conclusion we already intuited at the very beginning of the chapter — we cannot do it all at once and will instead have to improve our parameter values numerically, in an incremental fashion.

Let's try to think through how we might be able to pull this off.

Lunch money

At the very beginning of training, our model's parameters will simply have whatever values they were initialized with. Right off the bat, this presents us with the interesting and important problem of how to best initialize them, but we are trying to focus on the optimization aspect right now, so let's just assume we have drawn the initial parameter values from a standard normal distribution and add the initialization issue to our list of problems to revisit later:

For whatever initial parameter values we do end up with, we can readily calculate the corresponding loss. But once we have done that, how can we improve?

It seems that if we assume absolutely nothing about how the negative-likelihood loss behaves as a function of our model parameters, there isn't all that much we can do and we are essentially stuck with trial and error, that is, with simply trying different sets of parameter values and seeing which one results in the lowest loss. This is because if we don't assume any correlation between the loss values corresponding to different sets of parameter values, then the loss value affiliated with the current set of parameter values contains zero information regarding what other values might be better, and even if we have been finding, say, lower and lower loss values by increasing a parameter value, that does not imply at all that increasing it further would keep this trend going.

This is called the No-Free-Lunch Theorem and is easily visualized and/or proven by assuming we are dealing with a discrete loss function, i.e. both the parameter values and the loss values can only be varied in a discrete fashion (as is the case in the real world, since we can only store numbers with finite precision):

Single-parameter example: In the absence of any assumptions about the loss landscape, the loss at previously sampled parameter values contains no information about that at unsampled ones, and the minimum of the loss is just as likely to be found at $\theta_A$ as at $\theta_B$.

Single-parameter example: In the absence of any assumptions about the loss landscape, the loss at previously sampled parameter values contains no information about that at unsampled ones, and the minimum of the loss is just as likely to be found at $\theta_A$ as at $\theta_B$.

The no-free-lunch theorem implies that no learning algorithm is generally superior to any other when averaged across all thinkable loss landscapes, meaning that if an algorithm works better for one loss function, it must work more poorly for another. So in some sense, designing a good optimization algorithm is all about tailoring your algorithm to the character of the loss function at hand.

Generally, we have very little knowledge of what the loss function actually looks like in parameter space for models with millions or billions of parameters — the loss simply depends on way too many parameters to be mapped out computationally. This is why inventing better learning algorithms in machine learning, like most things in the field, ultimately tends to boil down to making educated guesses and seeing what works best empirically.

That being said, we do know something about the loss, namely its functional form, so we might be able to derive some reasonable assumptions from there.

For one, we can tell from the definition of our loss function that it is is continuous, meaning any infinitesimal change in the parameters will result in only an infinitesimal change in the loss. Now technically, there is some minimum size to any change we can make because the parameter values will be stored in whatever computer we use at finite precision, but unless we have reason to believe otherwise, it seems safe to assume that the scale at which we can adjust the parameters is smaller than the scale on which the loss varies significantly (and if this weren't the case, we could always increase the precision at which we store our parameters).

Now even if the loss-value changes are tiny if we take sufficiently small steps in parameter space, they could in principle still be unpredictable from one point to the next, for example if the loss had some sort of infinitesimal "jitter" on an even smaller scale than we can sample with our parameter values. But we wouldn't expect that either, of course — we'd intuitively expect the way the loss changes to be somewhat "steady" on the tiniest scale we can access. In other words, we'd expect not only the scale on which the loss itself varies but also the scale on which the loss's slope varies to be much larger than the smallest scale we can sample.

But if that is true, we should be able to use said slope at our initial point in parameter space to point us towards lower loss!

Of course, there is no guarantee that moving towards lower loss from our starting point will lead us to an actual global minimum of the loss (for example, the loss could just be decreasing towards a local minimum while the global one lies behind a "hump" we would be moving away from), but generally speaking, it certainly seems better to decrease the loss than to increase it.

Let's try to think this through in detail.

A 1D toy problem

To start simple, let's figure this out in 1D first. Assume we have a single-parameter model outputting a probability distribution and, as a result, a loss function $L(\theta)$ which depends on only a single parameter $\theta$.

If that parameter has been initialized to some value $\theta^{(0)}$, the initial value of the loss will simply be $L(\theta^{(0)})$.

The starting point.

The starting point.

We'd like to use the slope of the loss function at this point to figure out whether we should increase or decrease the value of $\theta$. For this 1D toy problem, we could do this numerically by giving the parameter value a tiny nudge and checking how much the loss increases, but we actually know the functional form of the loss $L(\theta)$, so assuming that the function is differentiable, we might as well calculate the slope exactly:

If the value of this derivative is positive (meaning the loss is increasing at our current parameter value), we would then conclude that we should decrease the value of our parameter, whereas if this derivative is negative (meaning the loss is decreasing at our current parameter value), we would conclude that we should increase the parameter value.

But by how much should we increase or decrease the value? The simplest choice would be to choose some fixed step size $\delta$ and to update our parameter to its new value $\theta^{(1)}$ as follows:

where

For example, if the slope at our initial parameter value in our earlier example were negative, this would look somewhat like this:

After the first learning update.

After the first learning update.

Once we have $\theta^{(1)}$, we could then either choose to apply the same rule again based only on the slope of the loss at the new parameter-value $\theta^{(1)}$, or we could try to make use of all the information we have about the slope of the loss and somehow take into account both $L'(\theta^{(0)})$ and $L'(\theta^{(1)})$ in deciding how to take the next step.

This seems like a pretty general question — should we only use local information in our learning updates or should we take into account previous updates as well? It's not immediately clear. The slope from the current parameter value certainly seems more relevant than the slope at any other point, but that doesn't mean the information from the other points has no value at all. For example, what if we have been consistently trending in one direction for many, many updates and now all of a sudden arrive at a point sloping the other way? Should we immediately turn around or keep going for a while because it might just be a bump in the road (or, in more mathematical terms, because we might have passed through a local minimum that is not the global minimum)?

As promised by the no-free-lunch theorem, the answer generally depends on the nature of the loss landscape and we will likely have to figure out what's best empirically down the line. So let's focus on finding a good local update rule for now and add the inclusion of information from previous updates to our growing list of things to revisit later:

Applying the above rule many times in a row will have us descend towards lower and lower loss in parameter space, with our steps in parameter space always being of the same size $\delta$ irrespective of the magnitude of the derivative.

This might feel like an intuitive start, but it is actually somewhat ill-motivated.

For one, even if our attention was to be equally careful in adjusting our model on every learning update, simply taking fixed-size steps in parameter space would be unlikely to achieve that. This is because the sensitivity of the model's output probability distribution on the parameter is likely to vary with the value of the parameter itself. So a fixed-size step in parameter space might correspond to smaller or larger changes in the model distribution depending on where in parameter space we are.

Secondly, fixed steps will never have us converge towards any parameter value at all – they would simply have us jump back and forth in parameter space as soon as the sign of the loss function's slope changes:

Taking fixed-size steps will have us jump back and forth across a minimum. No convergence.

Taking fixed-size steps will have us jump back and forth across a minimum. No convergence.

So in order to actually converge, we need to ensure that our steps somehow get smaller once we get closer to the minimum. Given that the only information we have is the slope at our current point in parameter space and that this slope should vanish towards the minimum, this seems to suggest that we should somehow translate a smaller slope into a smaller step size.

But how exactly should we do that?

Ultimately, what's best is once again bound to depend on what the loss landscape looks like, and it's hard to predict any general patterns for that far away from the minimum, meaning for the beginning of training when we have very high loss and our model is far from being optimized.

Closer to the minimum, however, we do know what the loss should look like solely based on what a minimum looks like — if we can get close enough to the minimum, the loss should be well approximated by a parabola.†This is visually intuitive, but it also follows algebraically from the Taylor expansion around the minimum.[note] So let's use this knowledge and see if we can derive more insights from there.

Sufficiently close to the minimum, the quadratic approximation yields

where $\theta_{\text{min}}$ is the target parameter value for which the loss is minimized, $C$ is a constant representing the loss function's curvature near the minimum and $L_{\text{min}}$ is the minimal loss that can be achieved.

The derivative of the loss at some parameter value $\theta_i$ close to the minimum is then

But wait! This is directly proportional to the "perfect step" $\theta_{\text{perfect}}$ that would immediately get us to the minimum starting from $\theta^{(i)}$:

This suggests that we should not just make our step in parameter space "smaller when the slope is smaller, and larger when the slope is larger", as we intuited above, but that at least close to the minimum, making the two directly proportional actually seems ideal.

Of course, we still need to be careful and make sure our steps aren't too large, since we do not know the value of the curvature $C$ and we do not want to be too bold far away from the minimum either. Given that we do not have any way to know the perfect scale, let's just multiply our derivative by some small factor $\epsilon$ whose size we will have to figure out empirically:

where

In machine learning, this number $\epsilon$ is called the Learning Rate, as it scales the size of an individual "learning step" in parameter space.

Choosing the right value of $\epsilon$ is critical. If we make it too small, learning takes a prohibitive amount of updates and our model cannot finish learning in a realistic amount of time. If we make it too large, some learning updates might have us go too far and take us way past a minimum into a region where the loss gets worse again, which usally results in an escalating situation and diverging loss.

For example, if we unknowingly happened to choose $\epsilon \approx \frac{1}{C}$, we would get $\Delta \theta = -\frac{1}{C}\, L'(\theta^{(i)}) \approx 2 \,\theta_{\text{perfect}}$ close to the minimum and jump all the way to the other side of the minimum, ending up about as far away as we started. If we accidentally chose $\epsilon > \frac{1}{C}$, we would overstep even harder and actually cause the loss to increase with each step as we keep jumping across the minimum, meaning we would never reach the minimum and might even cause the loss to diverge.

Here is a little widget making this more palpable:

The sensitivity to the learning-rate value illustrates that making our updates proportional to the slope of the loss, even with just a single parameter, is not withouts its risks. But the benefits seem to outweigh the cost. After all, if we know that close to the minimum the ideal step size is proportional to the derivative, our update rule should reflect that.

Plus, we are under no obligation to always use the same exact value for our learning rate $\epsilon$ throughout training, so if we are ever worried about this rule resulting in too large or too small updates far away from the minimum, we could compensate by ramping the learning rate up or down at various points of training. In the field, this is referred to as using a Learning Rate Schedule.

Let's keep this in mind for now and see what happens when we move to a model with more parameters.

More dimensions

If we have more than a single parameter, our parameter space is multidimensional.

This immediately presents us with a challenge — in the 1D case, there was one direction at any given point in which the loss decreased and one in which it increased, and we could simply use the slope to figure out which one was which. But as soon as we have more than one parameter, there are infinitely many directions we could move in, so there will be infinitely many directions along which the loss decreases, just like there will be infinitely many along which it increases.

A priori, the only thing that seems clear is that we should increase each parameter $\theta_i$ whose partial derivative $\pder{L}{\theta_i}$ is negative and decrease each parameter whose partial derivative $\pder{L}{\theta_i}$ is positive, but that still leaves infinitely many directions for us to choose from.

Two-parameter example: Many directions of loss decrease

Two-parameter example: Many directions of loss decrease

So before we consider the overall magnitude of the step we wish to take in each learning update, let's try to figure out what direction we would like to move in.

Ideally, we'd like to move in the direction of the global minimum, but we have no idea what direction that is. Assuming we are still disregarding the less immediately useful information from previous learning steps for now, all we have to base our decision on are the derivatives of the loss with respect to the various parameters at our current location in parameter space.

Given that this local slope information is all we have and that we ultimately want to minimize the loss, the next best thing seems to be to step in the direction along which the loss landscape falls off the most rapidly at our current point in parameter space, meaning the one for which the loss will decrease the most if we take a fixed-size step. This is commonly referred to as pursuing Steepest Descent in parameter space.

Two-parameter example: Directions of steepest descent at various points across a loss landscape

Two-parameter example: Directions of steepest descent at various points across a loss landscape

It is by no means guaranteed that steepest descent will be a successful, let alone optimal strategy.

Moving in the direction where the loss landscape drops off the most swiftly could completely mislead us in the long term depending on the loss landscape, and even if it does ultimately lead to the global minimum, it could lead us on a circuitous path through parameter space wasting lots of compute.

However, it's a very intuitive approach and we have already established that there cannot be any other approach which would always be optimal (no free lunch!), so let's attempt steepest descent for now and see if we can figure out how to determine the steepest-descent direction in parameter space based on our first derivatives.

To do so, let's consider a small step $\vct{\Delta \theta}$ in parameter space, consisting of some tiny change $\Delta\theta_1$ to the value of the first parameter $\theta_1$, some tiny change $\Delta\theta_2$ to the value of the second parameter $\theta_2$, and so on:

If all of these individual parameter-value changes are sufficiently small, simple linear approximation suggests that the change of the loss due to the update of the first parameter will be roughly $\pder{L}{\theta_1} \Delta\theta_1$, the change of the loss due to the update of the second parameter will be roughly $\pder{L}{\theta_2} \Delta\theta_2$, and so on. So the overall change of the loss as a result of our step in parameter space will be

This sum can be rewritten as the scalar product of the overall step $\vct{\Delta\theta}$ and a vector whose elements are the partial derivatives of $L(\vct{\theta})$ with respect to all the individual parameters $\theta_j$.

This vector of partial derivatives is called the Gradient of the loss with respect to the various parameters and is the natural extension of the derivative $L'(\theta)$ to higher dimensions. It would usually be written as $\gradL (\boldsymbol{\theta})$ in mathematical notation, but we will simply refer to it as $\vct{g}$ because it will pop up a lot and we are lazy:

Since a scalar product of two vectors is maximized when they point in the same direction and minimized when they point in opposite directions, the fact that the change in loss for a small step $\vct{\Delta\theta}$ can be written as

means that for a fixed-size step, the loss will be maximally increased if $\vct{\Delta\theta}$ points along $\vct{g}$ and maximally decreased if $\vct{\Delta\theta}$ points along $-\vct{g}$.

Two-parameter example: Gradients (grey) and negative gradients (teal) at a few points of a loss landscape.

Two-parameter example: Gradients (grey) and negative gradients (teal) at a few points of a loss landscape.

Thus, any learning algorithm that has us move in the direction of $-\vct{g}$ will have us move in the direction where the loss decreases most rapidly locally, i.e. the direction of steepest descent at our given point in parameter space.

For example, if we wanted to implement a simple fixed-step-size version of this, we could use the normalized gradient vector $\vct{\hat{g}} \coloneqq \vct{g}/\|\vct{g}\|$ and update our parameters using

But we already know from our single-parameter example that fixed-size updates fail to converge, so let's take inspiration from the single-parameter case and choose to make our step size proportional to the magnitude of the gradient instead, just as we chose to make it proportional to the magnitude of the derivative (= 1D gradient) in the single-parameter case. In this case, we do not even normalize $\vct{g}$ but simply multiply it by the learning rate $\epsilon$ directly, so the update rule simply becomes

This is known as Gradient Descent.

Here is what it might look like if things work well and we start from a specific point on the two-parameter example loss from the previous figures:

But wait! The elements of $\vct{g}$ are just the individual partial derivatives $\pder{L}{\theta_1}$, $\pder{L}{\theta_2}$, and so on, so this is basically the same result we would have gotten if we had simply applied our 1D derivative update rule to every single parameter and used the same learning rate $\epsilon$ for all parameters:

On the one hand, this is nice — the simplest possible extension of our 1D rule and the intuitive approach of moving in the direction of steepest descent in parameter space both lead us to the exact same result!

On the other hand, it might hint at trouble. After all, we observed that for the simple 1D case, the ideal learning rate close to the minimum was inversely proportional to the curvature of the loss, the intuition being that if we only use first-derivative information, we cannot know if the slope is large because we are far from the minimum (in which case we'd like to be greedy and take big steps) or because the loss is highly curved (in which case we'd like to be careful and take smaller ones).

In the single-parameter scenario, this was somewhat okay. It simply meant we'd have to choose a small-enough error rate to respect the loss surface's curvature along the single parameter dimension, and we would always have to find the learning-rate value empirically anyway since training is messy and the quadratic approximation only holds close to the minimum to begin with.

But once we have multiple parameters, the situation becomes trickier, as the curvature along different parameter-space dimensions might vary, even close to the minimum where the loss should be well-approximated by a quadratic. This means that, if we wanted to keep our update proportional to the "ideal step" in the direction of the minimum, we would have to use a different learning rate for each parameter since the mapping from "What's the slope here?" to "How far are we from the minimum?" will depend on the curvature along the direction in question.

Gradient descent, which uses the same learning rate for all parameters, does not do this. If we have one direction in parameter space along which the loss is highly curved, that curvature will limit the overall learning rate and slow down progress along the other directions. This degree to which such curvature disparities exist is referred to as the Conditioning of the optimization problem, with large disparities corresponding to poor conditioning.

In some sense, this is just a restatement of our earlier intuition that pursuing steepest descent will often fail to yield the shortest path to the minimum. For example, for asymmetric quadratic loss in a 2D parameter space, the steepest-descent direction will not point towards the minimum unless we happen to be on one of the principal axes of the loss, and descent will lead us onto a path containing unnecessarily many steps along the lower-curvature direction. This is exactly the sort of slow-down we are talking about, and, crucially, it cannot be helped by cranking up the learning rate because that would make training unstable along the high-curvature direction:

Unfortunately, to truly address this issue, we would have to take the different curvatures of the loss along different directions into account, meaning we would have to calculate all the second derivatives $\pdermix{L}{\theta_i}{\theta_j}$ of the loss with respect to the parameters. This is possible in theory, but for $n$ parameters, there will be $n^2$ such derivatives. For a model with billions of parameters, this is simply computationally prohibitive. In fact, it will be hard enough to even calculate the $n$ first derivatives forming the gradient.

Short of using second derivatives to fix the issue, the next-best thing we could try is to alleviate it.

One way to do this would be to tackle the problem at its root by trying to engineer a model whose loss exhibits better conditioning, i.e. exhibits a relatively narrow range of curvatures along different directions. Assuming we keep our NLL loss definition itself which we just tediously motivated, this translates into the vague goal of somehow tweaking our model architecture itself to ensure that it isn't far more sensitive to some parameters than others, as we might intuitively expect such sensitivity to translate into higher loss curvature along the corresponding direction in parameter space.

A second way of doing this would be to deviate from our steepest-descent strategy and to somehow adjust our algorithm to be greedier along suspected lower curvature directions, but there is no obvious way to determine those without using second derivatives and we have already established that improvements to learning algorithms will likely have to be made empirically based on some trial and error anyway.

So let's kick this can down the road and add these two thoughts to our list of problems for later:

In the meantime, let's stick with gradient descent, which, in spite of everything, has a lot going for it in that it is a highly intuitive approach guaranteed to converge if we can get close enough to the minimum to begin with.

To get a feel for gradient descent, conditioning and learning-rate sensitivity, here is a little widget showing various two-parameter loss landscapes near their minimum:

Of course, the loss even for a two-parameter model would not have to look anything like this at all far from the minimum, and we should generally expect the loss landscape to be extremely complicated. Also, the actual loss landscape for a large model would be completely unknown to us, meaning the real-life situation is much more akin to the "Hide Loss" scenario in the above widget, with the only information we have being the loss and derivative values at the exact points in parameter space we actually visited.

Essentially, optimizing a large model is like trying to descend a high-dimensional mountain of unknown shape, except the fog is so dense you can only see the space immediately surrounding your feet (bad) and you can teleport yourself instantly to wherever you like (good), at your own peril.

If that sounds difficult, that's because it is. Even having settled for gradient descent for now, pulling this off in practice is no small feat, especially when we are hoping to work with models that have billions of parameters and training datasets of trillions of tokens.

So before we can fully cross the "finding parameters" issue of our list of urgent problems, let's take a closer look at how we might be able to actually implement this GD approach in the next chapter.