Now that we have a model architecture which can process unknown contexts and generalize to new tokens on the output end, the obvious next step would be to scale up training, i.e. to train on much more data using a much larger context window.
The storage problem
Unfortunately, this is problematic.
Our recent progress hasn't changed the fact that for large context windows, each context in the training data (up to copies) will be unique, and the model still fundamentally has to store every single context+token pair from the training data in order to generalize from there.
Assuming we ultimately want to train our chatbot on much of what the internet has to offer and use very long context windows, this is strictly prohibitive. For example, for a dataset of about 10 trillion tokens and a context window of just 10,000 tokens, explicitly storing every single (unique) context+token pairing would require us to store about tokens, which is absurd, and even if we somehow came up with a way to avoid the redundancy of storing lots of overlapping contexts, the uniqueness of the contexts means we would at a bare minimum have to store the entirety of the training data if we wanted to keep all the n-gram information intact, which is neither smart nor feasible. And all of that isn't even taking into account the 10 trillion gigantic context embedding vectors we would either have to store or constantly construct in this scenario to keep pulling off our input-generalization trick.
In other words, the uniqueness of long contexts means the n-gram approach becomes completely infeasible for long context windows + large datasets. This is yet another outcome of the curse of dimensionality, which we will refer to as the Storage Problem in the following.
For long context windows, virtually all contexts in the training data become unique and the amount of information an n-gram model would have to store becomes proportional to the size of the dataset, making the n-gram approach infeasible for large context windows + datasets.
In some sense, this failure of the n-gram approach is unsurprising: The key idea underlying n-grams was to leverage simple statistical properties of the training data, and for long contexts, the statistics in questions become trivial, as each context tends to be followed by a single token only.
But regardless of how we frame the problem, we will have to overcome it, meaning we need to figure out how to build a model that requires less memory. At the same time, we still need our model to generate meaningful output, so we still need for the distribution it ends up using for any given context to be similar to the original one.
In other words, we are facing a Compression problem: We need to find a way for our model to produce approximately the right output (i.e. next-token distribution) for any given input (i.e. context) without storing so much stuff.
This seems like a pretty fundamental challenge and our specific situation of mapping text-form contexts to probability distributions over thousands of tokens is quite complicated, so keeping with our “always start simple” philosophy, let’s start with a simplified toy model instead and see if we can make some conceptual headway.
A toy model
Arguably, the simplest form the input to any model could take is that of a single number, and the simplest form the output could take is that of another number, so let’s start there and imagine we have a simple toy model that has learned to map some known real-number inputs to some other real numbers using a stored list — no contexts, no distributions, and no generalization strategies, just a finite number of known inputs and the corresponding outputs. We can visualize this as a simple table with two columns or as a set of discrete points in a graph:
A simple toy model mapping inputs to outputs .How can we redesign this toy model such that we don’t need to store all the individual ’s and ’s and it nonetheless approximates the desired well for any given input ?
The simplest way to reduce the amount of numbers stored would be to throw away some of the data points altogether, but this would cause the model to no longer recognize the input values of the discarded points and would defeat the entire purpose of gathering a lot of data to begin with. Not good.
Arguably, the next-simplest thing we could do would be to reduce the number of values we have to store by having some of the inputs share the same value, for example the mean value of the corresponding ’s. This would result in a bad approximation if the original ’s were very dissimilar but should work reasonably well if they were roughly equal to begin with. At the same time, it would save memory as long as whatever reference we need to store for each to associate it with the shared value takes up less space than the original 's (which would likely be the case, especially if the original 's require a high-precision data format to be stored accurately). So for any subset of data points with roughly equal values, we might save memory by compressing the individual ’s into a single number while still approximating the data well. Interesting.
A simple act of compression: Say we have some contiguous inputs with very similar outputs. If we choose to only store the average output value and assign it to all the inputs, we approximate the original data points while saving memory.Let’s stick with such a subset of data points for a moment. What about the values? Can we compress the corresponding inputs into fewer numbers as well? Our first instinct might be to employ the same strategy as for the values by assigning the same value to data points whose values are close to begin with, but this wouldn’t work — recall that the values represent the inputs our simple toy model recognizes, so collapsing several ’s into a single value would decrease the number of inputs the model recognizes.
Does this mean we strictly have to store all the individual ’s then?
Not necessarily. We don't actually care that the model remembers the specific inputs , we just care that it knows what to do with them, i.e. that it knows to assign . If there is no particular pattern the model can exploit to recognize an input as one of the ’s in question, having a stored list of all the ’s and explicitly checking if the input matches one of them is the only way to ensure this. But if there is a pattern and it isn’t too complicated, we can have the model exploit that pattern to recognize that a given input matches one of the ’s and save memory in the process.
For example, if our subset of input-output points with roughly equal values contains a group of contiguous data points falling on some interval , we could just have the model output whenever the input satisfies . With this method, the model will output when fed any of the ’s, but instead of storing the 20 individual values, we only have to store and the two interval boundaries and . We do accrue the additional computational cost of having to check whether falls between and , but given that our main concern is memory right now, this seems like a price worth paying.
Mathematically, this checking and mapping to can be thought of as having the model implement a piecewise function
While this is a somewhat minor adjustment codewise for this simple example, it is a seismic change on a conceptual level. Instead of having a model that is a glorified table and always looks up the outputs for specific inputs it recognizes, we now have a model that, at least on some interval, represents a continuous function.
A further improvement: Instead of assigning a constant value only to the specific it is based on, let's assign it to the corresponding interval of inputs , enabling the model to generalize to new inputs.This does not change the model output for the known inputs (which we were already assigning to anyway) on that interval, but it does change how the model treats unknown inputs — the model will now assign not only to the ’s we started out with, but also to any other inputs on . This might seem spooky at first, but it’s actually a boon, since would have been our best guess for the output some intermediate value on this interval should produce anyway!
In other words, defining a constant function over all possible inputs on an interval of near-constant data points like this automatically results in best-guess generalization to new inputs from that same interval. Whether this best guess is actually a good guess is a whole other question, of course — it entirely depends on whether the output should actually be independent of the input across the entire interval, or whether our "toy training data" (i.e. the specific input-output mappings that we are basing our approximation on) just happened to make it seem that way. Still, having a model that doesn’t flinch in the face of unknown inputs and simply produces a reasonable output guess seems pretty awesome.
That being said, using a constant function will only get us so far. What about intervals where the desired outputs aren’t roughly constant? Well, again, if the values are completely random and don't follow any pattern we can exploit, there is not much we can do. But if there is some pattern we can identify, even if it is more complicated than “the points are roughly equal”, we should be able to exploit it by finding an appropriate function.
Do the desired output values roughly fall on a line on some interval? Then we could perform simple linear regression and have the model use the line of best fit to calculate its output for that interval, which would only require storing a slope and an offset in addition to the two interval boundaries, no matter how many points we are approximating. Are there other intervals where the desired output values grow quadratically? Shrink exponentially? Oscillate sinusoidally? Then we could define parameterized versions of the corresponding functions, fit them to the data, and use the resulting fits to calculate the outputs for those intervals.
If we can consistently pull this off, we will be left with a compressed model that necessitates more computation but requires far less storage and comes with in-built generalization.
The best way to achieve compression + generalization is for our model to be a function which approximates the functional relationship underlying the data. For a tiny toy dataset, we can manually identify suitable functional forms + subintervals and fit our model piece-by-piece, but that's infeasible for larger datasets + models.Of course, identifying all these intervals and behaviors manually isn’t feasible for gigantic datasets, so the best we might hope to do is come up with an extremely flexible type of function that can approximate a broad range of other functions very well, and then somehow fit that to the data.
A simple way of doing this for our 1D toy example would be to divide the input space into intervals and to fit a line to the data points on each interval. This results in a function consisting of linear segments:
Given enough segments, this type of function should — in principle — be able to approximate any finite amount of data points (and even any continuous function on a finite interval) arbitrarily well.
A more generalizable approach: If we use enough intervals, we can approximate pretty much anything with a bunch of line segments. But using predetermined interval boundaries is inefficient, and discontinuous jumps hurt generalization.If we simply chose fixed interval boundaries ahead of time, each piece of this function would be completely independent, and we would be able to perform this fit in a fairly straightforward segment-by-segment fashion, without having to optimize a gazillion parameters at once.
However, using predetermined interval boundaries is clearly suboptimal in terms of compression. We would ideally like to use more of our segments (and, by extension, parameters) on parts of input space where the unknown target function (and, by extension, the training data) exhibits more intricate variations. This is a pretty fundamental piece of intuition we will return to down the line.
Memory-efficient function approximation requires a model that can flexibly allocate more of its parameter budget to regions of input space where the target function exhibits more complexity.
As a result, we have to give up the convenience of using predetermined interval boundaries, meaning we need to include the boundaries in the list of parameters the model needs to learn.
But the graph above reveals another issue that would persist even with smartly chosen interval boundaries: Fitting all the individual segments independently results in discontinuous jumps. This does not hurt how well our parameterized function matches the data, but it would likely hurt generalization. After all, we might intuitively expect whatever real-world input-output relationship gave rise to this data to be continuous, and even if it weren't, the exact positioning of each interval boundary between the surrounding data points according to our fit would be completely arbitrary and unlikely to match that of any real discontinuous jump.
To address this, we should define our model function in such a way that it is guaranteed to be continuous. For this type of piecewise linear function, this is easily done by ensuring we always start our line segment for the interval at the value where the previous segment left off:
A good approach: Defining a flexible, continuous, parameterized function, which can expend more parameters where the data + target function show more intricacy.The potential benefits of defining our model this way are enormous: If at least some of the data points fall on linear-ish segments and we don’t overdo it with the number of parameters, the resulting function should require far less storage than the individual points and would thus represent a compressed representation of the mappings , effectively addressing the storage problem. Even better, if these linear segments actually happen to capture an underlying input-output relationship that extends beyond the specific data points , using this function as our model will get us reasonable generalization to new inputs for free.
But crucially, neither of these benefits are exclusive to the linear-segment method or the 1D toy scenario we have been considering. As long as there is some non-random relationship underlying our training data, it should in principle be possible to exploit that and build a memory-efficient model which generalizes well, so long as we can figure out how to define a suitable parameterized function and how to have it learn the appropriate parameter values.
This approach is so powerful that it underlies virtually all of modern machine learning. It is known as Function Approximation.
New paradigm! Our model should be a parameterized function. The potential benefits are huge:
-
Compression: The size of our model will simply be determined by how we define our function, not by the size of our training data. And if our chosen parameterization efficiently exploits the patterns in the training data, it will be vastly more memory-efficient than an n-gram model.
-
Generalization: The model will automatically generalize to new inputs. If the model function approximates not just the training data but the functional relationship underlying the data, this generalization will actually be good.
Of course, these are some big ifs and it's not at all obvious that we should be able to pull this off.†In fact, virtually nobody expected LLMs to work as well as they do even 10-15 years ago, precisely because approximating the extremely complex unknown functional relationship underlying natural-language generation, based on very sparse data no less, seems incredibly daunting, to say the least.[note] Doing so is pretty much the central challenge of AI, and most efforts in machine learning are either about trying to come up with more ingenious ways of defining the parameterized model in question or about trying to come up with better ways of optimizing its parameter values.
Of course, we do not have to go all the way to the state of the art in one fell swoop, but if we want to make our function-approximation dreams a reality, we will have to put our other problems on hold and come up with at least half-decent solutions to both of these challenges:
Of these two issues, the more natural one to start with is how to define our parameterized model. Let's make some progress towards that by making the jump from our toy model to our actual chatbot scenario and seeing where that leaves us.
The LM scenario
Our chatbot scenario differs quite significantly from the simple toy scenario on both the input and the output end.
Instead of being fed a single number, it will have to process a context of (many) tokens as input, and instead of producing a single number as output, it will have to produce a probability distribution over all the tokens in the vocabulary.


Translating contexts into numbers
Let’s start with the former — contexts are not numbers. To use the approach we just outlined, we will have to somehow bridge that gap and translate our contexts into numbers that we can use as inputs of our parameterized function.
Of course, we already have a way of doing this, namely using our concatenated token embeddings. So the question is whether we should keep using that approach or try to come up with something else now that we are operating under a new paradigm.
Thinking about it from scratch, a minimal requirement our mapping from contexts to numerical inputs should meet is that it should never destroy information by mapping multiple contexts to the same thing. This could be achieved in many different ways, including by coming up with a clever way to map any possible context to a unique single number, in which case the input to our parameterized model would just be that: a single number, just like in our toy scenario.
However, mindlessly mapping contexts to numerical inputs while only preserving uniqueness will effectively lead to the contexts being randomly strewn across input space. This is exactly the scenario we said would be horrible for both compression and generalization, which inherently rely on the exploitation of patterns and benefit from nearby inputs being associated with similar outputs. After all, a smooth function is much easier to approximate than a chaotic one.
Randomly scrambled data points are a nightmare to approximate efficiently. If the data varies smoothly, it's much easier.So ideally, whatever embeddings we do use should ensure that our target function is reasonably smooth by mapping contexts requiring similar outputs to similar points in input space.
Embeddings map similar contexts to similar points in input space. Compared to more arbitrary mappings, this should make the function we are seeking to approximate simpler and benefit both compression and generalization.
Thus, even though we are no longer aiming to directly compute the dot-product similarity of different contexts, we still need similarity-based context embeddings.
Such similarity-capturing embeddings will necessarily have many elements, so this means one big difference to our toy scenario is that our parameterized model will have to take in not just a single number, but many numbers, namely however many elements our context embedding vectors have. As usual when moving to higher dimensions, this will make things harder to visualize, but it shouldn't be an issue otherwise. We will just have to define our paramterized function accordingly.
As for how to create these embeddings, recall that in the previous chapter, we decided to build them by concatenating the embeddings of the individual tokens, but we kind of punted on the question of how to get those token embeddings in the first place and simply assumed some semi-usable token embeddings had been found via manual scoring. Of course, manual scoring is neither optimal nor scaleable, which is why we have this item on our to-do list:
However, our new list already contains a big automization problem anyway, namely the challenge of automatically determining the parameter values of our function, and in an operational sense, there is no difference between the numbers that we have so far considered the parameters of that function and the feature scores that define the embeddings for all the individual tokens — they are all just numbers used to calculate the output of the model.
So we can catch two birds with one stone by simply thinking of the elements of our token embeddings as parameters of our model as well! If we can figure out how to optimize our model parameters, the model will then simply learn better embeddings while also learning to translate them into better outputs.
The elements of our token embeddings are just fixed numbers our model uses to calculate some output. Let's treat them just like any other function parameter and attempt to have our model learn their values.
Of course, simply saying we consider the token-embedding elements parameters does not mean we have truly solved the issue, since we have not yet come up with a suitable method to determine our model's parameter values. But that's a much bigger issue we have to solve anyway, so we have practically administered away our smaller problem by making it part of a bigger one:
We decided to make the token-embedding scores parameters, so if we can figure out a general method to find good parameter values, our model will learn them automatically.Interestingly, one thing we will lose if we can pull off this automated learning of embeddings is the interpretability of our features — the model will simply learn to assign each token an appropriate embedding vector by optimizing the individual elements of that vector. We will have to define how many elements that vector should have by picking the number of parameters used to embed each token, and we could still determine how similar the learned embeddings for various tokens are by looking at the alignment of these vectors, but we will no longer have a way of knowing the meaning of each individual vector component.
This might feel strange, but it is not a problem in terms of the model’s performance — the model didn’t know or care about the meaning of the individual feature scores even when we had come up with them manually. It only cares about the numbers themselves, as it is these numbers that impact its output.
With this context-embedding stage included, our updated model architecture looks like this:


As before, creating context embeddings by concatenating token embeddings does bring about the additional challenge that the size of the context embedding vectors depends on the length of the context. When we only used the embeddings for similarity-based generalization (see chapter 4), we said we could effectively deal with this by zero-padding or truncating the context embedding vectors, since the dot products in question would then automatically guarantee that we were effectively considering context embeddings of the appropriate length only.
Now that we want to feed the embeddings to some more complicated mathematical function , things are a bit trickier, as that function will presumably expect a fixed number of inputs and it is not clear that simply zero-padding the embeddings of shorter context to always reach the same, maximal length will not cause problems.
However, this issue does not feel as urgent as the others we are working on and even getting a function-approximation-based model to work well for a fixed context length would already be a huge leap forward, so let's assume we simply zero-pad our context embeddings at earlier-token positions to always get embeddings of the same size for now, and add this as another problem to our list of things to investigate more deeply down the road.
For now, let's move on to the output end of things and see if we can glean any more insights on how to define our language-model function from there.
Predicting next-token probabilities
In the case of our toy-model, the output simply took the form of a single number , but for our actual model, it must take the form of a probability distribution over all the tokens in the vocabulary, meaning for a vocabulary of size , it must consist of numbers on the interval which sum to .
Recall that for the embeddings-assisted n-gram model, we got these probabilities in two steps. First, we set each probability equal to the probability that the given token followed the context in the training data. Even though we didn’t think of it in those terms, this was essentially the equivalent of “fitting our parameterized model function to the data”, with the parameters of our model simply being the individual next-token probabilities and the "fitting" being the single step of setting them equal to the empirical training-data probabilities. Then, in order to generalize beyond the data and enable the chatbot to create original token sequences not contained in the training text, we used the token embeddings + temperature + softmax recipe to partially spread these initial probabilities to similar tokens across the vocabulary.
Switching to a function-approximation-based model does not affect the need for the first step. The training data is all we have to determine our parameter values, so we will still have to make our model learn to generate probabilities that reflect the empirical next-token probabilities from that data.
For the second step, the situation is slightly more subtle. We still want our model to be able to produce original output even when fed a training-data context, but unless our model function is able to fit the training data exactly, it will actually do so by default – even if the given context appears only once in the training data and there is a single token with empirical probability 1, a model which does not match this data point perfectly will assign some probability to other tokens. If we could somehow ensure that this generalization to new tokens due to imperfect fitting of the data ends up spreading probability in a reasonable fashion, we might not need an additional probability-spreading stage at all.
However, this sounds pretty complicated and demands further investigation, so let's just add this to our list of things to revisit down the road for now:
In the meantime, let's keep the token similarity + temperature + softmax strategy from earlier to spread our "fitted" model probabilities across the vocabulary.


The next thing we have to figure out is how to define the parameterized function producing those original probabilities, which we have already established should take in the context embedding vector and spit out next-token probabilities.
A conceptually simple approach would be to use different functions, each of which independently maps the context embedding vector to a token score signifying how well a given token matches the context, and to then use a softmax to map these scores to probabilities.
Unlike the softmax we use to spread probabilities across the vocabulary in the final stage, this first softmax shouldn't require a temperature hyperparameter, since all such a parameter effectively does is rescale the scores fed to the softmax. In the generalization stage, this rescaling is useful, because the token embeddings and resulting similarity scores are fixed, and we would like a manual, post-training dial that determines how aggresively the model spreads probability across the vocabulary while generating output. In the main part of our architecture (which will be subject to parameter optimization during training), it is redundant, because the model will have to learn appropriate scales for all the token scores anyway and whatever functionality a softmax-temperature hyperparameter would add should already be absorbed in the individual token-score functions to begin with.


But using such independent functions would be very inefficient — the probability that a specific token should be next necessarily depends on a large number of abstract, high-level features of the context, such as the meaning of the individual tokens comprising it, its grammatical structure, the literary style in which it is written, and so on. Even though a mathematical function mapping the context embedding vector to some next-token probability does not “think” in those specific terms, the calculation it represents has to somehow take into account all the information we think of in these terms in order to properly predict the next token.
Since this information is relevant to all the next-token probabilities, having separate functions for the individual probabilities would be wasteful, as all these functions would partially have to do (and learn how to do) the same thing. It should be much better to avoid this redundancy by using a single function that learns to extract all the necessary information from the context embeddings, and to only differentiate between various next-token candidates afterwards.†Note that if we wanted to, we could still think of this as having different functions predicting the individual probabilities, except that these functions now share most of their parameters and only differ in the final stage.[note]


It's not obvious at all what type of function would theoretically be optimal to calculate the token scores from the extracted context feature scores, but given that our vocabulary is huge and we are going to need one such function for every single token it contains, we desperately need these functions to be extremely simple in order to avoid a prohibitively large model, so let's simply choose a very basic type of function and hope that we can get away with that.
Perhaps the simplest non-trivial way of turning a bunch of numbers into a single one is via a weighted sum, i.e. to multiply each of the numbers by a Weight and to add up the results. Applying this to our context feature scores yields
Here, each weight is a parameter of our model whose value will have to be learned. If we think of the various context feature scores and weights as elements of vectors, this weighted sum can simply be rewritten as the dot product of the weight vector and the feature-score vector:
If our main function maps the context to features, each of the weight vectors will have parameters, leaving us with parameters total for this stage.
As it stands, this definition allows each token score to depend on all the various context features in distinct ways, but it does not allow the model any control over the context scores which is independent of those features. For example, there is no obvious way to make any token more probable regardless of the context features, even though we would intuitively say that some tokens should be more probable by default than others.
The simplest way to fix this is to add a Bias parameter , which will bias token scores to be larger (if positive) or smaller (if negative):
With this minimalistic mapping of context-feature scores to token scores, our overall architecture looks like this:


At this point, all that's left is the biggest challenge of all, namely to figure out how to define the main function at the heart of our model, which needs to be capable of mapping context embedding vectors to suitable high-level feature scores from which the next-token probabilities can be derived with relative ease.
We have no idea what the unknown target function pulling that off looks like, so we will need to come up with a type of parameterized function that can approximate an extremely broad range of functions for different sets of parameter values in order to have any hope that it might be able to get the job done, much like we had to come up with a very flexible type of function in our far simpler 1D toy scenario.
This seems like a massive challenge, so let's tackle it in the next chapter. Once we have settled on a suitable function definition, we can then move on to the remaining one of our two big challenges and try to figure out how to actually determine the parameter values.
